|
| 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