Support iam_instance_profile for AWS (#2365)

* Support iam_instance_profile for AWS

* Fix tests

Author: r4victor

docs/docs/concepts/backends.md

@@ -122,13 +122,44 @@ There are two ways to configure AWS: using an access key or using the default cr
                     "acm:ListCertificates"
                 ],
                 "Resource": "*"
+            },
+            {
+                "Effect": "Allow",
+                "Action": [
+                    "iam:GetInstanceProfile",
+                    "iam:GetRole",
+                    "iam:PassRole"
+                ],
+                "Resource": "*"
             }
         ]
     }
     ```
 
     The `elasticloadbalancing:*` and `acm:*` permissions are only needed for provisioning gateways with ACM (AWS Certificate Manager) certificates.
 
+    The `iam:*` permissions are only needed if you specify `iam_instance_profile` to assign to EC2 instances.
+
+    You can also limit permissions to specific resources in your account:
+    
+    ```
+    {
+        "Version": "2012-10-17",
+        "Statement": [
+            ...
+            {
+                "Effect": "Allow",
+                "Action": [
+                    "iam:GetInstanceProfile",
+                    "iam:GetRole",
+                    "iam:PassRole"
+                ],
+                "Resource": "arn:aws:iam::account-id:role/EC2-roles-for-XYZ-*"
+            }
+        ]
+    }
+    ```
+
 ??? info "VPC"
     By default, `dstack` uses the default VPC. It's possible to customize it:
 

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

@@ -219,7 +219,7 @@ def create_instance(
                         disk_size=disk_size,
                         image_id=image_id,
                         instance_type=instance_offer.instance.name,
-                        iam_instance_profile_arn=None,
+                        iam_instance_profile=self.config.iam_instance_profile,
                         user_data=get_user_data(authorized_keys=instance_config.get_public_keys()),
                         tags=aws_resources.make_tags(tags),
                         security_group_id=aws_resources.create_security_group(
@@ -264,6 +264,9 @@ def create_instance(
                 )
             except botocore.exceptions.ClientError as e:
                 logger.warning("Got botocore.exceptions.ClientError: %s", e)
+                if e.response["Error"]["Code"] == "InvalidParameterValue":
+                    msg = e.response["Error"].get("Message", "")
+                    raise ComputeError(f"Invalid AWS request: {msg}")
                 continue
         raise NoCapacityError()
 
@@ -380,7 +383,7 @@ def create_gateway(
                 disk_size=10,
                 image_id=aws_resources.get_gateway_image_id(ec2_client),
                 instance_type="t2.micro",
-                iam_instance_profile_arn=None,
+                iam_instance_profile=None,
                 user_data=get_gateway_user_data(configuration.ssh_key_pub),
                 tags=tags,
                 security_group_id=security_group_id,

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

@@ -131,7 +131,7 @@ def create_instances_struct(
     disk_size: int,
     image_id: str,
     instance_type: str,
-    iam_instance_profile_arn: Optional[str],
+    iam_instance_profile: Optional[str],
     user_data: str,
     tags: List[Dict[str, str]],
     security_group_id: str,
@@ -166,8 +166,8 @@ def create_instances_struct(
             },
         ],
     )
-    if iam_instance_profile_arn:
-        struct["IamInstanceProfile"] = {"Arn": iam_instance_profile_arn}
+    if iam_instance_profile:
+        struct["IamInstanceProfile"] = {"Name": iam_instance_profile}
     if spot:
         struct["InstanceMarketOptions"] = {
             "MarketType": "spot",

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

@@ -32,6 +32,7 @@ class AWSConfigInfo(CoreModel):
     vpc_ids: Optional[Dict[str, str]] = None
     default_vpcs: Optional[bool] = None
     public_ips: Optional[bool] = None
+    iam_instance_profile: Optional[str] = None
     tags: Optional[Dict[str, str]] = None
     os_images: Optional[AWSOSImageConfig] = None
 
@@ -70,6 +71,7 @@ class AWSConfigInfoWithCredsPartial(CoreModel):
     vpc_ids: Optional[Dict[str, str]]
     default_vpcs: Optional[bool]
     public_ips: Optional[bool]
+    iam_instance_profile: Optional[str]
     tags: Optional[Dict[str, str]]
     os_images: Optional["AWSOSImageConfig"]
 

src/dstack/_internal/server/services/backends/configurators/aws.py

@@ -2,6 +2,7 @@
 import json
 from typing import List
 
+import botocore.exceptions
 from boto3.session import Session
 
 from dstack._internal.core.backends.aws import AWSBackend, auth, compute, resources
@@ -35,6 +36,9 @@
     Configurator,
     raise_invalid_credentials_error,
 )
+from dstack._internal.utils.logging import get_logger
+
+logger = get_logger(__name__)
 
 REGIONS = [
     ("US East, N. Virginia", "us-east-1"),
@@ -137,7 +141,8 @@ def _get_regions_element(self, selected: List[str]) -> ConfigMultiElement:
 
     def _check_config(self, session: Session, config: AWSConfigInfoWithCredsPartial):
         self._check_tags_config(config)
-        self._check_vpc_config(session=session, config=config)
+        self._check_iam_instance_profile_config(session, config)
+        self._check_vpc_config(session, config)
 
     def _check_tags_config(self, config: AWSConfigInfoWithCredsPartial):
         if not config.tags:
@@ -151,6 +156,31 @@ def _check_tags_config(self, config: AWSConfigInfoWithCredsPartial):
         except BackendError as e:
             raise ServerClientError(e.args[0])
 
+    def _check_iam_instance_profile_config(
+        self, session: Session, config: AWSConfigInfoWithCredsPartial
+    ):
+        if config.iam_instance_profile is None:
+            return
+        try:
+            iam_client = session.client("iam")
+            iam_client.get_instance_profile(InstanceProfileName=config.iam_instance_profile)
+        except botocore.exceptions.ClientError as e:
+            if e.response["Error"]["Code"] == "NoSuchEntity":
+                raise ServerClientError(
+                    f"IAM instance profile {config.iam_instance_profile} not found"
+                )
+            logger.exception(
+                "Got botocore.exceptions.ClientError when checking iam_instance_profile"
+            )
+            raise ServerClientError(
+                f"Failed to check IAM instance profile {config.iam_instance_profile}"
+            )
+        except Exception:
+            logger.exception("Got exception when checking iam_instance_profile")
+            raise ServerClientError(
+                f"Failed to check IAM instance profile {config.iam_instance_profile}"
+            )
+
     def _check_vpc_config(self, session: Session, config: AWSConfigInfoWithCredsPartial):
         allocate_public_ip = config.public_ips if config.public_ips is not None else True
         use_default_vpcs = config.default_vpcs if config.default_vpcs is not None else True

src/dstack/_internal/server/services/config.py

@@ -107,6 +107,16 @@ class AWSConfig(CoreModel):
             )
         ),
     ] = None
+    iam_instance_profile: Annotated[
+        Optional[str],
+        Field(
+            description=(
+                "The name of the IAM instance profile to associate with EC2 instances."
+                " You can also specify the IAM role name for roles created via the AWS console."
+                " AWS automatically creates an instance profile and gives it the same name as the role"
+            )
+        ),
+    ] = None
     tags: Annotated[
         Optional[Dict[str, str]],
         Field(description="The tags that will be assigned to resources created by `dstack`"),
@@ -251,7 +261,7 @@ class GCPConfig(CoreModel):
         ),
     ] = None
     vm_service_account: Annotated[
-        Optional[str], Field(description="The service account associated with provisioned VMs")
+        Optional[str], Field(description="The service account to associate with provisioned VMs")
     ] = None
     tags: Annotated[
         Optional[Dict[str, str]],

src/tests/_internal/server/routers/test_backends.py

@@ -1335,6 +1335,7 @@ async def test_returns_config_info(self, test_db, session: AsyncSession, client:
             "vpc_ids": None,
             "default_vpcs": None,
             "public_ips": None,
+            "iam_instance_profile": None,
             "tags": None,
             "os_images": None,
             "creds": json.loads(backend.auth.plaintext),