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

[mohamedzeidan2021]

…eneous clusters

ModelTrainer injects code/sm_drivers/recipe channels that users cannot configure. On a heterogeneous cluster, SageMaker enforces an all-or-nothing rule: if any channel is assigned to instance groups, every channel must be. So setting instance_group_names on any user channel caused CreateTrainingJob to fail with "Some channels have assigned instance groups ... while others not: [sm_drivers, code]", making instance_group_names unusable with ModelTrainer.

When a user assigns instance groups to any channel, assign the SDK-managed channels to the full set of instance groups (code/drivers must be present on every node). Behavior is unchanged when instance groups are not in use.

Fixes #6089

Issue #, if available:

Description of changes:

By submitting this pull request, I confirm that you can use, modify, copy, and redistribute this contribution, under the terms of your choice.

[github-actions]

Claude finished @mohamedzeidan2021's task in 2m 42s —— View job

---

PR Review: assign SDK-managed channels to instance groups on heterogeneous clusters

• Read PR diff
• Analyze model_trainer.py changes for correctness
• Review test coverage
• Post review feedback

---

Overall this is a clean, well-scoped fix for issue #6089. The approach is correct: the SDK-injected channels (recipe/code/sm_drivers) are not user-configurable, so when the user assigns instance groups to any channel, propagating the full instance-group list onto the managed channels is the right way to satisfy SageMaker's all-or-nothing rule. Code/drivers legitimately need to be present on every node, so assigning the full set is appropriate.

What's done well

• Correct default preservation. Using instance_group_kwargs to conditionally pass the argument (rather than passing None) is important — S3DataSource.instance_group_names defaults to Unassigned() (sagemaker-core/.../shapes.py:1532), so passing None explicitly would change the serialized payload. The conditional keeps the field unset for the non-heterogeneous path. 👍
• Backward compatible. create_input_data_channel gains only an optional keyword arg defaulting to None, so existing callers are unaffected.
• Detection handles both wrapper types. _channel_has_instance_groups correctly unwraps Channel (→ data_source.s3_data_source) vs. InputData (→ data_source when it's an S3DataSource), and safely ignores string/FileSystemDataSource sources (which have no instance_group_names).
• Managed channels all flow through the URI/local-path branches. recipe_file_path, source_dir, and the temp drivers dir are all strings, so they hit the branches that actually apply instance_group_kwargs — the caller-supplied S3DataSource/FileSystemDataSource branches (which intentionally skip it) are never used for managed channels.
• Good test coverage. The three added tests cover the three meaningful states: user-assigns → managed assigned; user-unassigned on heterogeneous → managed stays unset; homogeneous cluster → never assigned. Regression reference to 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 #6089 is helpful.

Minor observations (non-blocking)

1. Partial user assignment is still the user's responsibility. If a user has multiple channels and assigns instance groups to only some of their own channels, SageMaker will still reject the request for the unassigned user channels — this PR only fixes the SDK-managed channels. That's the correct scope (the SDK can't guess intent for user channels), but it may be worth a one-line note in the create_input_data_channel docstring or user-facing docs so users understand they must assign all their own channels.

2. getattr defensiveness. getattr(self.compute, "instance_groups", None) and getattr(s3_data_source, "instance_group_names", None) are defensive against attribute absence. Since Compute and S3DataSource are typed shapes that always define these attributes, direct attribute access would also work — but the getattr form is harmless and arguably safer against future shape changes.

3. Test assertion depends on list equality with StrPipeVar. _instance_group_names guards with isinstance(names, list), which is a reasonable way to normalize Unassigned() → None. Just be aware that if the shape ever wraps values in PipelineVariable, the direct == comparison in the assertions could need adjustment. Not a concern for current behavior.

No correctness, security, or performance issues found. Nice work on the targeted fix and the regression tests.

[github-actions]

Claude finished @mohamedzeidan2021's task in 2m 45s —— View job

---

PR Review: assign SDK-managed channels to instance groups on heterogeneous clusters

• Read PR diff (/tmp/pr.diff)
• Verify _create_training_job_args wiring and channel creation flow
• Verify _resolve_managed_channel_instance_groups / detection logic
• Verify Unassigned default-preservation behavior
• Review test coverage

---

This is a clean, correctly-scoped fix for #6089. I traced the full flow and confirm the approach is sound.

Correctness — verified

• Default preservation is correct. S3DataSource.instance_group_names defaults to Unassigned(), and Unassigned.__bool__ returns False (sagemaker-core/.../utils.py:313). The instance_group_kwargs gate (model_trainer.py:974-979) keeps the field entirely absent from the constructor call when no groups are in use, so the serialized payload is unchanged on the homogeneous path. Passing None explicitly would have been wrong — this avoids it. 👍
• Detection reads only user channels. _resolve_managed_channel_instance_groups is computed at line 598 against final_input_data_config before the managed channels are appended, so the managed channels don't circularly influence the decision. Correct.
• Both wrapper types handled. _channel_has_instance_groups unwraps Channel (→ data_source.s3_data_source) vs. InputData (→ data_source when it's an S3DataSource), and the getattr(..., None) guards safely absorb Unassigned() and string/FileSystemDataSource sources.
• User assignments survive the round-trip. User InputData channels carrying an S3DataSource flow through the elif isinstance(data_source, S3DataSource) branch of create_input_data_channel (line 1068), which preserves the caller's object as-is — so the user's own instance_group_names are retained while the managed channels get the full set.
• All three managed channels covered. recipe, code, and sm_drivers are all built from string paths, so they hit the URI/local-path branches that actually apply instance_group_kwargs.

Non-blocking observations

1. Partial user assignment stays the user's responsibility. If a user assigns instance groups to only some of their own channels, SageMaker still rejects the request for their unassigned channels — this PR only fixes the SDK-managed ones. That's the right scope (the SDK can't infer intent for user channels). The docstring addition at create_input_data_channel helps; a one-line user-facing note that "you must assign all your own channels" would be a nice follow-up.
2. getattr defensiveness. getattr(self.compute, "instance_groups", None) and getattr(s3_data_source, "instance_group_names", None) are defensive against typed shapes that always define these attributes. Harmless and arguably safer against future shape changes.

Tests

The three added tests cover the meaningful states cleanly: user-assigns → managed assigned; heterogeneous but user-unassigned → managed stays unset; homogeneous → never assigned. Regression reference to #6089 is helpful. One coverage gap worth a follow-up: a test where the user assigns groups via a Channel (not InputData) wrapper — the detection branch for Channel at line 928 isn't directly exercised, though it's straightforward.

No correctness, security, or performance issues found. Nice, targeted fix. ✅

[mohamedzeidan2021]

checks passed from before- just removed comments: https://github.com/aws/sagemaker-python-sdk/actions/runs/30401366494