feat(hyperpod-slurm): add prepare_extensions.sh helper for OnInitComplete extensions

[aravneelaws]

Purpose

Adds a helper to prepare and upload the Extensions/ directory to S3 for use with the new HyperPod Slurm cluster creation flow (no-LCS). When creating a cluster and choosing Custom Setup → Lifecycle configuration → None, the console accepts an S3 location containing extension scripts. Assembling that bucket by hand is error-prone (per-extension config files, run_extensions.sh toggles, bucket naming/region rules, IAM), so this PR adds prepare_extensions.sh to automate it end-to-end.

Changes

• Add 1.architectures/5.sagemaker-hyperpod/Extensions/prepare_extensions.sh that:
  • Picks which extensions to include via --add-users / --observability flags.
  • Generates shared_users.txt from --users alice,bob,carol (auto-assigns UIDs from 2001) or from --users + --uids, or from an interactive prompt (matches the UX of LifecycleScripts/base-config/utils/create_users.sh). Also accepts a pre-made file via --users-file.
  • Prompts for (or accepts via --amp-url) the Amazon Managed Prometheus remote_write URL and patches observability/config.json.
  • Selects the console entrypoint based on selection:
    • --observability alone → uploads only observability/; entrypoint is observability/setup_observability.sh.
    • --add-users alone or with --observability → uploads detect-node/ + the selected extension dir(s) + a patched run_extensions.sh; entrypoint is run_extensions.sh.
  • Resolves an S3 bucket via --bucket (existing) or --create-bucket (new; sets versioning + public-access-block).
  • Validates bucket names client-side (S3 naming rules) and then performs a live status check that distinguishes owned / other-account / absent / wrong-region, using --expected-bucket-owner (needed because the 2024 HeadBucket API change now returns metadata for any locatable bucket regardless of ownership).
  • Supports --aws-profile threaded through every AWS call, plus --dry-run and --yes for non-interactive use.
  • Stages everything in a temp dir and uploads from there — no files in the repo are mutated.
  • Prints the exact s3://… entrypoint path to paste into the HyperPod console.
• Add 1.architectures/5.sagemaker-hyperpod/Extensions/README.md briefly describing each of the three extensions and documenting the helper script's flags and usage.

Test Plan

Environment:

• AWS Service: SageMaker HyperPod (Slurm, AMI-based, OnInitComplete)
• Instance type: ml.m5.2xlarge controller + ml.g5.8xlarge compute
• Number of nodes: 2 (1 controller, 1 compute)

Test commands:

# 1. Client-side bucket name validation (no AWS calls needed)
./Extensions/prepare_extensions.sh --add-users --users alice \
  --create-bucket "MyBucket" --region us-west-2 --dry-run --yes
# expected: rejects uppercase before any AWS call

# 2. Bucket edge cases against live S3
./Extensions/prepare_extensions.sh --add-users --users alice \
  --create-bucket elasticbeanstalk --region us-west-2 --yes
# expected: "already taken by another AWS account"

./Extensions/prepare_extensions.sh --add-users --users alice \
  --bucket "nonexistent-hp-$(date +%s)" --region us-west-2 --yes
# expected: "does not exist"

./Extensions/prepare_extensions.sh --add-users --users alice \
  --bucket sagemaker-sample-files --region us-west-2 --yes
# expected: "exists but is not accessible with this profile/account"

# 3. Wrong-region detection (owned by us, in us-east-1, asked for us-west-2)
aws s3api create-bucket --bucket hp-ext-regiontest-$$ --region us-east-1
./Extensions/prepare_extensions.sh --add-users --users alice \
  --bucket hp-ext-regiontest-$$ --region us-west-2 --yes
# expected: "in region 'us-east-1' but --region was 'us-west-2'"

# 4. End-to-end happy path -- add-users only
./Extensions/prepare_extensions.sh \
  --add-users \
  --users alice,bob \
  --create-bucket my-hyperpod-extensions-$(date +%s) \
  --region us-west-2 \
  --aws-profile my-profile

# 5. Use printed entrypoint in HyperPod console
# Custom setup -> Lifecycle configuration -> None
# Paste: s3://<bucket>/hyperpod-extensions/run_extensions.sh

Test Results

• Client-side name validation rejects invalid bucket names (uppercase, underscores, <3 or >63 chars, adjacent periods, IP-format, reserved prefixes/suffixes) before any AWS call.
• Bucket status checks produce actionable error messages for all four edge cases (globally taken, absent, exists-but-not-owned, wrong-region).
• End-to-end run with --add-users --users alice,bob uploaded the following to the target bucket:

  add-users/            (all scripts + generated shared_users.txt, sample files stripped)
  detect-node/
  run_extensions.sh     (ENABLE_ADD_USERS="true", ENABLE_OBSERVABILITY="false")
  

• HyperPod cluster provisioned with the printed entrypoint; users were created with the expected UIDs (2001, 2002), home directories set up on the shared filesystem, SSH keypairs generated, and Slurm accounting entries added on the controller.
• --aws-profile correctly routes every AWS call to the named profile; preflight prints the caller ARN and account ID so the user can verify before upload.
• --dry-run prints all aws invocations without executing them.

Checklist

• I have read the contributing guidelines.
• I am working against the latest main branch.
• I have searched existing open and recently merged PRs to confirm this is not a duplicate.
• The contribution is self-contained with documentation and scripts.
• External dependencies are pinned to a specific version or tag (no latest).
• A README is included or updated with prerequisites, instructions, and known issues.
• New test cases follow the expected directory structure. (N/A — this PR adds tooling, not a test case)