From a99cce41e6f73ff7f3c139339c1ce32031ef6300 Mon Sep 17 00:00:00 2001 From: Andrey Cheptsov Date: Sat, 15 Aug 2026 19:30:24 +0200 Subject: [PATCH] Use DO's AI/ML-ready image for AMD GPU droplets Co-Authored-By: Claude Fable 5 --- .../backends/digitalocean_base/compute.py | 32 ++++++---- .../backends/digitalocean_base/__init__.py | 0 .../digitalocean_base/test_compute.py | 64 +++++++++++++++++++ 3 files changed, 82 insertions(+), 14 deletions(-) create mode 100644 src/tests/_internal/core/backends/digitalocean_base/__init__.py create mode 100644 src/tests/_internal/core/backends/digitalocean_base/test_compute.py diff --git a/src/dstack/_internal/core/backends/digitalocean_base/compute.py b/src/dstack/_internal/core/backends/digitalocean_base/compute.py index 8f672480bc..4c43664bb4 100644 --- a/src/dstack/_internal/core/backends/digitalocean_base/compute.py +++ b/src/dstack/_internal/core/backends/digitalocean_base/compute.py @@ -1,4 +1,4 @@ -from typing import List, Optional +from typing import List, NamedTuple, Optional import gpuhunt from gpuhunt.providers.digitalocean import DigitalOceanProvider @@ -29,6 +29,13 @@ logger = get_logger(__name__) MAX_INSTANCE_NAME_LEN = 60 + + +class _Image(NamedTuple): + slug: str + ships_docker: bool + + DOCKER_INSTALL_COMMANDS = [ "export DEBIAN_FRONTEND=noninteractive", "mkdir -p /etc/apt/keyrings", @@ -86,10 +93,8 @@ def create_instance( ) size_slug = instance_offer.instance.name - if not instance_offer.instance.resources.gpus: - backend_specific_commands = DOCKER_INSTALL_COMMANDS - else: - backend_specific_commands = None + image = self._get_image_for_instance(instance_offer) + backend_specific_commands = None if image.ships_docker else DOCKER_INSTALL_COMMANDS project_id = None if self.config.project_name: @@ -100,7 +105,7 @@ def create_instance( "name": instance_name, "region": instance_offer.region, "size": size_slug, - "image": self._get_image_for_instance(instance_offer), + "image": image.slug, "ssh_keys": [ssh_key_id], "backups": False, "ipv6": False, @@ -148,28 +153,27 @@ def terminate_instance( ): self.api_client.delete_droplet(instance_id) - def _get_image_for_instance(self, instance_offer: InstanceOfferWithAvailability) -> str: + def _get_image_for_instance(self, instance_offer: InstanceOfferWithAvailability) -> _Image: if not instance_offer.instance.resources.gpus: - # No GPUs, use CPU image - return "ubuntu-24-04-x64" + return _Image("ubuntu-24-04-x64", ships_docker=False) gpu_count = len(instance_offer.instance.resources.gpus) gpu_vendor = instance_offer.instance.resources.gpus[0].vendor if gpu_vendor == gpuhunt.AcceleratorVendor.AMD: - # AMD GPU - return "digitaloceanai-rocmjupyter" + # Ubuntu 24.04, ROCm 7.14 (amdgpu 6.19.14). + return _Image("gpu-amd-base", ships_docker=False) else: # NVIDIA GPUs - DO only supports 1 and 8 GPU configurations. # DO says for single GPU plans using GPUs other than H100s use "gpu-h100x1-base". DO does not provide guidance for x8 GPUs so assuming the same applies. # See (https://docs.digitalocean.com/products/droplets/getting-started/recommended-gpu-setup/#aiml-ready-image) if gpu_count == 8: - return "gpu-h100x8-base" + return _Image("gpu-h100x8-base", ships_docker=True) elif gpu_count == 1: - return "gpu-h100x1-base" + return _Image("gpu-h100x1-base", ships_docker=True) else: # For Unsupported GPU count - use single GPU image and log warning logger.warning( f"Unsupported NVIDIA GPU count: {gpu_count}, using single GPU image" ) - return "gpu-h100x1-base" + return _Image("gpu-h100x1-base", ships_docker=True) diff --git a/src/tests/_internal/core/backends/digitalocean_base/__init__.py b/src/tests/_internal/core/backends/digitalocean_base/__init__.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/src/tests/_internal/core/backends/digitalocean_base/test_compute.py b/src/tests/_internal/core/backends/digitalocean_base/test_compute.py new file mode 100644 index 0000000000..3ac89c3c33 --- /dev/null +++ b/src/tests/_internal/core/backends/digitalocean_base/test_compute.py @@ -0,0 +1,64 @@ +from unittest.mock import Mock, patch + +import gpuhunt +import pytest + +from dstack._internal.core.backends.digitalocean_base.compute import BaseDigitalOceanCompute +from dstack._internal.core.models.backends.base import BackendType +from dstack._internal.core.models.instances import ( + Gpu, + InstanceOffer, + InstanceType, + Resources, +) + +pytestmark = pytest.mark.windows + + +def _offer(gpus: list[Gpu]) -> InstanceOffer: + return InstanceOffer( + backend=BackendType.AMDDEVCLOUD, + instance=InstanceType( + name="test-instance", + resources=Resources(cpus=8, memory_mib=65536, spot=False, gpus=gpus), + ), + region="atl1", + price=1.99, + ) + + +@pytest.fixture +def compute() -> BaseDigitalOceanCompute: + with patch("dstack._internal.core.backends.digitalocean_base.compute.DigitalOceanAPIClient"): + return BaseDigitalOceanCompute( + config=Mock(creds=Mock(api_key="key"), regions=None, project_name=None), + api_url="https://api-amd.digitalocean.com", + type=BackendType.AMDDEVCLOUD, + ) + + +class TestGetImageForInstance: + def test_amd_gpu_image_ships_no_docker(self, compute): + gpu = Gpu(vendor=gpuhunt.AcceleratorVendor.AMD, name="MI300X", memory_mib=196608) + + image = compute._get_image_for_instance(_offer([gpu])) + + assert image.slug == "gpu-amd-base" + assert image.ships_docker is False + + def test_nvidia_gpu_images_ship_docker(self, compute): + gpu = Gpu(vendor=gpuhunt.AcceleratorVendor.NVIDIA, name="H100", memory_mib=81920) + + single = compute._get_image_for_instance(_offer([gpu])) + eight = compute._get_image_for_instance(_offer([gpu] * 8)) + + assert single.slug == "gpu-h100x1-base" + assert eight.slug == "gpu-h100x8-base" + assert single.ships_docker is True + assert eight.ships_docker is True + + def test_cpu_image_ships_no_docker(self, compute): + image = compute._get_image_for_instance(_offer([])) + + assert image.slug == "ubuntu-24-04-x64" + assert image.ships_docker is False