Fix multiple output parameters in callback step properties. Additional unit tests

Author: Payton Staub

src/sagemaker/workflow/callback_step.py

@@ -105,12 +105,14 @@ def __init__(
 
         root_path = f"Steps.{name}"
         root_prop = Properties(path=root_path)
+
+        property_dict = {}
         for output in outputs:
-            property_dict = {}
             property_dict[output.output_name] = Properties(
                 f"{root_path}.OutputParameters['{output.output_name}']"
             )
-            root_prop.__dict__["Outputs"] = property_dict
+
+        root_prop.__dict__["Outputs"] = property_dict
         self._properties = root_prop
 
     @property

tests/unit/sagemaker/workflow/test_callback_step.py

@@ -12,9 +12,21 @@
 # language governing permissions and limitations under the License.
 from __future__ import absolute_import
 
-from sagemaker.workflow.parameters import ParameterInteger
+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
 
+from tests.unit.sagemaker.workflow.helpers import ordered
+
+@pytest.fixture
+def sagemaker_session_mock():
+    return Mock()
 
 def test_callback_step():
     param = ParameterInteger(name="MyInt")
@@ -39,3 +51,91 @@ def test_callback_step():
         ],
         "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"
+                        }
+                    ]
+                }
+            ]
+        }
+