Skip to content

Commit 0178e9b

Browse files
Bihanjvstme
andauthored
Add digital ocean provider (#172)
* Add DigitalOcean Provider * fix: Correct DigitalOcean env var name and add flavor parameter * fix: correct digitalocean provider miscategorization * rename env vars: TOKEN to API_KEY and flavor to API_URL * Update src/gpuhunt/providers/digitalocean.py Co-authored-by: jvstme <36324149+jvstme@users.noreply.github.com> * Update src/gpuhunt/providers/digitalocean.py Co-authored-by: jvstme <36324149+jvstme@users.noreply.github.com> * Update src/gpuhunt/providers/digitalocean.py Co-authored-by: jvstme <36324149+jvstme@users.noreply.github.com> * Resolve Review Comments * Resolve Review Comments --------- Co-authored-by: Bihan Rana Co-authored-by: jvstme <36324149+jvstme@users.noreply.github.com>
1 parent af4574e commit 0178e9b

7 files changed

Lines changed: 228 additions & 15 deletions

File tree

src/gpuhunt/__main__.py

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -16,6 +16,7 @@ def main():
1616
"cloudrift",
1717
"cudo",
1818
"datacrunch",
19+
"digitalocean",
1920
"gcp",
2021
"hotaisle",
2122
"lambdalabs",
@@ -54,6 +55,12 @@ def main():
5455
provider = DataCrunchProvider(
5556
os.getenv("DATACRUNCH_CLIENT_ID"), os.getenv("DATACRUNCH_CLIENT_SECRET")
5657
)
58+
elif args.provider == "digitalocean":
59+
from gpuhunt.providers.digitalocean import DigitalOceanProvider
60+
61+
provider = DigitalOceanProvider(
62+
api_key=os.getenv("DIGITAL_OCEAN_API_KEY"), api_url=os.getenv("DIGITAL_OCEAN_API_URL")
63+
)
5764
elif args.provider == "gcp":
5865
from gpuhunt.providers.gcp import GCPProvider
5966

src/gpuhunt/_internal/catalog.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -32,7 +32,7 @@
3232
"runpod",
3333
"cloudrift",
3434
]
35-
ONLINE_PROVIDERS = ["cudo", "hotaisle", "tensordock", "vastai", "vultr"]
35+
ONLINE_PROVIDERS = ["cudo", "digitalocean", "hotaisle", "tensordock", "vastai", "vultr"]
3636
RELOAD_INTERVAL = 15 * 60 # 15 minutes
3737

3838

src/gpuhunt/_internal/constraints.py

Lines changed: 13 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -102,6 +102,18 @@ def get_compute_capability(gpu_name: str) -> Optional[tuple[int, int]]:
102102
return None
103103

104104

