feat: add hyp migrate (Slurm cluster migration) capability

[allela-roy]

Adds the hyp migrate command for migrating a HyperPod Slurm cluster to a new Availability Zone or Region. Subcommands: snapshot, converge, validate, plan, provision. Bundles an Ansible project (identity + Slurm state capture/apply) and adds a migrate extra (ansible-core>=2.16).

What's changing and why?

SMHP Slurm customers provision P* capacity through training plans purchased for a fixed duration. When a plan expires they must buy a new one, and due to capacity constraints the new plan is frequently in a different AZ or Region. Several cluster resources are AZ/Region-scoped and cannot be moved: FSx for Lustre (/fsx) and FSx for OpenZFS (/home) are single-AZ, the S3 lifecycle-script (LCS) bucket is Region-scoped, and the cluster, its VPC/subnet/security group, and the capacity itself are AZ/Region-scoped.

Unless customers manage their clusters with Terraform (not the case ~90% of the time), they must rebuild from scratch with a manual, error-prone runbook. The most common misconfiguration is identity: FSx restores files by numeric owner, so if the new cluster assigns different UIDs/GIDs, every file — including each user's ~/.ssh — is owned by the wrong account and users cannot log in. The rebuild burns significant time and wastes paid training-plan hours while the new plan's clock is already running.

This change adds a staged, resumable hyp migrate command group that treats a cluster's configuration as a portable artifact and makes the rebuild fast, repeatable, and correct-by-construction:

• Infrastructure (VPC, subnets, security group, S3, IAM, cluster) — provisioned via CloudFormation, matching what the console/CLI already emit.
• Identity + Slurm state — captured and reproduced by a purpose-built Ansible layer that pins exact numeric UID/GID, enforces ~/.ssh permissions, and reapplies slurm.conf / gres.conf / accounting.
• File data (/fsx, /home) — the tool creates new FSx for Lustre and OpenZFS filesystems (matching source specs), restores /home from a native OpenZFS backup, and carries login-critical ~/.ssh data in the manifest. Bulk /fsx datasets remain the customer's responsibility via their existing Data Repository Associations (DRA).

