fix(train): assign SDK-managed channels to instance groups on heterog…

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",

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",
     [