105+
def get_gpu_vendor(gpu_name: Optional[str]) -> Optional[str]:
106+
if gpu_name is None:
107+
return None
108+
for gpu in KNOWN_NVIDIA_GPUS:
109+
if gpu.name.upper() == gpu_name.upper():
110+
return AcceleratorVendor.NVIDIA.value
111+
for gpu in KNOWN_AMD_GPUS:
112+
if gpu.name.upper() == gpu_name.upper():
113+
return AcceleratorVendor.AMD.value
114+
return None
115+
116+
105117
def correct_gpu_memory_gib(gpu_name: str, memory_mib: float) -> int:
106118
"""
107119
Round to whole number of gibibytes and attempt correcting the reported GPU
@@ -146,6 +158,7 @@ def is_nvidia_superchip(gpu_name: str) -> bool:
146158
NvidiaGPUInfo(name="A6000", memory=48, compute_capability=(8, 6)),
147159
NvidiaGPUInfo(name="H100", memory=80, compute_capability=(9, 0)),
148160
NvidiaGPUInfo(name="H100NVL", memory=94, compute_capability=(9, 0)),
161+
NvidiaGPUInfo(name="H200", memory=141, compute_capability=(9, 0)),
149162
NvidiaGPUInfo(name="L4", memory=24, compute_capability=(8, 9)),
150163
NvidiaGPUInfo(name="L40", memory=48, compute_capability=(8, 9)),
151164
NvidiaGPUInfo(name="L40S", memory=48, compute_capability=(8, 9)),

src/gpuhunt/_internal/default.py

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -24,6 +24,7 @@ def default_catalog() -> Catalog:
2424
("gpuhunt.providers.cudo", "CudoProvider"),
2525
("gpuhunt.providers.vultr", "VultrProvider"),
2626
("gpuhunt.providers.hotaisle", "HotAisleProvider"),
27+
("gpuhunt.providers.digitalocean", "DigitalOceanProvider"),
2728
]:
2829
try:
2930
module = importlib.import_module(module)
Lines changed: 103 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,103 @@
1+
import logging
2+
import os
3+
from typing import Optional
4+
5+
import requests
6+
7+
from gpuhunt._internal.constraints import get_gpu_vendor
8+
from gpuhunt._internal.models import QueryFilter, RawCatalogItem
9+
from gpuhunt.providers import AbstractProvider
10+
11+
logger = logging.getLogger(__name__)
12+
13+
# DigitalOcean Default API endpoints
14+
STANDARD_CLOUD_API_URL = "https://api.digitalocean.com"
15+
16+
17+
class DigitalOceanProvider(AbstractProvider):
18+
NAME = "digitalocean"
19+
20+
def __init__(self, api_key: Optional[str] = None, api_url: Optional[str] = None):
21+
self.api_key = api_key or os.getenv("DIGITAL_OCEAN_API_KEY")
22+
if not self.api_key:
23+
raise ValueError("Set the DIGITAL_OCEAN_API_KEY environment variable.")
24+
25+
self.api_url = api_url or os.getenv("DIGITAL_OCEAN_API_URL", STANDARD_CLOUD_API_URL)
26+
27+
def get(
28+
self, query_filter: Optional[QueryFilter] = None, balance_resources: bool = True
29+
) -> list[RawCatalogItem]:
30+
offers = self.fetch_offers()
31+
return sorted(offers, key=lambda i: i.price)
32+
33+
def fetch_offers(self) -> list[RawCatalogItem]:
34+
url = "/v2/sizes"
35+
response = self._make_request("GET", url)
36+
return convert_response_to_raw_catalog_items(response)
37+
38+
def _make_request(self, method: str, url: str):
39+
full_url = f"{self.api_url}{url}"
40+
params = {"per_page": 500}
41+
headers = {}
42+
if self.api_key:
43+
headers["Authorization"] = f"Bearer {self.api_key}"
44+
45+
response = requests.request(
46+
method=method, url=full_url, params=params, headers=headers, timeout=30
47+
)
48+
response.raise_for_status()
49+
return response
50+
51+
52+
def convert_response_to_raw_catalog_items(response) -> list[RawCatalogItem]:
53+
data = response.json()
54+
offers = []
55+
56+
for size in data["sizes"]:
57+
gpu_info = size.get("gpu_info")
58+
if gpu_info:
59+
gpu_count = gpu_info["count"]
60+
gpu_vram_info = gpu_info["vram"]
61+
gpu_memory = float(gpu_vram_info["amount"])
62+
gpu_model = gpu_info["model"]
63+
# gpu_model uses patterns like "amd_mi300x", "nvidia_h100", "nvidia_rtx6000_ada"
64+
model_parts = gpu_model.split("_")
65+
gpu_name = "".join(part.upper() for part in model_parts[1:])
66+
gpu_vendor = get_gpu_vendor(gpu_name)
67+
if gpu_vendor is None:
68+
logger.warning(
69+
f"Could not determine GPU vendor for model '{gpu_model}'. Skipping droplet '{size['slug']}'."
70+
)
71+
continue
72+
else:
73+
gpu_count = 0
74+
gpu_vendor = None
75+
gpu_name = ""
76+
gpu_memory = 0
77+
78+
total_disk_size = sum(
79+
float(disk["size"]["amount"]) for disk in size["disk_info"] if disk["type"] == "local"
80+
)
81+
82+
memory_gb = float(size["memory"]) / 1024 # MB -> GB
83+
84+
# Creates an offer for each available region.
85+
# If regions list is empty, instance type is not available.
86+
for region in size["regions"]:
87+
offer = RawCatalogItem(
88+
instance_name=size["slug"],
89+
location=region,
90+
price=size["price_hourly"],
91+
cpu=size["vcpus"],
92+
memory=memory_gb,
93+
gpu_vendor=gpu_vendor,
94+
gpu_count=gpu_count,
95+
gpu_name=gpu_name,
96+
gpu_memory=gpu_memory,
97+
spot=False,
98+
disk_size=total_disk_size,
99+
flags=[],
100+
)
101+
offers.append(offer)
102+
103+
return offers

src/gpuhunt/providers/vultr.py

Lines changed: 7 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -5,8 +5,13 @@
55
from requests import Response
66

77
from gpuhunt import QueryFilter, RawCatalogItem
8-
from gpuhunt._internal.constraints import KNOWN_AMD_GPUS, KNOWN_NVIDIA_GPUS, is_nvidia_superchip
9-
from gpuhunt._internal.models import AcceleratorVendor, CPUArchitecture
8+
from gpuhunt._internal.constraints import (
9+
KNOWN_AMD_GPUS,
10+
KNOWN_NVIDIA_GPUS,
11+
get_gpu_vendor,
12+
is_nvidia_superchip,
13+
)
14+
from gpuhunt._internal.models import CPUArchitecture
1015
from gpuhunt.providers import AbstractProvider
1116

1217
logger = logging.getLogger(__name__)
@@ -168,18 +173,6 @@ def get_gpu_memory(gpu_name: str) -> Optional[int]:
168173
return None
169174

170175

171-
def get_gpu_vendor(gpu_name: Optional[str]) -> Optional[str]:
172-
if gpu_name is None:
173-
return None
174-
for gpu in KNOWN_NVIDIA_GPUS:
175-
if gpu.name.upper() == gpu_name.upper():
176-
return AcceleratorVendor.NVIDIA.value
177-
for gpu in KNOWN_AMD_GPUS:
178-
if gpu.name.upper() == gpu_name.upper():
179-
return AcceleratorVendor.AMD.value
180-
return None
181-
182-
183176
def _make_request(method: str, path: str, data: Any = None) -> Response:
184177
response = requests.request(
185178
method=method,
Lines changed: 96 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,96 @@
1+
import gpuhunt._internal.catalog as internal_catalog
2+
from gpuhunt import Catalog
3+
from gpuhunt.providers.digitalocean import DigitalOceanProvider
4+
5+
sizes_response = {
6+
"sizes": [
7+
{
8+
"slug": "s-1vcpu-512mb-10gb",
9+
"memory": 512,
10+
"vcpus": 1,
11+
"disk": 10,
12+
"transfer": 0.5,
13+
"price_monthly": 4,
14+
"price_hourly": 0.00595,
15+
"regions": ["ams3", "fra1", "nyc1", "nyc2", "sfo2", "sfo3", "sgp1", "syd1"],
16+
"available": True,
17+
"description": "Basic",
18+
"networking_throughput": 2000,
19+
"disk_info": [{"type": "local", "size": {"amount": 10, "unit": "gib"}}],
20+
},
21+
{
22+
"slug": "gpu-h100x8-640gb",
23+
"memory": 1966080,
24+
"vcpus": 160,
25+
"disk": 2046,
26+
"transfer": 60.0,
27+
"price_monthly": 17796.48,
28+
"price_hourly": 23.92,
29+
"regions": ["nyc1"],
30+
"available": True,
31+
"description": "H100 GPU - 8X",
32+
"networking_throughput": 10000,
33+
"gpu_info": {
34+
"count": 8,
35+
"vram": {"amount": 640, "unit": "gib"},
36+
"model": "nvidia_h100",
37+
},
38+
"disk_info": [
39+
{"type": "local", "size": {"amount": 2046, "unit": "gib"}},
40+
{"type": "scratch", "size": {"amount": 40960, "unit": "gib"}},
41+
],
42+
},
43+
{
44+
"slug": "gpu-mi300x8-1536gb",
45+
"memory": 1966080,
46+
"vcpus": 160,
47+
"disk": 2046,
48+
"transfer": 60.0,
49+
"price_monthly": 11844.48,
50+
"price_hourly": 15.92,
51+
"regions": ["nyc1"],
52+
"available": True,
53+
"description": "AMD MI300X - 8X",
54+
"networking_throughput": 10000,
55+
"gpu_info": {
56+
"count": 8,
57+
"vram": {"amount": 1536, "unit": "gib"},
58+
"model": "amd_mi300x",
59+
},
60+
"disk_info": [
61+
{"type": "local", "size": {"amount": 2046, "unit": "gib"}},
62+
{"type": "scratch", "size": {"amount": 40960, "unit": "gib"}},
63+
],
64+
},
65+
],
66+
"links": {},
67+
"meta": {"total": 147},
68+
}
69+
70+
71+
def test_fetch_offers(requests_mock):
72+
requests_mock.get("https://api.digitalocean.com/v2/sizes", json=sizes_response)
73+
74+
provider = DigitalOceanProvider(api_key="test-token", api_url="https://api.digitalocean.com")
75+
offers = provider.fetch_offers()
76+
assert len(offers) == 10 # 8 CPU offers (8 regions) + 1 NVIDIA + 1 AMD (1 region each)
77+
catalog = Catalog(balance_resources=False, auto_reload=False)
78+
digitalocean = DigitalOceanProvider(
79+
api_key="test-token", api_url="https://api.digitalocean.com"
80+
)
81+
internal_catalog.ONLINE_PROVIDERS = ["digitalocean"]
82+
internal_catalog.OFFLINE_PROVIDERS = []
83+
catalog.add_provider(digitalocean)
84+
85+
# Test queries
86+
assert (
87+
len(catalog.query(provider=["digitalocean"], max_gpu_count=0)) == 8
88+
) # CPU only (8 regions)
89+
assert len(catalog.query(provider=["digitalocean"], min_gpu_count=1)) == 2 # GPU instances
90+
assert len(catalog.query(provider=["digitalocean"], gpu_vendor="nvidia")) == 1 # NVIDIA GPU
91+
assert len(catalog.query(provider=["digitalocean"], gpu_name="H100")) == 1 # Specific GPU
92+
assert len(catalog.query(provider=["digitalocean"], gpu_vendor="amd")) == 1 # AMD GPU
93+
assert len(catalog.query(provider=["digitalocean"], gpu_name="MI300X")) == 1 # AMD GPU
94+
assert (
95+
len(catalog.query(provider=["digitalocean"], min_gpu_memory=1000)) == 1
96+
) # MI300X: 1536GB

0 commit comments

Comments
 (0)