fix: [Bug] Pipeline parameters (ParameterInteger, ParameterString) fail in ModelTrain (5504) by aviruthen · Pull Request #5688 · aws/sagemaker-python-sdk

aviruthen · 2026-03-27T15:53:03Z

Description

No response from agent

Related Issue

Related issue: 5504

Changes Made

sagemaker-train/src/sagemaker/train/utils.py
sagemaker-train/src/sagemaker/train/model_trainer.py
sagemaker-train/tests/unit/train/test_model_trainer_pipeline_variable.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)

…il in ModelTrain (5504)

mufaddal-rohawala

🤖 AI Code Review

This PR fixes a bug where Pipeline parameters (PipelineVariable types) fail in ModelTrainer validation because truthiness checks on PipelineVariable objects don't behave like regular strings. The core fix is correct, but there are several issues: type annotations are removed instead of being updated, the import placement is suboptimal, and the test changes appear incomplete.

sagemaker-train/src/sagemaker/train/model_trainer.py

mufaddal-rohawala · 2026-03-27T15:57:36Z

sagemaker-train/src/sagemaker/train/utils.py



-def _get_repo_name_from_image(image: str) -> str:
+def _get_repo_name_from_image(image) -> str:


Same issue — don't remove the type annotation, update it:

def _get_repo_name_from_image(image: str | PipelineVariable) -> str | None:

Note the return type should also be updated to str | None since you now return None for PipelineVariable inputs.

mufaddal-rohawala · 2026-03-27T15:57:36Z

sagemaker-train/src/sagemaker/train/utils.py

