Run pre-commit for cudo

dstackai · Mar 13, 2024 · abf59ba · abf59ba
1 parent 9a64a2c
commit abf59ba
Show file tree

Hide file tree

Showing 3 changed files with 154 additions and 115 deletions.
diff --git a/src/gpuhunt/_internal/default.py b/src/gpuhunt/_internal/default.py
@@ -21,7 +21,7 @@ def default_catalog() -> Catalog:
     for module, provider in [
         ("gpuhunt.providers.tensordock", "TensorDockProvider"),
         ("gpuhunt.providers.vastai", "VastAIProvider"),
-        ("gpuhunt.providers.cudo", "CudoProvider")
+        ("gpuhunt.providers.cudo", "CudoProvider"),
     ]:
         try:
             module = importlib.import_module(module)

diff --git a/src/gpuhunt/providers/cudo.py b/src/gpuhunt/providers/cudo.py
@@ -27,19 +27,23 @@ class CudoProvider(AbstractProvider):
     NAME = "cudo"
 
     def get(
-            self, query_filter: Optional[QueryFilter] = None, balance_resources: bool = True
+        self, query_filter: Optional[QueryFilter] = None, balance_resources: bool = True
     ) -> List[RawCatalogItem]:
         offers = self.fetch_offers(query_filter, balance_resources)
         return sorted(offers, key=lambda i: i.price)
 
-    def fetch_offers(self, query_filter: Optional[QueryFilter], balance_resources) -> List[RawCatalogItem]:
+    def fetch_offers(
+        self, query_filter: Optional[QueryFilter], balance_resources
+    ) -> List[RawCatalogItem]:
         machine_types = self.list_vm_machine_types()
         if query_filter is not None:
             return self.optimize_offers(machine_types, query_filter, balance_resources)
         else:
             offers = []
             for machine_type in machine_types:
-                optimized_specs = optimize_offers_with_gpu(QueryFilter(), machine_type, balance_resources=False)
+                optimized_specs = optimize_offers_with_gpu(
+                    QueryFilter(), machine_type, balance_resources=False
+                )
                 raw_catalogs = [get_raw_catalog(machine_type, spec) for spec in optimized_specs]
                 offers.append(raw_catalogs)
             return list(chain.from_iterable(offers))
@@ -56,29 +60,43 @@ def list_vm_machine_types() -> list[dict]:
         resp.raise_for_status()
 
     @staticmethod
-    def optimize_offers(machine_types: list[dict], q: QueryFilter, balance_resource) -> List[RawCatalogItem]:
+    def optimize_offers(
+        machine_types: list[dict], q: QueryFilter, balance_resource
+    ) -> List[RawCatalogItem]:
         offers = []
-        if any(condition is not None for condition in
-               [q.min_gpu_count, q.max_gpu_count, q.min_total_gpu_memory, q.max_total_gpu_memory,
-                q.min_gpu_memory, q.max_gpu_memory, q.gpu_name]):
+        if any(
+            condition is not None
+            for condition in [
+                q.min_gpu_count,
+                q.max_gpu_count,
+                q.min_total_gpu_memory,
+                q.max_total_gpu_memory,
+                q.min_gpu_memory,
+                q.max_gpu_memory,
+                q.gpu_name,
+            ]
+        ):
             # filter offers with gpus
-            gpu_machine_types = [vm for vm in machine_types if vm['maxGpuFree'] != 0]
+            gpu_machine_types = [vm for vm in machine_types if vm["maxGpuFree"] != 0]
             for machine_type in gpu_machine_types:
                 machine_type["gpu_name"] = gpu_name(machine_type["gpuModel"])
                 machine_type["gpu_memory"] = get_memory(machine_type["gpu_name"])
