Add optional SageMaker Training Plan support for HyperPod compute instance groups

[newabdosheham]

This PR adds optional support for attaching a SageMaker Training Plan to a specific HyperPod instance group (e.g. compute/workers) in the hyperpod-slurm-tf Terraform module.

The change enables users to run HyperPod clusters under a Training Plan while preserving full backward compatibility for existing workflows.

---

What’s new

• Introduces optional variables to enable Training Plan usage:

  • use_training_plan

  • training_plan_arn

  • training_plan_instance_group_name

• Attaches training_plan_arn only to the configured instance group (default: compute)

• Supports custom instance group names (e.g. workers, compute-nodes)

• Adds optional safety validation for:

  • instance type

  • instance count
    via:

  • training_plan_expected_instance_type

  • training_plan_expected_instance_count

---

Behavior

Scenario | Result -- | -- use_training_plan = false | No change from current behavior use_training_plan = true + valid config | Training Plan attached to target group ARN missing | Terraform fails fast with clear error Group name not found | Terraform fails fast with clear error Instance type/count mismatch (when provided) | Terraform fails fast with clear error

Validation is implemented using Terraform precondition to avoid runtime cluster failures.

---

Backward compatibility

• The feature is fully optional and disabled by default

• Existing configurations continue to work without modification

• No behavior changes unless use_training_plan = true is explicitly set

---

Implementation details

• Training Plan is injected using a conditional merge() in instance_groups

• Target group is resolved dynamically from var.instance_groups map key

• Validation is implemented using Terraform 1.2+ lifecycle.precondition

• Default target group name is compute (matching existing examples)

---

Example usage

use_training_plan = true
training_plan_arn = "arn:aws:sagemaker:us-west-2:123456789012:training-plan/my-plan"
training_plan_instance_group_name = "compute-nodes"
training_plan_expected_instance_type  = "ml.trn1.32xlarge"

training_plan_expected_instance_count = 4

---

Testing

• terraform validate

• terraform plan

• Verified:

  • Training Plan attached only to target group

  • Validation triggers correctly on mismatch

  • No behavior change when feature is disabled

---

Why this matters

SageMaker Training Plans are increasingly used for:

• capacity reservation

• cost optimization

• scheduling control

This change allows users to adopt Training Plans without forking the module and keeps the official AWS sample aligned with current SageMaker capabilities.

[newabdosheham]

Hi maintainers,
This PR adds optional support for attaching a SageMaker Training Plan to a specific HyperPod instance group (e.g. compute/workers), with full backward compatibility.

Happy to adjust naming, defaults, or structure based on your feedback. Thanks for reviewing!

[bluecrayon52]

Convert instance_group to list of objects to preserve the user-defined ordering and embed the training_plan_arn parameter into the instance group object definition. See the HyperPod EKS implementation here for reference: https://github.com/aws-samples/awsome-distributed-training/blob/8fae1a909e91a9bc5288e0cb95006573b047f89b/1.architectures/7.sagemaker-hyperpod-eks/terraform-modules/hyperpod-eks-tf/modules/hyperpod_cluster/variables.tf#L42

Doing this should remove the need for the additional environment variables (use_training_plan, training_plan_instance_group_name). Given that the variables training_plan_expected_instance_type and training_plan_expected_instance_count are user defined, they provide no greater protection. This is essentially asking the user to verify their own inputs twice.

[bluecrayon52]

Move instance_group_name value into the object definition (config.name) of the instance group, see HyperPod EKS implementation here for reference: https://github.com/aws-samples/awsome-distributed-training/blob/8fae1a909e91a9bc5288e0cb95006573b047f89b/1.architectures/7.sagemaker-hyperpod-eks/terraform-modules/hyperpod-eks-tf/modules/hyperpod_cluster/main.tf#L13

Make training_plane_arn conditional based on config.training_plan_arn, see HyperPod EKS implementation here for reference:
https://github.com/aws-samples/awsome-distributed-training/blob/8fae1a909e91a9bc5288e0cb95006573b047f89b/1.architectures/7.sagemaker-hyperpod-eks/terraform-modules/hyperpod-eks-tf/modules/hyperpod_cluster/main.tf#L47

[bluecrayon52]

Left comments recommending changes. We want to standardize on a list of objects for instance_groups and embed the name and training_plan_arn into the object definition for an instance group.

[KeitaW]

Hi @newabdosheham — this PR was automatically closed by GitHub because the repository history was rewritten as part of a cleanup to reduce clone size (details in #959).

Your fork's branch still references the old (pre-rewrite) history, so it can't be reopened directly. To restore your PR, please rebase your fork onto the new main:

# Update your fork's main branch
cd your-fork
git remote add upstream https://github.com/aws-samples/awsome-distributed-training.git  # if not already set
git fetch upstream
git checkout main
git reset --hard upstream/main
git push --force origin main

# Rebase your feature branch
git checkout feature/hyperpod-training-plan
git rebase main
# Resolve any conflicts if needed
git push --force origin feature/hyperpod-training-plan

# Then open a new PR from your fork

Apologies for the inconvenience! If you run into any issues rebasing, feel free to ask for help.