aws / aws/sagemaker-python-sdk

ModelTrainer doesn't support heterogeneous clusters

Closed
#5,225 1 comment 0 reactions 0 assignees View on GitHub
component: training type: bug
Dominant language
Python
Stars
2.3k
Forks
1.3k
Avg merge
1d 22h
Merged PRs (30d)
35

Description

**Describe the bug**
Seems there is a bug with [ModelTrainer](https://sagemaker.readthedocs.io/en/stable/api/training/model_trainer.html) and [heterogeneous clusters](https://docs.aws.amazon.com/sagemaker/latest/dg/train-heterogeneous-cluster-configure.html). When I'm trying to run a SageMaker training job with an heterogeneous clusters, even if I'm configuring the instance_group in the compute, and in the input channels, I get the following error:

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

`sm_drivers` and `code` are private channels configured by the SDK.

**To reproduce**
A clear, step-by-step set of instructions to reproduce the bug.

```
from sagemaker.modules.configs import InstanceGroup
from sagemaker.modules.configs import InputData, S3DataSource
from sagemaker.modules.configs import (
Compute,
OutputDataConfig,
RemoteDebugConfig,
SourceCode,
StoppingCondition,
)
from sagemaker.modules.train import ModelTrainer

group_1 = InstanceGroup(
instance_type="ml.g5.xlarge",
instance_count=2,
instance_group_name="group_1",
)

group_2 = InstanceGroup(
instance_type="ml.t3.xlarge",
instance_count=1,
instance_group_name="group_2",
)

# Define the script to be run
source_code = SourceCode(
source_dir="./scripts",
requirements="requirements.txt",
command="python launcher.py -e train.py",
)

# Define the compute
compute_configs = Compute(
instance_groups=[group_1, group_2]
keep_alive_period_in_seconds=0,
)

job_name = "train-ray-processing-train"

output_path = f"s3://{bucket_name}/{job_name}"

model_trainer = ModelTrainer(
training_image=image_uri,
source_code=source_code,
base_job_name=job_name,
compute=compute_configs,
hyperparameters={
"epochs": 100,
"learning_rate": 0.001,
"batch_size": 100,
},
stopping_condition=StoppingCondition(max_runtime_in_seconds=18000),
output_data_config=OutputDataConfig(
s3_output_path=output_path, compression_type="NONE"
),
role=role,
)

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=["group_1", "group_2"]
),
)

data = [train_input]

model_trainer.train(input_data_config=data, wait=False)
```

**Expected behavior**
See Estimator behavior. `sm_drivers` and `code` are private channels configured by the SDK, so `instance_group_names` should be automatically set

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

**System information**
A description of your system. Please provide:
- **SageMaker Python SDK version**: 2.247.1
- **Framework name (eg. PyTorch) or algorithm (eg. KMeans)**: Any
- **Framework version**: Any
- **Python version**: 3.12
- **CPU or GPU**: CPU and GPU
- **Custom Docker image (Y/N)**: N

**Additional context**
Add any other context about the problem here.

Contributor guide

Open the contributing guide

Research direction

Start by running the provided ModelTrainer reproduction with SageMaker Python SDK 2.247.1 and inspect the ModelTrainer.train entry point and its generated input channels. Compare the generated configuration with Estimator behavior, especially the private sm_drivers and code channels. Done means a heterogeneous training job is accepted with instance groups assigned consistently to every channel.

Written by the indexing model from the issue text.

Assessment

Tech stack
aws, python
Domain
cloud, machine-learning
Issue type
Bug
Difficulty
4/5
Estimated time
3-5 days
Activity status
Stale
Clarity
Mostly clear
Newbie friendliness
30/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.