From 358492821808c1183843ce28e3b18c774588ff2b Mon Sep 17 00:00:00 2001 From: aviruthen <91846056+aviruthen@users.noreply.github.com> Date: Thu, 26 Mar 2026 16:44:18 -0400 Subject: [PATCH 1/3] fix: Invalid import (5637) --- .../src/sagemaker/mlops/workflow/steps.py | 2 +- .../tests/unit/workflow/test_steps.py | 61 +++++++++++++++++++ 2 files changed, 62 insertions(+), 1 deletion(-) diff --git a/sagemaker-mlops/src/sagemaker/mlops/workflow/steps.py b/sagemaker-mlops/src/sagemaker/mlops/workflow/steps.py index 5e0eb3dda3..76e90a5309 100644 --- a/sagemaker-mlops/src/sagemaker/mlops/workflow/steps.py +++ b/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): diff --git a/sagemaker-mlops/tests/unit/workflow/test_steps.py b/sagemaker-mlops/tests/unit/workflow/test_steps.py index 06cc729ceb..87f87100e9 100644 --- a/sagemaker-mlops/tests/unit/workflow/test_steps.py +++ b/sagemaker-mlops/tests/unit/workflow/test_steps.py @@ -386,6 +386,38 @@ def test_step_validate_json_get_function_with_property_file(): Step._validate_json_get_function(step, json_get, step_map) +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 + # Verify it's a class + assert isinstance(DelayedReturn, type) + + +def test_find_dependencies_in_step_arguments_with_delayed_return_uses_correct_import(): + """Verify _find_dependencies_in_step_arguments uses the mlops import path for DelayedReturn.""" + from sagemaker.mlops.workflow.steps import Step, StepTypeEnum + from sagemaker.mlops.workflow.function_step import DelayedReturn, _FunctionStep + from sagemaker.core.workflow.functions import JsonGet + from unittest.mock import patch, Mock + + # Create a mock function step + mock_func_step = Mock(spec=_FunctionStep) + mock_func_step.name = "func-step" + + # Create a DelayedReturn that acts as a PipelineVariable + delayed = Mock(spec=DelayedReturn) + delayed._referenced_steps = [mock_func_step] + delayed._step = mock_func_step + + mock_json_get = Mock(spec=JsonGet) + mock_json_get.property_file = None + delayed._to_json_get = Mock(return_value=mock_json_get) + + # Verify isinstance check works with the correct import + assert isinstance(delayed, DelayedReturn) + + def test_step_find_dependencies_in_step_arguments_with_json_get(): from unittest.mock import patch from sagemaker.mlops.workflow.steps import Step, StepTypeEnum @@ -415,6 +447,34 @@ def test_step_find_dependencies_in_step_arguments_with_json_get(): def test_step_find_dependencies_in_step_arguments_with_delayed_return(): + from unittest.mock import patch + from sagemaker.mlops.workflow.steps import Step, StepTypeEnum + from sagemaker.mlops.workflow.function_step import DelayedReturn + from sagemaker.core.workflow.functions import JsonGet + from sagemaker.core.helper.pipeline_variable import PipelineVariable + + step1 = Mock(spec=Step) + step1.name = "step1" + step1.step_type = StepTypeEnum.PROCESSING + step1.property_files = [] + step1.arguments = {} + + json_get = Mock(spec=JsonGet) + json_get.property_file = None + + delayed_return = Mock(spec=DelayedReturn) + delayed_return._referenced_steps = [step1] + delayed_return._to_json_get = Mock(return_value=json_get) + + step2 = Mock(spec=Step) + step2.name = "step2" + step2._validate_json_get_function = Mock() + step2._get_step_name_from_str = Step._get_step_name_from_str + + obj = {"key": delayed_return} + + dependencies = Step._find_dependencies_in_step_arguments(step2, obj, {"step1": step1}) + assert "step1" in dependencies from unittest.mock import patch from sagemaker.mlops.workflow.steps import Step, StepTypeEnum from sagemaker.core.workflow.functions import JsonGet @@ -450,6 +510,7 @@ def test_step_find_dependencies_in_step_arguments_with_delayed_return(): assert "step1" in dependencies + def test_step_find_dependencies_in_step_arguments_with_string_reference(): from unittest.mock import patch from sagemaker.mlops.workflow.steps import Step From 9ce5304e85de26e86d776a549ce5a9916cb831f8 Mon Sep 17 00:00:00 2001 From: aviruthen <91846056+aviruthen@users.noreply.github.com> Date: Thu, 26 Mar 2026 16:48:18 -0400 Subject: [PATCH 2/3] fix: address review comments (iteration #1) --- .../tests/unit/workflow/test_steps.py | 122 ++++++------------ 1 file changed, 40 insertions(+), 82 deletions(-) diff --git a/sagemaker-mlops/tests/unit/workflow/test_steps.py b/sagemaker-mlops/tests/unit/workflow/test_steps.py index 87f87100e9..bb812c7d74 100644 --- a/sagemaker-mlops/tests/unit/workflow/test_steps.py +++ b/sagemaker-mlops/tests/unit/workflow/test_steps.py @@ -252,7 +252,7 @@ def test_step_find_dependencies_in_depends_on_list_with_string(): def test_step_validate_json_get_property_file_reference_invalid_step_type(): from sagemaker.mlops.workflow.steps import Step, StepTypeEnum - from sagemaker.core.workflow.functions import JsonGet + step = Mock(spec=Step) step.name = "current-step" @@ -386,36 +386,37 @@ def test_step_validate_json_get_function_with_property_file(): Step._validate_json_get_function(step, json_get, step_map) -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 - # Verify it's a class - assert isinstance(DelayedReturn, type) -def test_find_dependencies_in_step_arguments_with_delayed_return_uses_correct_import(): - """Verify _find_dependencies_in_step_arguments uses the mlops import path for DelayedReturn.""" + + + + + + + + from sagemaker.mlops.workflow.steps import Step, StepTypeEnum - from sagemaker.mlops.workflow.function_step import DelayedReturn, _FunctionStep + from sagemaker.core.workflow.functions import JsonGet - from unittest.mock import patch, Mock - # Create a mock function step - mock_func_step = Mock(spec=_FunctionStep) - mock_func_step.name = "func-step" - # Create a DelayedReturn that acts as a PipelineVariable - delayed = Mock(spec=DelayedReturn) - delayed._referenced_steps = [mock_func_step] - delayed._step = mock_func_step - mock_json_get = Mock(spec=JsonGet) - mock_json_get.property_file = None - delayed._to_json_get = Mock(return_value=mock_json_get) - # Verify isinstance check works with the correct import - assert isinstance(delayed, DelayedReturn) + + + + + + + + + + + + + + def test_step_find_dependencies_in_step_arguments_with_json_get(): @@ -440,103 +441,60 @@ 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()}): - dependencies = Step._find_dependencies_in_step_arguments(step2, obj, {"step1": step1}) - assert "step1" in dependencies + dependencies = Step._find_dependencies_in_step_arguments(step2, obj, {"step1": step1}) + assert "step1" in dependencies def test_step_find_dependencies_in_step_arguments_with_delayed_return(): - from unittest.mock import patch from sagemaker.mlops.workflow.steps import Step, StepTypeEnum from sagemaker.mlops.workflow.function_step import DelayedReturn from sagemaker.core.workflow.functions import JsonGet - from sagemaker.core.helper.pipeline_variable import PipelineVariable - + step1 = Mock(spec=Step) step1.name = "step1" step1.step_type = StepTypeEnum.PROCESSING step1.property_files = [] step1.arguments = {} - + json_get = Mock(spec=JsonGet) json_get.property_file = None - + delayed_return = Mock(spec=DelayedReturn) delayed_return._referenced_steps = [step1] delayed_return._to_json_get = Mock(return_value=json_get) - + step2 = Mock(spec=Step) step2.name = "step2" step2._validate_json_get_function = Mock() step2._get_step_name_from_str = Step._get_step_name_from_str - + obj = {"key": delayed_return} - + dependencies = Step._find_dependencies_in_step_arguments(step2, obj, {"step1": step1}) assert "step1" in dependencies - from unittest.mock import patch - from sagemaker.mlops.workflow.steps import Step, StepTypeEnum - from sagemaker.core.workflow.functions import JsonGet - from sagemaker.core.helper.pipeline_variable import PipelineVariable - - step1 = Mock(spec=Step) - step1.name = "step1" - step1.step_type = StepTypeEnum.PROCESSING - step1.property_files = [] - step1.arguments = {} - - json_get = Mock(spec=JsonGet) - json_get.property_file = None - - delayed_return_class = type('DelayedReturn', (PipelineVariable,), {}) - delayed_return = Mock(spec=delayed_return_class) - delayed_return._referenced_steps = [step1] - delayed_return._to_json_get = Mock(return_value=json_get) - delayed_return.__class__ = delayed_return_class - - step2 = Mock(spec=Step) - step2.name = "step2" - step2._validate_json_get_function = Mock() - step2._get_step_name_from_str = Step._get_step_name_from_str - - obj = {"key": delayed_return} - - mock_module = Mock() - mock_module.DelayedReturn = delayed_return_class - - with patch.dict('sys.modules', {'sagemaker.core.workflow.function_step': mock_module}): - dependencies = Step._find_dependencies_in_step_arguments(step2, obj, {"step1": step1}) - assert "step1" in dependencies def test_step_find_dependencies_in_step_arguments_with_string_reference(): - from unittest.mock import patch from sagemaker.mlops.workflow.steps import Step from sagemaker.core.helper.pipeline_variable import PipelineVariable - + step1 = Mock(spec=Step) step1.name = "step1" - + pipeline_var = Mock(spec=PipelineVariable) pipeline_var._referenced_steps = ["step1"] - + step2 = Mock(spec=Step) step2.name = "step2" step2._get_step_name_from_str = Step._get_step_name_from_str - + obj = {"key": pipeline_var} - + step_map = {"step1": step1} - - delayed_return_class = type('DelayedReturn', (PipelineVariable,), {}) - mock_module = Mock() - mock_module.DelayedReturn = delayed_return_class - - with patch.dict('sys.modules', {'sagemaker.core.workflow.function_step': mock_module}): - dependencies = Step._find_dependencies_in_step_arguments(step2, obj, step_map) - assert "step1" in dependencies + + dependencies = Step._find_dependencies_in_step_arguments(step2, obj, step_map) + assert "step1" in dependencies def test_tuning_step_requires_step_args(): From dcdf2a0306a893d24936f1977366f8dd22519afd Mon Sep 17 00:00:00 2001 From: aviruthen <91846056+aviruthen@users.noreply.github.com> Date: Thu, 26 Mar 2026 16:52:02 -0400 Subject: [PATCH 3/3] fix: address review comments (iteration #2) --- sagemaker-mlops/tests/unit/workflow/test_steps.py | 8 +++++++- 1 file changed, 7 insertions(+), 1 deletion(-) diff --git a/sagemaker-mlops/tests/unit/workflow/test_steps.py b/sagemaker-mlops/tests/unit/workflow/test_steps.py index bb812c7d74..d4108abcff 100644 --- a/sagemaker-mlops/tests/unit/workflow/test_steps.py +++ b/sagemaker-mlops/tests/unit/workflow/test_steps.py @@ -383,7 +383,12 @@ def test_step_validate_json_get_function_with_property_file(): step_map = {"processing-step": processing_step} - Step._validate_json_get_function(step, json_get, step_map) + Step._validate_json_get_function(step, json_get, step_map) + + +def test_step_find_dependencies_in_step_arguments_with_json_get(): + from sagemaker.mlops.workflow.steps import Step, StepTypeEnum + from sagemaker.core.workflow.functions import JsonGet @@ -475,6 +480,7 @@ def test_step_find_dependencies_in_step_arguments_with_delayed_return(): + def test_step_find_dependencies_in_step_arguments_with_string_reference(): from sagemaker.mlops.workflow.steps import Step from sagemaker.core.helper.pipeline_variable import PipelineVariable