Address PR comments

Author: Payton Staub

src/sagemaker/workflow/callback_step.py

@@ -68,6 +68,8 @@ def _expr(cls, name, step_name):
 
         Args:
             name (str): The name of the callback output.
+            step_name (str): The name of the step the callback step associated
+                with this output belongs to.
         """
         return {"Get": f"Steps.{step_name}.OutputParameters['{name}']"}
 

src/sagemaker/workflow/pipeline.py

@@ -240,7 +240,8 @@ def definition(self) -> str:
         """Converts a request structure to string representation for workflow service calls."""
         request_dict = self.to_request()
         request_dict["PipelineExperimentConfig"] = interpolate(
-            request_dict["PipelineExperimentConfig"]
+            request_dict["PipelineExperimentConfig"],
+            {}
         )
         callback_output_to_step_map = _map_callback_outputs(self.steps)
         request_dict["Steps"] = interpolate(
@@ -266,7 +267,9 @@ def format_start_parameters(parameters: Dict[str, Any]) -> List[Dict[str, Any]]:
     return [{"Name": name, "Value": str(value)} for name, value in parameters.items()]
 
 
-def interpolate(request_obj: RequestType, **kwargs) -> RequestType:
+def interpolate(
+    request_obj: RequestType, callback_output_to_step_map: Dict[str, str]
+) -> RequestType:
     """Replaces Parameter values in a list of nested Dict[str, Any] with their workflow expression.
 
     Args:
@@ -276,13 +279,10 @@ def interpolate(request_obj: RequestType, **kwargs) -> RequestType:
         RequestType: The request dict with Parameter values replaced by their expression.
     """
     request_obj_copy = deepcopy(request_obj)
-    return _interpolate(
-        request_obj_copy,
-        callback_output_to_step_map=kwargs.get("callback_output_to_step_map", None),
-    )
+    return _interpolate(request_obj_copy, callback_output_to_step_map=callback_output_to_step_map)
 
 
-def _interpolate(obj: Union[RequestType, Any], **kwargs):
+def _interpolate(obj: Union[RequestType, Any], callback_output_to_step_map: Dict[str, str]):
     """Walks the nested request dict, replacing Parameter type values with workflow expressions.
 
     Args:
@@ -291,22 +291,14 @@ def _interpolate(obj: Union[RequestType, Any], **kwargs):
     if isinstance(obj, (Expression, Parameter, Properties)):
         return obj.expr
     if isinstance(obj, CallbackOutput):
-        callback_output_to_step_map = kwargs.get("callback_output_to_step_map", {})
         step_name = callback_output_to_step_map[obj.output_name]
         return obj.expr(step_name)
     if isinstance(obj, dict):
         new = obj.__class__()
         for key, value in obj.items():
-            new[key] = interpolate(
-                value, callback_output_to_step_map=kwargs.get("callback_output_to_step_map", None)
-            )
+            new[key] = interpolate(value, callback_output_to_step_map)
     elif isinstance(obj, (list, set, tuple)):
-        new = obj.__class__(
-            interpolate(
-                value, callback_output_to_step_map=kwargs.get("callback_output_to_step_map", None)
-            )
-            for value in obj
-        )
+        new = obj.__class__(interpolate(value, callback_output_to_step_map) for value in obj)
     else:
         return obj
     return new

tests/integ/test_workflow.py

@@ -739,6 +739,46 @@ def test_one_step_callback_pipeline(sagemaker_session, role, pipeline_name, regi
             pass
 
 
+def test_two_step_callback_pipeline_with_output_reference(
+    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_callback1 = CallbackStep(
+        name="callback-step1",
+        sqs_queue_url="https://sqs.us-east-2.amazonaws.com/123456789012/MyQueue",
+        inputs={"arg1": "foo"},
+        outputs=[outputParam1],
+    )
+
+    step_callback2 = CallbackStep(
+        name="callback-step2",
+        sqs_queue_url="https://sqs.us-east-2.amazonaws.com/123456789012/MyQueue",
+        inputs={"arg1": outputParam1},
+    )
+
+    pipeline = Pipeline(
+        name=pipeline_name,
+        parameters=[instance_count],
+        steps=[step_callback1, step_callback2],
+        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,
+        )
+    finally:
+        try:
+            pipeline.delete()
+        except Exception:
+            pass
+
+
 def test_conditional_pytorch_training_model_registration(
     sagemaker_session,
     role,