+        str: The repository name, or None if image is a PipelineVariable
    """
+    from sagemaker.core.helper.pipeline_variable import PipelineVariable
+


Move the import to the top of the module (or at minimum to the top of the function). Inline imports inside functions are acceptable for avoiding circular dependencies, but please add a comment explaining why it's done here:

# Import here to avoid circular dependency from sagemaker.core.helper.pipeline_variable import PipelineVariable

Also, is there actually a circular dependency risk? If not, this import should be at the module level with other imports.

sagemaker-train/src/sagemaker/train/utils.py

mufaddal-rohawala · 2026-03-27T15:57:36Z

sagemaker-train/tests/unit/train/test_model_trainer_pipeline_variable.py

    OutputDataConfig,
 )
-from sagemaker.train.defaults import DEFAULT_INSTANCE_TYPE
+from sagemaker.train.defaults import DEFAULT_INSTANCE_TYPE, TrainDefaults


TrainDefaults is imported but never used in the diff. Is this import used elsewhere in the file (not shown in the diff)? If not, this is an unused import that will fail linting. If it IS used in existing code not shown in the diff, please disregard this comment.

mufaddal-rohawala · 2026-03-27T15:57:36Z

sagemaker-train/tests/unit/train/test_model_trainer_pipeline_variable.py

    OutputDataConfig,
 )
-from sagemaker.train.defaults import DEFAULT_INSTANCE_TYPE
+from sagemaker.train.defaults import DEFAULT_INSTANCE_TYPE, TrainDefaults


The test changes seem insufficient for the scope of the fix. The PR modifies validation logic in _validate_training_image_and_algorithm_name and adds PipelineVariable handling in _get_repo_name_from_image, but the test diff only shows an import change. Where are the new test cases that:

Pass a PipelineVariable as training_image and verify validation passes?

Pass a PipelineVariable to _get_repo_name_from_image and verify it returns None?

Verify that both PipelineVariable for training_image AND algorithm_name still raises ValueError?

Verify that None for both still raises ValueError?

Please add explicit unit tests for the changed behavior. Target >90% coverage per SDK standards.

mufaddal-rohawala

🤖 AI Code Review

This PR fixes a bug where PipelineVariable objects (ParameterInteger, ParameterString) were failing in ModelTrainer validation because truthiness checks on these objects don't behave like regular strings/None. The fix correctly switches to is not None checks and adds a PipelineVariable guard in _get_repo_name_from_image. The approach is sound, but there are a few issues worth addressing.

mufaddal-rohawala · 2026-03-27T16:01:34Z

sagemaker-train/src/sagemaker/train/utils.py



-def _get_repo_name_from_image(image: str) -> str:
+def _get_repo_name_from_image(image: "str | PipelineVariable") -> "str | None":


The type annotation uses PEP 604 string-quoted syntax ("str | PipelineVariable") for the parameter but the module doesn't use from __future__ import annotations. While this works at runtime due to quoting, it's inconsistent with the return type which is also quoted. Per SDK conventions, prefer adding from __future__ import annotations at the top of the module and using unquoted str | PipelineVariable and str | None, or use Union[str, PipelineVariable] / Optional[str] consistently with the rest of the codebase (as done in model_trainer.py).

mufaddal-rohawala · 2026-03-27T16:01:34Z

sagemaker-train/src/sagemaker/train/utils.py

 from typing import Literal, Any

 from sagemaker.core.helper.session_helper import Session
+from sagemaker.core.helper.pipeline_variable import PipelineVariable


The import was changed from sagemaker.core.workflow.parameters.PipelineVariable to sagemaker.core.helper.pipeline_variable.PipelineVariable. Please verify these are the same class (or that the helper version is the correct base class that ParameterString/ParameterInteger inherit from). If sagemaker.core.workflow.parameters.PipelineVariable is a re-export of the same class, this is fine, but if they're different classes, the isinstance check on line 163 might not catch all pipeline variable types.

mufaddal-rohawala · 2026-03-27T16:01:34Z

sagemaker-train/src/sagemaker/train/model_trainer.py


    def _validate_training_image_and_algorithm_name(
-        self, training_image: Optional[str], algorithm_name: Optional[str]
+        self,


Good fix switching from truthiness (not training_image) to identity (is not None) checks. The truthiness check would evaluate a PipelineVariable object as truthy but could also fail for empty strings. The is not None approach is more correct and explicit.

Minor: The type annotation uses Union[str, PipelineVariable, None] — consider whether Optional[Union[str, PipelineVariable]] or str | PipelineVariable | None (with from __future__ import annotations) would be more consistent with the SDK's preference for newer syntax in new code.

mufaddal-rohawala · 2026-03-27T16:01:34Z

sagemaker-train/tests/unit/train/test_model_trainer_pipeline_variable.py

+    def test_safe_serialize_parameter_string_returns_variable(self):
+        """safe_serialize should return ParameterString as-is."""
+        param = ParameterString(name="MyParam", default_value="val")
+        result = safe_serialize(param)


The TestSafeSerialize tests are good additions, but I don't see any corresponding changes to the safe_serialize function in utils.py in this diff. If safe_serialize already handles PipelineVariable correctly, these tests are documenting existing behavior (which is fine). But if the bug also affects safe_serialize and it needs a fix, that change appears to be missing from this PR. Could you confirm that safe_serialize already has the PipelineVariable guard?

mufaddal-rohawala · 2026-03-27T16:01:34Z

sagemaker-train/tests/unit/train/test_model_trainer_pipeline_variable.py

+    """Tests for _validate_training_image_and_algorithm_name with PipelineVariable."""
+
+    def test_pipeline_variable_training_image_passes_validation(self):
+        """PipelineVariable as training_image should pass validation."""


Good test coverage for the validation fix. Consider also adding a test case where training_image is a ParameterString with an empty string default value (default_value="") to verify the is not None check doesn't regress for edge cases where the resolved value might be empty.

fix: [Bug] Pipeline parameters (ParameterInteger, ParameterString) fa…

04d768d

…il in ModelTrain (5504)

aviruthen temporarily deployed to auto-approve March 27, 2026 15:53 — with GitHub Actions Inactive

mufaddal-rohawala reviewed Mar 27, 2026

View reviewed changes

fix: address review comments (iteration #1)

a2f9e47

aviruthen temporarily deployed to auto-approve March 27, 2026 16:00 — with GitHub Actions Inactive

mufaddal-rohawala reviewed Mar 27, 2026

View reviewed changes

aviruthen closed this Mar 27, 2026

aviruthen deleted the fix/bug-pipeline-parameters-parameterinteger-5504 branch March 27, 2026 16:03

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

fix: [Bug] Pipeline parameters (ParameterInteger, ParameterString) fail in ModelTrain (5504)#5688

fix: [Bug] Pipeline parameters (ParameterInteger, ParameterString) fail in ModelTrain (5504)#5688
aviruthen wants to merge 2 commits intoaws:masterfrom
aviruthen:fix/bug-pipeline-parameters-parameterinteger-5504

aviruthen commented Mar 27, 2026

Uh oh!

mufaddal-rohawala left a comment

Uh oh!

Uh oh!

mufaddal-rohawala Mar 27, 2026

Uh oh!

mufaddal-rohawala Mar 27, 2026

Uh oh!

Uh oh!

mufaddal-rohawala Mar 27, 2026

Uh oh!

mufaddal-rohawala Mar 27, 2026

Uh oh!

mufaddal-rohawala left a comment

Uh oh!

mufaddal-rohawala Mar 27, 2026

Uh oh!

mufaddal-rohawala Mar 27, 2026

Uh oh!

mufaddal-rohawala Mar 27, 2026

Uh oh!

mufaddal-rohawala Mar 27, 2026

Uh oh!

mufaddal-rohawala Mar 27, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants



		def _get_repo_name_from_image(image: str) -> str:
		def _get_repo_name_from_image(image) -> str:



		def _get_repo_name_from_image(image: str) -> str:
		def _get_repo_name_from_image(image: "str \| PipelineVariable") -> "str \| None":

Conversation

aviruthen commented Mar 27, 2026

Description

Related Issue

Changes Made

AI-Generated PR

Merge Checklist

Uh oh!

mufaddal-rohawala left a comment

Choose a reason for hiding this comment

🤖 AI Code Review

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

mufaddal-rohawala left a comment

Choose a reason for hiding this comment

🤖 AI Code Review

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants