From 538e97da2ded463bdc3bd2e567c842e8c79c1225 Mon Sep 17 00:00:00 2001
From: allegroai <>
Date: Sun, 6 Mar 2022 18:47:44 +0200
Subject: [PATCH] Support IAM Instance Profile in AWS auto-scaler

---
 clearml/automation/aws_driver.py | 14 +++++++++++---
 1 file changed, 11 insertions(+), 3 deletions(-)

diff --git a/clearml/automation/aws_driver.py b/clearml/automation/aws_driver.py
index 0cdd0ea6..8dc25083 100644
--- a/clearml/automation/aws_driver.py
+++ b/clearml/automation/aws_driver.py
@@ -28,6 +28,8 @@ class AWSDriver(CloudDriver):
     aws_region = attr.ib(validator=instance_of(str), default='')
     use_credentials_chain = attr.ib(validator=instance_of(bool), default=False)
     use_iam_instance_profile = attr.ib(validator=instance_of(bool), default=False)
+    iam_arn = attr.ib(validator=instance_of(str), default='')
+    iam_name = attr.ib(validator=instance_of(str), default='')
 
     @classmethod
     def from_config(cls, config):
@@ -37,6 +39,8 @@ class AWSDriver(CloudDriver):
         obj.aws_region = config['hyper_params'].get('cloud_credentials_region')
         obj.use_credentials_chain = config['hyper_params'].get('use_credentials_chain', False)
         obj.use_iam_instance_profile = config['hyper_params'].get('use_iam_instance_profile', False)
+        obj.iam_arn = config['hyper_params'].get('iam_arn')
+        obj.iam_name = config['hyper_params'].get('iam_name')
         return obj
 
     def __attrs_post_init__(self):
@@ -71,10 +75,14 @@ class AWSDriver(CloudDriver):
             launch_specification["SecurityGroupIds"] = resource_conf[
                 "security_group_ids"
             ]
-        if resource_conf.get("iam_arn", None) and resource_conf.get("iam_name", None):
+        # Adding iam role - you can have Arn OR Name, not both, Arn getting priority
+        if self.iam_arn:
             launch_specification["IamInstanceProfile"] = {
-                'Arn': resource_conf["iam_arn"],
-                'Name': resource_conf["iam_name"]
+                'Arn': self.iam_arn,
+            }
+        elif self.iam_name:
+            launch_specification["IamInstanceProfile"] = {
+                'Name': self.iam_name
             }
 
         if resource_conf["is_spot"]: