feature: Add support for Callback steps in model building pipelines

src/sagemaker/workflow/callback_step.py

@@ -0,0 +1,137 @@
+# Copyright 2020 Amazon.com, Inc. or its affiliates. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License"). You
+# may not use this file except in compliance with the License. A copy of
+# the License is located at
+#
+#     http://aws.amazon.com/apache2.0/
+#
+# or in the "license" file accompanying this file. This file is
+# distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF
+# ANY KIND, either express or implied. See the License for the specific
+# language governing permissions and limitations under the License.
+"""The step definitions for workflow."""
+from __future__ import absolute_import
+
+from typing import List, Dict
+from enum import Enum
+
+import attr
+
+from sagemaker.workflow.entities import (
+    RequestType,
+)
+from sagemaker.workflow.properties import (
+    Properties,
+)
+from sagemaker.workflow.entities import (
+    DefaultEnumMeta,
+)
+from sagemaker.workflow.steps import Step, StepTypeEnum, CacheConfig
+
+
+class CallbackOutputTypeEnum(Enum, metaclass=DefaultEnumMeta):
+    """CallbackOutput type enum."""
+
+    String = "String"
+    Integer = "Integer"
+    Boolean = "Boolean"
+    Float = "Float"
+
+
+@attr.s
+class CallbackOutput:
+    """Output for a callback step.
+
+    Attributes:
+        output_name (str): The output name
+        output_type (CallbackOutputTypeEnum): The output type
+    """
+
+    output_name: str = attr.ib(default=None)
+    output_type: CallbackOutputTypeEnum = attr.ib(default=CallbackOutputTypeEnum.String.value)
+
+    def to_request(self) -> RequestType:
+        """Get the request structure for workflow service calls."""
+        return {
+            "OutputName": self.output_name,
+            "OutputType": self.output_type.value,
+        }
+
+    @property
+    def expr(self) -> Dict[str, str]:
+        """The 'Get' expression dict for a `Parameter`."""
+        return CallbackOutput._expr(self.output_name)
+
+    @classmethod
+    def _expr(cls, name):
+        """An internal classmethod for the 'Get' expression dict for a `CallbackOutput`.
+
+        Args:
+            name (str): The name of the callback output.
+        """
+        return {"Get": f"Steps.{name}.OutputParameters['{name}']"}
+
+
+class CallbackStep(Step):
+    """Callback step for workflow."""
+
+    def __init__(
+        self,
+        name: str,
+        sqs_queue_url: str,
+        inputs: dict,
+        outputs: List[CallbackOutput],
+        cache_config: CacheConfig = None,
+        depends_on: List[str] = None,
+    ):
+        """Constructs a CallbackStep.
+
+        Args:
+            name (str): The name of the callback step.
+            sqs_queue_url (str): An SQS queue URL for receiving callback messages.
+            inputs (dict): Input arguments that will be provided
+                in the SQS message body of callback messages.
+            outputs (List[CallbackOutput]): Outputs that can be provided when completing a callback.
+            cache_config (CacheConfig):  A `sagemaker.workflow.steps.CacheConfig` instance.
+            depends_on (List[str]): A list of step names this `sagemaker.workflow.steps.TransformStep`
+                depends on
+        """
+        super(CallbackStep, self).__init__(name, StepTypeEnum.CALLBACK, depends_on)
+        self.sqs_queue_url = sqs_queue_url
+        self.outputs = outputs
+        self.cache_config = cache_config
+        self.inputs = inputs
+
+        root_path = f"Steps.{name}"
+        root_prop = Properties(path=root_path)
+
+        property_dict = {}
+        for output in outputs:
+            property_dict[output.output_name] = Properties(
+                f"{root_path}.OutputParameters['{output.output_name}']"
+            )
+
+        root_prop.__dict__["Outputs"] = property_dict
+        self._properties = root_prop
+
+    @property
+    def arguments(self) -> RequestType:
+        """The arguments dict that is used to define the callback step."""
+        return self.inputs
+
+    @property
+    def properties(self):
+        """A Properties object representing the output parameters of the callback step."""
+        return self._properties
+
+    def to_request(self) -> RequestType:
+        """Updates the dictionary with cache configuration."""
+        request_dict = super().to_request()
+        if self.cache_config:
+            request_dict.update(self.cache_config.config)
+
+        request_dict["SqsQueueUrl"] = self.sqs_queue_url
+        request_dict["OutputParameters"] = list(map(lambda op: op.to_request(), self.outputs))
+
+        return request_dict

src/sagemaker/workflow/pipeline.py

