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

Aberta Para iniciantes
#6,577 0 comentários 0 reações 0 responsáveis Ver no GitHub

Mantenedores costumam responder em até 1 dia

Ninguém assumiu esta issue ainda.

Avaliação

Dificuldade
2/5
Tempo estimado
1-3 horas
Facilidade para iniciantes
82/100
Tipo de issue
Bug
Clareza
Claramente especificada
Status de atividade
Ativa
Stack de tecnologia
aws, python

Direção de pesquisa

Leia sagemaker-core/src/sagemaker/core/tools/constants.py e o gerador de shapes em sagemaker/core/tools/shapes_codegen.py. Adicione a sobrescrita da variável de pipeline TransformResources.InstanceCount descrita na issue, gere shapes.py novamente e verifique se a anotação gerada aceita IntPipeVar. A tarefa estará concluída quando a reprodução construir uma definição de pipeline com InstanceCount fazendo referência a MyInstanceCount.

Escrita pelo modelo de indexação a partir do texto da issue.

Descrição

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
Linguagem predominante
Python
Estrelas
2.3k
Forks
1.3k
Merge médio
2d 10h
PRs com merge (30d)
65

Preparar o ambiente

Primeiros passos

  1. Leia a issue inteira e depois o guia de contribuição do projeto.
  2. Comente na issue dizendo que vai assumir — evita que duas pessoas façam o mesmo trabalho.
  3. Faça um fork do repositório e trabalhe em uma branch.
  4. Abra um pull request que referencie o número da issue.

Mais de aws/sagemaker-python-sdk

Todas as issues de aws/sagemaker-python-sdk

Issues semelhantes

Mais issues de Python

Receba novas issues na sua caixa de entrada

Um resumo curto de issues do GitHub para quem está começando.