[CZID-8390] Add sqs step notifications (#116)

* add current sqs queue to sqs notification step

* run black

* terraform fmt

* typing

* change to using sns

* add formatting

* add options to turn off step notifications

Author: rzlim08

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/sns_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

main.tf

@@ -64,6 +64,7 @@ module "sfn" {
   stage_vcpu_defaults           = var.stage_vcpu_defaults
   extra_env_vars                = var.extra_env_vars
   sqs_queues                    = var.sqs_queues
+  step_notifications            = var.step_notifications
   call_cache                    = var.call_cache
   output_status_json_files      = var.output_status_json_files
   tags                          = var.tags

miniwdl-plugins/sns_notification/README.md

@@ -0,0 +1 @@
+# sns_notifications

miniwdl-plugins/sns_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="sns_notification",
+    version="0.0.1",
+    description="miniwdl plugin for notification of task completion to Amazon SQS",
+    url="https://github.com/chanzuckerberg/swipe",
+    project_urls={},
+    long_description=long_description,
+    long_description_content_type="text/markdown",
+    author="",
+    py_modules=["sns_notification"],
+    python_requires=">=3.6",
+    setup_requires=["reentry"],
+    install_requires=["boto3"],
+    reentry_register=True,
+    entry_points={
+        "miniwdl.plugin.task": ["sns_notification_task = sns_notification:task"],
+        "miniwdl.plugin.workflow": [
+            "sns_notification_workflow = sns_notification:workflow"
+        ],
+    },
+)

miniwdl-plugins/sns_notification/sns_notification.py

@@ -0,0 +1,80 @@
+"""
+Send SNS notifications after each miniwdl step
+"""
+
+import os
+import json
+from typing import Dict
+from datetime import datetime
+from WDL import values_to_json
+from WDL._util import StructuredLogMessage as _
+
+import boto3
+
+sns_client = boto3.client("sns", endpoint_url=os.getenv("AWS_ENDPOINT_URL"))
+topic_arn = os.getenv('STEP_NOTIFICATION_TOPIC_ARN')
+
+
+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 SNS"""
+    sns_resp = sns_client.publish(
+        TopicArn=topic_arn,
+        Message=body,
+        MessageAttributes=attr,
+    )
+    return sns_resp
+
+
+def task(cfg, logger, run_id, run_dir, task, **recv):
+    """
+    on completion of any task sends a message to sns with the output files
+    """
+    log = logger.getChild("sns_step_notification")
+
+    # ignore inputs
+    recv = yield recv
+    # ignore command/runtime/container
+    recv = yield recv
+
+    log.info(_("sending message to sns"))
+
+    if topic_arn:
+        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):
+    log = logger.getChild("sns_step_notification")
+
+    # ignore inputs
+    recv = yield recv
+
+    log.info(_("ignores workflow calls"))
+    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]),
+    "STEP_NOTIFICATION_TOPIC_ARN"               = var.sfn_notification_topic_arn,
     "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,9 @@ variable "docker_network" {
   type        = string
   default     = ""
 }
+
+variable "sfn_notification_topic_arn" {
+  description = "ARN of notification sns topic"
+  type        = string
+}
+

terraform/modules/swipe-sfn/main.tf

@@ -32,19 +32,20 @@ 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_topic_arn = length(var.sqs_queues) > 0 && var.step_notifications ? aws_sns_topic.sfn_notifications_topic[0].arn : ""
+  tags                       = var.tags
 }
 
 module "sfn_io_helper" {

terraform/modules/swipe-sfn/notifications.tf

@@ -64,9 +64,10 @@ data "aws_iam_policy_document" "sfn_notifications_topic_policy_document" {
 resource "aws_sns_topic_subscription" "sfn_notifications_sqs_target" {
   for_each = var.sqs_queues
 
-  topic_arn = aws_sns_topic.sfn_notifications_topic[0].arn
-  protocol  = "sqs"
-  endpoint  = aws_sqs_queue.sfn_notifications_queue[each.key].arn
+  topic_arn            = aws_sns_topic.sfn_notifications_topic[0].arn
+  protocol             = "sqs"
+  endpoint             = aws_sqs_queue.sfn_notifications_queue[each.key].arn
+  raw_message_delivery = true
 }
 
 resource "aws_sqs_queue" "sfn_notifications_queue" {

terraform/modules/swipe-sfn/variables.tf

@@ -116,3 +116,10 @@ variable "metrics_schedule" {
   type        = string
   default     = "rate(1 minute)"
 }
+
+variable "step_notifications" {
+  description = "Boolean to determine whether or not to use send step notifications with SNS"
+  type        = bool
+  default     = false
+}
+

test/terraform/moto/main.tf

@@ -32,7 +32,7 @@ module "swipetest" {
     "Two" : { "spot" : 12800, "on_demand" : 256000 },
   }
 
-  workspace_s3_prefixes = ["swipe-test"]
-
+  workspace_s3_prefixes    = ["swipe-test"]
   output_status_json_files = true
+  step_notifications       = true
 }