fix: Invalid import (5637)

Author: aviruthen

sagemaker-mlops/src/sagemaker/mlops/workflow/steps.py

@@ -205,7 +205,7 @@ def _find_dependencies_in_step_arguments(
                 else:
                     dependencies.add(self._get_step_name_from_str(referenced_step, step_map))
 
-            from sagemaker.core.workflow.function_step import DelayedReturn
+            from sagemaker.mlops.workflow.function_step import DelayedReturn
 
             # TODO: we can remove the if-elif once move the validators to JsonGet constructor
             if isinstance(pipeline_variable, JsonGet):

sagemaker-mlops/tests/unit/workflow/test_steps.py

@@ -409,7 +409,7 @@ def test_step_find_dependencies_in_step_arguments_with_json_get():
     obj = {"key": json_get}
 
     with patch('sagemaker.mlops.workflow.steps.TYPE_CHECKING', False):
-        with patch.dict('sys.modules', {'sagemaker.core.workflow.function_step': Mock()}):
+        with patch.dict('sys.modules', {'sagemaker.mlops.workflow.function_step': Mock()}):
             dependencies = Step._find_dependencies_in_step_arguments(step2, obj, {"step1": step1})
             assert "step1" in dependencies
 
@@ -445,7 +445,7 @@ def test_step_find_dependencies_in_step_arguments_with_delayed_return():
     mock_module = Mock()
     mock_module.DelayedReturn = delayed_return_class
 
-    with patch.dict('sys.modules', {'sagemaker.core.workflow.function_step': mock_module}):
+    with patch.dict('sys.modules', {'sagemaker.mlops.workflow.function_step': mock_module}):
         dependencies = Step._find_dependencies_in_step_arguments(step2, obj, {"step1": step1})
         assert "step1" in dependencies
 
@@ -473,11 +473,18 @@ def test_step_find_dependencies_in_step_arguments_with_string_reference():
     mock_module = Mock()
     mock_module.DelayedReturn = delayed_return_class
 
-    with patch.dict('sys.modules', {'sagemaker.core.workflow.function_step': mock_module}):
+    with patch.dict('sys.modules', {'sagemaker.mlops.workflow.function_step': mock_module}):
         dependencies = Step._find_dependencies_in_step_arguments(step2, obj, step_map)
         assert "step1" in dependencies
 
 
+def test_delayed_return_import_from_correct_module():
+    """Verify that DelayedReturn can be imported from sagemaker.mlops.workflow.function_step."""
+    from sagemaker.mlops.workflow.function_step import DelayedReturn
+    assert DelayedReturn is not None
+    assert hasattr(DelayedReturn, '_to_json_get')
+
+
 def test_tuning_step_requires_step_args():
     from sagemaker.mlops.workflow.steps import TuningStep