Initial commit moves code from private repo to public repo (#1)

.dockerignore

@@ -0,0 +1,4 @@
+**/__pycache__/
+**/*.py[cod]
+**/*$py.class
+.env

.env.sample

@@ -0,0 +1,6 @@
+NAME=
+LAUNCH_TEMPLATE_NAME=
+SUBNET_ID=
+AWS_DEFAULT_REGION=
+AWS_ACCESS_KEY_ID=
+AWS_SECRET_ACCESS_KEY=

.gitignore

@@ -0,0 +1,4 @@
+__pycache__/
+*.py[cod]
+*$py.class
+.env

Dockerfile

@@ -0,0 +1,11 @@
+FROM continuumio/miniconda:4.5.4
+
+SHELL ["/bin/bash", "-c"]
+WORKDIR /opt/app
+
+COPY environment.yml .
+RUN conda env create --file environment.yml
+
+COPY . .
+
+ENTRYPOINT ["./entrypoint.sh"]

README.md

@@ -0,0 +1,58 @@
+# Run on EC2
+
+Executable Docker image for running commands on custom ephemeral EC2 instances defined by launch templates.
+
+## Usage
+
+Docker image creates a Key Pair and initializes an EC2 instance both with the `NAME` variable. The command is run on the EC2 instance via SSH using the Key Pair. The Key Pair is deleted and EC2 instance terminated upon success, kill, or failure of the command.
+
+### Environment Variables
+
+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.
+- `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.
+- `AWS_SECRET_ACCESS_KEY` - AWS secret access key used to create EC2 Key Pair and Instance.
+
+### Example
+
+```sh
+$ cp .env.sample .env # Fill out .env file
+$ docker build -t run-on-ec2 .
+$ docker run --rm -it --env-file=.env run-on-ec2 echo \"hello world\"
+Creating key pair run-on-ec2-eb5f9910-1635-40e1-b120-0e08b06a60ce...
+Launching instance run-on-ec2-eb5f9910-1635-40e1-b120-0e08b06a60ce...
+Waiting for instance run-on-ec2-eb5f9910-1635-40e1-b120-0e08b06a60ce to be ready...
+Running command echo "hello world" on 10.128.130.201...
+hello world
+Terminating instance run-on-ec2-eb5f9910-1635-40e1-b120-0e08b06a60ce...
+Deleting key pair run-on-ec2-eb5f9910-1635-40e1-b120-0e08b06a60ce...
+```
+
+## Development
+
+### Setup Environment
+
+Use [Miniconda](https://conda.io/miniconda.html) to setup Python virtual environment.
+
+```sh
+$ conda env create -f environment.yml
+$ source activate run-on-ec2
+```
+
+### Run Locally
+
+```sh
+$ cp .env.sample .env # Fill out .env file
+$ env $(cat .env | xargs) python main.py echo \"hello world\"
+Creating key pair run-on-ec2-eb5f9910-1635-40e1-b120-0e08b06a60ce...
+Launching instance run-on-ec2-eb5f9910-1635-40e1-b120-0e08b06a60ce...
+Waiting for instance run-on-ec2-eb5f9910-1635-40e1-b120-0e08b06a60ce to be ready...
+Running command echo "hello world" on 10.128.130.201...
+hello world
+Terminating instance run-on-ec2-eb5f9910-1635-40e1-b120-0e08b06a60ce...
+Deleting key pair run-on-ec2-eb5f9910-1635-40e1-b120-0e08b06a60ce...
+```

entrypoint.sh

@@ -0,0 +1,4 @@
+#!/bin/bash
+set -e
+source activate run-on-ec2
+python main.py "$@"

environment.yml

@@ -0,0 +1,9 @@
+name: run-on-ec2
+channels:
+  - conda-forge
+  - defaults
+dependencies:
+  - boto3=1.9.31
+  - python=3.7.0
+  - fabric=2.4.0
+

lib/ec2.py

@@ -0,0 +1,86 @@
+import sys
+import boto3
+
+EC2_CLIENT = boto3.client('ec2')
+EC2_RESOURCE = boto3.resource('ec2')
+
+
+class TempKeyPair:
+    """
+    Create a temporary EC2 Key Pair.
+
+    On enter a Key Pair is created. The private key is returned.
+
+    On exit the Key Pair is deleted.
+
+    :param name: Name of Key Pair to create.
+    :type name: str
+    """
+
+    def __init__(self, name):
+        self.name = name
+
+    def __enter__(self):
+        print(f'Creating key pair {self.name}...')
+        key_pair = EC2_CLIENT.create_key_pair(KeyName=self.name)
+        pem = key_pair['KeyMaterial']
+        return pem
+
+    def __exit__(self, type, value, traceback):
+        print(f'Deleting key pair {self.name}...')
+        EC2_CLIENT.delete_key_pair(KeyName=self.name)
+
+
+class TempInstance():
+    """
+    Create a temporary EC2 Instance from a launch template.
+
+    On enter an Instance is created and tagged with a name. The Instance is
+    returned after it is successfully running.
+
+    On exit the Instance is terminated.
+
+    :param name: Name to tag Instance with.
+    :type name: str
+    :param launch_template_name: Name of launch template to launch Instance with.
+    :type launch_template_name: str
+    :param subnet_id: ID for subnet to launch Instance in.
+    :type subnet_id: str
+    """
+
+    def __init__(self, name, launch_template_name, subnet_id):
+        self.name = name
+        self.launch_template_name = launch_template_name
+        self.subnet_id = subnet_id
+
+    def __enter__(self):
+        print(f'Launching instance {self.name}...')
+        self.instance = EC2_RESOURCE.create_instances(
+            LaunchTemplate={'LaunchTemplateName': self.launch_template_name},
+            KeyName=self.name,
+            MinCount=1,
+            MaxCount=1,
+            SubnetId=self.subnet_id,
+            TagSpecifications=[
+                {
+                    'ResourceType': 'instance',
+                    'Tags': [
+                        {
+                            'Key': 'Name',
+                            'Value': self.name,
+                        },
+                    ],
+                },
+            ],
+        )[0]
+
+        try:
+            self.instance.wait_until_running()
+            return self.instance
+        except BaseException:
+            self.__exit__(*sys.exc_info())
+            raise
+
+    def __exit__(self, type, value, traceback):
+        print(f'Terminating instance {self.name}...')
+        self.instance.terminate()

lib/ssh.py

@@ -0,0 +1,54 @@
+from io import StringIO
+from time import sleep
+from fabric import Connection
+from paramiko import RSAKey
+from paramiko.ssh_exception import NoValidConnectionsError
+
+
+class SSH():
+    """
+    SSH context manager for creating an SSH connection.
+
+    On enter an SSH connection is attempted every 5 seconds until successful.
+    An exception is raised after 60 attempts.
+
+    On exit the connection is closed.
+
+    :param host: Host to connect to.
+    :type host: str
+    :param user: User to connect with.
+    :type user: str
+    :param private_key: RSA private key.
+    :type private: str
+    """
+
+    def __init__(self, host, user, private_key):
+        self.host = host
+        self.user = user
+        self.private_key = RSAKey.from_private_key(StringIO(private_key))
+        self.wait_count = 0
+
+    def __enter__(self):
+        self.connection = Connection(
+            host=self.host,
+            user=self.user,
+            connect_kwargs={'pkey': self.private_key},
+        )
+        print(f'Waiting for SSH to become available on {self.host}...')
+        self.wait_for_ssh()
+        return self.connection
+
+    def __exit__(self, type, value, traceback):
+        print(f'Closing SSH connection to {self.host}...')
+        self.connection.close()
+
+    def wait_for_ssh(self):
+        self.wait_count += 1
+        try:
+            self.connection.open()
+        except NoValidConnectionsError:
+            if self.wait_count >= 60:
+                raise
+
+            sleep(5)
+            self.wait_for_ssh()

main.py

@@ -0,0 +1,26 @@
+from contextlib import ExitStack
+import os
+import sys
+from uuid import uuid4
+from lib.ec2 import TempKeyPair, TempInstance
+from lib.ssh import SSH
+
+if __name__ == '__main__':
+    name = os.environ['NAME']
+    unique_name = f'{name}-{uuid4()}'
+    launch_template_name = os.environ['LAUNCH_TEMPLATE_NAME']
+    subnet_id = os.environ['SUBNET_ID']
+    command = ' '.join(sys.argv[1:])
+
+    with ExitStack() as stack:
+        private_key = stack.enter_context(TempKeyPair(unique_name))
+        instance = stack.enter_context(TempInstance(
+            unique_name,
+            launch_template_name,
+            subnet_id,
+        ))
+        host = instance.private_ip_address
+        connection = stack.enter_context(SSH(host, 'ec2-user', private_key))
+
+        print(f'Running command {command} on {host}...')
+        connection.run(command)