`TransformResources.instance_count` rejects `ParameterInteger`, breaking parameterized instance counts in V3 Pipelines

Leggi sagemaker-core/src/sagemaker/core/tools/constants.py e il generatore di shapes in sagemaker/core/tools/shapes_codegen.py. Aggiungi l’override della variabile della pipeline TransformResources.InstanceCount descritto nell’issue, rigenera shapes.py e verifica che l’annotazione generata accetti IntPipeVar. Il lavoro è completato quando la riproduzione costruisce una definizione di pipeline con InstanceCount che fa riferimento a MyInstanceCount.

TransformResources.instance_count rejects ParameterInteger, breaking parameterized instance counts in V3 Pipelines

Describe the bug

In V3, a batch transform step inside a SageMaker Pipeline cannot have its instance count
driven by a pipeline parameter. Passing a ParameterInteger to
Transformer(instance_count=...) raises a pydantic ValidationError when the pipeline
definition is built.

The generated TransformResources shape annotates instance_count as a bare int, so
pydantic rejects any PipelineVariable:

class TransformResources(Base):
    instance_type: StrPipeVar      # pipeline-variable aware
    instance_count: int            # <-- not IntPipeVar
    volume_kms_key_id: Optional[StrPipeVar] = Unassigned()
    transform_ami_version: Optional[StrPipeVar] = Unassigned()

instance_type is StrPipeVar (Union[str, PipelineVariable]) and resolves correctly;
only the count is affected.

This appears to be an omission rather than an intentional restriction. shapes.py is
generated by sagemaker/core/tools/shapes_codegen.py, and the fields that should accept
pipeline variables are listed explicitly in PIPE_VAR_OVERRIDES in
sagemaker-core/src/sagemaker/core/tools/constants.py:

PIPE_VAR_OVERRIDES = {
    "ResourceConfig": {
        "InstanceCount": "IntPipeVar",
        "VolumeSizeInGB": "IntPipeVar",
        "KeepAlivePeriodInSeconds": "IntPipeVar",
        "SelectedInstanceCount": "IntPipeVar",
    },
    "InstancePreference": {"InstanceCount": "IntPipeVar"},
    "ProcessingClusterConfig": {
        "InstanceCount": "IntPipeVar",
        "SelectedInstanceCount": "IntPipeVar",
    },
    "ProcessingInstancePreference": {"InstanceCount": "IntPipeVar"},
    # RandomSeed accepts a pipeline variable (e.g. a ParameterInteger) so tuning pipelines
    # can parameterize reproducibility (issue #5614 / #6171).
    "HyperParameterTuningJobConfig": {"RandomSeed": "IntPipeVar"},
}

Training (ResourceConfig) and processing (ProcessingClusterConfig) both have their
InstanceCount listed, but TransformResources has no entry. The net effect is that
within a single pipeline you can parameterize the instance count of a processing step or
a training step, but not of a transform step.

To reproduce

import boto3
from sagemaker.core.transformer import Transformer
from sagemaker.core.workflow.parameters import ParameterInteger, ParameterString
from sagemaker.core.workflow.pipeline_context import PipelineSession
from sagemaker.mlops.workflow.pipeline import Pipeline
from sagemaker.mlops.workflow.steps import TransformStep

session = PipelineSession(
    default_bucket="example-bucket",
    boto_session=boto3.Session(region_name="us-east-1"),
)

instance_count = ParameterInteger(name="MyInstanceCount", default_value=5)
instance_type = ParameterString(name="MyInstanceType", default_value="ml.c6i.8xlarge")

transformer = Transformer(
    model_name="my-model",
    base_transform_job_name="my-transform-job",
    instance_count=instance_count,   # ParameterInteger -> ValidationError
    instance_type=instance_type,     # ParameterString  -> accepted
    sagemaker_session=session,
    output_path="s3://example-bucket/out",
)

step = TransformStep(
    name="MyTransformStep",
    step_args=transformer.transform(
        data="s3://example-bucket/in", data_type="S3Prefix"
    ),
)

print(
    Pipeline(
        name="p",
        parameters=[instance_count, instance_type],
        steps=[step],
        sagemaker_session=session,
    ).definition()
)

Result:

File "sagemaker/core/transformer.py", line 732, in _get_transform_request
    "TransformResources": TransformResources(**resource_config),
pydantic_core._pydantic_core.ValidationError: 1 validation error for TransformResources
instance_count
  Input should be a valid integer [type=int_type,
  input_value=ParameterInteger(name='MyInstanceCount', ..., default_value=5),
  input_type=ParameterInteger]

Replacing instance_count=instance_count with instance_count=5 succeeds, and confirms
that instance_type resolves correctly as a pipeline reference while the count is forced
to a literal:

{"InstanceType": {"Get": "Parameters.MyInstanceType"}, "InstanceCount": 5}

Expected behavior

The same definition that V2 produces, where both fields resolve to pipeline parameter
references. Running the equivalent V2 script on sagemaker==2.251.1:

{"InstanceCount": {"Get": "Parameters.MyInstanceCount"},
 "InstanceType": {"Get": "Parameters.MyInstanceType"}}

In V2, Transformer.__init__ annotates the parameter as
typing.Union[int, sagemaker.workflow.entities.PipelineVariable], so this is a V2 → V3
regression.

Suggested fix

Add a TransformResources entry to PIPE_VAR_OVERRIDES in
sagemaker-core/src/sagemaker/core/tools/constants.py, alongside the existing training
and processing entries, and regenerate the shapes:

    "TransformResources": {"InstanceCount": "IntPipeVar"},

This mirrors the fix already applied for HyperParameterTuningJobConfig.RandomSeed in
response to issues #5614 and #6171, which look like the same class of problem.

I verified the outcome by patching the generated annotation to IntPipeVar in an
installed 3.29.0 and re-running the repro above. It then produces a definition matching
V2 exactly:

{"InstanceType": {"Get": "Parameters.MyInstanceType"},
 "InstanceCount": {"Get": "Parameters.MyInstanceCount"}}

I have not sent a PR since the change is to codegen input and requires regenerating
shapes.py, but happy to do so if that is useful.

Impact

This blocks migrating any pipeline that exposes transform instance count as a tunable
parameter. We have several such parameters, set per-run by operators to scale batch
transform jobs to the size of the input dataset. The only workaround is hardcoding the
counts into the pipeline definition, which removes per-run tuning and would require
redeploying the pipeline to change capacity.

Every other part of the V2 → V3 port works: once the counts are hardcoded, our pipelines
build on 3.29.0 and produce definitions with identical step names, step types, parameter
counts, and identical S3 artifact destinations.

System information

• SageMaker Python SDK version: 3.29.0 (sagemaker-core 2.29.0, sagemaker-mlops 1.29.0)
• Version that works: 2.251.1
• Framework / algorithm: batch transform (Transformer) inside Pipeline
• Python version: 3.12.15
• pydantic version: 2.14.3
• Is the issue with a custom Docker image: no