Expose SSH key and Instance IAM Role config options in EC2Cluster (#208)

* Expose SSH key and Instance IAM Role config options in EC2Cluster

* Remove accidental deletion

* Add SSH example to docstring

dask_cloudprovider/aws/ec2.py

@@ -45,6 +45,8 @@ def __init__(
         subnet_id=None,
         security_groups=None,
         filesystem_size=None,
+        key_name=None,
+        iam_instance_profile=None,
         **kwargs,
     ):
         super().__init__(*args, **kwargs)
@@ -62,6 +64,8 @@ def __init__(
         self.subnet_id = subnet_id
         self.security_groups = security_groups
         self.filesystem_size = filesystem_size
+        self.key_name = key_name
+        self.iam_instance_profile = iam_instance_profile
 
     async def create_vm(self):
         """
@@ -86,8 +90,8 @@ async def create_vm(self):
                 "099720109477",  # Canonical
             )
 
-            response = await client.run_instances(
-                BlockDeviceMappings=[
+            vm_kwargs = {
+                "BlockDeviceMappings": [
                     {
                         "DeviceName": "/dev/sda1",
                         "VirtualName": "sda1",
@@ -99,20 +103,20 @@ async def create_vm(self):
                         },
                     }
                 ],
-                ImageId=self.ami,
-                InstanceType=self.instance_type,
-                MaxCount=1,
-                MinCount=1,
-                Monitoring={"Enabled": False},
-                UserData=self.cluster.render_cloud_init(
+                "ImageId": self.ami,
+                "InstanceType": self.instance_type,
+                "MaxCount": 1,
+                "MinCount": 1,
+                "Monitoring": {"Enabled": False},
+                "UserData": self.cluster.render_cloud_init(
                     image=self.docker_image,
                     command=self.command,
                     gpu_instance=self.gpu_instance,
                     bootstrap=self.bootstrap,
                     env_vars=self.env_vars,
                 ),
-                InstanceInitiatedShutdownBehavior="terminate",
-                NetworkInterfaces=[
+                "InstanceInitiatedShutdownBehavior": "terminate",
+                "NetworkInterfaces": [
                     {
                         "AssociatePublicIpAddress": True,
                         "DeleteOnTermination": True,
@@ -122,7 +126,15 @@ async def create_vm(self):
                         "SubnetId": self.subnet_id,
                     }
                 ],
-            )
+            }
+
+            if self.key_name:
+                vm_kwargs["KeyName"] = self.key_name
+
+            if self.iam_instance_profile:
+                vm_kwargs["IamInstanceProfile"] = self.iam_instance_profile
+
+            response = await client.run_instances(**vm_kwargs)
             [self.instance] = response["Instances"]
             await client.create_tags(
                 Resources=[self.instance["InstanceId"]],
@@ -236,6 +248,16 @@ class EC2Cluster(VMCluster):
         The instance filesystem size in GB.
 
         Defaults to ``40``.
+    key_name: str (optional)
+        The SSH key name to assign to all instances created by the cluster manager.
+        You can list your existing key pair names with
+        ``aws ec2 describe-key-pairs  --query 'KeyPairs[*].KeyName' --output text``.
+
+        NOTE: You will need to ensure your security group allows access on port 22. If ``security_groups``
+        is not set the default group will not contain this rule and you will need to add it manually.
+    iam_instance_profile: dict (optional)
+        An IAM profile to assign to VMs. This can be used for allowing access to other AWS resources such as S3.
+        See https://docs.aws.amazon.com/AWSEC2/latest/UserGuide/iam-roles-for-amazon-ec2.html.
     n_workers: int
         Number of workers to initialise the cluster with. Defaults to ``0``.
     worker_module: str
@@ -316,6 +338,18 @@ class EC2Cluster(VMCluster):
                              worker_module="dask_cuda.cli.dask_cuda_worker",
                              bootstrap=False,
                              filesystem_size=120)
+
+    Enable SSH for debugging
+
+    >>> from dask_cloudprovider.aws import EC2Cluster
+    >>> cluster = EC2Cluster(key_name="myawesomekey",
+                             # Security group which allows ports 22, 8786, 8787 and all internal traffic
+                             security_groups=["sg-aabbcc112233"])
+
+    # You can now SSH to an instance with `ssh ubuntu@public_ip`
+
+    >>> cluster.close()
+
     """
 
     def __init__(
@@ -329,6 +363,8 @@ def __init__(
         subnet_id=None,
         security_groups=None,
         filesystem_size=None,
+        key_name=None,
+        iam_instance_profile=None,
         docker_image=None,
         **kwargs,
     ):
@@ -367,6 +403,14 @@ def __init__(
             else self.config.get("filesystem_size")
         )
 
+        self.key_name = (
+            key_name if key_name is not None else self.config.get("key_name")
+        )
+        self.iam_instance_profile = (
+            iam_instance_profile
+            if iam_instance_profile is not None
+            else self.config.get("iam_instance_profile")
+        )
         self.options = {
             "cluster": self,
             "config": self.config,
@@ -380,6 +424,8 @@ def __init__(
             "subnet_id": self.subnet_id,
             "security_groups": self.security_groups,
             "filesystem_size": self.filesystem_size,
+            "key_name": self.key_name,
+            "iam_instance_profile": self.iam_instance_profile,
         }
         self.scheduler_options = {**self.options}
         self.worker_options = {**self.options}

dask_cloudprovider/cloudprovider.yaml

@@ -46,6 +46,10 @@ cloudprovider:
     # subnet_id: "" # Subnet ID for instances to. Defaults to all subnets in default VPC.
     # security_groups: [] # Security groups for instances. Will create a minimal Dask security group by default.
     filesystem_size: 40 # Default root filesystem size for scheduler and worker VMs in GB
+    key_name: null # SSH Key name to assign to instances
+    iam_instance_profile: {} # Iam role to assign to instances
+      # Arn: 'string'
+      # Name: 'string'
 
   azure:
     location: null # The Azure location to launch your cluster

requirements_test.txt

@@ -1,3 +1,3 @@
 pytest
 pytest-asyncio
-pytest-timeout
+pytest-timeout