diff --git a/MANIFEST.in b/MANIFEST.in new file mode 100644 index 00000000..72662ba2 --- /dev/null +++ b/MANIFEST.in @@ -0,0 +1,2 @@ +# Ship the vendored migration Ansible assets used by `hyp migrate`. +recursive-include src/sagemaker/hyperpod/cli/migrate/ansible *.yml *.yaml *.cfg *.ini *.py *.txt *.md diff --git a/README.md b/README.md index 8d596ed7..e915b5ec 100644 --- a/README.md +++ b/README.md @@ -17,6 +17,9 @@ Note: Old `hyperpod`CLI V2 has been moved to `release_v2` branch. Please refer [ - [Getting Started](#getting-started) - [CLI](#cli) - [Cluster Management](#cluster-management) + - [Slurm Cluster Migration](#slurm-cluster-migration) + - [Training plans](#training-plans) + - [Cross-Region migrations](#cross-region-migrations) - [Training](#training) - [Inference](#inference) - [Jumpstart Endpoint](#jumpstart-endpoint-creation) @@ -235,6 +238,100 @@ Reset configuration to default values: hyp reset ``` +### Slurm Cluster Migration + +Migrate a HyperPod Slurm cluster to a new Availability Zone or Region — for +example when a training plan expires and new capacity is only available +elsewhere. `hyp migrate` captures the source cluster's user identity, `~/.ssh`, +and Slurm configuration (including accounting configuration: accounts, users, +associations, QOS) into a versioned S3 manifest, then reproduces that state on a +target cluster with **pinned numeric UID/GID** and a **fail-closed validation +gate**. The target FSx for Lustre / OpenZFS filesystems are created and `/home` +is restored; bulk `/fsx` datasets are re-attached by the customer via their FSx +Data Repository Associations (DRA). + +**Prerequisites** (installed separately from the base CLI): + +```bash +pip install "sagemaker-hyperpod[migrate]" # ansible-core +ansible-galaxy collection install amazon.aws community.aws ansible.posix +# plus the AWS Session Manager plugin (session-manager-plugin) +``` + +`hyp migrate` provides five subcommands: + +| Subcommand | Purpose | +|------------|---------| +| `snapshot` | Capture identity + Slurm state from the SOURCE cluster (read-only) → S3 | +| `plan` | Inspect a target training plan (reserved instance type, AZ, capacity) | +| `provision`| Generate the target cluster instance-group config (training-plan aware) | +| `converge` | Reproduce identity + Slurm state on the TARGET cluster | +| `validate` | Fail-closed gate: assert TARGET identity/`~/.ssh` match the source | + +Use an S3 bucket named `sagemaker-*` (both the runner and the node execution +role must be able to read/write it). `snapshot`/`converge`/`validate` share a +`--run-id`: + +```bash +# 1) Snapshot the SOURCE cluster (read-only) +hyp migrate snapshot --cluster my-source --region us-east-2 \ + --bucket s3://sagemaker-my-artifacts --run-id 20260101-migration + +# 2) Converge the TARGET cluster (recreate ID-pinned identity + Slurm state) +hyp migrate converge --cluster my-target --region us-west-1 \ + --bucket s3://sagemaker-my-artifacts --run-id 20260101-migration + +# 3) Validate the TARGET before cutover (fails if identity does not match) +hyp migrate validate --cluster my-target --region us-west-1 \ + --bucket s3://sagemaker-my-artifacts --run-id 20260101-migration +``` + +Add `--groups controller-machine,login-group` to target specific instance +groups, or `--debug` for verbose Ansible output. Ansible reaches HyperPod nodes +over the SSM `aws_ssm` connection plugin; no SSH keys or bastion are required. + +#### Training plans + +When the target cluster is backed by a **training plan**, capacity is AZ-pinned +and bound per instance group. Inspect the plan, then generate a cluster config +that binds the plan to the **worker group** (controller/login stay on-demand): + +```bash +# Inspect the target plan — reports the reserved instance type + AZ +hyp migrate plan --plan-arn --region us-west-1 + +# Generate the CreateCluster instance-group config (worker on the plan) +hyp migrate provision \ + --bucket s3://sagemaker-my-tgt-artifacts \ + --execution-role-arn \ + --plan-arn \ + --region us-west-1 --out cluster-config.json +``` + +Place the target cluster **subnet in the plan's reserved AZ** (reported by +`hyp migrate plan`) or the reserved capacity will not attach. The execution role +must include EC2/VPC + ENI permissions and read/write on the `sagemaker-*` +artifact bucket. + +#### Cross-Region migrations + +The SSM transport bucket is Region-local, so for a cross-Region move copy the +snapshot manifest into a `sagemaker-*` bucket in the **target Region** before +`converge`, and grant the target node role read access to it: + +```bash +aws s3 cp s3://sagemaker-src-artifacts/RUN/ s3://sagemaker-tgt-artifacts/RUN/ \ + --recursive --source-region us-east-2 --region us-west-1 +``` + +Then run `converge`/`validate` against the target-Region bucket. + +> **Storage note:** the tool creates the target FSx filesystems and restores +> `/home` (OpenZFS) plus login-critical `~/.ssh`. Bulk `/fsx` datasets are +> re-attached by the customer via their FSx Data Repository Associations (DRA); +> users whose home is on `/fsx` have their identity migrated, and their `~/.ssh` +> arrives once `/fsx` is hydrated. + ### Training #### **Option 1**: Create Pytorch job through init experience diff --git a/setup.cfg b/setup.cfg index e883c540..f902bd32 100644 --- a/setup.cfg +++ b/setup.cfg @@ -20,6 +20,9 @@ exclude = [options.package_data] hyperpod_cli = py.typed +sagemaker.hyperpod.cli.migrate = + ansible/* + ansible/**/* [options.entry_points] diff --git a/setup.py b/setup.py index 4f452095..c31ee3ca 100644 --- a/setup.py +++ b/setup.py @@ -96,6 +96,14 @@ "hyperpod-cluster-stack-template>=1.0.0, <2.0.0", "hyperpod_space_template>=1.0.0, <2.0.0" ], + extras_require={ + # Required only for `hyp migrate` (Slurm cluster migration). + # Also needs the AWS Session Manager plugin and the collections: + # ansible-galaxy collection install amazon.aws community.aws ansible.posix + "migrate": [ + "ansible-core>=2.16", + ], + }, entry_points={ "console_scripts": [ "hyp=sagemaker.hyperpod.cli.hyp_cli:cli", diff --git a/src/sagemaker/hyperpod/cli/commands/migrate.py b/src/sagemaker/hyperpod/cli/commands/migrate.py new file mode 100644 index 00000000..d6775d47 --- /dev/null +++ b/src/sagemaker/hyperpod/cli/commands/migrate.py @@ -0,0 +1,417 @@ +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"). You +# may not use this file except in compliance with the License. A copy of +# the License is located at +# +# http://aws.amazon.com/apache2.0/ +# +# or in the "license" file accompanying this file. This file is +# distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF +# ANY KIND, either express or implied. See the License for the specific +# language governing permissions and limitations under the License. +""" +`hyp migrate` — capture a HyperPod Slurm cluster's identity + Slurm state and +reproduce it on an equivalent cluster in a different AZ/Region. + +Wraps the vendored migration Ansible playbooks (snapshot / converge / validate) +and drives them over the SSM `aws_ssm` connection plugin against a runtime- +generated inventory built from `list-cluster-nodes`. See the module README for +the full architecture. + +Note: this migrates identity (users, groups, pinned UID/GID, ~/.ssh perms) and +Slurm configuration + accounting *configuration* (accounts/users/associations/ +QOS). Bulk datasets on /fsx remain the customer's responsibility via their FSx +Data Repository Associations; historical job records are out of scope. +""" +import datetime +import json +import os +import shutil +import subprocess +from pathlib import Path +from typing import List, Optional + +import boto3 +import click + +from sagemaker.hyperpod.cli.utils import setup_logger + +logger = setup_logger(__name__) + +CONTROLLER_GROUP_HINTS = ("controller", "head", "login") + + +# --------------------------------------------------------------------------- # +# Helpers +# --------------------------------------------------------------------------- # +def _ansible_root() -> Path: + """Locate the vendored ansible assets shipped with the package.""" + # Assets live in the sibling `migrate` package: cli/migrate/ansible. + root = Path(__file__).resolve().parent.parent / "migrate" / "ansible" + if not root.is_dir(): + raise click.ClickException( + f"Bundled ansible assets not found at {root}. " + "The package may be installed incorrectly." + ) + return root + + +def _require_tooling() -> None: + """Fail early with a clear message if runtime prerequisites are missing.""" + missing = [] + if shutil.which("ansible-playbook") is None: + missing.append("ansible-core (provides ansible-playbook)") + if shutil.which("session-manager-plugin") is None: + missing.append("session-manager-plugin (AWS Session Manager plugin)") + if missing: + raise click.ClickException( + "Missing required tooling for `hyp migrate`:\n - " + + "\n - ".join(missing) + + "\n\nInstall with:\n" + " pip install 'sagemaker-hyperpod[migrate]'\n" + " ansible-galaxy collection install amazon.aws community.aws ansible.posix\n" + " # plus the AWS Session Manager plugin" + ) + + +def _resolve_cluster(cluster_name: str, region: str): + sm = boto3.client("sagemaker", region_name=region) + desc = sm.describe_cluster(ClusterName=cluster_name) + cluster_id = desc["ClusterArn"].rsplit("/", 1)[-1] + nodes = sm.list_cluster_nodes(ClusterName=cluster_name).get( + "ClusterNodeSummaries", [] + ) + return cluster_id, nodes + + +def _build_inventory( + cluster_name: str, + region: str, + groups: Optional[str], + transport_bucket: str, + out_path: Path, +) -> int: + """Generate an Ansible INI inventory targeting the controller over aws_ssm. + + Returns the number of controller hosts resolved. + """ + cluster_id, nodes = _resolve_cluster(cluster_name, region) + wanted = {g.strip() for g in (groups or "").split(",") if g.strip()} + + controllers: List[str] = [] + bucket = transport_bucket.replace("s3://", "").split("/")[0] + for n in nodes: + if n.get("InstanceStatus", {}).get("Status") != "Running": + continue + group = n["InstanceGroupName"] + if wanted and group not in wanted: + continue + low = group.lower() + if wanted or any(k in low for k in CONTROLLER_GROUP_HINTS): + iid = n["InstanceId"] + target = f"sagemaker-cluster:{cluster_id}_{group}-{iid}" + controllers.append( + f"{group}-{iid} ansible_host={target} " + f"ansible_connection=community.aws.aws_ssm " + f"ansible_aws_ssm_region={region} " + f"ansible_aws_ssm_instance_id={target} " + f"ansible_aws_ssm_bucket_name={bucket}" + ) + + if not controllers: + raise click.ClickException( + f"No running controller/login nodes found for cluster '{cluster_name}'. " + "Pass --groups to target specific instance groups." + ) + + out_path.write_text( + "[controller]\n" + + "\n".join(controllers) + + "\n\n[all:vars]\nansible_python_interpreter=/usr/bin/python3\n" + ) + return len(controllers) + + +def _run_playbook( + playbook: str, + inventory: Path, + extra_vars: dict, + debug: bool, +) -> None: + root = _ansible_root() + cmd = [ + "ansible-playbook", + str(root / "playbooks" / playbook), + "-i", + str(inventory), + ] + for k, v in extra_vars.items(): + cmd += ["-e", f"{k}={v}"] + if debug: + cmd.append("-vvv") + + env = dict(os.environ) + env.setdefault("ANSIBLE_HOST_KEY_CHECKING", "False") + # Run with the vendored ansible.cfg as CWD so roles_path resolves. + logger.info("Running: %s", " ".join(cmd)) + proc = subprocess.run(cmd, cwd=str(root), env=env) + if proc.returncode != 0: + raise click.ClickException( + f"Ansible playbook '{playbook}' failed (exit {proc.returncode}). " + "Re-run with --debug for detail." + ) + + +def _default_run_id() -> str: + return datetime.datetime.utcnow().strftime("%Y%m%d-%H%M%S") + + +def _describe_training_plan(plan_arn: str, region: str) -> dict: + """Resolve a training plan's reserved capacity: instance type, AZ, count. + + Training-plan capacity is AZ-pinned; the target cluster subnet MUST be + placed in the plan's reserved AZ or the capacity will not attach. + """ + sm = boto3.client("sagemaker", region_name=region) + name = plan_arn.rsplit("/", 1)[-1] + tp = sm.describe_training_plan(TrainingPlanName=name) + caps = tp.get("ReservedCapacitySummaries", []) + if not caps: + raise click.ClickException( + f"Training plan '{name}' has no reserved capacity summaries." + ) + cap = caps[0] + return { + "arn": tp["TrainingPlanArn"], + "status": tp.get("Status"), + "instance_type": cap.get("InstanceType"), + "availability_zone": cap.get("AvailabilityZone"), + "availability_zone_id": cap.get("AvailabilityZoneId"), + "total_instances": cap.get("TotalInstanceCount"), + "available_instances": tp.get("AvailableInstanceCount"), + } + + +def _build_cluster_config( + bucket: str, + exec_role_arn: str, + plan_arn: Optional[str], + worker_instance: str, + worker_count: int, + ondemand_instance: str, +) -> list: + """Build a HyperPod instance-group config. + + The worker group binds the training-plan capacity (TrainingPlanArn is a + per-instance-group field); controller and login run on-demand, since a plan + typically reserves only the accelerated worker instances. + """ + uri = f"s3://{bucket.replace('s3://', '').split('/')[0]}/LifecycleScripts/base-config/" + lcs = {"SourceS3Uri": uri, "OnCreate": "on_create.sh"} + + def group(name, instance, count, on_plan=False): + g = { + "InstanceGroupName": name, + "InstanceType": instance, + "InstanceCount": count, + "LifeCycleConfig": lcs, + "ExecutionRole": exec_role_arn, + "ThreadsPerCore": 1, + } + if on_plan and plan_arn: + g["TrainingPlanArn"] = plan_arn + return g + + return [ + group("controller-machine", ondemand_instance, 1), + group("login-group", ondemand_instance, 1), + group("worker-group", worker_instance, worker_count, on_plan=bool(plan_arn)), + ] + + +# --------------------------------------------------------------------------- # +# Command group +# --------------------------------------------------------------------------- # +@click.group("migrate") +def migrate(): + """Migrate a HyperPod Slurm cluster to a new AZ/Region. + + Captures user identity, ~/.ssh, and Slurm configuration from a SOURCE + cluster into a versioned S3 manifest, then reproduces that state on a + TARGET cluster with pinned numeric UID/GID and a fail-closed validation + gate. Storage: the target FSx/OpenZFS filesystems are created and /home is + restored; bulk /fsx datasets are re-attached by the customer via FSx DRA. + + Typical flow: + + \b + hyp migrate snapshot --cluster SRC --region us-east-2 --bucket s3://sagemaker-...-artifacts --run-id RUN + hyp migrate plan --plan-arn --region us-west-1 + hyp migrate provision --bucket s3://sagemaker-...-tgt --execution-role-arn --plan-arn --region us-west-1 + hyp migrate converge --cluster TGT --region us-west-1 --bucket s3://sagemaker-...-artifacts --run-id RUN + hyp migrate validate --cluster TGT --region us-west-1 --bucket s3://sagemaker-...-artifacts --run-id RUN + + Training plans: capacity is AZ-pinned and bound per instance group. The + worker group draws from the plan (TrainingPlanArn); controller/login run + on-demand. Cross-Region migrations require copying the manifest to a + sagemaker-* bucket in the target Region before converge. + """ + pass + + +def _common_options(f): + f = click.option("--cluster", "cluster_name", required=True, + help="HyperPod cluster name.")(f) + f = click.option("--region", required=True, help="AWS region of the cluster.")(f) + f = click.option("--bucket", required=True, + help="S3 artifact bucket for the manifest + SSM transport " + "(must be named sagemaker-*).")(f) + f = click.option("--run-id", default=None, + help="Migration run id (shared across snapshot/converge/" + "validate). Defaults to a UTC timestamp for snapshot.")(f) + f = click.option("--groups", default=None, + help="Comma-separated instance groups to target " + "(default: auto-detect controller/login).")(f) + f = click.option("--debug", is_flag=True, help="Enable verbose Ansible output.")(f) + return f + + +@migrate.command("snapshot") +@_common_options +def snapshot(cluster_name, region, bucket, run_id, groups, debug): + """Capture identity + Slurm state from the SOURCE cluster (read-only). + + Writes users/groups/ssh_inventory/identity_model/shape and the Slurm config + + accounting-configuration dump to s3:////. Nothing on the + source is modified. + """ + _require_tooling() + run_id = run_id or _default_run_id() + inv = _ansible_root() / "inventory.snapshot.ini" + count = _build_inventory(cluster_name, region, groups, bucket, inv) + logger.info("Resolved %d controller node(s) for '%s'.", count, cluster_name) + _run_playbook( + "snapshot.yml", + inv, + {"hpm_bucket": bucket, "hpm_run_id": run_id}, + debug, + ) + click.echo(f"\nSnapshot complete. Artifacts: {bucket.rstrip('/')}/{run_id}/") + click.echo(f"Use --run-id {run_id} for the converge/validate steps.") + + +@migrate.command("converge") +@_common_options +def converge(cluster_name, region, bucket, run_id, groups, debug): + """Reproduce the snapshot on the TARGET cluster (pins UID/GID, ssh, Slurm). + + Reads the manifest from s3://// and recreates groups/users + with their exact numeric IDs, enforces ~/.ssh permissions, and applies the + Slurm configuration. Idempotent. + """ + if not run_id: + raise click.ClickException("--run-id is required for converge (use the id from snapshot).") + _require_tooling() + inv = _ansible_root() / "inventory.converge.ini" + count = _build_inventory(cluster_name, region, groups, bucket, inv) + logger.info("Resolved %d controller node(s) for '%s'.", count, cluster_name) + _run_playbook( + "converge.yml", + inv, + {"hpm_bucket": bucket, "hpm_run_id": run_id}, + debug, + ) + click.echo("\nConverge complete. Run `hyp migrate validate` before cutover.") + + +@migrate.command("validate") +@_common_options +def validate(cluster_name, region, bucket, run_id, groups, debug): + """Fail-closed gate: assert TARGET identity + ~/.ssh perms match the source. + + Any mismatch fails the command and should block cutover. + """ + if not run_id: + raise click.ClickException("--run-id is required for validate (use the id from snapshot).") + _require_tooling() + inv = _ansible_root() / "inventory.validate.ini" + _build_inventory(cluster_name, region, groups, bucket, inv) + _run_playbook( + "validate.yml", + inv, + {"hpm_bucket": bucket, "hpm_run_id": run_id}, + debug, + ) + click.echo("\nValidation passed — identity matches source. Safe to proceed to cutover.") + + +@migrate.command("plan") +@click.option("--plan-arn", required=True, help="Target training-plan ARN.") +@click.option("--region", required=True, help="Region of the training plan.") +def plan(plan_arn, region): + """Inspect a target training plan (instance type, reserved AZ, capacity). + + Training-plan capacity is AZ-pinned: the target cluster subnet MUST be in the + plan's reserved AZ. Use the reported AZ when provisioning target networking. + """ + info = _describe_training_plan(plan_arn, region) + click.echo(json.dumps(info, indent=2)) + click.echo( + f"\nPlace the target cluster subnet in AZ '{info['availability_zone']}' " + f"({info['availability_zone_id']}). Worker instance: {info['instance_type']}, " + f"{info['available_instances']} available." + ) + + +@migrate.command("provision") +@click.option("--bucket", required=True, + help="Target LCS/artifact bucket (sagemaker-*).") +@click.option("--execution-role-arn", required=True, + help="SageMaker cluster execution role ARN (needs EC2/VPC+ENI and " + "artifact-bucket read/write).") +@click.option("--plan-arn", default=None, + help="Training-plan ARN to bind to the worker group (optional; " + "omit for on-demand-only clusters).") +@click.option("--region", required=True, help="Target region.") +@click.option("--worker-instance", default="ml.p5.48xlarge", show_default=True) +@click.option("--worker-count", type=int, default=1, show_default=True) +@click.option("--ondemand-instance", default="ml.m5.xlarge", show_default=True, + help="Instance type for controller/login (on-demand).") +@click.option("--out", default="cluster-config.json", show_default=True, + help="Where to write the generated instance-group config.") +def provision(bucket, execution_role_arn, plan_arn, region, worker_instance, + worker_count, ondemand_instance, out): + """Generate the target cluster instance-group config (training-plan aware). + + Emits a CreateCluster --instance-groups document in which the worker group + binds the training plan (TrainingPlanArn) and controller/login run + on-demand. When --plan-arn is given, the command also prints the plan's + reserved AZ so the cluster subnet can be placed correctly. + """ + if plan_arn: + info = _describe_training_plan(plan_arn, region) + if info["instance_type"] and info["instance_type"] != worker_instance: + logger.warning( + "Plan reserves %s but --worker-instance is %s; using the plan's type.", + info["instance_type"], worker_instance, + ) + worker_instance = info["instance_type"] + click.echo( + f"Training plan capacity: {worker_instance} in AZ " + f"{info['availability_zone']} ({info['availability_zone_id']}). " + f"Ensure the target subnet is in this AZ." + ) + + config = _build_cluster_config( + bucket, execution_role_arn, plan_arn, + worker_instance, worker_count, ondemand_instance, + ) + Path(out).write_text(json.dumps(config, indent=2)) + click.echo(f"\nWrote instance-group config to {out}.") + click.echo( + "Create the cluster with:\n" + f" aws sagemaker create-cluster --cluster-name " + f"--instance-groups file://{out} --vpc-config file:// " + f"--region {region}" + ) diff --git a/src/sagemaker/hyperpod/cli/hyp_cli.py b/src/sagemaker/hyperpod/cli/hyp_cli.py index 4b2107c0..2dcc19a6 100644 --- a/src/sagemaker/hyperpod/cli/hyp_cli.py +++ b/src/sagemaker/hyperpod/cli/hyp_cli.py @@ -54,6 +54,7 @@ space_template_update, ) from sagemaker.hyperpod.cli.commands.space_access import space_access_create +from sagemaker.hyperpod.cli.commands.migrate import migrate from sagemaker.hyperpod.cli.commands.init import ( init, @@ -205,6 +206,7 @@ def exec(): cli.add_command(reset) cli.add_command(configure) cli.add_command(validate) +cli.add_command(migrate) create.add_command(pytorch_create) # create.add_command(create_recipe_job_interactive) diff --git a/src/sagemaker/hyperpod/cli/migrate/__init__.py b/src/sagemaker/hyperpod/cli/migrate/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/src/sagemaker/hyperpod/cli/migrate/ansible/.ansible-lint b/src/sagemaker/hyperpod/cli/migrate/ansible/.ansible-lint new file mode 100644 index 00000000..0bbb1435 --- /dev/null +++ b/src/sagemaker/hyperpod/cli/migrate/ansible/.ansible-lint @@ -0,0 +1,16 @@ +--- +# .ansible-lint +profile: production + +exclude_paths: + - .venv/ + - inventory.generated.ini + +# The runtime inventory is generated; role vars come from S3-loaded snapshots, +# so some vars are legitimately not defined at lint time. +warn_list: + - var-naming[no-role-prefix] # roles share hpm_* vars by design + - no-changed-when # some aws s3 cp are inherently side-effecting + +skip_list: + - fqcn[action-core] # allow builtin short names in a few spots diff --git a/src/sagemaker/hyperpod/cli/migrate/ansible/.yamllint b/src/sagemaker/hyperpod/cli/migrate/ansible/.yamllint new file mode 100644 index 00000000..98bcf56d --- /dev/null +++ b/src/sagemaker/hyperpod/cli/migrate/ansible/.yamllint @@ -0,0 +1,22 @@ +--- +extends: default + +rules: + line-length: + max: 160 + level: warning + truthy: + allowed-values: ["true", "false"] + check-keys: false + comments: + min-spaces-from-content: 1 + comments-indentation: disable + braces: + max-spaces-inside: 1 + octal-values: + forbid-implicit-octal: true + forbid-explicit-octal: false + +ignore: | + inventory.generated.ini + .venv/ diff --git a/src/sagemaker/hyperpod/cli/migrate/ansible/README.md b/src/sagemaker/hyperpod/cli/migrate/ansible/README.md new file mode 100644 index 00000000..1d72d88d --- /dev/null +++ b/src/sagemaker/hyperpod/cli/migrate/ansible/README.md @@ -0,0 +1,200 @@ +# hpmigrate-ansible + +Purpose-built Ansible for the **SageMaker HyperPod Slurm cluster migration +utility** (`hpmigrate`). This is a **fresh, migration-specific** codebase — it is +NOT the general node-config playbook (`hyperpod-ansible/site.yml`). Its only job +is to **snapshot** identity + Slurm state from a source cluster and **converge** +an equivalent, ID-pinned state onto a target cluster during an AZ/Region +migration. + +See the design doc `hyperpod-slurm-migration-design.md` (§9 data, §10 identity, +§11 user management) for the "what". This repo is the "how". + +--- + +## 1. Why a separate Ansible codebase + +| `hyperpod-ansible` (existing) | `hpmigrate-ansible` (this repo) | +|---|---| +| Day-2 config mgmt of a **running** cluster | One-shot **migration** between two clusters | +| Push desired users/packages | **Capture** live state, then **reproduce** it elsewhere | +| No notion of source vs target | Explicit **source** (read-only) and **target** (converge) roles | +| IDs are whatever the file says | **Pins exact numeric UID/GID** captured from source (correctness gate) | +| Single `site.yml` | `snapshot.yml` + `converge.yml` with distinct role sets | + +The migration has a hard correctness requirement the day-2 playbook doesn't: +DRA import / OpenZFS restore preserve files by **numeric owner**, so the target +must reproduce identical UID/GID or `~/.ssh` (and all data) is owned by the wrong +account and logins break. This codebase is organized around enforcing that. + +--- + +## 2. Transport (inherited lesson, re-implemented) + +HyperPod nodes are **not** SSM-managed instances — no `describe-instance-information`, +`send-command` is rejected, no `mi-*`. The only programmatic entry is **SSM +Session Manager** via the target `sagemaker-cluster:_-`. + +We therefore drive Ansible with **`--connection=community.aws.aws_ssm`** (SSM as +the connection plugin), using a **dynamically generated inventory** built at run +time from `aws sagemaker list-cluster-nodes`. This gives real Ansible semantics +(idempotency, handlers, `slurp`/`fetch`, check mode) over the only channel +HyperPod exposes — instead of scraping a raw session shell. + +- Inventory is generated by `bin/gen_inventory.py` (see §5), so scaling / node + replacement is handled at run time. +- Only the **controller/login group** is targeted for identity + Slurm state + (that is where `getent`, `slurm.conf`, `sacctmgr` live). Compute nodes are + converged by the LCS bootstrap, not by this fan-out. +- An **S3 artifact bucket** is the durable store for the captured snapshot and + run results (survives session drops; readable by both source and target + contexts). + +--- + +## 3. Layout + +``` +hpmigrate-ansible/ + ansible.cfg + README.md # this file (architecture) + bin/ + gen_inventory.py # AWS -> Ansible inventory over SSM target + hpmigrate-ansible.sh # thin wrapper: gen inventory + run a playbook + playbooks/ + snapshot.yml # SOURCE (read-only): capture identity + slurm -> S3 + converge.yml # TARGET: recreate ID-pinned users, ssh perms, slurm + validate.yml # TARGET: assert UID/GID + .ssh perms == source + group_vars/ + all.yml # bucket, paths, toggles (shape/model aware) + roles/ + discover/ # detect storage shape + identity model + homes + identity_capture/ # getent passwd/group -> structured snapshot -> S3 + slurm_capture/ # slurm.conf, gres.conf, sacctmgr dump -> S3 + identity_apply/ # groups+users with PINNED numeric UID/GID + ssh_perms/ # enforce ~/.ssh 700 / keys 600 / correct owner + slurm_apply/ # place slurm.conf/gres.conf, sacctmgr load + validate_identity/ # compare target vs captured; fail on mismatch +``` + +--- + +## 4. The two flows + +### 4.1 Snapshot (source cluster, read-only) +`playbooks/snapshot.yml` → roles `discover`, `identity_capture`, `slurm_capture`. +Produces, under `s3:////`: +``` +identity/ + identity_model.json # local | sssd (detected) + users.json # [{name, uid, gid, groups[], shell, home, home_fs}] + groups.json # [{name, gid, members[]}] incl. fsx-users + ssh_inventory.json # per-user .ssh path, owner, mode (pre-migration truth) +slurm/ + slurm.conf gres.conf cgroup.conf topology.conf + sacctmgr_dump.cfg # accounts/users/assoc/QOS +storage/ + shape.json # fsx_only | fsx_openzfs +``` +Nothing on the source is modified. + +### 4.2 Converge (target cluster) +`playbooks/converge.yml` → roles `identity_apply`, `ssh_perms`, `slurm_apply`. +Reads the snapshot from S3 and **pins** every group/user to the captured numeric +UID/GID, recreates `~/.ssh` perms, and places Slurm config. Idempotent, so the +same playbook is also what `lcs/on_create.sh`-style bootstrap runs to self-heal +replaced nodes. Model 2 (SSSD/AD) skips local user creation and only ensures the +directory join. + +### 4.3 Validate (target cluster, gate) +`playbooks/validate.yml` → role `validate_identity`. Re-reads live `getent` on +the target, diffs against the captured snapshot, and **fails** if any UID/GID +differs or if a user's `~/.ssh` owner/mode does not match the **captured source** +(not a hardcoded policy — a successful migration preserves the source's perms, +even non-standard ones). `.ssh` is only enforced for homes the tool migrates +(`/home`); `/fsx` homes ride customer DRA and are checked after hydration. For +`fsx+OpenZFS` clusters it also asserts `/home` is actually NFS-mounted. This is +the cutover gate from design §11.4. + +### 4.4 Training plans and cross-Region +- **Training plans:** capacity is AZ-pinned and bound per instance group. The + CLI (`hyp migrate plan` / `provision`) binds the plan's `TrainingPlanArn` to + the **worker group** and pins the cluster subnet to the plan's reserved AZ; + controller/login run on-demand. +- **Cross-Region:** the `aws_ssm` transport bucket is Region-local, so the + snapshot must be copied to a `sagemaker-*` bucket in the target Region before + `converge` (`aws s3 cp --source-region … --region …`), and the target node + role needs read on it. + +--- + +## 5. Usage + +```bash +# 0. Prereqs: awscli v2, session-manager-plugin, ansible-core, +# community.aws + amazon.aws collections, boto3. + +export HPM_BUCKET=s3://my-hpmigrate-artifacts +export HPM_RUN_ID=$(date +%Y%m%d-%H%M%S) + +# 1. Snapshot the SOURCE cluster (read-only): +./bin/hpmigrate-ansible.sh snapshot \ + --cluster src-cluster --region us-west-2 --group controller-machine + +# 2. Converge the TARGET cluster (recreate ID-pinned identity + slurm): +./bin/hpmigrate-ansible.sh converge \ + --cluster tgt-cluster --region us-east-1 --group controller-machine + +# 3. Gate before cutover: +./bin/hpmigrate-ansible.sh validate \ + --cluster tgt-cluster --region us-east-1 --group controller-machine +``` + +The wrapper generates the SSM inventory, then invokes the matching playbook with +`-e hpm_bucket=... -e hpm_run_id=...`. `hpmigrate` (the parent CLI) calls this +wrapper at the identity/slurm steps of `provision`/`validate`. + +--- + +## 6. IAM + +**Runner principal** (CI / operator host): +- `sagemaker:DescribeCluster`, `sagemaker:ListClusterNodes` +- `ssm:StartSession` on `arn:aws:sagemaker:*:*:cluster/*`, `ssm:TerminateSession` +- `s3:PutObject/GetObject/ListBucket` on the artifact bucket + +**Node execution role** (already on HyperPod nodes): +- read/write the artifact bucket prefix (for `slurp`/`fetch` staging + results) +- standard SSM agent permissions HyperPod already grants + +--- + +## 7. Guarantees + +- **Read-only source**: `snapshot.yml` uses only `getent`, `slurp`, `sacctmgr + dump` — never mutates. +- **ID correctness**: `identity_apply` always sets explicit `-u/-g`; it never + lets `useradd` auto-assign. +- **Shape/model aware**: `discover` gates OpenZFS vs fsx-only and local vs SSSD, + so a `fsx_only` cluster is never told to touch `/home`, and an SSSD cluster + never gets local users created. +- **Fail-closed validation**: `validate.yml` is the hard gate; cutover does not + proceed on any mismatch. + +--- + +## 8. Testing + +- **Offline / CI** — `yamllint`, `ansible-lint`, and `--syntax-check` on every + playbook (no AWS). Runs via `.github/workflows/ci.yml`. +- **Live** — end-to-end over SSM against real clusters. The authoritative + pre-flight (incl. the critical `aws_ssm`-through-HyperPod-shell-wrapper sanity + check) is in `TESTING.md`. + +```bash +pip install -r requirements.txt +ansible-galaxy collection install -r requirements.yml +yamllint . && ansible-lint +for p in playbooks/*.yml; do ansible-playbook "$p" --syntax-check -i localhost, \ + -e hpm_bucket=s3://ci-noop -e hpm_run_id=ci; done +``` diff --git a/src/sagemaker/hyperpod/cli/migrate/ansible/TESTING.md b/src/sagemaker/hyperpod/cli/migrate/ansible/TESTING.md new file mode 100644 index 00000000..9949fdad --- /dev/null +++ b/src/sagemaker/hyperpod/cli/migrate/ansible/TESTING.md @@ -0,0 +1,114 @@ +# TESTING — hpmigrate-ansible + +Two layers of testing: + +1. **Offline / CI** — syntax and lint. No AWS, no HyperPod. Runs on every push + (see `.github/workflows/ci.yml`). +2. **Live** — end-to-end against real HyperPod clusters over SSM. Manual, + gated, and destructive-adjacent (it creates users on a *target*). The + checklist below is the authoritative pre-flight. + +--- + +## 1. Offline / CI (fast, safe) + +```bash +python3 -m pip install -r requirements.txt +ansible-galaxy collection install -r requirements.yml + +# a) YAML lint + ansible-lint +yamllint . +ansible-lint + +# b) Playbook syntax (no connection made) +ansible-playbook playbooks/snapshot.yml --syntax-check -i localhost, \ + -e hpm_bucket=s3://ci-noop -e hpm_run_id=ci +ansible-playbook playbooks/converge.yml --syntax-check -i localhost, \ + -e hpm_bucket=s3://ci-noop -e hpm_run_id=ci +ansible-playbook playbooks/validate.yml --syntax-check -i localhost, \ + -e hpm_bucket=s3://ci-noop -e hpm_run_id=ci +``` + +--- + +## 2. Live test checklist (HyperPod, over SSM) + +> Do this first on a **throwaway pair** of small clusters, never on production +> capacity. The converge/validate steps must run against a **target** you are +> willing to mutate. + +### 2.0 Prereqs on the runner host +- [ ] `awscli v2`, `session-manager-plugin`, `ansible-core`, collections from + `requirements.yml`, `boto3` installed. +- [ ] Runner principal IAM: `sagemaker:DescribeCluster`, `sagemaker:ListClusterNodes`, + `ssm:StartSession`/`TerminateSession` on `arn:aws:sagemaker:*:*:cluster/*`, + and `s3:*Object`/`ListBucket` on the artifact bucket. +- [ ] Node execution role can read/write the artifact bucket prefix. +- [ ] `export HPM_BUCKET=s3://... HPM_RUN_ID=$(date +%Y%m%d-%H%M%S)` + +### 2.1 Transport sanity (the риск area — validate FIRST) +The `community.aws.aws_ssm` plugin uses its own SSM document + S3 file transfer, +**not** an interactive shell. HyperPod wraps the login shell (drops into a nested +root `bash`), which historically broke stdout-scraping approaches. Confirm the +plugin tunnels cleanly through that wrapper before trusting the playbooks: +- [ ] `python3 bin/gen_inventory.py --cluster SRC --region R --group controller-machine --bucket $HPM_BUCKET` + prints a `[controller]` host with `ansible_connection=community.aws.aws_ssm`. +- [ ] Raw ping over SSM: + `ansible -i inventory.generated.ini controller -m ping` + → expect `pong`. **If this hangs or errors**, the wrapper/plugin + interaction is the culprit — fall back to the shell-based fan-out from + `hyperpod-ansible/ci/run_ansible.sh` and file a note in the README. +- [ ] `ansible -i inventory.generated.ini controller -m command -a 'id'` + returns `uid=0(root)` (confirms privilege + real command execution). +- [ ] `ansible -i inventory.generated.ini controller -m command -a 'aws s3 ls $HPM_BUCKET'` + works from the node (confirms node role S3 access used by capture). + +### 2.2 Snapshot (SOURCE, read-only) +- [ ] `./bin/hpmigrate-ansible.sh snapshot --cluster SRC --region R --group controller-machine` +- [ ] `aws s3 ls $HPM_BUCKET/$HPM_RUN_ID/ --recursive` shows + `identity/{users,groups,ssh_inventory,identity_model}.json`, + `slurm/{slurm.conf,gres.conf,sacctmgr_dump.cfg}`, `storage/shape.json`. +- [ ] `users.json` UID/GIDs match `getent passwd` on the source (spot-check a + few, incl. members of `fsx-users`). +- [ ] `identity_model.json` correctly reports `local` vs `sssd`. +- [ ] `shape.json` correctly reports `fsx_only` vs `fsx_openzfs`. +- [ ] **Source unchanged**: re-run `getent passwd | wc -l` on source before/after + → identical (no accidental writes). + +### 2.3 Converge (TARGET) +- [ ] Run BEFORE data-move/mount so ownership lines up: + `./bin/hpmigrate-ansible.sh converge --cluster TGT --region R --group controller-machine` +- [ ] `getent passwd ` on target shows **identical UID/GID** to source. +- [ ] `getent group fsx-users` GID matches source. +- [ ] Idempotency: run converge **again** → `changed=0`. +- [ ] SSSD case: on a directory-managed source, converge **skips** local user + creation (check the debug line) and does not create local accounts. +- [ ] fsx_only case: no task attempts to touch `/home`. + +### 2.4 Post-data-move .ssh perms +After DRA-import (`/fsx`) and, in Shape B, OpenZFS-restore (`/home`): +- [ ] Re-run converge (or just the `ssh_perms` role) so perms are enforced on the + now-populated homes. +- [ ] For a sample user: `ls -ld ~/.ssh` → `700`, owned by the user; + `ls -l ~/.ssh/authorized_keys` → `600`, owned by the user. +- [ ] Actually SSH in as a migrated user with their existing key → **succeeds**. + +### 2.5 Validate (the gate) +- [ ] `./bin/hpmigrate-ansible.sh validate --cluster TGT --region R --group controller-machine` + → `IDENTITY VALIDATION PASSED`. +- [ ] Negative test: on a scratch target, manually `usermod -u `, + re-run validate → it **fails closed** with the mismatch message. Restore. + +### 2.6 Slurm +- [ ] `sinfo` on target shows the source partitions. +- [ ] `sacctmgr show assoc` matches source associations/QOS. +- [ ] `gres.conf` customization (e.g. per-GPU booking) is present on target. + +--- + +## 3. What is intentionally NOT tested here +- **Bulk data movement** (DRA export/import, OpenZFS backup/restore) — owned by + the parent `hpmigrate` CLI, not Ansible. Ansible only enforces ownership/perms + on the resulting files and validates identity. +- **Full slurmdbd `mysqldump`** — parent CLI, not Ansible. +- **Cluster/VPC/FSx provisioning** — CloudFormation via the parent CLI. diff --git a/src/sagemaker/hyperpod/cli/migrate/ansible/ansible.cfg b/src/sagemaker/hyperpod/cli/migrate/ansible/ansible.cfg new file mode 100644 index 00000000..97505dd7 --- /dev/null +++ b/src/sagemaker/hyperpod/cli/migrate/ansible/ansible.cfg @@ -0,0 +1,12 @@ +[defaults] +inventory = ./inventory.generated.ini +host_key_checking = False +retry_files_enabled = False +stdout_callback = default +roles_path = ./roles +gathering = explicit +# HyperPod nodes are reached over SSM Session Manager, not SSH. +# The inventory sets ansible_connection=community.aws.aws_ssm per host. + +[ssh_connection] +pipelining = True diff --git a/src/sagemaker/hyperpod/cli/migrate/ansible/bin/gen_inventory.py b/src/sagemaker/hyperpod/cli/migrate/ansible/bin/gen_inventory.py new file mode 100755 index 00000000..e08de6ac --- /dev/null +++ b/src/sagemaker/hyperpod/cli/migrate/ansible/bin/gen_inventory.py @@ -0,0 +1,94 @@ +#!/usr/bin/env python3 +""" +gen_inventory.py — build an Ansible inventory for a HyperPod cluster reachable +only over SSM Session Manager. + +HyperPod nodes are not SSM-managed instances, so we cannot use a normal dynamic +inventory keyed on instance IDs. Instead each node is reachable via the special +Session target: sagemaker-cluster:_- + +We emit an INI inventory where each host uses the community.aws.aws_ssm +connection plugin with that target as ansible_host. + +Usage: + gen_inventory.py --cluster NAME --region REGION [--group G1,G2] \ + [--bucket s3://...] > inventory.generated.ini +""" +import argparse +import json +import subprocess +import sys + + +def aws(region, *args): + out = subprocess.check_output( + ["aws", "--region", region, "--output", "json", *args] + ) + return json.loads(out) + + +def main(): + ap = argparse.ArgumentParser() + ap.add_argument("--cluster", required=True) + ap.add_argument("--region", required=True) + ap.add_argument("--group", default="", help="comma-separated instance groups") + ap.add_argument("--bucket", default="", help="s3 bucket for aws_ssm transport") + args = ap.parse_args() + + desc = aws(args.region, "sagemaker", "describe-cluster", + "--cluster-name", args.cluster) + cluster_arn = desc["ClusterArn"] + cluster_id = cluster_arn.rsplit("/", 1)[-1] + + nodes = aws(args.region, "sagemaker", "list-cluster-nodes", + "--cluster-name", args.cluster).get("ClusterNodeSummaries", []) + + wanted = {g.strip() for g in args.group.split(",") if g.strip()} + + lines = [] + lines.append("[controller]") + controllers = [] + workers = [] + for n in nodes: + if n.get("InstanceStatus", {}).get("Status") != "Running": + continue + group = n["InstanceGroupName"] + if wanted and group not in wanted: + continue + iid = n["InstanceId"] + target = f"sagemaker-cluster:{cluster_id}_{group}-{iid}" + host = f"{group}-{iid}" + entry = ( + f"{host} ansible_host={target} " + f"ansible_connection=community.aws.aws_ssm " + f"ansible_aws_ssm_region={args.region} " + f"ansible_aws_ssm_instance_id={target}" + ) + if args.bucket: + b = args.bucket.replace("s3://", "") + entry += f" ansible_aws_ssm_bucket_name={b}" + # Heuristic: controller/head/login groups host identity + slurm state. + low = group.lower() + if any(k in low for k in ("controller", "head", "login")): + controllers.append(entry) + else: + workers.append(entry) + + if not controllers and not workers: + sys.stderr.write("ERROR: no running nodes matched.\n") + sys.exit(1) + + # If nothing matched the controller heuristic, treat all matched as controller + # (caller likely passed --group controller-machine explicitly). + if not controllers: + controllers = workers + workers = [] + + out = ["[controller]", *controllers, "", "[workers]", *workers, "", + "[all:vars]", + "ansible_python_interpreter=/usr/bin/python3"] + print("\n".join(out)) + + +if __name__ == "__main__": + main() diff --git a/src/sagemaker/hyperpod/cli/migrate/ansible/bin/hpmigrate-ansible.sh b/src/sagemaker/hyperpod/cli/migrate/ansible/bin/hpmigrate-ansible.sh new file mode 100755 index 00000000..706d38ea --- /dev/null +++ b/src/sagemaker/hyperpod/cli/migrate/ansible/bin/hpmigrate-ansible.sh @@ -0,0 +1,48 @@ +#!/bin/bash +# hpmigrate-ansible.sh — generate the SSM inventory then run a migration playbook. +# +# Subcommands: snapshot | converge | validate +# +# Example: +# HPM_BUCKET=s3://my-artifacts HPM_RUN_ID=20260101-1200 \ +# ./bin/hpmigrate-ansible.sh snapshot --cluster src --region us-west-2 \ +# --group controller-machine +set -euo pipefail + +HERE="$(cd "$(dirname "$0")/.." && pwd)" +cd "${HERE}" + +SUB="${1:-}"; shift || true +case "${SUB}" in + snapshot) PLAYBOOK=playbooks/snapshot.yml ;; + converge) PLAYBOOK=playbooks/converge.yml ;; + validate) PLAYBOOK=playbooks/validate.yml ;; + *) echo "Usage: $0 {snapshot|converge|validate} --cluster N --region R [--group G]"; exit 1 ;; +esac + +CLUSTER=""; REGION=""; GROUP="" +while [ $# -gt 0 ]; do + case "$1" in + --cluster) CLUSTER="$2"; shift 2 ;; + --region) REGION="$2"; shift 2 ;; + --group) GROUP="$2"; shift 2 ;; + *) echo "unknown arg: $1"; exit 1 ;; + esac +done +[ -n "${CLUSTER}" ] && [ -n "${REGION}" ] || { echo "need --cluster and --region"; exit 1; } + +: "${HPM_BUCKET:?set HPM_BUCKET=s3://...}" +: "${HPM_RUN_ID:=$(date +%Y%m%d-%H%M%S)}" +export HPM_BUCKET HPM_RUN_ID + +echo ">> generating SSM inventory for ${CLUSTER} (${REGION})" +python3 bin/gen_inventory.py --cluster "${CLUSTER}" --region "${REGION}" \ + --group "${GROUP}" --bucket "${HPM_BUCKET}" > inventory.generated.ini +cat inventory.generated.ini + +echo ">> running ${PLAYBOOK} (run_id=${HPM_RUN_ID})" +ansible-playbook "${PLAYBOOK}" \ + -e "hpm_bucket=${HPM_BUCKET}" \ + -e "hpm_run_id=${HPM_RUN_ID}" \ + -e "target_cluster=${CLUSTER}" \ + -e "target_region=${REGION}" diff --git a/src/sagemaker/hyperpod/cli/migrate/ansible/group_vars/all.yml b/src/sagemaker/hyperpod/cli/migrate/ansible/group_vars/all.yml new file mode 100644 index 00000000..2a6dc000 --- /dev/null +++ b/src/sagemaker/hyperpod/cli/migrate/ansible/group_vars/all.yml @@ -0,0 +1,47 @@ +--- +# group_vars/all.yml — shared config for snapshot / converge / validate. +# Most values are provided at runtime by bin/hpmigrate-ansible.sh via -e, but +# defaults live here so playbooks are runnable standalone for debugging. + +# S3 artifact store (durable snapshot + results). Overridden by HPM_BUCKET. +hpm_bucket: "{{ lookup('env', 'HPM_BUCKET') | default('s3://CHANGE-ME-hpmigrate-artifacts', true) }}" +hpm_run_id: "{{ lookup('env', 'HPM_RUN_ID') | default('dev-run', true) }}" + +# Derived S3 prefixes for the snapshot artifacts. +hpm_prefix: "{{ hpm_bucket }}/{{ hpm_run_id }}" +hpm_identity_prefix: "{{ hpm_prefix }}/identity" +hpm_slurm_prefix: "{{ hpm_prefix }}/slurm" +hpm_storage_prefix: "{{ hpm_prefix }}/storage" + +# Local staging dir on the controller node during capture/converge. +hpm_workdir: /opt/hpmigrate + +# ---- Detection outputs (filled by the discover role) ---------------------- +# storage_shape: "fsx_only" | "fsx_openzfs" +# identity_model: "local" | "sssd" +# These are auto-detected; set here only to force a value for testing. +storage_shape: "" +identity_model: "" + +# ---- Identity policy ------------------------------------------------------ +# Never auto-assign IDs. If a captured user is missing a numeric UID/GID the +# converge run FAILS rather than guessing (correctness over convenience). +require_numeric_ids: true + +# UID/GID floor: skip system accounts below this when capturing human users. +# HyperPod human users are typically >= 1000; fsx-users group is captured +# regardless via the explicit system_groups_keep list. +uid_min: 1000 +gid_min: 1000 +system_groups_keep: + - fsx-users + +# ---- SSH perms enforcement ------------------------------------------------ +ssh_dir_mode: "0700" +ssh_privkey_mode: "0600" +ssh_authkeys_mode: "0600" + +# ---- Slurm ---------------------------------------------------------------- +slurm_conf_dir: /opt/slurm/etc # base-config default location +capture_accounting_dump: true # sacctmgr dump (assoc/qos) +# Full DB mysqldump is handled by the parent hpmigrate CLI, not Ansible. diff --git a/src/sagemaker/hyperpod/cli/migrate/ansible/playbooks/converge.yml b/src/sagemaker/hyperpod/cli/migrate/ansible/playbooks/converge.yml new file mode 100644 index 00000000..f5b190ce --- /dev/null +++ b/src/sagemaker/hyperpod/cli/migrate/ansible/playbooks/converge.yml @@ -0,0 +1,34 @@ +--- +# converge.yml — TARGET cluster. +# Reads the snapshot from S3 and reproduces an ID-pinned identity + Slurm state. +# Idempotent: safe to re-run and usable as the LCS bootstrap convergence step. +# +# Order matters: identity (groups+users with pinned numeric IDs) must exist +# BEFORE data is populated/mounted so DRA-import / OpenZFS-restore ownership +# lines up. hpmigrate sequences this playbook before the data-move for that +# reason. + +- name: Converge HyperPod target cluster (identity + ssh + slurm) + hosts: controller + gather_facts: true + become: true + vars: + hpm_mode: converge + pre_tasks: + - name: Load captured snapshot from S3 + ansible.builtin.include_role: + name: discover + tasks_from: load_snapshot.yml + + roles: + - role: identity_apply # groups + users with EXACT numeric UID/GID + - role: ssh_perms # ~/.ssh 700, keys 600, correct owner + - role: slurm_apply # place slurm.conf/gres.conf, sacctmgr load + + post_tasks: + - name: Converge summary + ansible.builtin.debug: + msg: + - "identity_model = {{ identity_model }}" + - "users applied = {{ (identity_model == 'sssd') | ternary('skipped (directory-managed)', captured_users | length) }}" + - "run local `validate.yml` before cutover" diff --git a/src/sagemaker/hyperpod/cli/migrate/ansible/playbooks/seed_source.yml b/src/sagemaker/hyperpod/cli/migrate/ansible/playbooks/seed_source.yml new file mode 100644 index 00000000..e5a08b5a --- /dev/null +++ b/src/sagemaker/hyperpod/cli/migrate/ansible/playbooks/seed_source.yml @@ -0,0 +1,61 @@ +--- +# seed_source.yml — populate the SOURCE cluster with representative state so the +# migration has something meaningful to capture and reproduce. +# Creates: fsx-users group (gid 5000), two users with pinned UID/GID whose homes +# live on /fsx, their .ssh dirs + authorized_keys, and a marker file. +- name: Seed source cluster identity state + hosts: controller + become: true + vars: + seed_users: + - { name: researcher1, uid: 5001, gid: 5000 } + - { name: researcher2, uid: 5002, gid: 5000 } + tasks: + - name: Create fsx-users group (gid 5000) + ansible.builtin.group: + name: fsx-users + gid: 5000 + state: present + + - name: Create seed users with pinned UID/GID, home on /fsx + ansible.builtin.user: + name: "{{ item.name }}" + uid: "{{ item.uid }}" + group: "{{ item.gid }}" + home: "/fsx/{{ item.name }}" + shell: /bin/bash + create_home: true + loop: "{{ seed_users }}" + loop_control: { label: "{{ item.name }} uid={{ item.uid }}" } + + - name: Create .ssh dirs + ansible.builtin.file: + path: "/fsx/{{ item.name }}/.ssh" + state: directory + owner: "{{ item.uid }}" + group: "{{ item.gid }}" + mode: "0700" + loop: "{{ seed_users }}" + loop_control: { label: "{{ item.name }}" } + + - name: Install an authorized_keys per user + ansible.builtin.copy: + dest: "/fsx/{{ item.name }}/.ssh/authorized_keys" + content: "ssh-ed25519 AAAAC3NzaC1SEEDKEY {{ item.name }}@source\n" + owner: "{{ item.uid }}" + group: "{{ item.gid }}" + mode: "0600" + loop: "{{ seed_users }}" + loop_control: { label: "{{ item.name }}" } + + - name: Show resulting identity state + ansible.builtin.shell: | + echo "=== passwd ==="; getent passwd researcher1 researcher2 + echo "=== group ==="; getent group fsx-users + echo "=== ssh ==="; ls -ln /fsx/researcher1/.ssh /fsx/researcher2/.ssh + register: seed_out + changed_when: false + + - name: Seed summary + ansible.builtin.debug: + var: seed_out.stdout_lines diff --git a/src/sagemaker/hyperpod/cli/migrate/ansible/playbooks/snapshot.yml b/src/sagemaker/hyperpod/cli/migrate/ansible/playbooks/snapshot.yml new file mode 100644 index 00000000..4b2c512c --- /dev/null +++ b/src/sagemaker/hyperpod/cli/migrate/ansible/playbooks/snapshot.yml @@ -0,0 +1,27 @@ +--- +# snapshot.yml — SOURCE cluster, READ-ONLY. +# Captures identity + Slurm state and stores it in S3 for later converge. +# Runs on the controller/head/login node(s) only. +# +# Guarantee: nothing on the source is modified. Only getent / slurp / +# `sacctmgr dump` (read) are used. + +- name: Snapshot HyperPod source cluster (identity + slurm) + hosts: controller + gather_facts: true + become: true + vars: + hpm_mode: snapshot + roles: + - role: discover # storage_shape + identity_model + per-user homes + - role: identity_capture # getent -> users.json / groups.json / ssh_inventory.json + - role: slurm_capture # slurm.conf, gres.conf, sacctmgr dump + + post_tasks: + - name: Snapshot summary + ansible.builtin.debug: + msg: + - "storage_shape = {{ storage_shape }}" + - "identity_model = {{ identity_model }}" + - "users captured = {{ captured_users | length }}" + - "artifacts -> {{ hpm_prefix }}/" diff --git a/src/sagemaker/hyperpod/cli/migrate/ansible/playbooks/validate.yml b/src/sagemaker/hyperpod/cli/migrate/ansible/playbooks/validate.yml new file mode 100644 index 00000000..0750e891 --- /dev/null +++ b/src/sagemaker/hyperpod/cli/migrate/ansible/playbooks/validate.yml @@ -0,0 +1,18 @@ +--- +# validate.yml — TARGET cluster. The hard cutover GATE (design §11.4). +# Re-reads live getent + .ssh perms and diffs against the captured snapshot. +# FAILS the run on any UID/GID or .ssh owner/mode mismatch. + +- name: Validate target identity matches source snapshot + hosts: controller + gather_facts: true + become: true + vars: + hpm_mode: validate + pre_tasks: + - name: Load captured snapshot from S3 + ansible.builtin.include_role: + name: discover + tasks_from: load_snapshot.yml + roles: + - role: validate_identity diff --git a/src/sagemaker/hyperpod/cli/migrate/ansible/requirements.txt b/src/sagemaker/hyperpod/cli/migrate/ansible/requirements.txt new file mode 100644 index 00000000..fb9e5ca4 --- /dev/null +++ b/src/sagemaker/hyperpod/cli/migrate/ansible/requirements.txt @@ -0,0 +1,8 @@ +# Python deps for the runner host (CI or operator laptop). +# pip install -r requirements.txt +ansible-core>=2.16,<2.19 +boto3>=1.34.0 +botocore>=1.34.0 +# Dev / CI only: +ansible-lint>=24.0.0 +yamllint>=1.35.0 diff --git a/src/sagemaker/hyperpod/cli/migrate/ansible/requirements.yml b/src/sagemaker/hyperpod/cli/migrate/ansible/requirements.yml new file mode 100644 index 00000000..1397c271 --- /dev/null +++ b/src/sagemaker/hyperpod/cli/migrate/ansible/requirements.yml @@ -0,0 +1,10 @@ +--- +# Ansible collection dependencies for hpmigrate-ansible. +# Install with: ansible-galaxy collection install -r requirements.yml +collections: + - name: amazon.aws + version: ">=8.0.0" + - name: community.aws + version: ">=8.0.0" # provides the aws_ssm connection plugin + - name: ansible.posix + version: ">=1.5.0" # authorized_key, mount, etc. diff --git a/src/sagemaker/hyperpod/cli/migrate/ansible/roles/discover/tasks/load_snapshot.yml b/src/sagemaker/hyperpod/cli/migrate/ansible/roles/discover/tasks/load_snapshot.yml new file mode 100644 index 00000000..a35cad17 --- /dev/null +++ b/src/sagemaker/hyperpod/cli/migrate/ansible/roles/discover/tasks/load_snapshot.yml @@ -0,0 +1,65 @@ +--- +# discover/tasks/load_snapshot.yml — pull the captured snapshot from S3 into +# facts on the TARGET, for converge.yml / validate.yml. + +- name: Ensure workdir exists + ansible.builtin.file: + path: "{{ hpm_workdir }}" + state: directory + mode: "0755" + +- name: Download identity + storage snapshot from S3 + ansible.builtin.command: >- + aws s3 cp {{ item.src }} {{ hpm_workdir }}/{{ item.dest }} + loop: + - { src: "{{ hpm_identity_prefix }}/users.json", dest: users.json } + - { src: "{{ hpm_identity_prefix }}/groups.json", dest: groups.json } + - { src: "{{ hpm_identity_prefix }}/ssh_inventory.json", dest: ssh_inventory.json } + - { src: "{{ hpm_identity_prefix }}/identity_model.json", dest: identity_model.json } + - { src: "{{ hpm_storage_prefix }}/shape.json", dest: shape.json } + changed_when: false + +# The snapshot files live on the REMOTE node (aws s3 cp ran there), so read them +# with slurp (remote) rather than lookup('file') (which reads the control host). +- name: Slurp snapshot files from the remote node + ansible.builtin.slurp: + src: "{{ hpm_workdir }}/{{ item }}" + loop: + - users.json + - groups.json + - ssh_inventory.json + - identity_model.json + - shape.json + register: _snap + +- name: Index slurped snapshot content by filename + ansible.builtin.set_fact: + _snap_map: "{{ dict(_snap.results | map(attribute='item') | list + | zip(_snap.results | map(attribute='content') | list)) }}" + +- name: Load captured users + ansible.builtin.set_fact: + captured_users: "{{ (_snap_map['users.json'] | b64decode) | from_json }}" + +- name: Load captured groups + ansible.builtin.set_fact: + captured_groups: "{{ (_snap_map['groups.json'] | b64decode) | from_json }}" + +- name: Load captured ssh inventory + ansible.builtin.set_fact: + captured_ssh: "{{ (_snap_map['ssh_inventory.json'] | b64decode) | from_json }}" + +- name: Load identity model + storage shape + ansible.builtin.set_fact: + identity_model: "{{ ((_snap_map['identity_model.json'] | b64decode) | from_json).identity_model }}" + storage_shape: "{{ ((_snap_map['shape.json'] | b64decode) | from_json).storage_shape }}" + +- name: Fail if numeric IDs missing and required + ansible.builtin.assert: + that: + - item.uid is defined and item.uid | int > 0 + - item.gid is defined and item.gid | int > 0 + fail_msg: "captured user {{ item.name }} lacks a numeric UID/GID; refusing to auto-assign" + loop: "{{ captured_users }}" + loop_control: { label: "{{ item.name }}" } + when: require_numeric_ids | bool and identity_model == 'local' diff --git a/src/sagemaker/hyperpod/cli/migrate/ansible/roles/discover/tasks/main.yml b/src/sagemaker/hyperpod/cli/migrate/ansible/roles/discover/tasks/main.yml new file mode 100644 index 00000000..8a0a9fdd --- /dev/null +++ b/src/sagemaker/hyperpod/cli/migrate/ansible/roles/discover/tasks/main.yml @@ -0,0 +1,61 @@ +--- +# discover/tasks/main.yml — detect storage shape, identity model, and per-user +# home locations on the SOURCE cluster. Read-only. + +- name: Ensure workdir exists + ansible.builtin.file: + path: "{{ hpm_workdir }}" + state: directory + mode: "0755" + +# ---- Storage shape: fsx_only vs fsx_openzfs -------------------------------- +- name: Detect mounts + ansible.builtin.command: findmnt -rno TARGET,SOURCE,FSTYPE + register: _mounts + changed_when: false + +- name: Determine storage shape + ansible.builtin.set_fact: + storage_shape: >- + {{ 'fsx_openzfs' if (_mounts.stdout is search('/home') + and (_mounts.stdout is search('nfs') or _mounts.stdout is search('zfs'))) + else 'fsx_only' }} + +- name: Detect FSx Lustre mount (/fsx) + ansible.builtin.set_fact: + has_fsx_lustre: "{{ _mounts.stdout is search('/fsx') }}" + +# ---- Identity model: local (shared_users.txt) vs SSSD/AD ------------------- +- name: Check for SSSD + ansible.builtin.stat: + path: /etc/sssd/sssd.conf + register: _sssd + +- name: Check sssd service + ansible.builtin.command: systemctl is-active sssd + register: _sssd_svc + changed_when: false + failed_when: false + +- name: Determine identity model + ansible.builtin.set_fact: + identity_model: >- + {{ 'sssd' if (_sssd.stat.exists or _sssd_svc.stdout == 'active') + else 'local' }} + +- name: Announce detection + ansible.builtin.debug: + msg: "shape={{ storage_shape }} identity_model={{ identity_model }} fsx={{ has_fsx_lustre }}" + +- name: Persist storage shape artifact + ansible.builtin.copy: + dest: "{{ hpm_workdir }}/shape.json" + content: "{{ {'storage_shape': storage_shape, 'has_fsx_lustre': has_fsx_lustre} | to_nice_json }}" + mode: "0644" + when: hpm_mode == 'snapshot' + +- name: Upload shape + identity_model to S3 + ansible.builtin.command: >- + aws s3 cp {{ hpm_workdir }}/shape.json {{ hpm_storage_prefix }}/shape.json + changed_when: true + when: hpm_mode == 'snapshot' diff --git a/src/sagemaker/hyperpod/cli/migrate/ansible/roles/identity_apply/tasks/main.yml b/src/sagemaker/hyperpod/cli/migrate/ansible/roles/identity_apply/tasks/main.yml new file mode 100644 index 00000000..c5993762 --- /dev/null +++ b/src/sagemaker/hyperpod/cli/migrate/ansible/roles/identity_apply/tasks/main.yml @@ -0,0 +1,50 @@ +--- +# identity_apply/tasks/main.yml — TARGET. +# Recreate groups + users with the EXACT numeric UID/GID captured from source. +# This is the correctness core: never auto-assign. +# +# For identity_model == 'sssd' the whole local-user creation is skipped; users +# and IDs come from the directory. We only ensure the SSSD join is present +# (handled by the parent hpmigrate CLI re-running setup_sssd.py) and manage +# nothing locally here. + +- name: Skip local identity for directory-managed (SSSD/AD) clusters + ansible.builtin.debug: + msg: "identity_model=sssd -> users/IDs come from directory; skipping local user creation" + when: identity_model == 'sssd' + +- name: Recreate groups with pinned GIDs + ansible.builtin.group: + name: "{{ item.name }}" + gid: "{{ item.gid }}" + state: present + loop: "{{ captured_groups }}" + loop_control: { label: "{{ item.name }} (gid={{ item.gid }})" } + when: identity_model == 'local' + +- name: Recreate users with pinned UID/GID and same home path + ansible.builtin.user: + name: "{{ item.name }}" + uid: "{{ item.uid }}" + group: "{{ item.gid }}" # primary group by GID + home: "{{ item.home }}" + shell: "{{ item.shell | default('/bin/bash') }}" + comment: "{{ item.gecos | default('') }}" + create_home: false # home comes from DRA-import / OpenZFS-restore + state: present + loop: "{{ captured_users }}" + loop_control: { label: "{{ item.name }} (uid={{ item.uid }} gid={{ item.gid }})" } + when: identity_model == 'local' + +- name: Restore supplementary group memberships (per user) + ansible.builtin.user: + name: "{{ item.name }}" + # all groups (other than primary) whose member list includes this user + groups: >- + {{ captured_groups + | selectattr('members', 'contains', item.name) + | map(attribute='name') | list }} + append: true + loop: "{{ captured_users }}" + loop_control: { label: "{{ item.name }}" } + when: identity_model == 'local' diff --git a/src/sagemaker/hyperpod/cli/migrate/ansible/roles/identity_capture/tasks/main.yml b/src/sagemaker/hyperpod/cli/migrate/ansible/roles/identity_capture/tasks/main.yml new file mode 100644 index 00000000..adac1517 --- /dev/null +++ b/src/sagemaker/hyperpod/cli/migrate/ansible/roles/identity_capture/tasks/main.yml @@ -0,0 +1,94 @@ +--- +# identity_capture/tasks/main.yml — SOURCE, read-only. +# Turns live getent state into structured JSON artifacts. getent is the source +# of record: it captures ALL users regardless of how they were created +# (shared_users.txt, SSSD, or ad-hoc useradd). + +- name: Capture passwd entries + ansible.builtin.getent: + database: passwd + # populates ansible_facts.getent_passwd + +- name: Capture group entries + ansible.builtin.getent: + database: group + +- name: Build human user list (uid >= uid_min), with real home + which fs it is on + ansible.builtin.set_fact: + captured_users: >- + {{ captured_users | default([]) + [ { + 'name': item.key, + 'uid': item.value[1] | int, + 'gid': item.value[2] | int, + 'gecos': item.value[3], + 'home': item.value[4], + 'shell': item.value[5], + 'home_fs': ('lustre' if item.value[4] is match('^/fsx/') + else ('openzfs' if item.value[4] is match('^/home/') + else 'other')) + } ] }} + loop: "{{ ansible_facts.getent_passwd | dict2items }}" + loop_control: { label: "{{ item.key }}" } + when: + - (item.value[1] | int) >= (uid_min | int) + - (item.value[1] | int) < 65000 # exclude nobody/system ceiling accounts + +- name: Build group list (keep human groups + explicitly-kept system groups) + ansible.builtin.set_fact: + captured_groups: >- + {{ captured_groups | default([]) + [ { + 'name': item.key, + 'gid': item.value[1] | int, + 'members': (item.value[2] | default('')) | split(',') | reject('equalto','') | list + } ] }} + loop: "{{ ansible_facts.getent_group | dict2items }}" + loop_control: { label: "{{ item.key }}" } + when: >- + (item.value[1] | int) >= (gid_min | int) + or item.key in system_groups_keep + +# ---- Per-user .ssh inventory (truth before migration) ---------------------- +- name: Stat each user's .ssh dir + ansible.builtin.stat: + path: "{{ item.home }}/.ssh" + loop: "{{ captured_users }}" + loop_control: { label: "{{ item.name }}" } + register: _ssh_stats + +- name: Build ssh inventory + ansible.builtin.set_fact: + captured_ssh: >- + {{ captured_ssh | default([]) + [ { + 'name': item.item.name, + 'home': item.item.home, + 'home_fs': item.item.home_fs, + 'ssh_dir': item.item.home ~ '/.ssh', + 'exists': item.stat.exists, + 'owner_uid': (item.stat.uid | default(omit)), + 'mode': (item.stat.mode | default(omit)) + } ] }} + loop: "{{ _ssh_stats.results }}" + loop_control: { label: "{{ item.item.name }}" } + +# ---- Write + upload artifacts --------------------------------------------- +- name: Write identity artifacts locally + ansible.builtin.copy: + dest: "{{ hpm_workdir }}/{{ item.file }}" + content: "{{ item.data | to_nice_json }}" + mode: "0600" + loop: + - { file: users.json, data: "{{ captured_users }}" } + - { file: groups.json, data: "{{ captured_groups }}" } + - { file: ssh_inventory.json, data: "{{ captured_ssh }}" } + - { file: identity_model.json, data: "{{ {'identity_model': identity_model} }}" } + loop_control: { label: "{{ item.file }}" } + +- name: Upload identity artifacts to S3 + ansible.builtin.command: >- + aws s3 cp {{ hpm_workdir }}/{{ item }} {{ hpm_identity_prefix }}/{{ item }} + changed_when: true + loop: + - users.json + - groups.json + - ssh_inventory.json + - identity_model.json diff --git a/src/sagemaker/hyperpod/cli/migrate/ansible/roles/slurm_apply/handlers/main.yml b/src/sagemaker/hyperpod/cli/migrate/ansible/roles/slurm_apply/handlers/main.yml new file mode 100644 index 00000000..ca321440 --- /dev/null +++ b/src/sagemaker/hyperpod/cli/migrate/ansible/roles/slurm_apply/handlers/main.yml @@ -0,0 +1,8 @@ +--- +- name: reload slurm + ansible.builtin.systemd: + name: "{{ item }}" + state: restarted + loop: + - slurmctld + failed_when: false diff --git a/src/sagemaker/hyperpod/cli/migrate/ansible/roles/slurm_apply/tasks/main.yml b/src/sagemaker/hyperpod/cli/migrate/ansible/roles/slurm_apply/tasks/main.yml new file mode 100644 index 00000000..edc1551a --- /dev/null +++ b/src/sagemaker/hyperpod/cli/migrate/ansible/roles/slurm_apply/tasks/main.yml @@ -0,0 +1,41 @@ +--- +# slurm_apply/tasks/main.yml — TARGET. +# Place captured slurm config + reload accounting. Config files are captured +# verbatim; hpmigrate (parent CLI) re-templates node/host-specific values before +# this playbook runs, so here we place-and-reload only. + +- name: Ensure slurm conf dir exists + ansible.builtin.file: + path: "{{ slurm_conf_dir }}" + state: directory + mode: "0755" + +- name: Download slurm artifacts from S3 + ansible.builtin.command: >- + aws s3 cp {{ hpm_slurm_prefix }}/ {{ hpm_workdir }}/slurm/ --recursive + changed_when: false + +- name: Place slurm config files + ansible.builtin.copy: + src: "{{ hpm_workdir }}/slurm/{{ item }}" + dest: "{{ slurm_conf_dir }}/{{ item }}" + remote_src: true + owner: slurm + group: slurm + mode: "0644" + loop: + - slurm.conf + - gres.conf + - cgroup.conf + - topology.conf + register: _placed + failed_when: false + notify: reload slurm + +- name: Load accounting dump (associations / users / QOS) + ansible.builtin.command: >- + sacctmgr --immediate load file={{ hpm_workdir }}/slurm/sacctmgr_dump.cfg + register: _load + changed_when: "'Nothing new added' not in _load.stdout | default('')" + failed_when: false + when: capture_accounting_dump | bool diff --git a/src/sagemaker/hyperpod/cli/migrate/ansible/roles/slurm_capture/tasks/main.yml b/src/sagemaker/hyperpod/cli/migrate/ansible/roles/slurm_capture/tasks/main.yml new file mode 100644 index 00000000..a5762370 --- /dev/null +++ b/src/sagemaker/hyperpod/cli/migrate/ansible/roles/slurm_capture/tasks/main.yml @@ -0,0 +1,56 @@ +--- +# slurm_capture/tasks/main.yml — SOURCE, read-only. +# Stage Slurm config files + accounting dump into the remote workdir, then push +# to S3. All work happens ON the node (remote_src copies), consistent with +# identity_capture; nothing is pulled to the control host. + +- name: Ensure remote slurm workdir exists + ansible.builtin.file: + path: "{{ hpm_workdir }}/slurm" + state: directory + mode: "0755" + +- name: Locate slurm config files + ansible.builtin.stat: + path: "{{ slurm_conf_dir }}/{{ item }}" + loop: + - slurm.conf + - gres.conf + - cgroup.conf + - topology.conf + register: _slurm_files + +- name: Stage existing slurm config files into workdir (remote copy) + ansible.builtin.copy: + src: "{{ item.stat.path }}" + dest: "{{ hpm_workdir }}/slurm/{{ item.item }}" + remote_src: true + mode: "0644" + loop: "{{ _slurm_files.results }}" + loop_control: { label: "{{ item.item }}" } + when: item.stat.exists + +- name: Dump slurm accounting (associations / users / QOS) + ansible.builtin.shell: >- + set -o pipefail; + ( command -v sacctmgr >/dev/null 2>&1 && sacctmgr --noheader dump ) + || ( [ -x /opt/slurm/bin/sacctmgr ] && /opt/slurm/bin/sacctmgr --noheader dump ) + || echo "SACCTMGR_UNAVAILABLE" + args: + executable: /bin/bash + register: _sacctmgr + changed_when: false + failed_when: false + when: capture_accounting_dump | bool + +- name: Save sacctmgr dump + ansible.builtin.copy: + dest: "{{ hpm_workdir }}/slurm/sacctmgr_dump.cfg" + content: "{{ _sacctmgr.stdout | default('') }}" + mode: "0644" + when: capture_accounting_dump | bool and _sacctmgr.stdout is defined + +- name: Upload slurm artifacts to S3 + ansible.builtin.command: >- + aws s3 cp {{ hpm_workdir }}/slurm/ {{ hpm_slurm_prefix }}/ --recursive + changed_when: true diff --git a/src/sagemaker/hyperpod/cli/migrate/ansible/roles/ssh_perms/tasks/main.yml b/src/sagemaker/hyperpod/cli/migrate/ansible/roles/ssh_perms/tasks/main.yml new file mode 100644 index 00000000..8b2a1e4d --- /dev/null +++ b/src/sagemaker/hyperpod/cli/migrate/ansible/roles/ssh_perms/tasks/main.yml @@ -0,0 +1,48 @@ +--- +# ssh_perms/tasks/main.yml — TARGET. +# DRA-import / OpenZFS-restore bring the .ssh dirs and keys with their numeric +# ownership. This role ENFORCES correct owner + mode so sshd accepts the keys. +# It operates on whatever filesystem holds each user's $HOME (Shape A: /fsx, +# Shape B: /fsx and/or /home) — driven entirely by the captured home path. + +- name: Build set of users whose source .ssh existed + ansible.builtin.set_fact: + ssh_users: >- + {{ captured_users | selectattr('name', 'in', + captured_ssh | selectattr('exists') | map(attribute='name') | list) + | list }} + +- name: Ensure ~/.ssh has correct ownership and mode + ansible.builtin.file: + path: "{{ item.home }}/.ssh" + state: directory + owner: "{{ item.uid }}" + group: "{{ item.gid }}" + mode: "{{ ssh_dir_mode }}" + loop: "{{ ssh_users }}" + loop_control: { label: "{{ item.name }} -> {{ item.home }}/.ssh" } + when: identity_model == 'local' + failed_when: false + +- name: Fix authorized_keys mode/owner if present + ansible.builtin.file: + path: "{{ item.home }}/.ssh/authorized_keys" + owner: "{{ item.uid }}" + group: "{{ item.gid }}" + mode: "{{ ssh_authkeys_mode }}" + loop: "{{ ssh_users }}" + loop_control: { label: "{{ item.name }}" } + when: identity_model == 'local' + failed_when: false + +- name: Normalize private key perms (id_* files, not .pub) + ansible.builtin.shell: >- + if [ -d "{{ item.home }}/.ssh" ]; then find "{{ item.home }}/.ssh" -maxdepth 1 -type f -name 'id_*' ! -name '*.pub' -exec chown {{ item.uid }}:{{ item.gid }} {} + -exec chmod {{ ssh_privkey_mode }} {} +; fi + args: + executable: /bin/bash + loop: "{{ ssh_users }}" + loop_control: { label: "{{ item.name }}" } + when: + - identity_model == 'local' + - item.home is match('^/(fsx|home)/') + changed_when: false diff --git a/src/sagemaker/hyperpod/cli/migrate/ansible/roles/validate_identity/tasks/main.yml b/src/sagemaker/hyperpod/cli/migrate/ansible/roles/validate_identity/tasks/main.yml new file mode 100644 index 00000000..17af5a32 --- /dev/null +++ b/src/sagemaker/hyperpod/cli/migrate/ansible/roles/validate_identity/tasks/main.yml @@ -0,0 +1,90 @@ +--- +# validate_identity/tasks/main.yml — TARGET. The hard cutover gate. +# Re-read live getent + .ssh perms and diff against the captured snapshot. +# FAIL on any UID/GID or .ssh owner/mode mismatch (design §11.4). + +- name: Skip user-ID diff for SSSD (directory guarantees stable IDs) + ansible.builtin.debug: + msg: "identity_model=sssd -> IDs are directory-managed; validating .ssh perms only" + when: identity_model == 'sssd' + +- name: Read live passwd on target + ansible.builtin.getent: + database: passwd + +- name: Assert each captured user exists with identical UID/GID + ansible.builtin.assert: + that: + - item.name in ansible_facts.getent_passwd + - (ansible_facts.getent_passwd[item.name][1] | int) == (item.uid | int) + - (ansible_facts.getent_passwd[item.name][2] | int) == (item.gid | int) + fail_msg: >- + IDENTITY MISMATCH for {{ item.name }}: + expected uid={{ item.uid }} gid={{ item.gid }}, + got {{ ansible_facts.getent_passwd[item.name] | default('MISSING') }}. + Cutover blocked — files (incl. ~/.ssh) would be owned by the wrong account. + success_msg: "{{ item.name }} uid/gid OK" + loop: "{{ captured_users }}" + loop_control: { label: "{{ item.name }}" } + when: identity_model == 'local' + +# ---- .ssh ownership + perms (both models) ---------------------------------- +# Validate against the CAPTURED SOURCE state (owner + mode), not a hardcoded +# policy — the goal is "target matches source". Only enforce for users whose +# home is on a filesystem the tool migrates (/home OpenZFS). Homes on /fsx ride +# the customer's DRA and are validated after the customer hydrates /fsx. +- name: Stat target .ssh dirs + ansible.builtin.stat: + path: "{{ item.home }}/.ssh" + loop: "{{ captured_users }}" + loop_control: { label: "{{ item.name }}" } + register: _tgt_ssh + +- name: Assert .ssh dir matches captured source (owner + mode) + ansible.builtin.assert: + that: + - item.0.stat.exists + - item.0.stat.uid | int == item.1.uid | int + - item.0.stat.mode == (_src_ssh.mode | regex_replace('^0', '')) + or item.0.stat.mode == _src_ssh.mode + fail_msg: >- + SSH PERMS MISMATCH for {{ item.1.name }} at {{ item.1.home }}/.ssh: + target owner={{ item.0.stat.uid | default('?') }} mode={{ item.0.stat.mode | default('?') }}, + source owner={{ item.1.uid }} mode={{ _src_ssh.mode | default('?') }}. Login would fail. + success_msg: "{{ item.1.name }} .ssh matches source (mode {{ _src_ssh.mode }})" + vars: + _src_ssh: "{{ captured_ssh | selectattr('name','equalto', item.1.name) | first | default({}) }}" + loop: "{{ _tgt_ssh.results | zip(captured_users) | list }}" + loop_control: { label: "{{ item.1.name }}" } + # Only enforce where the source .ssh existed AND home is on a tool-migrated + # filesystem (/home). /fsx homes are customer-DRA — skipped here. + when: >- + (captured_ssh | selectattr('name','equalto', item.1.name) + | map(attribute='exists') | first | default(false)) + and (item.1.home is match('^/home/')) + +# ---- Storage: /home must actually be mounted for Shape B (fsx+OpenZFS) ------ +# OpenZFS /home mounts only when BOTH enable_fsx_openzfs and fsx_openzfs_dns_name +# are set; otherwise the cluster comes up InService with no /home. Assert it. +- name: Check /home mount when source was fsx+OpenZFS + ansible.builtin.command: findmnt --noheadings /home + register: _home_mnt + changed_when: false + failed_when: false + when: storage_shape == 'fsx_openzfs' + +- name: Assert /home is mounted (OpenZFS) + ansible.builtin.assert: + that: + - _home_mnt.rc == 0 + - _home_mnt.stdout is search('nfs') + fail_msg: >- + Source was fsx+OpenZFS but /home is not mounted on the target. Check that + enable_fsx_openzfs=True (config.py) AND fsx_openzfs_dns_name + (provisioning_parameters.json) are both set, then re-provision. + success_msg: "/home (OpenZFS) is mounted" + when: storage_shape == 'fsx_openzfs' + +- name: Validation passed + ansible.builtin.debug: + msg: "IDENTITY VALIDATION PASSED — safe to proceed to cutover." diff --git a/test/unit_tests/cli/test_migrate.py b/test/unit_tests/cli/test_migrate.py new file mode 100644 index 00000000..29836cbf --- /dev/null +++ b/test/unit_tests/cli/test_migrate.py @@ -0,0 +1,146 @@ +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"). You +# may not use this file except in compliance with the License. A copy of +# the License is located at +# +# http://aws.amazon.com/apache2.0/ +# +# or in the "license" file accompanying this file. This file is +# distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF +# ANY KIND, either express or implied. See the License for the specific +# language governing permissions and limitations under the License. +"""Unit tests for `hyp migrate` (snapshot / converge / validate).""" +from unittest import mock + +import pytest +from click.testing import CliRunner + +from sagemaker.hyperpod.cli.commands.migrate import ( + migrate, + _ansible_root, + _build_inventory, + _build_cluster_config, +) + + +@pytest.fixture +def runner(): + return CliRunner() + + +def test_migrate_group_has_expected_subcommands(runner): + result = runner.invoke(migrate, ["--help"]) + assert result.exit_code == 0 + for sub in ("snapshot", "converge", "validate", "plan", "provision"): + assert sub in result.output + + +def test_cluster_config_binds_plan_to_worker_only(): + cfg = _build_cluster_config( + bucket="s3://sagemaker-bkt", + exec_role_arn="arn:aws:iam::1:role/r", + plan_arn="arn:aws:sagemaker:us-west-1:1:training-plan/p", + worker_instance="ml.p5.48xlarge", + worker_count=1, + ondemand_instance="ml.m5.xlarge", + ) + by_name = {g["InstanceGroupName"]: g for g in cfg} + # worker bound to the plan + assert by_name["worker-group"]["TrainingPlanArn"].endswith("training-plan/p") + assert by_name["worker-group"]["InstanceType"] == "ml.p5.48xlarge" + # controller/login on-demand, NOT on the plan + assert "TrainingPlanArn" not in by_name["controller-machine"] + assert "TrainingPlanArn" not in by_name["login-group"] + assert by_name["controller-machine"]["InstanceType"] == "ml.m5.xlarge" + + +def test_cluster_config_no_plan_is_ondemand_only(): + cfg = _build_cluster_config( + bucket="s3://sagemaker-bkt", exec_role_arn="arn:aws:iam::1:role/r", + plan_arn=None, worker_instance="ml.m5.xlarge", worker_count=1, + ondemand_instance="ml.m5.xlarge", + ) + assert all("TrainingPlanArn" not in g for g in cfg) + + +def test_provision_writes_config(runner, tmp_path): + out = tmp_path / "cfg.json" + result = runner.invoke( + migrate, + ["provision", "--bucket", "s3://sagemaker-bkt", + "--execution-role-arn", "arn:aws:iam::1:role/r", + "--region", "us-west-1", "--out", str(out)], + ) + assert result.exit_code == 0, result.output + import json as _json + cfg = _json.loads(out.read_text()) + assert {g["InstanceGroupName"] for g in cfg} == { + "controller-machine", "login-group", "worker-group"} + + +@pytest.mark.parametrize("sub", ["snapshot", "converge", "validate"]) +def test_subcommand_help_renders(runner, sub): + result = runner.invoke(migrate, [sub, "--help"]) + assert result.exit_code == 0 + assert "--cluster" in result.output + assert "--bucket" in result.output + + +def test_vendored_ansible_assets_present(): + root = _ansible_root() + assert (root / "playbooks" / "snapshot.yml").is_file() + assert (root / "playbooks" / "converge.yml").is_file() + assert (root / "playbooks" / "validate.yml").is_file() + assert (root / "roles" / "identity_apply").is_dir() + assert (root / "ansible.cfg").is_file() + + +def test_converge_requires_run_id(runner): + with mock.patch( + "sagemaker.hyperpod.cli.commands.migrate._require_tooling" + ): + result = runner.invoke( + migrate, + ["converge", "--cluster", "c", "--region", "us-east-1", + "--bucket", "s3://sagemaker-x"], + ) + assert result.exit_code != 0 + assert "run-id" in result.output.lower() + + +def test_build_inventory_targets_controller(tmp_path): + nodes = [ + {"InstanceGroupName": "controller-machine", "InstanceId": "i-1", + "InstanceStatus": {"Status": "Running"}}, + {"InstanceGroupName": "worker-group", "InstanceId": "i-2", + "InstanceStatus": {"Status": "Running"}}, + ] + with mock.patch( + "sagemaker.hyperpod.cli.commands.migrate._resolve_cluster", + return_value=("clabc123", nodes), + ): + out = tmp_path / "inv.ini" + count = _build_inventory("src", "us-east-1", None, + "s3://sagemaker-bkt", out) + text = out.read_text() + assert count == 1 + assert "sagemaker-cluster:clabc123_controller-machine-i-1" in text + assert "worker-group-i-2" not in text + assert "community.aws.aws_ssm" in text + assert "ansible_aws_ssm_bucket_name=sagemaker-bkt" in text + + +def test_build_inventory_errors_when_no_controller(tmp_path): + nodes = [ + {"InstanceGroupName": "worker-group", "InstanceId": "i-2", + "InstanceStatus": {"Status": "Running"}}, + ] + import click + with mock.patch( + "sagemaker.hyperpod.cli.commands.migrate._resolve_cluster", + return_value=("clabc123", nodes), + ): + with pytest.raises(click.ClickException): + _build_inventory("src", "us-east-1", None, + "s3://sagemaker-bkt", tmp_path / "inv.ini")