ModelTrainer: heterogeneous-cluster training fails when instance_group_names is set on a user channel — SDK-managed `code`/`sm_drivers` channels have no instance-group assignment

[brunopistone]

PySDK Version

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

Describe the bug

When using a heterogeneous cluster (multiple InstanceGroups) with ModelTrainer, setting instance_group_names on the S3DataSource of a user-defined input channel makes CreateTrainingJob fail with a ValidationException.

SageMaker enforces an all-or-nothing rule: if any channel is assigned to instance groups, every channel must be. But ModelTrainer.train() automatically injects its own input channels — code (the packaged source_code) and sm_drivers (the SDK driver scripts) — and these are not user-managed: there is no public API to assign instance_group_names to them. As a result they never get an instance-group assignment, the user channel does, and the API rejects the mismatch:

ClientError: An error occurred (ValidationException) when calling the CreateTrainingJob operation: Some channels have assigned instance groups: [processing] while others not: [sm_drivers, code]

The net effect: instance_group_names on S3DataSource is effectively unusable with ModelTrainer, because the user cannot satisfy the "all channels" requirement for the SDK-created channels.

To reproduce

Complete, runnable script (creates a minimal source_dir and a dummy data file, then submits the job):

import os

# --- minimal source code (this is what becomes the SDK-managed "code" channel) ---
os.makedirs("scripts", exist_ok=True)
with open("scripts/train.py", "w") as f:
    f.write("print('hello from training')\n")

# --- minimal input data ---
os.makedirs("data", exist_ok=True)
with open("data/data.csv", "w") as f:
    f.write("a,b,c\n1,2,3\n")
input_data = sagemaker_session.upload_data("./data/data.csv", key_prefix="repro-het/input")

from sagemaker.core import image_uris
from sagemaker.core.helper.session_helper import Session, get_execution_role
from sagemaker.train.model_trainer import ModelTrainer
from sagemaker.train.configs import (
    Compute,
    InstanceGroup,
    SourceCode,
    StoppingCondition,
    InputData,
    S3DataSource,
)

sagemaker_session = Session()
region = sagemaker_session.boto_session.region_name
role = get_execution_role()

# --- heterogeneous cluster: two instance groups ---
instance_groups = [
    InstanceGroup(
        instance_group_name="head-instance-group",
        instance_type="ml.t3.large",
        instance_count=1,
    ),
    InstanceGroup(
        instance_group_name="worker-instance-group-1",
        instance_type="ml.m5.2xlarge",
        instance_count=2,
    ),
]

image_uri = image_uris.retrieve(
    framework="pytorch",
    region=region,
    version="2.8.0",
    instance_type="ml.m5.2xlarge",
    image_scope="training",
)

model_trainer = ModelTrainer(
    training_image=image_uri,
    source_code=SourceCode(source_dir="./scripts", command="python train.py"),
    base_job_name="repro-heterogeneous",
    compute=Compute(instance_groups=instance_groups),
    stopping_condition=StoppingCondition(max_runtime_in_seconds=3600),
    role=role,
)

# Scope the user channel to the instance groups (all of them, in this example)
train_input = InputData(
    channel_name="processing",
    data_source=S3DataSource(
        s3_data_type="S3Prefix",
        s3_uri=input_data,
        s3_data_distribution_type="FullyReplicated",
        instance_group_names=[
            "head-instance-group",
            "worker-instance-group-1",
        ],
    ),
)

# Fails at CreateTrainingJob with the ValidationException
model_trainer.train(input_data_config=[train_input], wait=False)

Expected behavior

train() should submit successfully. When the user assigns instance_group_names to any channel in a heterogeneous cluster, the SDK should automatically assign its internally-managed channels (code, sm_drivers, and recipe when present) to the full set of instance groups, so the CreateTrainingJob all-or-nothing validation is satisfied. Alternatively, expose a supported way to set instance groups on those SDK-managed channels. Users have no way to do this today because those channels are created internally.

Screenshots or logs

ClientError: An error occurred (ValidationException) when calling the CreateTrainingJob operation: Some channels have assigned instance groups: [processing] while others not: [sm_drivers, code]

System information

• SageMaker Python SDK version: sagemaker 3.7.0 (sagemaker-core 2.7.0, sagemaker-train 1.7.0)
• Framework name (eg. PyTorch) or algorithm (eg. KMeans): PyTorch (SageMaker training image via image_uris.retrieve)
• Framework version: 2.8.0
• Python version: 3.12.11
• CPU or GPU: CPU (ml.t3.large + ml.m5.2xlarge)
• Custom Docker image (Y/N): N