-                if not is_between(machine_type["gpu_memory"], q.min_gpu_memory,
-                                  q.max_total_gpu_memory):
+                if not is_between(
+                    machine_type["gpu_memory"], q.min_gpu_memory, q.max_total_gpu_memory
+                ):
                     continue
                 if q.gpu_name is not None and machine_type["gpu_name"].lower() not in q.gpu_name:
                     continue
                 cc = get_compute_capability(machine_type["gpu_name"])
-                if not cc or not is_between(cc, q.min_compute_capability, q.max_compute_capability):
+                if not cc or not is_between(
+                    cc, q.min_compute_capability, q.max_compute_capability
+                ):
                     continue
                 optimized_specs = optimize_offers_with_gpu(q, machine_type, balance_resource)
                 raw_catalogs = [get_raw_catalog(machine_type, spec) for spec in optimized_specs]
                 offers.append(raw_catalogs)
         else:
-            cpu_only_machine_types = [vm for vm in machine_types if vm['maxGpuFree'] == 0]
+            cpu_only_machine_types = [vm for vm in machine_types if vm["maxGpuFree"] == 0]
             for machine_type in cpu_only_machine_types:
                 optimized_specs = optimize_offers_no_gpu(q, machine_type, balance_resource)
                 raw_catalogs = [get_raw_catalog(machine_type, spec) for spec in optimized_specs]
@@ -103,11 +121,11 @@ def get_raw_catalog(machine_type, spec):
         instance_name=machine_type["machineType"],
         location=machine_type["dataCenterId"],
         spot=False,
-        price=(round(float(machine_type["vcpuPriceHr"]["value"]), 5) * spec["cpu"]) +
-              (round(float(machine_type["memoryGibPriceHr"]["value"]), 5) * spec["memory"]) +
-              (round(float(machine_type["gpuPriceHr"]["value"]), 5) * spec.get("gpu", 0)) +
-              (round(float(machine_type["minStorageGibPriceHr"]["value"]), 5) * spec["disk_size"]) +
-              (round(float(machine_type["ipv4PriceHr"]["value"]), 5)),
+        price=(round(float(machine_type["vcpuPriceHr"]["value"]), 5) * spec["cpu"])
+        + (round(float(machine_type["memoryGibPriceHr"]["value"]), 5) * spec["memory"])
+        + (round(float(machine_type["gpuPriceHr"]["value"]), 5) * spec.get("gpu", 0))
+        + (round(float(machine_type["minStorageGibPriceHr"]["value"]), 5) * spec["disk_size"])
+        + (round(float(machine_type["ipv4PriceHr"]["value"]), 5)),
         cpu=spec["cpu"],
         memory=spec["memory"],
         gpu_count=spec.get("gpu", 0),
@@ -124,15 +142,17 @@ def optimize_offers_with_gpu(q: QueryFilter, machine_type, balance_resources) ->
     gpu_range = get_gpu_range(q.min_gpu_count, q.max_gpu_count, machine_type["maxGpuFree"])
     memory_range = get_memory_range(q.min_memory, q.max_memory, machine_type["maxMemoryGibFree"])
     min_vcpu_per_memory_gib = machine_type.get("minVcpuPerMemoryGib", 0)
-    max_vcpu_per_memory_gib = machine_type.get("maxVcpuPerMemoryGib", float('inf'))
+    max_vcpu_per_memory_gib = machine_type.get("maxVcpuPerMemoryGib", float("inf"))
     min_vcpu_per_gpu = machine_type.get("minVcpuPerGpu", 0)
-    max_vcpu_per_gpu = machine_type.get("maxVcpuPerGpu", float('inf'))
+    max_vcpu_per_gpu = machine_type.get("maxVcpuPerGpu", float("inf"))
     unbalanced_specs = []
     for cpu in cpu_range:
         for gpu in gpu_range:
             for memory in memory_range:
                 # Check CPU/memory constraints
