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']}" + )