Conversation
sagemaker-bot
left a comment
sagemaker-bot
left a comment
There was a problem hiding this comment.
🤖 AI Code Review
The fix correctly changes the import path for DelayedReturn from the non-existent sagemaker.core.workflow.function_step to the correct sagemaker.mlops.workflow.function_step. However, the test file has several issues: a duplicate function body in test_step_find_dependencies_in_step_arguments_with_delayed_return, trailing whitespace, and some tests that don't actually exercise the code path they claim to test.
| Step._validate_json_get_function(step, json_get, step_map) | ||
|
|
||
|
|
||
| def test_delayed_return_import_from_correct_module(): |
There was a problem hiding this comment.
Low-value test. This test (test_delayed_return_import_from_correct_module) only verifies that the import statement works, which is already implicitly tested by the other tests that import and use DelayedReturn. Consider removing it — the real value is in testing that _find_dependencies_in_step_arguments works correctly with the fixed import, which the other tests cover.
| assert isinstance(DelayedReturn, type) | ||
|
|
||
|
|
||
| def test_find_dependencies_in_step_arguments_with_delayed_return_uses_correct_import(): |
There was a problem hiding this comment.
Test doesn't exercise the fix. test_find_dependencies_in_step_arguments_with_delayed_return_uses_correct_import creates mocks and asserts isinstance(delayed, DelayedReturn), but it never calls _find_dependencies_in_step_arguments. The isinstance check on a Mock(spec=DelayedReturn) will always pass regardless of where DelayedReturn is imported from. This test doesn't actually verify the bug fix. Consider removing it in favor of the more comprehensive test at line 447 that actually calls the method.
| @@ -450,6 +510,7 @@ def test_step_find_dependencies_in_step_arguments_with_delayed_return(): | |||
| assert "step1" in dependencies | |||
There was a problem hiding this comment.
Extra blank line. There's a double blank line here. PEP 8 expects one blank line between methods inside a test module at this level.
sagemaker-bot
left a comment
sagemaker-bot
left a comment
There was a problem hiding this comment.
🤖 AI Code Review
This PR fixes a genuine bug where DelayedReturn was imported from a non-existent module path (sagemaker.core.workflow.function_step) and corrects it to the actual location (sagemaker.mlops.workflow.function_step). The fix is minimal and correct. The test changes appropriately simplify tests by removing unnecessary mocking of the broken import path. However, there are some issues with the test file changes.
| Step._validate_json_get_function(step, json_get, step_map) | ||
|
|
||
|
|
||
|
|
There was a problem hiding this comment.
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.
| 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}) |
There was a problem hiding this comment.
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.
|
|
||
| def test_step_find_dependencies_in_step_arguments_with_delayed_return(): | ||
| from unittest.mock import patch | ||
| from sagemaker.mlops.workflow.steps import Step, StepTypeEnum |
There was a problem hiding this comment.
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.
| 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 |
There was a problem hiding this comment.
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.
2 participants
Description
The file
sagemaker-mlops/src/sagemaker/mlops/workflow/steps.pyhas an invalid import at line 208 inside the_find_dependencies_in_step_argumentsmethod. It importsDelayedReturnfromsagemaker.core.workflow.function_step, but that module does not exist in thesagemaker-corepackage (there is nofunction_step.pyinsagemaker-core/src/sagemaker/core/workflow/). TheDelayedReturnclass is actually defined insagemaker-mlops/src/sagemaker/mlops/workflow/function_step.py, so the import should befrom sagemaker.mlops.workflow.function_step import DelayedReturn.Related Issue
Fixes GitHub issue 5637
Changes Made
sagemaker-mlops/src/sagemaker/mlops/workflow/steps.pysagemaker-mlops/tests/unit/workflow/test_steps.pyAI-Generated PR
This PR was automatically generated by the PySDK Issue Agent.
Merge Checklist
prefix: descriptionformat