A fail-closed validation gate asserts that target UID/GID and ~/.ssh permissions match the source before cutover, converting "files copied" into "users can actually log in." The tool is shape- and identity-aware (auto-detects fsx-only vs fsx+OpenZFS, and local shared_users.txt vs SSSD/AD) and training-plan aware (TrainingPlanArn binds per instance group, so the worker group draws from the plan while controller/login run on-demand, with the subnet pinned to the plan's reserved AZ).

The snapshot capability also doubles as a cost-saving mechanism: capture full cluster state, tear the cluster down to stop compute charges, and recreate it later (same identity, Slurm config, and mounts) when needed.

How the Ansible layer reaches nodes

Cluster-internal state (users, groups, ~/.ssh, slurm.conf, gres.conf, accounting) only exists on the live controller. HyperPod nodes aren't standard SSM managed instances, but are reachable through SSM Session Manager via the sagemaker-cluster:<clusterId>_<group>-<instanceId> target. The CLI drives Ansible over the amazon.aws.aws_ssm connection plugin against a runtime-generated inventory built from list-cluster-nodes, giving full idempotent Ansible semantics instead of fragile shell-scraping. Only the controller/login group is targeted; compute nodes self-heal from lifecycle scripts.

Before/After UX

Before: No first-class migration path. Customers hand-run a multi-step runbook (recreate filesystems, copy data preserving POSIX ownership, re-upload LCS, recreate the cluster, re-apply Slurm config, clean up). Identity drift silently breaks logins, and reserved training-plan hours are wasted during the rebuild.

After: A five-subcommand workflow inside the existing HyperPod CLI:

# 1) Snapshot the SOURCE cluster (read-only) to S3
hyp migrate snapshot --cluster <source> --region us-east-2 \
    --bucket s3://<artifacts-bucket> --run-id 20260101-migration

# 2) (training plans) Inspect the target plan — reports reserved instance type + AZ
hyp migrate plan --plan-arn <target-training-plan-arn> --region us-west-1

# 3) (training plans) Generate the target instance-group config (worker on the plan)
hyp migrate provision --plan-arn <target-training-plan-arn> --region us-west-1 ...

# 4) Converge the TARGET (recreate ID-pinned identity + Slurm state)
hyp migrate converge --cluster <target> --region us-west-1 \
    --bucket s3://<artifacts-bucket> --run-id 20260101-migration

# 5) Validate the TARGET before cutover (fails if identity does not match)
hyp migrate validate --cluster <target> --region us-west-1 \
    --bucket s3://<artifacts-bucket> --run-id 20260101-migration

Install the optional dependency set with pip install "sagemaker-hyperpod[migrate]" (adds ansible-core); also requires the AWS Session Manager plugin and the collections amazon.aws community.aws ansible.posix.

How was this change tested?

Validated end-to-end in two scenarios of increasing difficulty. The SSM transport was confirmed in every Region tested (us-east-1, us-east-2, us-west-1): Ansible drives HyperPod controllers over the aws_ssm connection plugin on the first try despite HyperPod's login-shell wrapper.

Test A — cross-AZ, dual filesystem (us-east-1). Deployed a Slurm cluster with FSx Lustre + OpenZFS in use1-az1, then migrated to use1-az2. Both healthy; /fsx (Lustre) and /home (OpenZFS / NFS 4.2) mounted; Slurm operational.

Test B — cross-Region training-plan migration (us-east-2 → us-west-1). The original intended use case: a Slurm cluster backed by a purchased training plan (reserved P* capacity), migrated to a different Region where a new plan provides capacity.

• Source (cluster A): us-east-2, plan AZ us-east-2a (use2-az1).
• Target (cluster B): us-west-1, plan AZ us-west-1a (usw1-az3).
• Both plans: 1× ml.p5.48xlarge, instance-type reserved capacity, ~22h40m window.
• Topology: worker (1× p5.48xlarge) on the training plan; controller + login (ml.m5.xlarge) on-demand (each plan reserves only 1 instance). Storage: FSx Lustre (/fsx) + FSx OpenZFS (/home).

Result — cross-Region training-plan migration SUCCEEDED:

• Source came up on the us-east-2 plan; worker was a real p5.48xlarge (96 vCPU, ~205 GB) drawn from reserved capacity, idle in Slurm partition dev. /fsx (Lustre) and /home (OpenZFS, NFS 4.2) both mounted.
• snapshot captured identity + Slurm state (shape=fsx_openzfs, identity_model=local, 3 users) to the source-Region S3.
• Manifest copied cross-Region (aws s3 cp --source-region us-east-2 --region us-west-1).
• Target came up on the us-west-1 plan; converge reproduced identity exactly (researcher1=5001, researcher2=5002, fsx-users=5000, ubuntu=1000) and validate PASSED (after a validator logic fix, below).

Are unit tests added?

Yes. test/unit_tests/cli/test_migrate.py (11 tests) covers the migrate command group and its subcommands (snapshot, converge, validate, plan, provision) — group wiring, expected subcommand set, and argument handling. All 11 pass locally against current main.

Are integration tests added?

No automated integration tests are added to CI (they require live HyperPod clusters, training-plan capacity, and multi-Region FSx). The workflow was validated manually end-to-end (Tests A and B above), including the cross-Region training-plan migration on ml.p5.48xlarge. Local run results are summarized in the testing section.

Reviewer Guidelines

‼️ Merge Requirements: PRs with failing integration tests cannot be merged without justification. One of the following must be true:

• All automated PR checks pass
• Failed tests include local run results/screenshots proving they work — manual end-to-end validation results are documented above (cross-AZ and cross-Region training-plan migrations succeeded; 11 unit tests pass locally)
• Changes are documentation-only