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

Ouverte Adaptée aux débutants
  • Quelque chose ne va pas ?
#6,577 0 commentaires 0 réactions 0 personnes assignées Voir sur GitHub

Les mainteneurs répondent en général sous 1 jour

Personne n'a encore pris cette issue.

Évaluation

Difficulté
2/5
Temps estimé
1-3 heures
Accessibilité débutants
82/100
Type d'issue
Bug
Clarté
Clairement spécifiée
Activité
Active
Stack technique
aws, python
Domaine
backend, machine-learning

Piste de recherche

Lisez sagemaker-core/src/sagemaker/core/tools/constants.py et le générateur de shapes dans sagemaker/core/tools/shapes_codegen.py. Ajoutez la surcharge de la variable de pipeline TransformResources.InstanceCount décrite dans l’issue, régénérez shapes.py et vérifiez que l’annotation générée accepte IntPipeVar. Le travail est terminé lorsque la reproduction construit une définition de pipeline où InstanceCount fait référence à MyInstanceCount.

Rédigé par le modèle d'indexation à partir du texte de l'issue.

Description

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
Langage dominant
Python
Étoiles
2.3k
Forks
1.3k
Merge moyen
2 j 10 h
PR mergées (30 j)
65

Préparer son environnement

Par où commencer

  1. Lisez l'issue en entier, puis le guide de contribution du projet.
  2. Signalez en commentaire que vous la prenez — cela évite que deux personnes fassent le même travail.
  3. Forkez le dépôt et travaillez sur une branche.
  4. Ouvrez une pull request qui référence le numéro de l'issue.

Autres issues de aws/sagemaker-python-sdk

Toutes les issues de aws/sagemaker-python-sdk

Issues similaires

Plus d'issues Python

Recevez les nouvelles issues par e-mail

Un résumé court des issues GitHub adaptées aux débutants.