-                if not is_between(cpu, memory * min_vcpu_per_memory_gib, memory * max_vcpu_per_memory_gib):
+                if not is_between(
+                    cpu, memory * min_vcpu_per_memory_gib, memory * max_vcpu_per_memory_gib
+                ):
                     continue
 
                 # Check CPU/GPU constraints
@@ -145,29 +165,38 @@ def optimize_offers_with_gpu(q: QueryFilter, machine_type, balance_resources) ->
 
     # If resource balancing is required, filter combinations to meet the balanced memory requirement
     if balance_resources:
-        memory_balanced = [spec for spec in unbalanced_specs
-                           if spec["memory"] ==
-                           get_balanced_memory(spec["gpu"], machine_type["gpu_memory"], q.max_memory)]
+        memory_balanced = [
+            spec
+            for spec in unbalanced_specs
+            if spec["memory"]
+            == get_balanced_memory(spec["gpu"], machine_type["gpu_memory"], q.max_memory)
+        ]
         balanced_specs = memory_balanced
         # Add disk
-        balanced_specs = [{"cpu": spec["cpu"],
-                           "memory": spec["memory"],
-                           "gpu": spec["gpu"],
-                           "disk_size": get_balanced_disk_size(machine_type["maxStorageGibFree"],
-                                                               spec["memory"],
-                                                               spec["gpu"] * machine_type["gpu_memory"],
-                                                               q.max_disk_size, q.min_disk_size)}
-                          for spec in balanced_specs]
+        balanced_specs = [
+            {
+                "cpu": spec["cpu"],
+                "memory": spec["memory"],
+                "gpu": spec["gpu"],
+                "disk_size": get_balanced_disk_size(
+                    machine_type["maxStorageGibFree"],
+                    spec["memory"],
+                    spec["gpu"] * machine_type["gpu_memory"],
+                    q.max_disk_size,
+                    q.min_disk_size,
+                ),
+            }
+            for spec in balanced_specs
+        ]
         # Return balanced combinations if any; otherwise, return all combinations
         return balanced_specs
 
     disk_size = q.min_disk_size if q.min_disk_size is not None else MIN_DISK_SIZE
     # Add disk
-    unbalanced_specs = [{"cpu": spec["cpu"],
-                         "memory": spec["memory"],
-                         "gpu": spec["gpu"],
-                         "disk_size": disk_size}
-                        for spec in unbalanced_specs]
+    unbalanced_specs = [
+        {"cpu": spec["cpu"], "memory": spec["memory"], "gpu": spec["gpu"], "disk_size": disk_size}
+        for spec in unbalanced_specs
+    ]
     return unbalanced_specs
 
 
@@ -178,75 +207,82 @@ def optimize_offers_no_gpu(q: QueryFilter, machine_type, balance_resource) -> Li
 
     # Cudo Specific Constraints
     min_vcpu_per_memory_gib = machine_type.get("minVcpuPerMemoryGib", 0)
-    max_vcpu_per_memory_gib = machine_type.get("maxVcpuPerMemoryGib", float('inf'))
+    max_vcpu_per_memory_gib = machine_type.get("maxVcpuPerMemoryGib", float("inf"))
 
     unbalanced_specs = []
     for cpu in cpu_range:
         for memory in memory_range:
             # Check CPU/memory constraints
-            if not is_between(cpu, memory * min_vcpu_per_memory_gib, memory * max_vcpu_per_memory_gib):
+            if not is_between(
+                cpu, memory * min_vcpu_per_memory_gib, memory * max_vcpu_per_memory_gib
+            ):
                 continue
             # If all constraints are met, append this combination
             unbalanced_specs.append({"cpu": cpu, "memory": memory})
 
     # If resource balancing is required, filter combinations to meet the balanced memory requirement
     if balance_resource:
-        cpu_balanced = [spec for spec in unbalanced_specs
-                        if spec["cpu"] ==
-                        get_balanced_cpu(spec["memory"], q.max_memory)]
+        cpu_balanced = [
+            spec
+            for spec in unbalanced_specs
+            if spec["cpu"] == get_balanced_cpu(spec["memory"], q.max_memory)
+        ]
 
         balanced_specs = cpu_balanced
         # Add disk
         disk_size = q.min_disk_size if q.min_disk_size is not None else MIN_DISK_SIZE
-        balanced_specs = [{"cpu": spec["cpu"],
-                           "memory": spec["memory"],
-                           "disk_size": disk_size}
-                          for spec in balanced_specs]
+        balanced_specs = [
+            {"cpu": spec["cpu"], "memory": spec["memory"], "disk_size": disk_size}
+            for spec in balanced_specs
+        ]
         # Return balanced combinations if any; otherwise, return all combinations
         return balanced_specs
 
     disk_size = q.min_disk_size if q.min_disk_size is not None else MIN_DISK_SIZE
     # Add disk
-    unbalanced_specs = [{"cpu": spec["cpu"],
-                         "memory": spec["memory"],
-                         "gpu": 0,
-                         "disk_size": min_none(machine_type["maxStorageGibFree"], disk_size)}
-                        for spec in unbalanced_specs]
+    unbalanced_specs = [
+        {
+            "cpu": spec["cpu"],
+            "memory": spec["memory"],
+            "gpu": 0,
+            "disk_size": min_none(machine_type["maxStorageGibFree"], disk_size),
+        }
+        for spec in unbalanced_specs
+    ]
     return unbalanced_specs
 
 
 def get_cpu_range(min_cpu, max_cpu, max_cpu_free):
     cpu_range = range(
         min_cpu if min_cpu is not None else MIN_CPU,
-        min(max_cpu if max_cpu is not None else max_cpu_free,
-            max_cpu_free) + 1
+        min(max_cpu if max_cpu is not None else max_cpu_free, max_cpu_free) + 1,
     )
     return cpu_range
 
 
 def get_gpu_range(min_gpu_count, max_gpu_count, max_gpu_free):
     gpu_range = range(
         min_gpu_count if min_gpu_count is not None else 1,
-        min(max_gpu_count if max_gpu_count is not None else max_gpu_free,
-            max_gpu_free) + 1
+        min(max_gpu_count if max_gpu_count is not None else max_gpu_free, max_gpu_free) + 1,
     )
     return gpu_range
 
 
 def get_memory_range(min_memory, max_memory, max_memory_gib_free):
     memory_range = range(
         int(min_memory) if min_memory is not None else MIN_MEMORY,
-        min(int(max_memory) if max_memory is not None else max_memory_gib_free,
-            max_memory_gib_free) + 1
+        min(
+            int(max_memory) if max_memory is not None else max_memory_gib_free, max_memory_gib_free
+        )
+        + 1,
     )
     return memory_range
 
 
 def get_balanced_memory(gpu_count, gpu_memory, max_memory):
     return min_none(
-        round_up(
-            RAM_PER_VRAM * gpu_memory * gpu_count, RAM_DIV),
-        round_down(max_memory, RAM_DIV))
+        round_up(RAM_PER_VRAM * gpu_memory * gpu_count, RAM_DIV), round_down(max_memory, RAM_DIV)
+    )
 
 
 def get_balanced_cpu(memory, max_cpu):
@@ -262,7 +298,9 @@ def get_balanced_disk_size(available_disk, memory, total_gpu_memory, max_disk_si
             available_disk,
             max(memory, total_gpu_memory),
             max_disk_size,
-        ), min_disk_size)
+        ),
+        min_disk_size,
+    )
 
 
 def gpu_name(name: str) -> Optional[str]:
@@ -311,5 +349,5 @@ def max_none(*args: Optional[T]) -> T:
     "RTX A6000": "A6000",
     "NVIDIA A40": "A40",
     "NVIDIA V100": "V100",
-    "RTX 3080": "RTX3080"
-}
+    "RTX 3080": "RTX3080",
+}