From abf59baa5ab817fd16956d3937e0efffa5c4a9fb Mon Sep 17 00:00:00 2001
From: Bihan  Rana <bihan@Bihans-MacBook-Pro.local>
Date: Wed, 13 Mar 2024 12:28:20 +0545
Subject: [PATCH] Run pre-commit for cudo

---
 src/gpuhunt/_internal/default.py |   2 +-
 src/gpuhunt/providers/cudo.py    | 162 +++++++++++++++++++------------
 src/tests/providers/test_cudo.py | 105 ++++++++++----------
 3 files changed, 154 insertions(+), 115 deletions(-)

diff --git a/src/gpuhunt/_internal/default.py b/src/gpuhunt/_internal/default.py
index d437789..3ea430b 100644
--- 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
index 5112d18..b891aa5 100644
--- 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,48 +207,55 @@ 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
 
@@ -227,8 +263,7 @@ def get_cpu_range(min_cpu, max_cpu, max_cpu_free):
 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
 
@@ -236,17 +271,18 @@ def get_gpu_range(min_gpu_count, max_gpu_count, max_gpu_free):
 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"
-}
\ No newline at end of file
+    "RTX 3080": "RTX3080",
+}
diff --git a/src/tests/providers/test_cudo.py b/src/tests/providers/test_cudo.py
index 9d0a8ef..852079e 100644
--- a/src/tests/providers/test_cudo.py
+++ b/src/tests/providers/test_cudo.py
@@ -1,4 +1,3 @@
-from itertools import chain
 from typing import List
 
 import pytest
@@ -17,54 +16,46 @@
 
 @pytest.fixture
 def machine_types() -> List[dict]:
-    return [{
-        "dataCenterId": "br-saopaulo-1",
-        "machineType": "cascade-lake",
-        "cpuModel": "Cascadelake-Server-noTSX",
-        "gpuModel": "RTX 3080",
-        "gpuModelId": "nvidia-rtx-3080",
-        "minVcpuPerMemoryGib": 0.25,
-        "maxVcpuPerMemoryGib": 1,
-        "minVcpuPerGpu": 1,
-        "maxVcpuPerGpu": 13,
-        "vcpuPriceHr": {
-            "value": "0.002500"
-        },
-        "memoryGibPriceHr": {
-            "value": "0.003800"
-        },
-        "gpuPriceHr": {
-            "value": "0.05"
-        },
-        "minStorageGibPriceHr": {
-            "value": "0.00013"
-        },
-        "ipv4PriceHr": {
-            "value": "0.005500"
-        },
-        "maxVcpuFree": 76,
-        "totalVcpuFree": 377,
-        "maxMemoryGibFree": 227,
-        "totalMemoryGibFree": 1132,
-        "maxGpuFree": 5,
-        "totalGpuFree": 24,
-        "maxStorageGibFree": 42420,
-        "totalStorageGibFree": 42420
-    }]
+    return [
+        {
+            "dataCenterId": "br-saopaulo-1",
+            "machineType": "cascade-lake",
+            "cpuModel": "Cascadelake-Server-noTSX",
+            "gpuModel": "RTX 3080",
+            "gpuModelId": "nvidia-rtx-3080",
+            "minVcpuPerMemoryGib": 0.25,
+            "maxVcpuPerMemoryGib": 1,
+            "minVcpuPerGpu": 1,
+            "maxVcpuPerGpu": 13,
+            "vcpuPriceHr": {"value": "0.002500"},
+            "memoryGibPriceHr": {"value": "0.003800"},
+            "gpuPriceHr": {"value": "0.05"},
+            "minStorageGibPriceHr": {"value": "0.00013"},
+            "ipv4PriceHr": {"value": "0.005500"},
+            "maxVcpuFree": 76,
+            "totalVcpuFree": 377,
+            "maxMemoryGibFree": 227,
+            "totalMemoryGibFree": 1132,
+            "maxGpuFree": 5,
+            "totalGpuFree": 24,
+            "maxStorageGibFree": 42420,
+            "totalStorageGibFree": 42420,
+        }
+    ]
 
 
 def test_get_offers_with_query_filter():
     cudo = CudoProvider()
     offers = cudo.get(QueryFilter(min_gpu_count=1, max_gpu_count=1), balance_resources=True)
