[CZID-8390] Add sqs step notifications

Dockerfile

@@ -46,6 +46,7 @@ RUN apt-get -q update && apt-get -q install -y \
 #   upgrade because of this issue https://github.com/chanzuckerberg/miniwdl/issues/607 in miniwdl
 RUN pip3 install importlib-metadata==4.13.0
 RUN pip3 install miniwdl==${MINIWDL_VERSION}
+RUN pip3 install urllib3==1.26.16
 
 RUN curl -Ls https://github.com/chanzuckerberg/s3parcp/releases/download/v1.0.1/s3parcp_1.0.1_linux_amd64.tar.gz | tar -C /usr/bin -xz s3parcp
 
@@ -62,6 +63,7 @@ ADD miniwdl-plugins miniwdl-plugins
 RUN pip install miniwdl-plugins/s3upload
 RUN pip install miniwdl-plugins/sfn_wdl
 RUN pip install miniwdl-plugins/s3parcp_download
+RUN pip install miniwdl-plugins/sqs_notification
 
 RUN cd /usr/bin; curl -O https://amazon-ecr-credential-helper-releases.s3.amazonaws.com/0.4.0/linux-amd64/docker-credential-ecr-login
 RUN chmod +x /usr/bin/docker-credential-ecr-login

miniwdl-plugins/sqs_notification/README.md

@@ -0,0 +1 @@
+# sqs_notifications

miniwdl-plugins/sqs_notification/setup.py

@@ -0,0 +1,29 @@
+#!/usr/bin/env python3
+from setuptools import setup
+from os import path
+
+this_directory = path.abspath(path.dirname(__file__))
+with open(path.join(path.dirname(__file__), "README.md")) as f:
+    long_description = f.read()
+
+setup(
+    name="sqs_notification",
+    version="0.0.1",
+    description="miniwdl plugin for notification of task completion to Amazon SQS",
+    url="https://github.com/chanzuckerberg/miniwdl-s3upload",
+    project_urls={},
+    long_description=long_description,
+    long_description_content_type="text/markdown",
+    author="",
+    py_modules=["sqs_notification"],
+    python_requires=">=3.6",
+    setup_requires=["reentry"],
+    install_requires=["boto3"],
+    reentry_register=True,
+    entry_points={
+        "miniwdl.plugin.task": ["sqs_notification_task = sqs_notification:task"],
+        "miniwdl.plugin.workflow": [
+            "sqs_notification_workflow = sqs_notification:workflow"
+        ],
+    },
+)

miniwdl-plugins/sqs_notification/sqs_notification.py

@@ -0,0 +1,83 @@
+"""
+TODO
+"""
+
+import os
+import json
+from typing import Dict
+from datetime import datetime
+from WDL import values_to_json
+
+import boto3
+
+sqs_client = boto3.client("sqs", endpoint_url=os.getenv("AWS_ENDPOINT_URL"))
+queue_url = os.getenv("AWS_STEP_NOTIFICATION_PLUGIN")
+
+
+def process_outputs(outputs: Dict):
+    """process outputs dict into string to be passed into SQS"""
+    # only stringify for now
+    return json.dumps(outputs)
+
+
+def send_message(attr, body):
+    """send message to SQS, eventually wrap this in a try catch to deal with throttling"""
+    sqs_resp = sqs_client.send_message(
+        QueueUrl=queue_url,
+        DelaySeconds=0,
+        MessageAttributes=attr,
+        MessageBody=body,
+    )
+    return sqs_resp
+
+
+def task(cfg, logger, run_id, run_dir, task, **recv):
+    """
+    on completion of any task, upload its output files to S3, and record the S3 URI corresponding
+    to each local file (keyed by inode) in _uploaded_files
+    """
+    logger = logger.getChild("s3_progressive_upload")
+
+    # ignore inputs
+    recv = yield recv
+
+    # ignore command/runtime/container
+    recv = yield recv
+
+    message_attributes = {
+        "WorkflowName": {"DataType": "String", "StringValue": run_id[0]},
+        "TaskName": {"DataType": "String", "StringValue": run_id[-1]},
+        "ExecutionId": {
+            "DataType": "String",
+            "StringValue": "execution_id_to_be_passed_in",
+        },
+    }
+
+    outputs = process_outputs(values_to_json(recv["outputs"]))
+    message_body = {
+        "version": "0",
+        "id": "0",
+        "detail-type": "Step Functions Execution Step Notification",
+        "source": "aws.batch",
+        "account": "",
+        "time": datetime.now().strftime("%Y-%m-%dT%H:%M:%SZ"),
+        "resources": [],
+        "detail": outputs,
+    }
+    send_message(message_attributes, json.dumps(message_body))
+
+    yield recv
+
+
+def workflow(cfg, logger, run_id, run_dir, workflow, **recv):
+    """
+    on workflow completion, add a file outputs.s3.json to the run directory, which is outputs.json
+    with local filenames rewritten to the uploaded S3 URIs (as previously recorded on completion of
+    each task).
+    """
+    logger = logger.getChild("s3_progressive_upload")
+
+    # ignore inputs
+    recv = yield recv
+
+    yield recv