@@ -24,6 +24,7 @@
 
 from sagemaker._studio import _append_project_tags
 from sagemaker.session import Session
+from sagemaker.workflow.callback_step import CallbackOutput
 from sagemaker.workflow.entities import (
     Entity,
     Expression,
@@ -281,7 +282,7 @@ def _interpolate(obj: Union[RequestType, Any]):
     Args:
         obj (Union[RequestType, Any]): The request dict.
     """
-    if isinstance(obj, (Expression, Parameter, Properties)):
+    if isinstance(obj, (Expression, Parameter, Properties, CallbackOutput)):
         return obj.expr
     if isinstance(obj, dict):
         new = obj.__class__()

src/sagemaker/workflow/steps.py

@@ -54,6 +54,7 @@ class StepTypeEnum(Enum, metaclass=DefaultEnumMeta):
     REGISTER_MODEL = "RegisterModel"
     TRAINING = "Training"
     TRANSFORM = "Transform"
+    CALLBACK = "Callback"
 
 
 @attr.s

tests/integ/test_workflow.py

@@ -45,6 +45,7 @@
 from sagemaker.spark.processing import PySparkProcessor, SparkJarProcessor
 from sagemaker.workflow.conditions import ConditionGreaterThanOrEqualTo, ConditionIn
 from sagemaker.workflow.condition_step import ConditionStep
+from sagemaker.workflow.callback_step import CallbackStep, CallbackOutput, CallbackOutputTypeEnum
 from sagemaker.wrangler.processing import DataWranglerProcessor
 from sagemaker.dataset_definition.inputs import DatasetDefinition, AthenaDatasetDefinition
 from sagemaker.workflow.execution_variables import ExecutionVariables
@@ -698,6 +699,46 @@ def test_one_step_sparkjar_processing_pipeline(
             pass
 
 
+def test_one_step_callback_pipeline(sagemaker_session, role, pipeline_name, region_name):
+    instance_count = ParameterInteger(name="InstanceCount", default_value=2)
+
+    outputParam1 = CallbackOutput(output_name="output1", output_type=CallbackOutputTypeEnum.String)
+    step_callback = CallbackStep(
+        name="callback-step",
+        sqs_queue_url="https://sqs.us-east-2.amazonaws.com/123456789012/MyQueue",
+        inputs={"arg1": "foo"},
+        outputs=[outputParam1],
+    )
+
+    pipeline = Pipeline(
+        name=pipeline_name,
+        parameters=[instance_count],
+        steps=[step_callback],
+        sagemaker_session=sagemaker_session,
+    )
+
+    try:
+        response = pipeline.create(role)
+        create_arn = response["PipelineArn"]
+        assert re.match(
+            fr"arn:aws:sagemaker:{region_name}:\d{{12}}:pipeline/{pipeline_name}",
+            create_arn,
+        )
+
+        pipeline.parameters = [ParameterInteger(name="InstanceCount", default_value=1)]
+        response = pipeline.update(role)
+        update_arn = response["PipelineArn"]
+        assert re.match(
+            fr"arn:aws:sagemaker:{region_name}:\d{{12}}:pipeline/{pipeline_name}",
+            update_arn,
+        )
+    finally:
+        try:
+            pipeline.delete()
+        except Exception:
+            pass
+
+
 def test_conditional_pytorch_training_model_registration(
     sagemaker_session,
     role,

tests/unit/sagemaker/workflow/test_callback_step.py

@@ -0,0 +1,128 @@
+# Copyright 2020 Amazon.com, Inc. or its affiliates. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License"). You
+# may not use this file except in compliance with the License. A copy of
+# the License is located at
+#
+#     http://aws.amazon.com/apache2.0/
+#
+# or in the "license" file accompanying this file. This file is
+# distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF
+# ANY KIND, either express or implied. See the License for the specific
+# language governing permissions and limitations under the License.
+from __future__ import absolute_import
+
+import json
+
+import pytest
+
+from mock import Mock
+
+from sagemaker.workflow.parameters import ParameterInteger, ParameterString
+from sagemaker.workflow.pipeline import Pipeline
+from sagemaker.workflow.callback_step import CallbackStep, CallbackOutput, CallbackOutputTypeEnum
+
+
+@pytest.fixture
+def sagemaker_session_mock():
+    return Mock()
+
+
+def test_callback_step():
+    param = ParameterInteger(name="MyInt")
+    outputParam1 = CallbackOutput(output_name="output1", output_type=CallbackOutputTypeEnum.String)
+    outputParam2 = CallbackOutput(output_name="output2", output_type=CallbackOutputTypeEnum.Boolean)
+    cb_step = CallbackStep(
+        name="MyCallbackStep",
+        depends_on=["TestStep"],
+        sqs_queue_url="https://sqs.us-east-2.amazonaws.com/123456789012/MyQueue",
+        inputs={"arg1": "foo", "arg2": 5, "arg3": param},
+        outputs=[outputParam1, outputParam2],
+    )
+    cb_step.add_depends_on(["SecondTestStep"])
+    assert cb_step.to_request() == {
+        "Name": "MyCallbackStep",
+        "Type": "Callback",
+        "DependsOn": ["TestStep", "SecondTestStep"],
+        "SqsQueueUrl": "https://sqs.us-east-2.amazonaws.com/123456789012/MyQueue",
+        "OutputParameters": [
+            {"OutputName": "output1", "OutputType": "String"},
+            {"OutputName": "output2", "OutputType": "Boolean"},
+        ],
+        "Arguments": {"arg1": "foo", "arg2": 5, "arg3": param},
+    }
+
+
+def test_callback_step_output_expr():
+    param = ParameterInteger(name="MyInt")
+    outputParam1 = CallbackOutput(output_name="output1", output_type=CallbackOutputTypeEnum.String)
+    outputParam2 = CallbackOutput(output_name="output2", output_type=CallbackOutputTypeEnum.Boolean)
+    cb_step = CallbackStep(
+        name="MyCallbackStep",
+        depends_on=["TestStep"],
+        sqs_queue_url="https://sqs.us-east-2.amazonaws.com/123456789012/MyQueue",
+        inputs={"arg1": "foo", "arg2": 5, "arg3": param},
+        outputs=[outputParam1, outputParam2],
+    )
+
+    assert cb_step.properties.Outputs["output1"].expr == {
+        "Get": "Steps.MyCallbackStep.OutputParameters['output1']"
+    }
+    assert cb_step.properties.Outputs["output2"].expr == {
+        "Get": "Steps.MyCallbackStep.OutputParameters['output2']"
+    }
+
+
+def test_pipeline_interpolates_callback_outputs():
+    parameter = ParameterString("MyStr")
+    outputParam1 = CallbackOutput(output_name="output1", output_type=CallbackOutputTypeEnum.String)
+    outputParam2 = CallbackOutput(output_name="output2", output_type=CallbackOutputTypeEnum.String)
+    cb_step1 = CallbackStep(
+        name="MyCallbackStep1",
+        depends_on=["TestStep"],
+        sqs_queue_url="https://sqs.us-east-2.amazonaws.com/123456789012/MyQueue",
+        inputs={"arg1": "foo"},
+        outputs=[outputParam1],
+    )
+    cb_step2 = CallbackStep(
+        name="MyCallbackStep2",
+        depends_on=["TestStep"],
+        sqs_queue_url="https://sqs.us-east-2.amazonaws.com/123456789012/MyQueue",
+        inputs={"arg1": cb_step1.properties.Outputs["output1"]},
+        outputs=[outputParam2],
+    )
+
+    pipeline = Pipeline(
+        name="MyPipeline",
+        parameters=[parameter],
+        steps=[cb_step1, cb_step2],
+        sagemaker_session=sagemaker_session_mock,
+    )
+
+    assert json.loads(pipeline.definition()) == {
+        "Version": "2020-12-01",
+        "Metadata": {},
+        "Parameters": [{"Name": "MyStr", "Type": "String"}],
+        "PipelineExperimentConfig": {
+            "ExperimentName": {"Get": "Execution.PipelineName"},
+            "TrialName": {"Get": "Execution.PipelineExecutionId"},
+        },
+        "Steps": [
+            {
+                "Name": "MyCallbackStep1",
+                "Type": "Callback",
+                "Arguments": {"arg1": "foo"},
+                "DependsOn": ["TestStep"],
+                "SqsQueueUrl": "https://sqs.us-east-2.amazonaws.com/123456789012/MyQueue",
+                "OutputParameters": [{"OutputName": "output1", "OutputType": "String"}],
+            },
+            {
+                "Name": "MyCallbackStep2",
+                "Type": "Callback",
+                "Arguments": {"arg1": {"Get": "Steps.MyCallbackStep1.OutputParameters['output1']"}},
+                "DependsOn": ["TestStep"],
+                "SqsQueueUrl": "https://sqs.us-east-2.amazonaws.com/123456789012/MyQueue",
+                "OutputParameters": [{"OutputName": "output2", "OutputType": "String"}],
+            },
+        ],
+    }