Fix outdated ROCm driver on DigitalOcean AMD GPU droplets

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)

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