diff --git a/sagemaker-train/src/sagemaker/train/model_trainer.py b/sagemaker-train/src/sagemaker/train/model_trainer.py index ea3d62d145..48a80af03b 100644 --- a/sagemaker-train/src/sagemaker/train/model_trainer.py +++ b/sagemaker-train/src/sagemaker/train/model_trainer.py @@ -594,6 +594,11 @@ def _create_training_job_args( final_input_data_config = list(existing_channels.values()) + new_channels + # Assign SDK-managed channels to instance groups when the user uses them (see helper). + managed_channel_instance_group_names = self._resolve_managed_channel_instance_groups( + final_input_data_config + ) + if self._is_nova_recipe or self._is_llmft_recipe: for input_data in final_input_data_config: if input_data.channel_name == SM_RECIPE: @@ -606,6 +611,7 @@ def _create_training_job_args( channel_name=SM_RECIPE, data_source=recipe_file_path, key_prefix=input_data_key_prefix, + instance_group_names=managed_channel_instance_group_names, ) final_input_data_config.append(recipe_channel) if self._is_nova_recipe or self._is_llmft_recipe: @@ -664,6 +670,7 @@ def _create_training_job_args( data_source=self.source_code.source_dir, key_prefix=input_data_key_prefix, ignore_patterns=self.source_code.ignore_patterns, + instance_group_names=managed_channel_instance_group_names, ) final_input_data_config.append(source_code_channel) @@ -686,6 +693,7 @@ def _create_training_job_args( data_source=self._temp_code_dir.name, key_prefix=input_data_key_prefix, ignore_patterns=self.source_code.ignore_patterns, + instance_group_names=managed_channel_instance_group_names, ) final_input_data_config.append(sm_drivers_channel) @@ -891,12 +899,52 @@ def _resolve_staging_bucket(self) -> tuple[str,str]: return default_bucket, None + def _resolve_managed_channel_instance_groups( + self, input_data_config: List[Union[Channel, InputData]] + ) -> Optional[List[str]]: + """Resolve the instance groups to assign to SDK-managed channels. + + SageMaker requires that, on a heterogeneous cluster, either all channels are + assigned to instance groups or none are. The SDK injects channels (recipe, code, + sm_drivers) that users cannot configure, so when a user assigns instance groups to + any of their channels, the managed channels must be assigned too. We assign them to + the full set of instance groups defined on the compute config, since the code and + drivers must be present on every node regardless of instance group. + + Args: + input_data_config (List[Union[Channel, InputData]]): The user-provided channels. + + Returns: + Optional[List[str]]: The full list of instance group names to assign to + managed channels, or ``None`` if instance groups are not in use. + """ + instance_groups = getattr(self.compute, "instance_groups", None) if self.compute else None + if not instance_groups: + return None + + def _channel_has_instance_groups(channel: Union[Channel, InputData]) -> bool: + # Channel nests its S3DataSource under data_source; InputData may hold it directly. + s3_data_source = None + if isinstance(channel, Channel): + if channel.data_source: + s3_data_source = channel.data_source.s3_data_source + elif isinstance(channel, InputData): + if isinstance(channel.data_source, S3DataSource): + s3_data_source = channel.data_source + return bool(getattr(s3_data_source, "instance_group_names", None)) + + if not any(_channel_has_instance_groups(channel) for channel in input_data_config): + return None + + return [group.instance_group_name for group in instance_groups] + def create_input_data_channel( self, channel_name: str, data_source: DataSourceType, key_prefix: Optional[str] = None, ignore_patterns: Optional[List[str]] = None, + instance_group_names: Optional[List[str]] = None, ) -> Channel: """Create an input data channel for the training job. @@ -915,9 +963,21 @@ def create_input_data_channel( ignore_patterns: (Optional[List[str]]) : The ignore patterns to ignore specific files/folders when uploading to S3. If not specified, default to: ['.env', '.git', '__pycache__', '.DS_Store', '.cache', '.ipynb_checkpoints']. + instance_group_names: (Optional[List[str]]) : + The names of the instance groups (for heterogeneous clusters) that this + channel's data should be assigned to. Only applied when the channel is + built from a URI/local-path data source (not a caller-supplied + ``S3DataSource``/``FileSystemDataSource``, which the caller controls). """ from sagemaker.core.helper.pipeline_variable import PipelineVariable - + + # Pass the field only when provided, so it stays unset (Unassigned) by default. + instance_group_kwargs = ( + {"instance_group_names": instance_group_names} + if instance_group_names is not None + else {} + ) + channel = None if isinstance(data_source, PipelineVariable): channel = Channel( @@ -927,6 +987,7 @@ def create_input_data_channel( s3_data_type="S3Prefix", s3_uri=data_source, s3_data_distribution_type="FullyReplicated", + **instance_group_kwargs, ), ), input_mode="File", @@ -940,6 +1001,7 @@ def create_input_data_channel( s3_data_type="S3Prefix", s3_uri=data_source, s3_data_distribution_type="FullyReplicated", + **instance_group_kwargs, ), ), input_mode="File", @@ -996,6 +1058,7 @@ def create_input_data_channel( s3_data_type="S3Prefix", s3_uri=s3_uri, s3_data_distribution_type="FullyReplicated", + **instance_group_kwargs, ), ), input_mode="File", diff --git a/sagemaker-train/tests/unit/train/test_model_trainer.py b/sagemaker-train/tests/unit/train/test_model_trainer.py index b2da7cd2a4..ce5d208bbc 100644 --- a/sagemaker-train/tests/unit/train/test_model_trainer.py +++ b/sagemaker-train/tests/unit/train/test_model_trainer.py @@ -73,6 +73,7 @@ Channel, DataSource, MetricDefinition, + InstanceGroup, ) from sagemaker.train.distributed import Torchrun, SMP, MPI from sagemaker.train.sm_recipes.utils import _load_recipes_cfg, _is_nova_recipe, _get_args_from_nova_recipe @@ -472,6 +473,153 @@ def test_create_input_data_channel(mock_default_bucket, mock_upload_data, model_ assert channel.data_source.s3_data_source.s3_uri == expected_s3_uri +def test_create_input_data_channel_with_instance_group_names(model_trainer): + """instance_group_names is propagated onto the channel's S3DataSource.""" + channel = model_trainer.create_input_data_channel( + channel_name="code", + data_source=f"s3://{DEFAULT_BUCKET}/{DEFAULT_BASE_NAME}-job/input/code", + instance_group_names=["head-instance-group", "worker-instance-group-1"], + ) + assert channel.data_source.s3_data_source.instance_group_names == [ + "head-instance-group", + "worker-instance-group-1", + ] + + +HETEROGENEOUS_INSTANCE_GROUPS = [ + InstanceGroup( + instance_type="ml.t3.large", instance_count=1, instance_group_name="head-instance-group" + ), + InstanceGroup( + instance_type="ml.m5.2xlarge", + instance_count=2, + instance_group_name="worker-instance-group-1", + ), +] + + +def _instance_group_names(channel): + """Return the channel's assigned instance_group_names as a list (or None).""" + s3_data_source = channel.data_source.s3_data_source if channel.data_source else None + names = getattr(s3_data_source, "instance_group_names", None) if s3_data_source else None + return names if isinstance(names, list) else None + + +def _managed_channel_names(input_data_config): + return { + channel.channel_name: _instance_group_names(channel) for channel in input_data_config + } + + +@patch("sagemaker.train.model_trainer.Session.upload_data") +@patch("sagemaker.train.model_trainer.Session.default_bucket") +def test_managed_channels_assigned_instance_groups_when_user_channel_assigns( + mock_default_bucket, mock_upload_data +): + """Regression test for issue #6089. + + On a heterogeneous cluster, when a user assigns instance_group_names to any of their + channels, the SDK-managed ``code``/``sm_drivers`` channels must also be assigned to the + full set of instance groups, otherwise CreateTrainingJob fails validation with + "Some channels have assigned instance groups ... while others not". + """ + mock_upload_data.return_value = f"s3://{DEFAULT_BUCKET}/{DEFAULT_BASE_NAME}-job/input/code" + mock_default_bucket.return_value = DEFAULT_BUCKET + expected_names = ["head-instance-group", "worker-instance-group-1"] + + trainer = ModelTrainer( + training_image=DEFAULT_IMAGE, + role=DEFAULT_ROLE, + source_code=DEFAULT_SOURCE_CODE, + compute=Compute(instance_groups=HETEROGENEOUS_INSTANCE_GROUPS), + stopping_condition=DEFAULT_STOPPING_CONDITION, + output_data_config=DEFAULT_OUTPUT_DATA_CONFIG, + ) + + user_channel = InputData( + channel_name="processing", + data_source=S3DataSource( + s3_data_type="S3Prefix", + s3_uri=f"s3://{DEFAULT_BUCKET}/data/", + s3_data_distribution_type="FullyReplicated", + instance_group_names=expected_names, + ), + ) + + args = trainer._create_training_job_args(input_data_config=[user_channel]) + channels = _managed_channel_names(args["input_data_config"]) + + assert channels["processing"] == expected_names + assert channels["code"] == expected_names + assert channels["sm_drivers"] == expected_names + + +@patch("sagemaker.train.model_trainer.Session.upload_data") +@patch("sagemaker.train.model_trainer.Session.default_bucket") +def test_managed_channels_not_assigned_when_user_channel_unassigned( + mock_default_bucket, mock_upload_data +): + """Managed channels stay unassigned when the user does not use instance groups. + + Even on a heterogeneous cluster, if no user channel assigns instance_group_names, the + SDK must not assign them to managed channels (which would itself violate the + all-or-nothing rule against the unassigned user channel). + """ + mock_upload_data.return_value = f"s3://{DEFAULT_BUCKET}/{DEFAULT_BASE_NAME}-job/input/code" + mock_default_bucket.return_value = DEFAULT_BUCKET + + trainer = ModelTrainer( + training_image=DEFAULT_IMAGE, + role=DEFAULT_ROLE, + source_code=DEFAULT_SOURCE_CODE, + compute=Compute(instance_groups=HETEROGENEOUS_INSTANCE_GROUPS), + stopping_condition=DEFAULT_STOPPING_CONDITION, + output_data_config=DEFAULT_OUTPUT_DATA_CONFIG, + ) + + user_channel = InputData( + channel_name="processing", + data_source=S3DataSource( + s3_data_type="S3Prefix", + s3_uri=f"s3://{DEFAULT_BUCKET}/data/", + s3_data_distribution_type="FullyReplicated", + ), + ) + + args = trainer._create_training_job_args(input_data_config=[user_channel]) + channels = _managed_channel_names(args["input_data_config"]) + + assert channels["code"] is None + assert channels["sm_drivers"] is None + + +@patch("sagemaker.train.model_trainer.Session.upload_data") +@patch("sagemaker.train.model_trainer.Session.default_bucket") +def test_managed_channels_not_assigned_on_homogeneous_cluster( + mock_default_bucket, mock_upload_data +): + """No instance groups configured -> managed channels are never assigned.""" + mock_upload_data.return_value = f"s3://{DEFAULT_BUCKET}/{DEFAULT_BASE_NAME}-job/input/code" + mock_default_bucket.return_value = DEFAULT_BUCKET + + trainer = ModelTrainer( + training_image=DEFAULT_IMAGE, + role=DEFAULT_ROLE, + source_code=DEFAULT_SOURCE_CODE, + compute=DEFAULT_COMPUTE_CONFIG, + stopping_condition=DEFAULT_STOPPING_CONDITION, + output_data_config=DEFAULT_OUTPUT_DATA_CONFIG, + ) + + user_channel = InputData(channel_name="train", data_source=f"s3://{DEFAULT_BUCKET}/train/") + + args = trainer._create_training_job_args(input_data_config=[user_channel]) + channels = _managed_channel_names(args["input_data_config"]) + + assert channels["code"] is None + assert channels["sm_drivers"] is None + + @pytest.mark.parametrize( "test_case", [