Skip to content

fix: Invalid import (5637) #5686

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
aviruthen wants to merge 3 commits into from
+63 −38
Closed

fix: Invalid import (5637) #5686

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
3584928
fix: Invalid import (5637)
aviruthen Mar 26, 2026
9ce5304
fix: address review comments (iteration #1)
aviruthen Mar 26, 2026
dcdf2a0
fix: address review comments (iteration #2)
aviruthen Mar 26, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 1 addition & 1 deletion sagemaker-mlops/src/sagemaker/mlops/workflow/steps.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Copy link
Copy Markdown
Collaborator

@sagemaker-bot sagemaker-bot Mar 26, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The fix is correct. The DelayedReturn class lives in sagemaker.mlops.workflow.function_step, not sagemaker.core.workflow.function_step. This is a lazy import (inside the method) which is appropriate here to avoid circular imports between steps.py and function_step.py.


# TODO: we can remove the if-elif once move the validators to JsonGet constructor
if isinstance(pipeline_variable, JsonGet):
Expand Down
99 changes: 62 additions & 37 deletions sagemaker-mlops/tests/unit/workflow/test_steps.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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"
Expand Down Expand Up @@ -383,7 +383,45 @@ 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



Copy link
Copy Markdown
Collaborator

@sagemaker-bot sagemaker-bot Mar 26, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

There are ~30 lines of seemingly orphaned code between test_step_validate_json_get_function_with_property_file and test_step_find_dependencies_in_step_arguments_with_json_get that contain bare import statements (from sagemaker.mlops.workflow.steps import Step, StepTypeEnum and from sagemaker.core.workflow.functions import JsonGet) outside of any function or class definition. These imports execute at module load time and appear to be leftover from an incomplete edit. They should either be placed inside a proper test function or removed entirely. As-is, they are dead code that will execute on import and could mask import errors.










from sagemaker.mlops.workflow.steps import Step, StepTypeEnum

from sagemaker.core.workflow.functions import JsonGet




















def test_step_find_dependencies_in_step_arguments_with_json_get():
Expand All @@ -408,74 +446,61 @@ 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})
Copy link
Copy Markdown
Collaborator

@sagemaker-bot sagemaker-bot Mar 26, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Good improvement — removing the patch.dict('sys.modules', ...) workaround now that the import path is fixed. The test is much cleaner and actually validates the real import path works correctly.

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
Copy link
Copy Markdown
Collaborator

@sagemaker-bot sagemaker-bot Mar 26, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Good: directly importing DelayedReturn from the corrected path and using Mock(spec=DelayedReturn) is much better than the previous approach of dynamically creating a fake class. This properly validates that the import works and that the isinstance check in the source code will function correctly.

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_class = type('DelayedReturn', (PipelineVariable,), {})
delayed_return = Mock(spec=delayed_return_class)

delayed_return = Mock(spec=DelayedReturn)
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

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():
Expand Down
Loading