Refactored SU Flavor logic

The Flavor object will no longer require an SU name. The SU type will be determined only from PCI requests.

src/openstack_billing_db/billing.py

@@ -107,20 +107,11 @@ def collect_invoice_data_from_openstack(database, billing_start, billing_end, ra
                 su = i.service_units
                 su_hours = runtime_hours * su
 
-                if i.service_unit_type == "CPU":
-                    invoice.cpu_su_hours += su_hours
-                elif i.service_unit_type == "GPU A100SXM4":
-                    invoice.gpu_a100sxm4_su_hours += su_hours
-                elif i.service_unit_type == "GPU A100":
-                    invoice.gpu_a100_su_hours += su_hours
-                elif i.service_unit_type == "GPU V100":
-                    invoice.gpu_v100_su_hours += su_hours
-                elif i.service_unit_type == "GPU K80":
-                    invoice.gpu_k80_su_hours += su_hours
-                elif i.service_unit_type == "GPU A2":
-                    invoice.gpu_a2_su_hours += su_hours
+                su_hour_attr = f"{i.service_unit_type}_su_hours"
+                if hasattr(invoice, su_hour_attr):
+                    invoice.__setattr__(su_hour_attr, invoice.__getattribute__(su_hour_attr) + su_hours)
                 else:
-                    raise Exception("Invalid flavor.")
+                    raise Exception(f"Invalid flavor {i.service_unit_type}.")
 
         invoices.append(invoice)
     return invoices

src/openstack_billing_db/model.py

@@ -14,47 +14,21 @@
 @dataclass()
 class Flavor(object):
     id: int
-    name: str
+    service_unit_type: str
     vcpus: int
     memory: int
     storage: int
+    gpu_count: int = 0
 
     @property
     def service_units(self):
-        if "gpu" not in self.name:
-            # 1 CPU SU = 0 GPU, 1 CPU, 4 GB RAM, 20 GB
-            return int(
+        # 1 CPU SU = 0 GPU, 1 CPU, 4 GB RAM, 20 GB
+        return self.gpu_count or int(
                 max(
                     self.vcpus,
                     self.memory / 4096,
                 )
             )
-        else:
-            # The flavor for 2 SUs of V100 is inconsistent with previous
-            # naming scheme.
-            if self.name == "gpu-su-v100.1m":
-                return 2
-
-            split = self.name.split(".")
-            return int(split[-1])
-
-    @property
-    def service_unit_type(self):
-        if "gpu" not in self.name:
-            return "CPU"
-        elif "a100-sxm4" in self.name:
-            return "GPU A100SXM4"
-        elif "a100" in self.name and "sxm4" not in self.name:
-            return "GPU A100"
-        elif "v100" in self.name:
-            return "GPU V100"
-        elif "k80" in self.name:
-            return "GPU K80"
-        elif "gpu-su-a2" in self.name:
-            return "GPU A2"
-        else:
-            # New GPU type that we need to take into account.
-            raise Exception()
 
 
 @dataclass()
@@ -240,7 +214,7 @@ def get_instances(self, project) -> list[Instance]:
                 logger.warning(
                     f"Could not parse pci requests from instance {instance}."
                 )
-            su_name = "cpu"
+            su_type = "cpu"
             if pci_info:
                 # The PCI Requests column of the database contains a JSON
                 # object with the below format. If the instance has an
@@ -266,14 +240,19 @@ def get_instances(self, project) -> list[Instance]:
                     raise Exception(f"Invalid pci_name {pci_name}.")
 
                 count = pci_info[0]["count"]
-                su_name = f"gpu-{pci_name}.{count}"
+                if count == "1m":
+                    count = 2
+                else: 
+                    count = int(count)
+                su_type = f"gpu_{pci_name}"
 
             flavor = Flavor(
                 id=instance["instance_type_id"],
-                name=su_name,
+                service_unit_type=su_type,
                 vcpus=instance["vcpus"],
                 memory=instance["memory_mb"],
                 storage=20,
+                gpu_count=count
             )
 
             i = Instance(

src/openstack_billing_db/tests/unit/test_instance.py

@@ -3,7 +3,7 @@
 
 from openstack_billing_db.model import Instance, InstanceEvent, Flavor
 
-FLAVORS = {1: Flavor(id=1, name="TestFlavor", vcpus=1, memory=4096, storage=10)}
+FLAVORS = {1: Flavor(id=1, service_unit_type="TestFlavor", vcpus=1, memory=4096, storage=10)}
 
 MINUTE = 60
 HOUR = 60 * MINUTE