Fix referencing of CallbackOutputs in other steps

Author: Payton Staub

src/sagemaker/workflow/callback_step.py

@@ -13,7 +13,7 @@
 """The step definitions for workflow."""
 from __future__ import absolute_import
 
-from typing import List
+from typing import List, Dict
 from enum import Enum
 
 import attr
@@ -59,6 +59,20 @@ 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)
+
+    @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."""
@@ -91,7 +105,12 @@ def __init__(
 
         root_path = f"Steps.{name}"
         root_prop = Properties(path=root_path)
-        root_prop.__dict__["OutputParameters"] = Properties(f"{root_path}.OutputParameters")
+        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
         self._properties = root_prop
 
     @property

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__()