-    print(f'{len(offers)} offers found')
-    assert len(offers) >= 1, f'No offers found'
+    print(f"{len(offers)} offers found")
+    assert len(offers) >= 1, "No offers found"
 
 
 def test_get_offers_no_query_filter():
     cudo = CudoProvider()
     offers = cudo.get(balance_resources=True)
-    print(f'{len(offers)} offers found')
-    assert len(offers) >= 1, f'No offers found'
+    print(f"{len(offers)} offers found")
+    assert len(offers) >= 1, "No offers found"
 
 
 def test_optimize_offers(machine_types):
@@ -85,24 +76,34 @@ def test_optimize_offers(machine_types):
         min_cpus_for_memory = machine_type["minVcpuPerMemoryGib"] * config["memory"]
         max_cpus_for_memory = machine_type["maxVcpuPerMemoryGib"] * config["memory"]
         min_cpus_for_gpu = machine_type["minVcpuPerGpu"] * config["gpu"]
-        assert config["cpu"] >= min_cpus_for_memory, \
-            f"VM config does not meet the minimum CPU:Memory requirement. Required minimum CPUs: " \
+        assert config["cpu"] >= min_cpus_for_memory, (
+            f"VM config does not meet the minimum CPU:Memory requirement. Required minimum CPUs: "
             f"{min_cpus_for_memory}, Found: {config['cpu']}"
-        assert config["cpu"] <= max_cpus_for_memory, \
-            f"VM config exceeds the maximum CPU:Memory allowance. Allowed maximum CPUs: " \
+        )
+        assert config["cpu"] <= max_cpus_for_memory, (
+            f"VM config exceeds the maximum CPU:Memory allowance. Allowed maximum CPUs: "
             f"{max_cpus_for_memory}, Found: {config['cpu']}"
-        assert config["cpu"] >= min_cpus_for_gpu, \
-            f"VM config does not meet the minimum CPU:GPU requirement. " \
+        )
+        assert config["cpu"] >= min_cpus_for_gpu, (
+            f"VM config does not meet the minimum CPU:GPU requirement. "
             f"Required minimum CPUs: {min_cpus_for_gpu}, Found: {config['cpu']}"
+        )
         # Perform the balance resource checks if balance_resource is True
         if balance_resource:
-            expected_memory = get_balanced_memory(config['gpu'], gpu_memory, max_memory)
-            expected_disk_size = get_balanced_disk_size(available_disk, config['memory'], config["gpu"] * gpu_memory,
-                                                        max_disk_size, min_disk_size)
+            expected_memory = get_balanced_memory(config["gpu"], gpu_memory, max_memory)
+            expected_disk_size = get_balanced_disk_size(
+                available_disk,
+                config["memory"],
+                config["gpu"] * gpu_memory,
+                max_disk_size,
+                min_disk_size,
+            )
 
-            assert config['memory'] == expected_memory, \
-                f"Memory allocation does not match the expected balanced memory. " \
+            assert config["memory"] == expected_memory, (
+                f"Memory allocation does not match the expected balanced memory. "
                 f"Expected: {expected_memory}, Found: {config['memory']} in config {config}"
-            assert config['disk_size'] == expected_disk_size, \
-                f"Disk size allocation does not match the expected balanced disk size. " \
+            )
+            assert config["disk_size"] == expected_disk_size, (
+                f"Disk size allocation does not match the expected balanced disk size. "
                 f"Expected: {expected_disk_size}, Found: {config['disk_size']}"
+            )