fix: Invalid import (5637)

[aviruthen]

Description

Line 208 in sagemaker-mlops/src/sagemaker/mlops/workflow/steps.py has an invalid import: from sagemaker.core.workflow.function_step import DelayedReturn. The module sagemaker.core.workflow.function_step does not exist — there is no function_step.py in the sagemaker-core/src/sagemaker/core/workflow/ directory. The DelayedReturn class is defined in sagemaker-mlops/src/sagemaker/mlops/workflow/function_step.py, so the import should be from sagemaker.mlops.workflow.function_step import DelayedReturn. This is a straightforward one-line import path fix.

Related Issue

Related issue: 5637

Changes Made

• sagemaker-mlops/src/sagemaker/mlops/workflow/steps.py
• sagemaker-mlops/tests/unit/workflow/test_steps.py

AI-Generated PR

This PR was automatically generated by the PySDK Issue Agent.

• Confidence score: 85%
• Classification: bug
• SDK version target: V3

Merge Checklist

• Changes are backward compatible
• Commit message follows prefix: description format
• Unit tests added/updated
• Integration tests added (if applicable)
• Documentation updated (if applicable)

[mufaddal-rohawala]

🤖 AI Code Review

This PR fixes a clearly broken import path in steps.py, changing sagemaker.core.workflow.function_step to sagemaker.mlops.workflow.function_step. The fix is correct and the test updates are consistent with the change. A few minor observations on the new test.

[mufaddal-rohawala]

This test performs an actual import of DelayedReturn from the real module, which makes it more of a smoke/integration test than a unit test. If sagemaker.mlops.workflow.function_step has heavy dependencies that aren't available in the unit test environment, this could fail unexpectedly.

Also, the assertion hasattr(DelayedReturn, '_to_json_get') couples this test to a private implementation detail. If the private method is renamed or removed, this test will break even though the import fix is still valid. Consider either:

1. Removing the hasattr check and just verifying the import succeeds, or
2. Checking a more stable public attribute if one exists.

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