Merge pull request #20 from aelzeiny/ami-helper

Ami helper

.github/workflows/main.yml

@@ -1,7 +1,7 @@
 # This workflow will install Python dependencies, run tests and lint with a variety of Python versions
 # For more information see: https://help.github.com/actions/language-and-framework-guides/using-python-with-github-actions
 
-name: AWS Airflow Executor Tester
+name: AWS Airflow Executors
 
 on:
   push:

batch_ami_helper.py

@@ -0,0 +1,42 @@
+"""
+It turns out that the #1 source of grief is setting up AWS correctly.
+This file will help folks debug their AWS Permissions
+"""
+import logging
+import time
+import unittest
+from airflow_aws_executors.batch_executor import AwsBatchExecutor
+from airflow.utils.state import State
+
+
+class BatchAMIHelper(unittest.TestCase):
+    def setUp(self):
+        self._log = logging.getLogger("FargateAMIHelper")
+        self.executor = AwsBatchExecutor()
+        self.executor.start()
+
+    def test_boto_submit_job(self):
+        self.executor._submit_job(None, ['airflow', 'version'], None)
+
+    def test_boto_describe_job(self):
+        job_id = self.executor._submit_job(None, ['airflow', 'version'], None)
+        self.executor._describe_tasks([job_id])
+
+    def test_boto_terminate_job(self):
+        job_id = self.executor._submit_job(None, ['airflow', 'version'], None)
+        self.executor.batch.terminate_job(
+            jobId=job_id,
+            reason='Testing AMI permissions'
+        )
+
+    def test_sample_airflow_task(self):
+        job_id = self.executor._submit_job(None, ['airflow', 'version'], None)
+        job = None
+        while job is None or job.get_job_state() == State.QUEUED:
+            responses = self.executor._describe_tasks([job_id])
+            assert responses, 'No response received'
+            job = responses[0]
+            time.sleep(1)
+
+        self.assertEqual(job.get_job_state(), State.SUCCESS, 'AWS Batch Job did not run successfully!')
+

getting_started_batch.md

@@ -48,38 +48,48 @@ Build this image and upload it to a private repository. You may want to use Dock
     the container. *BE SURE YOU HAVE THIS*!  It's heavily important that commands like 
     `["airflow", "run", <dag_id>, <task_id>, <execution_date>]` are accepted by your container's entrypoint script.
 
- 4. Run through the AWS Batch Creation Wizard on AWS Console. The executor does not have any
+4. Run through the AWS Batch Creation Wizard on AWS Console. The executor does not have any
  prerequisites to how you create your Job Queue or Compute Environment. Go nuts; have at it. I'll refer you to the 
  [AWS Docs' Getting Started with Batch](https://docs.aws.amazon.com/batch/latest/userguide/Batch_GetStarted.html).
  You will need to assign the right IAM roles for the remote S3 logging. 
  Also, your dynamically provisioned EC2 instances do not need to be connected to the public internet, 
  private subnets in private VPCs are encouraged. However, be sure that all instances has access to your Airflow MetaDB.
 
- 5. When creating a Job Definition choose the 'container' type and point to the private repo. The 'commands' array is
+5. When creating a Job Definition choose the 'container' type and point to the private repo. The 'commands' array is
  optional on the task-definition level. At runtime, Airflow commands will be injected here by the AwsBatchExecutor!
 
- 6. Let's go back to that machine in step #1 that's running the Scheduler. We'll use the same docker container as 
+6. Let's go back to that machine in step #1 that's running the Scheduler. We'll use the same docker container as 
  before; except we'll do something like `docker run ... airflow webserver` and `docker run ... airflow scheduler`. 
  Here are the minimum IAM roles that the executor needs to launch tasks, feel free to tighten the resources around the 
  job-queues and compute environments that you'll use.
-    ```json
-    {
-        "Version": "2012-10-17",
-        "Statement": [
-            {
-                "Sid": "AirflowBatchRole",
-                "Effect": "Allow",
-                "Action": [
-                    "batch:SubmitJob",
-                    "batch:DescribeJobs",
-                    "batch:TerminateJob"
-                ],
-                "Resource": "*"
-            }
-        ]
-    }
-    ```
- 7. You're done. Configure & launch your scheduler. However, maybe you did something real funky with your AWS Batch compute
+   ```json
+   {
+       "Version": "2012-10-17",
+       "Statement": [
+           {
+               "Sid": "AirflowBatchRole",
+               "Effect": "Allow",
+               "Action": [
+                   "batch:SubmitJob",
+                   "batch:DescribeJobs",
+                   "batch:TerminateJob"
+               ],
+               "Resource": "*"
+           }
+       ]
+   }
+   ```
+7. Let's test your AMI configuration before we launch the webserver or scheduler. 
+[Copy & Run this python file somewhere on the machine that has your scheduler](batch_ami_helper.py).
+To simplify, this file will make sure that your scheduler and your cluster have correct AMI permissions. Your scheduler
+should be able to launch, describe, plus terminate Batch jobs, and connect to your Airflow Meta DB of choice. 
+Meanwhile, your cluster should have the ability to pull the docker container, and also 
+connect to your Airflow MetaDB of choice.
+```bash
+python3 -m unittest batch_ami_helper.py
+```
+
+9. You're done. Configure & launch your scheduler. However, maybe you did something real funky with your AWS Batch compute
  environment. The good news is that you have full control over how the executor submits jobs. 
  See the [#Extensibility](./readme.md) section in the readme. Thank you for taking the time to set this up!