PS-547 - Added Support for Trying Multiple Different Launch Templates (…

…#12)

* added support for trying multiple different launch templates

* fix name of parameter

* added error handling for boto3 clientError

* added instance to the TempInstance class

* updated README and added split on comma

* updated docstring and Tempinstance description

* updated README.md with more detailed description and example

* updated example in readme

.env.sample

@@ -1,5 +1,5 @@
 NAME=
-LAUNCH_TEMPLATE_NAME=
+LAUNCH_TEMPLATE_NAMES=
 SUBNET_ID=
 AWS_DEFAULT_REGION=
 AWS_ACCESS_KEY_ID=

README.md

@@ -11,7 +11,12 @@ Docker image creates a Key Pair and initializes an EC2 instance both with the `N
 These variables can be [passed into the Docker run](https://docs.docker.com/engine/reference/commandline/run/#set-environment-variables--e---env---env-file) using the `-e` or `--env-file` flags:
 
 - `NAME` - Used to name EC2 Key Pair and Instance. UUID is appended to the end.
-- `LAUNCH_TEMPLATE_NAME` - Name of launch template to launch EC2 Instance.
+- `LAUNCH_TEMPLATE_NAMES` - A comma separated list of launch templates to be used to
+create EC2 instances. The launch templates will be used in the order they are passed in.
+If creating an EC2 instance with a launch template fails, then the next launch template given
+in the list will be used. If creating an EC2 instance with every launch templates fails,
+then the program will exit.
+Example: `launch-template-1,launch-template-2`
 - `SUBNET_ID` - ID for subnet to launch EC2 Instance in.
 - `AWS_DEFAULT_REGION` - AWS region to launch EC2 Instance in.
 - `AWS_ACCESS_KEY_ID` - AWS access key ID used to create EC2 Key Pair and Instance.

lib/ec2.py

@@ -1,7 +1,9 @@
 import logging
 import sys
+from typing import List
 
 import boto3
+from botocore.exceptions import ClientError
 
 EC2_CLIENT = boto3.client('ec2')
 EC2_RESOURCE = boto3.resource('ec2')
@@ -37,7 +39,11 @@ def __exit__(self, type, value, traceback):
 
 class TempInstance():
     """
-    Create a temporary EC2 Instance from a launch template.
+    Create a temporary EC2 Instance from a list of launch templates.
+    If creating an EC2 instance with a launch template fails,
+    then the next launch template given in the list will be used.
+    If creating an EC2 instance with every launch templates fails,
+    then the program will exit.
 
     On enter an Instance is created and tagged with a name. The Instance is
     returned after it is successfully running.
@@ -46,53 +52,71 @@ class TempInstance():
 
     Arguments:
         name: Name to tag Instance with.
-        launch_template_name: Name of launch template to launch Instance with.
+        launch_template_names: A list of launch templates to be
+                               used to create EC2 instances.
+                               The launch templates will be used
+                               in the order they are passed in.
         key_name: Name of Key Pair to associate Instance with.
         subnet_id: ID for subnet to launch Instance in.
     """
 
     def __init__(
         self,
         name: str,
-        launch_template_name: str,
+        launch_template_names: List[str],
         key_name: str,
         subnet_id: str,
     ):
         self.name = name
-        self.launch_template_name = launch_template_name
+        self.launch_template_names = launch_template_names
         self.key_name = key_name
         self.subnet_id = subnet_id
+        self.instance = None
 
-    def __enter__(self):
-        logger.info(f'Launching instance {self.name}...')
-        self.instance = EC2_RESOURCE.create_instances(
-            LaunchTemplate={'LaunchTemplateName': self.launch_template_name},
-            KeyName=self.key_name,
-            MinCount=1,
-            MaxCount=1,
-            SubnetId=self.subnet_id,
-            TagSpecifications=[
-                {
-                    'ResourceType': 'instance',
-                    'Tags': [
-                        {
-                            'Key': 'Name',
-                            'Value': self.name,
-                        },
-                    ],
-                },
-            ],
-        )[0]
-        logger.info(f'Waiting for instance {self.instance.instance_id} to be '
-                    'ready...')
+    def __launch_instance(self, index: int = 0):
+        if index == len(self.launch_template_names):
+            raise Exception('No instance can be launched.')
+
+        try:
+            logger.info(f'Launching instance {self.name} using template: {self.launch_template_names[index]}')
+            self.instance = EC2_RESOURCE.create_instances(
+                LaunchTemplate={'LaunchTemplateName': self.launch_template_names[index]},
+                KeyName=self.key_name,
+                MinCount=1,
+                MaxCount=1,
+                SubnetId=self.subnet_id,
+                TagSpecifications=[
+                    {
+                        'ResourceType': 'instance',
+                        'Tags': [
+                            {
+                                'Key': 'Name',
+                                'Value': self.name,
+                            },
+                        ],
+                    },
+                ],
+            )[0]
+        except ClientError as e:
+            error_code = e.response['Error']['Code']
+            if error_code == 'InsufficientInstanceCapacity' or error_code == 'SpotMaxPriceTooLow':
+                logger.error(f'The use of launch template: {self.launch_template_names[index]} failed.')
+                return self.__launch_instance(index+1)
+            else:
+                raise
+
+        logger.info(f'Waiting for instance {self.instance.instance_id} to be ready...')
+
+        self.instance.wait_until_running()
+        return self.instance
 
+    def __enter__(self):
         try:
-            self.instance.wait_until_running()
-            return self.instance
-        except BaseException:
+            return self.__launch_instance()
+        finally:
             self.__exit__(*sys.exc_info())
-            raise
 
     def __exit__(self, type, value, traceback):
-        logger.info(f'Terminating instance {self.instance.instance_id}...')
-        self.instance.terminate()
+        if self.instance:
+            logger.info(f'Terminating instance {self.instance.instance_id}...')
+            self.instance.terminate()

main.py

@@ -16,15 +16,15 @@ def main():
 
     name = os.environ['NAME']
     key_name = f'{name}-{uuid4()}'
-    launch_template_name = os.environ['LAUNCH_TEMPLATE_NAME']
+    launch_template_names = os.environ['LAUNCH_TEMPLATE_NAMES'].split(',')
     subnet_id = os.environ['SUBNET_ID']
     command = ' '.join(sys.argv[1:])
 
     with ExitStack() as stack:
         private_key = stack.enter_context(TempKeyPair(key_name))
         instance = stack.enter_context(TempInstance(
             name,
-            launch_template_name,
+            launch_template_names,
             key_name,
             subnet_id,
         ))