Merge pull request #1 from moee/0.0.1

0.0.1

Author: moee

.gitignore

@@ -87,3 +87,5 @@ ENV/
 
 # Rope project settings
 .ropeproject
+
+*.swp

README.md

@@ -1,2 +1,45 @@
 # ecs_task_runner
 Helper tools to run a task on AWS ECS and follow the logs
+
+# Why?
+The use case why I created this script is that we want to run ECS tasks in Jenkins and directly see the output of the tasks that are being run. This is very cumbersome to achieve using the `aws cli`. Hence this library.
+
+# Installation
+Until this package is distributed as pip package, you have to install it directly from this repository:
+
+```
+pip install git+https://github.com/moee/[email protected]
+```
+
+# Usage
+## Example 1: Jenkins Integration
+
+```sh
+#!/bin/sh
+pip install git+https://github.com/moee/[email protected]
+
+python << END
+import ecstaskrunner, sys, logging
+
+logging.basicConfig()
+logging.getLogger('ecstaskrunner').setLevel(logging.INFO)
+
+sys.exit(
+    ecstaskrunner.run_task(
+        cluster="YOUR-CLUSTER-NAME",
+        taskDefinition='YOUR-TASK-DEFINITION',
+    )
+)
+END
+```
+This runs the task named `YOUR-TASK-DEFINITION` on the cluster `YOUR-CLUSTER-NAME`, displays all the output (Note: this only works if the container definition uses the awslogs driver) and waits for the task to stop. Only if all containers have stopped and exited with `0` the job will be marked as success.
+
+## Example 2: Get the log output of a task
+
+```python
+import ecstaskrunner
+task = ecstaskrunner.task.Task(cluster='YOUR-CLUSTER-NAME', taskId='YOUR-TASK-ID')
+for container in task.containers:
+    for line in task.containers[container].get_log_events():
+        print "%s: %s" % (container, line)
+```

ecstaskrunner/__init__.py

@@ -0,0 +1,45 @@
+import boto3
+import sys
+import time
+import logging
+import botocore.errorfactory
+import ecstaskrunner
+import ecstaskrunner.task
+
+def run_task(**kwargs):
+    client = boto3.client('ecs')
+
+    response = client.run_task(**kwargs)
+
+    logger = logging.getLogger('ecstaskrunner')
+
+    taskResponse = response['tasks'][0]
+
+    taskId = taskResponse['taskArn'][taskResponse['taskArn'].rfind("/")+1:]
+    logger.debug(taskId)
+
+    task = ecstaskrunner.task.Task(taskResponse['clusterArn'], taskId)
+
+    logger.debug("task is pending")
+    while task.isPending():
+        time.sleep(1)
+
+    while task.isRunning():
+        for container in task.containers:
+            for log_event in task.containers[container].get_log_events():
+                logger.info("%s: %s" % (container, log_event))
+    logger.debug("task status: %s" % task.getLastStatus())
+
+    exitCode = 0
+
+    for container in task.describeTask()['containers']:
+        if 'reason' in container:
+            logger.warn("%s failed: %s" % (container['name'], container['reason']))
+            exitCode = 1
+            continue
+            
+        logger.info("%s exited with code %d" % (container['name'], container['exitCode']))
+        if container['exitCode'] != 0:
+            exitCode = 2
+
+    return exitCode

ecstaskrunner/container.py

@@ -0,0 +1,61 @@
+import boto3
+import logging
+import datetime
+
+class Container:
+    def __init__(self, taskId, config, definition):
+        self.taskId = taskId
+        self.config = config
+        self.definition = definition
+        self.name = config['name']
+        self.startTime = None
+        self.logger = logging.getLogger("Container")
+
+    def get_log_events(self):
+        logConfig = self.definition['logConfiguration']
+        tasklogger = logging.getLogger(self.name)
+
+        logs = boto3.client('logs')
+
+        logStreamName = '%s/%s/%s' % (
+            logConfig['options']['awslogs-stream-prefix'],
+            self.name,
+            self.taskId
+        )
+        nextToken = False 
+        while True:
+            a = {
+                'logGroupName': logConfig['options']['awslogs-group'],
+                'logStreamName': logStreamName,
+                'startFromHead': True,
+                'limit': 2
+            }
+
+            if nextToken:
+                a['nextToken'] = nextToken
+            else:    
+                if self.startTime:
+                    a['startTime'] = self.startTime
+
+            try:
+                response = logs.get_log_events(**a)
+            except Exception as e:
+                # todo not sure why i cannot check for the class directly
+                if e.__class__.__name__ == 'ResourceNotFoundException':
+                    self.logger.warn(e) 
+                    return
+                raise e
+            for event in response['events']:
+                yield "[%s] %s" % (
+                    datetime.datetime.fromtimestamp(
+                        event['timestamp']/1000
+                    ),
+                    event['message'])
+                self.startTime = event['timestamp']+1
+
+            if len(response['events']) != a['limit']:
+                self.logger.debug("[EOS]")
+                break
+
+            nextToken = response['nextForwardToken']
+

ecstaskrunner/task.py

@@ -0,0 +1,46 @@
+import boto3
+import logging
+from ecstaskrunner.container import Container
+
+class Task:
+    def __init__(self, cluster, taskId):
+        self.cluster = cluster
+        self.taskId = taskId
+        self.containers = {}
+        self.client = boto3.client('ecs')
+        self.logger = logging.getLogger('Task')
+       
+        task = self.describeTask()
+    
+        if not task:
+            self.logger.warn("Task with id %s does not exist" % taskId) 
+            return
+
+        self.taskDefinitionArn = task['taskDefinitionArn']
+        containers = task['containers']
+        response = self.client.describe_task_definition(taskDefinition=self.taskDefinitionArn)
+        containerDefinitions = response['taskDefinition']['containerDefinitions']
+
+        self.logger.debug("task definition arn: %s" % self.taskDefinitionArn)
+
+        for container in self.describeTask()['containers']:
+            self.containers[container['name']] = Container(
+                self.taskId,
+                container,
+                [x for x in containerDefinitions if x['name'] == container['name']][0]
+            )
+
+    def describeTask(self):
+        response = self.client.describe_tasks(cluster=self.cluster, tasks=[self.taskId])
+        if len(response['tasks']) != 1:
+            return None
+        return response['tasks'][0]
+
+    def getLastStatus(self):
+        return self.describeTask()['lastStatus']
+
+    def isRunning(self):
+        return self.getLastStatus() == "RUNNING"
+
+    def isPending(self):
+        return self.getLastStatus() == "PENDING"

setup.py

@@ -0,0 +1,14 @@
+#!/usr/bin/env python
+
+from distutils.core import setup
+
+setup(name='ecstaskrunner',
+    version='0.0.1',
+    description='Helper tools to run a task on AWS ECS and follow the logs',
+    author='Michael Osl',
+    author_email='[email protected]',
+    url='https://github.com/moee/ecstaskrunner',
+    packages=['ecstaskrunner'],
+    install_requires=['boto3']
+)
+