Skip to content

Extend PipelineVariable support in ModelTrainer to match Estimator (V2 -> V3 migration blocker) #5524

@admivsn

Description

@admivsn

PySDK Version

  • PySDK V2 (2.x)
  • PySDK V3 (3.x)

Describe the bug
In SageMaker V2 we had Estimator which had vast support for pipeline variables. For example image_uri:

class Estimator(EstimatorBase):
    """A generic Estimator to train using any supplied algorithm.

    This class is designed for use with algorithms that don't have their own, custom class.
    """

    def __init__(
        self,
        image_uri: Union[str, PipelineVariable],

In SageMaker V3, we replace Estimator with ModelTrainer, however the support for pipeline variables is lacking. For example:

class ModelTrainer(BaseModel):
    """Class that trains a model using AWS SageMaker.

    ...
    """
    ...
    training_image: Optional[str] = None

Please help us to migrate to SageMaker V3 by adding equivalent support for pipeline variables to ModelTrainer, or let us know when we could expect the support to be added.

Alternatively, happy to work on this given some guidance.

To reproduce

from sagemaker.core.workflow.parameters import ParameterString
from sagemaker.core.workflow.pipeline_context import PipelineSession
from sagemaker.train import ModelTrainer

pipeline_session = PipelineSession(default_bucket="my-bucket")
training_image = ParameterString(name="training_image")

trainer = ModelTrainer(
    training_image=training_image,
    sagemaker_session=pipeline_session
)
ValidationError: 1 validation error for ModelTrainer
training_image
  Input should be a valid string [type=string_type, input_value=ParameterString(name='tra...g'>, default_value=None), input_type=ParameterString]

Expected behavior
ModelTrainer should accept PipelineVariable for training_image (and other relevant fields), matching V2 Estimator behaviour.

Metadata

Metadata

Assignees

No one assigned

    Labels

    No labels
    No labels

    Type

    No type

    Projects

    No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions