diff --git a/pyproject.toml b/pyproject.toml index bba8433701..0e08aacdcd 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -33,7 +33,7 @@ dependencies = [ "python-multipart>=0.0.16", "filelock", "psutil", - "gpuhunt==0.1.32", + "gpuhunt==0.1.33", "argcomplete>=3.5.0", "ignore-python>=0.2.0", "apscheduler<4", diff --git a/src/dstack/_internal/core/backends/vultr/compute.py b/src/dstack/_internal/core/backends/vultr/compute.py index 8afcb79369..9b85e71eb0 100644 --- a/src/dstack/_internal/core/backends/vultr/compute.py +++ b/src/dstack/_internal/core/backends/vultr/compute.py @@ -2,7 +2,9 @@ import re from typing import List, Optional +import gpuhunt import requests +from gpuhunt.providers.vultr import VultrProvider from dstack._internal.core.backends.base.backend import Compute from dstack._internal.core.backends.base.compute import ( @@ -46,6 +48,10 @@ def __init__(self, config: VultrConfig): super().__init__() self.config = config self.api_client = VultrApiClient(config.creds.api_key) + self.catalog = gpuhunt.Catalog(balance_resources=False, auto_reload=False) + self.catalog.add_provider( + VultrProvider(api_key=config.creds.api_key, regions=config.regions) + ) def get_all_offers_with_availability( self, unallocated_resources: bool @@ -55,6 +61,7 @@ def get_all_offers_with_availability( requirements=None, locations=self.config.regions or None, extra_filter=_supported_instances, + catalog=self.catalog, ) offers = [ offer.with_availability(availability=InstanceAvailability.AVAILABLE) diff --git a/src/tests/_internal/core/backends/vultr/test_compute.py b/src/tests/_internal/core/backends/vultr/test_compute.py new file mode 100644 index 0000000000..530f43ec1e --- /dev/null +++ b/src/tests/_internal/core/backends/vultr/test_compute.py @@ -0,0 +1,55 @@ +from gpuhunt.providers.vultr import API_URL + +from dstack._internal.core.backends.vultr.compute import VultrCompute +from dstack._internal.core.backends.vultr.models import VultrConfig, VultrCreds + + +class TestGetAllOffersWithAvailability: + def test_uses_backend_credentials_for_gpu_availability(self, requests_mock, monkeypatch): + monkeypatch.setenv("VULTR_API_KEY", "environment-key") + requests_mock.get(f"{API_URL}/plans-metal?per_page=500", json={"plans_metal": []}) + requests_mock.get( + f"{API_URL}/plans?type=all&per_page=500", + json={ + "plans": [ + { + "id": "vc2-2c-4gb", + "type": "vc2", + "vcpu_count": 2, + "ram": 4096, + "disk": 80, + "hourly_cost": 0.027, + "locations": ["blr", "fra"], + }, + { + "id": "vcg-a40-24c-120g-48vram", + "type": "vdm", + "vcpu_count": 24, + "ram": 122880, + "disk": 1400, + "hourly_cost": 1.712, + "deploy_ondemand": True, + "locations": [], + }, + ], + }, + ) + requests_mock.get( + f"{API_URL}/regions/blr/availability?type=vdm", + request_headers={"Authorization": "Bearer backend-key"}, + json={ + "available_plans": ["vcg-a40-24c-120g-48vram"], + "available_vpc_only_plans": [], + }, + ) + compute = VultrCompute( + VultrConfig(creds=VultrCreds(api_key="backend-key"), regions=["blr"]) + ) + + offers = compute.get_all_offers_with_availability(unallocated_resources=False) + + assert {offer.instance.name for offer in offers} == { + "vc2-2c-4gb", + "vcg-a40-24c-120g-48vram", + } + assert all(offer.region == "blr" for offer in offers)