Use DLAMI on AWS (#2782)

Author: r4victor

src/dstack/_internal/core/backends/aws/compute.py

@@ -231,6 +231,7 @@ def create_instance(
                 image_id, username = aws_resources.get_image_id_and_username(
                     ec2_client=ec2_client,
                     cuda=len(instance_offer.instance.resources.gpus) > 0,
+                    instance_type=instance_offer.instance.name,
                     image_config=self.config.os_images,
                 )
                 response = ec2_resource.create_instances(

src/dstack/_internal/core/backends/aws/resources.py

@@ -12,11 +12,13 @@
 logger = get_logger(__name__)
 
 DSTACK_ACCOUNT_ID = "142421590066"
+DLAMI_OWNER_ACCOUNT_ID = "898082745236"
 
 
 def get_image_id_and_username(
     ec2_client: botocore.client.BaseClient,
     cuda: bool,
+    instance_type: str,
     image_config: Optional[AWSOSImageConfig] = None,
 ) -> tuple[str, str]:
     if image_config is not None:
@@ -27,6 +29,11 @@ def get_image_id_and_username(
         image_name = image.name
         image_owner = image.owner
         username = image.user
+    elif _supported_by_dlami(instance_type):
+        # TODO: Update DLAMI image version from time to time
+        image_name = "Deep Learning Base OSS Nvidia Driver GPU AMI (Ubuntu 22.04) 20250516"
+        image_owner = DLAMI_OWNER_ACCOUNT_ID
+        username = "ubuntu"
     else:
         image_name = (
             f"dstack-{version.base_image}" if not cuda else f"dstack-cuda-{version.base_image}"
@@ -628,6 +635,25 @@ def _is_private_subnet_with_internet_egress(
     return False
 
 
+def _supported_by_dlami(instance_type: str) -> bool:
+    # Currently only p3. instances are not supported by DLAMI among GPU instances.
+    return any(
+        instance_type.startswith(family)
+        for family in [
+            "g4dn.",
+            "g5.",
+            "g6.",
+            "gr6.",
+            "g6e.",
+            "p4d.",
+            "p4de.",
+            "p5.",
+            "p5e.",
+            "p6-b200.",
+        ]
+    )
+
+
 def get_reservation(
     ec2_client: botocore.client.BaseClient,
     reservation_id: str,

src/tests/_internal/core/backends/aws/test_resources.py

@@ -116,7 +116,11 @@ def test_returns_the_latest_available(self, ec2_client_mock: Mock):
                 },
             ]
         }
-        image_id, username = get_image_id_and_username(ec2_client_mock, cuda=False)
+        image_id, username = get_image_id_and_username(
+            ec2_client_mock,
+            cuda=False,
+            instance_type="some",
+        )
         assert image_id == "ami-00000000000000003"
         assert username == "ubuntu"
 
@@ -138,7 +142,11 @@ def test_raises_resource_not_found_if_none_available(
             ]
         }
         with pytest.raises(ComputeResourceNotFoundError):
-            get_image_id_and_username(ec2_client_mock, cuda=False)
+            get_image_id_and_username(
+                ec2_client_mock,
+                cuda=False,
+                instance_type="some",
+            )
         assert "image 'dstack-0.0' not found" in caplog.text
 
     @pytest.mark.parametrize(
@@ -152,7 +160,11 @@ def test_uses_dstack_image_name_and_account_id_if_image_config_not_provided(
         self, monkeypatch: pytest.MonkeyPatch, ec2_client_mock: Mock, cuda: bool, expected: str
     ):
         monkeypatch.setattr("dstack.version.base_image", "0.0")
-        _, username = get_image_id_and_username(ec2_client_mock, cuda)
+        _, username = get_image_id_and_username(
+            ec2_client_mock,
+            cuda=cuda,
+            instance_type="some",
+        )
         assert username == "ubuntu"
         ec2_client_mock.describe_images.assert_called_once_with(
             Filters=[{"Name": "name", "Values": [expected]}], Owners=["142421590066"]
@@ -184,7 +196,12 @@ def test_uses_image_config_if_provided(
                 user="dstack",
             ),
         )
-        _, username = get_image_id_and_username(ec2_client_mock, cuda, image_config)
+        _, username = get_image_id_and_username(
+            ec2_client_mock,
+            cuda=cuda,
+            instance_type="some",
+            image_config=image_config,
+        )
         assert username == expected_username
         ec2_client_mock.describe_images.assert_called_once_with(
             Filters=[{"Name": "name", "Values": [expected_name]}],
@@ -202,5 +219,10 @@ def test_raises_resource_not_found_if_image_config_property_not_set(
             ),
         )
         with pytest.raises(ComputeResourceNotFoundError):
-            get_image_id_and_username(ec2_client_mock, cuda=False, image_config=image_config)
+            get_image_id_and_username(
+                ec2_client_mock,
+                cuda=False,
+                instance_type="some",
+                image_config=image_config,
+            )
         assert "cpu image not configured" in caplog.text