Correctly interpolate Callback output parameters that are passed by reference

Author: Payton Staub

src/sagemaker/workflow/callback_step.py

@@ -58,19 +58,18 @@ def to_request(self) -> RequestType:
             "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)
+    def expr(self, step_name) -> Dict[str, str]:
+        """The 'Get' expression dict for a `CallbackOutput`."""
+        return CallbackOutput._expr(self.output_name, step_name)
 
     @classmethod
-    def _expr(cls, name):
+    def _expr(cls, name, step_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}']"}
+        return {"Get": f"Steps.{step_name}.OutputParameters['{name}']"}
 
 
 class CallbackStep(Step):

src/sagemaker/workflow/pipeline.py

@@ -24,7 +24,7 @@
 
 from sagemaker._studio import _append_project_tags
 from sagemaker.session import Session
-from sagemaker.workflow.callback_step import CallbackOutput
+from sagemaker.workflow.callback_step import CallbackOutput, CallbackStep
 from sagemaker.workflow.entities import (
     Entity,
     Expression,
@@ -242,7 +242,10 @@ def definition(self) -> str:
         request_dict["PipelineExperimentConfig"] = interpolate(
             request_dict["PipelineExperimentConfig"]
         )
-        request_dict["Steps"] = interpolate(request_dict["Steps"])
+        callback_output_to_step_map = _map_callback_outputs(self.steps)
+        request_dict["Steps"] = interpolate(
+            request_dict["Steps"], callback_output_to_step_map=callback_output_to_step_map
+        )
 
         return json.dumps(request_dict)
 
@@ -263,7 +266,7 @@ 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) -> RequestType:
+def interpolate(request_obj: RequestType, **kwargs) -> RequestType:
     """Replaces Parameter values in a list of nested Dict[str, Any] with their workflow expression.
 
     Args:
@@ -273,28 +276,59 @@ def interpolate(request_obj: RequestType) -> RequestType:
         RequestType: The request dict with Parameter values replaced by their expression.
     """
     request_obj_copy = deepcopy(request_obj)
-    return _interpolate(request_obj_copy)
+    return _interpolate(
+        request_obj_copy,
+        callback_output_to_step_map=kwargs.get("callback_output_to_step_map", None),
+    )
 
 
-def _interpolate(obj: Union[RequestType, Any]):
+def _interpolate(obj: Union[RequestType, Any], **kwargs):
     """Walks the nested request dict, replacing Parameter type values with workflow expressions.
 
     Args:
         obj (Union[RequestType, Any]): The request dict.
     """
-    if isinstance(obj, (Expression, Parameter, Properties, CallbackOutput)):
+    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)
+            new[key] = interpolate(
+                value, callback_output_to_step_map=kwargs.get("callback_output_to_step_map", None)
+            )
     elif isinstance(obj, (list, set, tuple)):
-        new = obj.__class__(interpolate(value) for value in obj)
+        new = obj.__class__(
+            interpolate(
+                value, callback_output_to_step_map=kwargs.get("callback_output_to_step_map", None)
+            )
+            for value in obj
+        )
     else:
         return obj
     return new
 
 
+def _map_callback_outputs(steps: List[Step]):
+    """Iterate over the provided steps, building a map of callback output parameters to step names.
+
+    Args:
+        step (List[Step]): The steps list.
+    """
+
+    callback_output_map = {}
+    for step in steps:
+        if isinstance(step, CallbackStep):
+            if step.outputs:
+                for output in step.outputs:
+                    callback_output_map[output.output_name] = step.name
+
+    return callback_output_map
+
+
 def update_args(args: Dict[str, Any], **kwargs):
     """Updates the request arguments dict with a value, if populated.
 

tests/unit/sagemaker/workflow/test_callback_step.py

@@ -88,7 +88,7 @@ def test_pipeline_interpolates_callback_outputs():
         name="MyCallbackStep2",
         depends_on=["TestStep"],
         sqs_queue_url="https://sqs.us-east-2.amazonaws.com/123456789012/MyQueue",
-        inputs={"arg1": cb_step1.properties.Outputs["output1"]},
+        inputs={"arg1": outputParam1},
         outputs=[outputParam2],
     )