SageMaker Core TrainingJob Incorrectly expects a resource Tag instead of shape Tag

[benieric]

Describe the bug
A clear and concise description of what the bug is.

When creating a TrainingJob and passing in the a Tag shape get an error like:

ValidationError: 1 validation error for create
tags.0
  Input should be a valid dictionary or instance of Tag [type=model_type, input_value=Tag(key='key', 
value='value'), input_type=Tag]
    For further information visit https://errors.pydantic.dev/2.9/v/model_type

Example of failing code:

from sagemaker.modules import Session
from sagemaker_core.resources import TrainingJob
from sagemaker_core.shapes import AlgorithmSpecification, ResourceConfig, StoppingCondition, Tag, OutputDataConfig

import time

session = Session()
training_job_name = f"base-name-{int(time.time())}"
role = session.get_caller_identity_arn()
output_bucket = session.default_bucket()


training_job = TrainingJob.create(
    training_job_name=training_job_name,
    role_arn=role,
    algorithm_specification=AlgorithmSpecification(
        training_image=image,
        training_input_mode="File",
    ),
    resource_config=ResourceConfig(
        instance_count=1,
        instance_type="ml.m5.large",
        volume_size_in_gb=30,
    ),
    output_data_config=OutputDataConfig(
        s3_output_path=f"s3://{output_bucket}/output",
    ),
    stopping_condition=StoppingCondition(
        max_runtime_in_seconds=86400,
    ),
    tags=[Tag(key="key", value="value")],

This is because there are duplicate class names for Tag in the shape module and resource module.

• Resource Tag - https://github.com/aws/sagemaker-core/blob/1918e57dce20ad44aff0ceb8b73e31ce497ff94e/src/sagemaker_core/main/resources.py#L27596
• Shape Tag -

  sagemaker-core/src/sagemaker_core/main/shapes.py

  Line 497 in 1918e57

  class Tag(Base):

To reproduce
A clear, step-by-step set of instructions to reproduce the bug.
The provided code need to be complete and runnable, if additional data is needed, please include them in the issue.

Expected behavior
A clear and concise description of what you expected to happen.

Screenshots or logs
If applicable, add screenshots or logs to help explain your problem.

Bug information
A description of your system. Please provide:

• SageMaker Core version:
• Python version:
• Is the issue with autogen code or with generate code ?:

Additional context
Add any other context about the problem here.

[MattFriedman]

In ModelTrainer pydantic is demanding one version of Tag and when you call model_trainer.train() it is expecting a different version of Tag. So if you are using ModelTrainer and use [Tag] it will break.

Maintainers need to pick a consistent version of Tag and stick with it.

There are at least 2 versions of Tag in the codebase.

[MattFriedman]

For instance if you set your tags like this:

tags=[Tag(key='AppCode', value='FOO')]

model_trainer = ModelTrainer.from_recipe(
        tags=tags,
        networking=networking,
        role=ROLE_ARN,
        sagemaker_session=helper_session,
        training_image=image_uri,
        training_recipe=path_to_recipe,
        recipe_overrides=recipe_overrides,
        requirements="requirements.txt",
        base_job_name=output_model_name,
        compute=compute_config,
        stopping_condition=StoppingCondition(max_runtime_in_seconds=7200),
        checkpoint_config=CheckpointConfig(
            s3_uri=f"{checkpoint_s3_path}"
        ),
        output_data_config=output_data_config
    )

You will then get a pydantic validation exception when you run:

model_trainer.train(input_data_config=data, wait=True)

NOTE: the schema for using a dict as your tags is highly restrictive and forced us to try to use Tags, but we also failed to use Tags due this bug.

In that case we comment out the validation annotation on the .train() method (I think). But this is a hack just to make it work.

[benieric]

Closing since the fix has been released