terraform/modules/swipe-sfn-batch-job/main.tf

@@ -34,6 +34,7 @@ locals {
     "MINIWDL__DOWNLOAD_CACHE__DISABLE_PATTERNS" = "[\"s3://swipe-samples-*/*\"]",
     "DOWNLOAD_CACHE_MAX_GB"                     = "500",
     "WDL_PASSTHRU_ENVVARS"                      = join(" ", [for k, v in var.extra_env_vars : k]),
+    "AWS_STEP_NOTIFICATION_PLUGIN"              = var.sfn_notification_queue_urls[0],
     "OUTPUT_STATUS_JSON_FILES"                  = tostring(var.output_status_json_files)
   })
   container_env_vars     = { "environment" : [for k in sort(keys(local.batch_env_vars)) : { "name" : k, "value" : local.batch_env_vars[k] }] }

terraform/modules/swipe-sfn-batch-job/variables.tf

@@ -70,3 +70,13 @@ variable "docker_network" {
   type        = string
   default     = ""
 }
+
+variable "sfn_notification_queue_arns" {
+  description = "ARNs of notification SQS queues"
+  type        = list(string)
+}
+
+variable "sfn_notification_queue_urls" {
+  description = "URLs of notification SQS queues"
+  type        = list(string)
+}

terraform/modules/swipe-sfn/main.tf

@@ -32,19 +32,21 @@ resource "aws_iam_role_policy_attachment" "swipe_sfn_service" {
 }
 
 module "batch_job" {
-  source                    = "../swipe-sfn-batch-job"
-  app_name                  = var.app_name
-  batch_job_docker_image    = var.batch_job_docker_image
-  batch_job_timeout_seconds = var.batch_job_timeout_seconds
-  miniwdl_dir               = var.miniwdl_dir
-  workspace_s3_prefixes     = var.workspace_s3_prefixes
-  wdl_workflow_s3_prefix    = var.wdl_workflow_s3_prefix
-  job_policy_arns           = var.job_policy_arns
-  extra_env_vars            = var.extra_env_vars
-  docker_network            = var.docker_network
-  call_cache                = var.call_cache
-  output_status_json_files  = var.output_status_json_files
-  tags                      = var.tags
+  source                      = "../swipe-sfn-batch-job"
+  app_name                    = var.app_name
+  batch_job_docker_image      = var.batch_job_docker_image
+  batch_job_timeout_seconds   = var.batch_job_timeout_seconds
+  miniwdl_dir                 = var.miniwdl_dir
+  workspace_s3_prefixes       = var.workspace_s3_prefixes
+  wdl_workflow_s3_prefix      = var.wdl_workflow_s3_prefix
+  job_policy_arns             = var.job_policy_arns
+  extra_env_vars              = var.extra_env_vars
+  docker_network              = var.docker_network
+  call_cache                  = var.call_cache
+  output_status_json_files    = var.output_status_json_files
+  sfn_notification_queue_arns = [for name, queue in aws_sqs_queue.sfn_notifications_queue : queue.arn]
+  sfn_notification_queue_urls = [for name, queue in aws_sqs_queue.sfn_notifications_queue : queue.url]
+  tags                        = var.tags
 }
 
 module "sfn_io_helper" {