Wave 20: SageMaker GRPO smoke artifacts (F3) — runnable g5.2xlarge GSM8K

The F3 runnable-now RL smoke, with two live-verified corrections to the
design doc's recipe:
- DLC tag is cu126 + requires the -v1.25 build suffix (NOT the cu124 bare
tag in the AWS docs page; resolved against the live ECR registry).
- sagemaker SDK pinned <3: v3 is an API rewrite that dropped
sagemaker.estimator.Estimator / sagemaker.pytorch.

Artifacts:
- examples/gsm8k_grpo/run_sagemaker.py: GPU+vLLM-colocate variant of run.py
(same GSM8K RLVR reward, plain GRPO alpha=beta=0), CLI/JSON hyperparams,
writes to SM_MODEL_DIR.
- examples/gsm8k_grpo/run_sagemaker_launch.py: sagemaker-v2 Estimator driver.
Default --image dlc uses the stock PyTorch DLC + source_dir (no local
15GB cross-arch build); --image <ecr> uses a prebuilt baked image.
Stages a minimal source_dir (framework + script + requirements).
- examples/gsm8k_grpo/requirements.txt: trl + vllm==0.8.5 (torch-2.6 line) +
RL stack, installed at job start atop the DLC's torch 2.6/cu126.
- docker/Dockerfile.sagemaker + scripts/build_and_push_ecr.sh: the
repeatable baked-image path (corrected DLC tag, cross-arch note).
- docs/AWS_SAGEMAKER_QUICKSTART.md: live facts + run instructions + gotchas.
- pyproject [aws] extra: add sagemaker>=2.200,<3.

Live preflight verified: g5.2xlarge training quota=1, exec role trusts
sagemaker, output bucket writable, 0 jobs running. Smoke job submitted
(composer-grpo-smoke-2026-06-09-18-35-31-353); result recorded separately.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

docker/Dockerfile.sagemaker

@@ -0,0 +1,25 @@
+# Baked SageMaker training image for the Composer-replication RL stack (F3 §3.2).
+#
+# The repeatable path: bake trl+vllm+the framework into an image so jobs don't
+# pip-install at startup (saves ~5-10 min/job and removes a flaky failure
+# surface). The one-shot smoke can instead use the stock DLC + source_dir
+# (run_sagemaker_launch.py --image dlc), which needs no local build.
+#
+# Base: AWS PyTorch DLC, tag RESOLVED LIVE against the us-west-2 registry
+# 2026-06-09 — it is cu126 (NOT the cu124 some docs list) and the -v1.25 build
+# suffix is required (no bare floating 2.6.0-gpu tag exists).
+FROM 763104351884.dkr.ecr.us-west-2.amazonaws.com/pytorch-training:2.6.0-gpu-py312-cu126-ubuntu22.04-sagemaker-v1.25
+
+# RL stack baked in. torch 2.6 + CUDA 12.6 already in the DLC — do NOT reinstall
+# torch. vllm 0.8.5 is the torch-2.6 line (pin to avoid a wheel/CUDA mismatch).
+RUN pip install --no-cache-dir \
+      "trl>=1.5,<2" "peft>=0.13" "accelerate>=1.0" "datasets>=3.0" \
+      "vllm==0.8.5" "fsspec>=2024.6" "s3fs>=2024.6" "hf_transfer>=0.1.6"
+
+# The framework itself (train + serverless extras → trainer, loss, executors,
+# replica_entrypoint, s3fs all present).
+COPY . /opt/composer_replication
+RUN pip install --no-cache-dir -e "/opt/composer_replication[train,serverless]"
+
+ENV HF_HOME=/opt/ml/input/hf_cache \
+    HF_HUB_ENABLE_HF_TRANSFER=1

docs/AWS_SAGEMAKER_QUICKSTART.md

@@ -0,0 +1,89 @@
+# AWS SageMaker Quickstart — the runnable-now GRPO smoke
+
+The minimum path to running the Composer-replication RL inner loop on a real
+GPU, end-to-end, for **under $1**. Implements F3 (`research/design-F3-rl-sagemaker.md`).
+
+## Live account facts (verified 2026-06-09, acct 386931836011, us-west-2)
+
+| Fact | Value |
+|---|---|
+| `ml.g5.2xlarge` training-job quota | **1** (live, code `L-2D6DEB3C`) → no quota ticket needed |
+| Execution role | `arn:aws:iam::386931836011:role/service-role/AmazonSageMaker-ExecutionRole-20250725T133247` |
+| Bucket (rendezvous + output) | `amazon-sagemaker-386931836011-us-west-2-7597bf4d9a3d` |
+| PyTorch DLC base image | `763104351884.dkr.ecr.us-west-2.amazonaws.com/pytorch-training:2.6.0-gpu-py312-cu126-ubuntu22.04-sagemaker-v1.25` |
+
+> **Two corrections vs the AWS DLC docs page** (found by querying the live ECR
+> registry, not the docs): the tag is **cu126**, not cu124, and the **`-v1.25`
+> build suffix is required** — there is no bare floating `2.6.0-gpu...` tag.
+> Always resolve the tag against the registry:
+> ```bash
+> aws ecr describe-images --registry-id 763104351884 \
+>   --repository-name pytorch-training --region us-west-2 \
+>   --query "reverse(sort_by(imageDetails,&imagePushedAt))[].imageTags" --output text \
+>   | tr '\t' '\n' | grep -E '^2.6.0-gpu-py312-cu126-.*-sagemaker-v[0-9.]+$' | head -1
+> ```
+
+> **SDK pin:** the smoke launcher uses the **sagemaker SDK v2** Estimator API.
+> SDK **v3 is an API rewrite** that dropped `sagemaker.estimator.Estimator` and
+> `sagemaker.pytorch` — install `pip install 'sagemaker>=2.200,<3'`.
+
+## Run it (no local Docker build)
+
+```bash
+pip install 'sagemaker>=2.200,<3'
+export AWS_REGION=us-west-2
+python examples/gsm8k_grpo/run_sagemaker_launch.py --max-steps 20
+```
+
+This uses the **stock PyTorch DLC directly** as the training image and ships the
+framework + entry script via `source_dir`; `examples/gsm8k_grpo/requirements.txt`
+(trl + vllm==0.8.5 + the RL stack) installs at job start. No 15 GB local image
+build, no ECR push. The script trains `Qwen/Qwen2.5-0.5B-Instruct` with GRPO +
+the GSM8K `#### NUMBER` RLVR reward, vLLM **colocated** in-process
+(`vllm_gpu_memory_utilization=0.3` on the 24 GB A10G).
+
+Flags: `--no-wait` (submit + poll later), `--spot` (managed spot, quota=1 too),
+`--no-vllm` (fall back to `model.generate` rollout if a vLLM/CUDA wheel mismatch
+appears), `--image <ecr-uri>` (use a prebuilt baked image instead of the DLC).
+
+**Cost:** `ml.g5.2xlarge` ≈ $1.52/hr on-demand; a 20-step 0.5B smoke is
+~15–25 min ⇒ **well under $1**. Spot ≈ $0.45–0.60/hr ⇒ pennies.
+
+## The repeatable path (baked image)
+
+For runs where the ~5–10 min per-job pip-install is unwanted (and for the
+DiLoCo N-replica `SageMakerExecutor`, which passes `ContainerEntrypoint` and
+needs the framework baked in), build the image once:
+
+```bash
+scripts/build_and_push_ecr.sh            # creates ECR repo, builds, pushes composer-rl:smoke
+python examples/gsm8k_grpo/run_sagemaker_launch.py \
+  --image 386931836011.dkr.ecr.us-west-2.amazonaws.com/composer-rl:smoke
+```
+
+On an Apple-Silicon host the build cross-compiles (`--platform linux/amd64`) the
+~15 GB GPU DLC under emulation — slow; prefer a linux/amd64 host or CodeBuild.
+
+## Gotchas (load-bearing)
+
+- **`EnableNetworkIsolation` stays False** (the default) so the container can
+  reach `huggingface.co` (model + GSM8K download) and S3.
+- **`vllm_gpu_memory_utilization=0.3`** is the load-bearing knob on a 24 GB
+  A10G: too high → OOM when the policy + grads also need the GPU; too low → tiny
+  KV cache. Use `--no-vllm` if a vLLM wheel/CUDA mismatch surfaces.
+- **Warm pools are off** — `g5 training warm pool usage` quota is 0 in this
+  account, so each job pays ~3–6 min cold-start. Request a warm-pool quota bump
+  for iterative dev, or move the long inner loop to HyperPod (persistent).
+- **"Waiting for capacity"** in `SecondaryStatus` is transient g5 capacity
+  contention in the region, not an error — the job proceeds when capacity frees.
+
+## Next: DiLoCo N-replica (the `SageMakerExecutor` path)
+
+`examples/diloco_sagemaker/run.py` (driver, F3 §4.3) drives N independent
+single-instance Training Jobs sharing one `s3://.../rendezvous/` prefix via
+`ObjectStoreAllReduce` — no cross-job NCCL. N=1 runs today; N=2–4 needs a
+`ml.g5.2xlarge for training job usage` quota increase. The DiLoCo math, loss,
+trainer, and `ObjectStoreAllReduce` are unchanged from the smoke — the S3
+rendezvous is the entire portability contract (validated `file://` and live
+`s3://`; see `test_serverless_local.py::test_s3_rendezvous_allreduce_across_replicas`).
+```

examples/gsm8k_grpo/requirements.txt

@@ -0,0 +1,14 @@
+# Installed by SageMaker at job start, layered on the PyTorch DLC 2.6.0
+# (torch 2.6 + CUDA 12.6 already present — do NOT reinstall torch).
+# F3 §3.2: baking these into a Dockerfile.sagemaker is the repeatable path;
+# this requirements.txt is the no-local-build path for the one-shot smoke.
+trl>=1.5,<2
+peft>=0.13
+accelerate>=1.0
+datasets>=3.0
+# vLLM must match the DLC's torch 2.6 / cu126. 0.8.x is the torch-2.6 line;
+# pinning avoids a silent wheel/CUDA mismatch at colocate time (F3 §7).
+vllm==0.8.5
+fsspec>=2024.6
+s3fs>=2024.6
+hf_transfer>=0.1.6

examples/gsm8k_grpo/run_sagemaker.py

@@ -0,0 +1,179 @@
+"""GPU + vLLM-colocated variant of run.py, for a SageMaker Training Job.
+
+F3 §3.3: the runnable-NOW GRPO smoke. Same GSM8K RLVR reward and same
+`ComposerReplicationTrainer(alpha_sdpo=0, beta_replay=0)` (plain GRPO, channels
+2/3 off) as the CPU example, lifted to one real GPU with vLLM colocated in the
+training process. Proves: container builds, trainer runs on GPU, vLLM rollout
+works, reward fires, checkpoint lands in S3.
+
+This script runs INSIDE the SageMaker container. SageMaker conventions used:
+  * Hyperparameters arrive as CLI args ``--key value`` (the Estimator's
+    ``hyperparameters=`` dict). We also read /opt/ml/input/config/
+    hyperparameters.json as a fallback.
+  * The final model must be written to ``/opt/ml/model`` (SM_MODEL_DIR);
+    SageMaker tars it to ``OutputDataConfig.S3OutputPath`` on exit.
+  * stdout/stderr stream to CloudWatch ``/aws/sagemaker/TrainingJobs/<job>``.
+
+Run via examples/gsm8k_grpo/run_sagemaker_launch.py (the Estimator driver).
+"""
+from __future__ import annotations
+
+import argparse
+import json
+import logging
+import os
+import re
+import sys
+
+import torch
+from datasets import load_dataset
+from transformers import AutoModelForCausalLM, AutoTokenizer
+
+# ---------------------------------------------------------------------------
+# Reward — identical RLVR `#### NUMBER` regex to the CPU example (run.py)
+# ---------------------------------------------------------------------------
+
+_ANSWER_RE = re.compile(r"####\s*(-?\d+(?:\.\d+)?)")
+
+
+def _extract_answer(text: str) -> str | None:
+    matches = _ANSWER_RE.findall(text or "")
+    return matches[-1].strip() if matches else None
+
+
+def gsm8k_reward(completions, **kwargs):
+    gold = kwargs.get("gold_answer")
+    if gold is None:
+        return [0.0] * len(completions)
+    rewards: list[float] = []
+    for completion, gold_ans in zip(completions, gold, strict=False):
+        if isinstance(completion, list):
+            text = "\n".join(m.get("content", "") for m in completion)
+        else:
+            text = str(completion)
+        pred = _extract_answer(text)
+        rewards.append(1.0 if (pred is not None and pred == str(gold_ans).strip()) else 0.0)
+    return rewards
+
+
+SYSTEM_PROMPT = (
+    "You are a math tutor. Solve the problem step by step. "
+    "End your answer with `#### N` where N is the final numeric answer."
+)
+
+
+def build_dataset(n_rows: int):
+    raw = load_dataset("openai/gsm8k", "main", split=f"train[:{n_rows}]")
+
+    def _format(row):
+        gold = _extract_answer(row["answer"]) or ""
+        return {
+            "prompt": [
+                {"role": "system", "content": SYSTEM_PROMPT},
+                {"role": "user", "content": row["question"]},
+            ],
+            "gold_answer": gold,
+        }
+
+    return raw.map(_format, remove_columns=raw.column_names)
+
+
+# ---------------------------------------------------------------------------
+# Hyperparameters — SageMaker passes them as CLI args; JSON fallback.
+# ---------------------------------------------------------------------------
+
+
+def _parse_hyperparameters() -> argparse.Namespace:
+    p = argparse.ArgumentParser()
+    p.add_argument("--model", default="Qwen/Qwen2.5-0.5B-Instruct")
+    p.add_argument("--n_train_rows", type=int, default=100)
+    p.add_argument("--max_steps", type=int, default=20)
+    p.add_argument("--num_generations", type=int, default=8)
+    p.add_argument("--per_device_train_batch_size", type=int, default=8)
+    p.add_argument("--max_completion_length", type=int, default=256)
+    p.add_argument("--learning_rate", type=float, default=1e-5)
+    p.add_argument("--beta", type=float, default=0.04)
+    p.add_argument("--vllm_gpu_memory_utilization", type=float, default=0.3)
+    p.add_argument("--use_vllm", type=lambda s: str(s).lower() != "false", default=True)
+    # SageMaker model output dir (env SM_MODEL_DIR, default /opt/ml/model).
+    p.add_argument("--model_dir", default=os.environ.get("SM_MODEL_DIR", "/opt/ml/model"))
+    args, _unknown = p.parse_known_args()
+    return args
+
+
+def main() -> int:
+    logging.basicConfig(
+        level=logging.INFO,
+        format="%(asctime)s [%(levelname)s] %(name)s: %(message)s",
+        handlers=[logging.StreamHandler(sys.stdout)],
+    )
+    log = logging.getLogger("gsm8k_grpo_sagemaker")
+    args = _parse_hyperparameters()
+
+    log.info("=" * 64)
+    log.info("GRPO + GSM8K + %s (SageMaker GPU, vLLM=%s)", args.model, args.use_vllm)
+    log.info("=" * 64)
+    log.info("hyperparameters: %s", json.dumps(vars(args), indent=2))
+
+    cuda = torch.cuda.is_available()
+    log.info("CUDA available: %s | device: %s", cuda,
+             torch.cuda.get_device_name(0) if cuda else "cpu")
+
+    log.info("[1/4] Loading model + tokenizer ...")
+    tokenizer = AutoTokenizer.from_pretrained(args.model)
+    if tokenizer.pad_token_id is None:
+        tokenizer.pad_token = tokenizer.eos_token
+    model = AutoModelForCausalLM.from_pretrained(
+        args.model, torch_dtype=torch.bfloat16 if cuda else torch.float32
+    )
+
+    log.info("[2/4] Loading %d GSM8K rows ...", args.n_train_rows)
+    dataset = build_dataset(args.n_train_rows)
+
+    log.info("[3/4] Building GRPOConfig + ComposerReplicationTrainer ...")
+    from trl import GRPOConfig
+
+    from composer_replication import ComposerReplicationTrainer
+
+    config = GRPOConfig(
+        output_dir=args.model_dir,
+        per_device_train_batch_size=args.per_device_train_batch_size,
+        num_generations=args.num_generations,
+        max_completion_length=args.max_completion_length,
+        learning_rate=args.learning_rate,
+        max_steps=args.max_steps,
+        logging_steps=1,
+        save_strategy="no",
+        report_to=[],
+        bf16=cuda,
+        beta=args.beta,
+        # vLLM colocated in-process on the same GPU (F3 §3.3 / §5).
+        use_vllm=bool(args.use_vllm and cuda),
+        vllm_mode="colocate",
+        vllm_gpu_memory_utilization=args.vllm_gpu_memory_utilization,
+        vllm_tensor_parallel_size=1,
+        seed=42,
+    )
+
+    trainer = ComposerReplicationTrainer(
+        model=model,
+        processing_class=tokenizer,
+        reward_funcs=[gsm8k_reward],
+        train_dataset=dataset,
+        args=config,
+        alpha_sdpo=0.0,   # channels 2/3 off — plain GRPO smoke
+        beta_replay=0.0,
+    )
+
+    log.info("[4/4] Training for %d steps ...", args.max_steps)
+    result = trainer.train()
+    log.info("Training complete: %s", result.metrics)
+
+    # Persist to SM_MODEL_DIR → SageMaker uploads to OutputDataConfig.
+    trainer.save_model(args.model_dir)
+    log.info("Model saved to %s (SageMaker will upload to S3).", args.model_dir)
+    return 0
+
+
+if __name__ == "__main__":
+    sys.exit(main())

examples/gsm8k_grpo/run_sagemaker_launch.py

@@ -0,0 +1,138 @@
+"""Launch the GSM8K GRPO smoke as a SageMaker Training Job (F3 §3.4).
+
+Two ways to run, selected by --image:
+  * --image dlc   (default): use the AWS PyTorch DLC directly as the training
+    image and ship the framework + script via source_dir; deps install at job
+    start from requirements.txt. No local Docker build — the lowest-friction
+    path for a one-shot smoke (F3's "source_dir on top of the baked image",
+    here on top of the *stock* DLC).
+  * --image <ecr-uri>: use a prebuilt baked image (docker/Dockerfile.sagemaker
+    pushed to ECR via scripts/build_and_push_ecr.sh) — the repeatable path that
+    avoids the ~5-10 min startup pip-install (F3 §3.2).
+
+Live facts (verified 2026-06-09, acct 386931836011, us-west-2):
+  * ml.g5.2xlarge training-job quota = 1 → runs with no quota ticket.
+  * DLC tag resolved live: 2.6.0-gpu-py312-cu126-ubuntu22.04-sagemaker-v1.25
+    (NOTE: cu126, not the cu124 in some docs; bare floating tag does not exist,
+    the -v1.25 build suffix is required).
+  * role: AmazonSageMaker-ExecutionRole-20250725T133247
+  * bucket: amazon-sagemaker-386931836011-us-west-2-7597bf4d9a3d
+
+Usage (from the laptop / Studio, sagemaker SDK v2):
+    pip install 'sagemaker>=2.200,<3'
+    python examples/gsm8k_grpo/run_sagemaker_launch.py --max-steps 20
+    python examples/gsm8k_grpo/run_sagemaker_launch.py --no-wait   # fire and poll later
+"""
+from __future__ import annotations
+
+import argparse
+import os
+import shutil
+import sys
+import tempfile
+
+REGION = "us-west-2"
+ACCOUNT = "386931836011"
+ROLE = f"arn:aws:iam::{ACCOUNT}:role/service-role/AmazonSageMaker-ExecutionRole-20250725T133247"
+BUCKET = f"amazon-sagemaker-{ACCOUNT}-{REGION}-7597bf4d9a3d"
+# DLC tag resolved live against the 763104351884 us-west-2 registry.
+DLC_IMAGE = (
+    "763104351884.dkr.ecr.us-west-2.amazonaws.com/"
+    "pytorch-training:2.6.0-gpu-py312-cu126-ubuntu22.04-sagemaker-v1.25"
+)
+
+_HERE = os.path.dirname(os.path.abspath(__file__))
+_REPO_ROOT = os.path.abspath(os.path.join(_HERE, "..", ".."))
+
+
+def _stage_source() -> str:
+    """Build a minimal source_dir: the entry script + requirements.txt + the
+    composer_replication package (importable as a local package from
+    /opt/ml/code). Keeps the S3 upload tiny — no .venv/research/docs/.git.
+    Returns the staging dir path (caller cleans up)."""
+    staging = tempfile.mkdtemp(prefix="composer-sm-src-")
+    shutil.copy2(os.path.join(_HERE, "run_sagemaker.py"), staging)
+    shutil.copy2(os.path.join(_HERE, "requirements.txt"), staging)
+    shutil.copytree(
+        os.path.join(_REPO_ROOT, "composer_replication"),
+        os.path.join(staging, "composer_replication"),
+        ignore=shutil.ignore_patterns("__pycache__", "*.pyc", "tests", "*.egg-info"),
+    )
+    return staging
+
+
+def main() -> int:
+    ap = argparse.ArgumentParser()
+    ap.add_argument("--image", default="dlc",
+                    help="'dlc' (stock PyTorch DLC + source_dir) or a prebuilt ECR image URI")
+    ap.add_argument("--instance-type", default="ml.g5.2xlarge")
+    ap.add_argument("--max-steps", type=int, default=20)
+    ap.add_argument("--n-train-rows", type=int, default=100)
+    ap.add_argument("--model", default="Qwen/Qwen2.5-0.5B-Instruct")
+    ap.add_argument("--no-vllm", action="store_true",
+                    help="disable colocated vLLM (use model.generate rollout) — safer fallback")
+    ap.add_argument("--spot", action="store_true", help="use managed spot (quota=1 too)")
+    ap.add_argument("--no-wait", action="store_true", help="submit and return; poll later")
+    ap.add_argument("--max-run", type=int, default=3600)
+    args = ap.parse_args()
+
+    import sagemaker
+    from sagemaker.estimator import Estimator
+
+    sess = sagemaker.Session(default_bucket=BUCKET)
+    image = DLC_IMAGE if args.image == "dlc" else args.image
+    print(f"[launch] region={REGION} image={image}")
+    print(f"[launch] role={ROLE}")
+    print(f"[launch] instance={args.instance_type} max_steps={args.max_steps} "
+          f"vllm={not args.no_vllm} spot={args.spot}")
+
+    staging = _stage_source()
+    print(f"[launch] staged source_dir at {staging}")
+
+    hyperparameters = {
+        "model": args.model,
+        "n_train_rows": args.n_train_rows,
+        "max_steps": args.max_steps,
+        "use_vllm": "false" if args.no_vllm else "true",
+    }
+
+    spot_kwargs = {}
+    if args.spot:
+        spot_kwargs = {"use_spot_instances": True, "max_wait": args.max_run + 3600}
+
+    est = Estimator(
+        image_uri=image,
+        role=ROLE,
+        instance_type=args.instance_type,
+        instance_count=1,
+        volume_size=100,
+        max_run=args.max_run,
+        sagemaker_session=sess,
+        output_path=f"s3://{BUCKET}/composer-rl/smoke/output",
+        base_job_name="composer-grpo-smoke",
+        entry_point="run_sagemaker.py",
+        source_dir=staging,
+        hyperparameters=hyperparameters,
+        environment={
+            "HF_HUB_ENABLE_HF_TRANSFER": "1",
+            # DLC sagemaker-training-toolkit installs requirements.txt from source_dir.
+        },
+        # EnableNetworkIsolation MUST be False (default) so the container can reach
+        # huggingface.co (model + GSM8K) and S3 (F3 §3.5).
+        keep_alive_period_in_seconds=0,  # warm-pool quota=0 in this acct → leave off
+        **spot_kwargs,
+    )
+
+    try:
+        est.fit(wait=not args.no_wait, logs=("All" if not args.no_wait else None))
+    finally:
+        shutil.rmtree(staging, ignore_errors=True)
+
+    print(f"[launch] job name: {est.latest_training_job.name}")
+    if not args.no_wait:
+        print(f"[launch] model artifact: {est.model_data}")
+    return 0
+
+
+if __name__ == "__main__":
+    sys.exit(main())

pyproject.toml

@@ -75,8 +75,13 @@ eks = [
     "kubernetes>=29",
 ]
 # Amazon SageMaker training-job executor (SageMakerExecutor, per ADR-005).
+# boto3: the executor uses raw create_training_job. sagemaker (<3): the v2
+# Estimator API the GSM8K smoke launcher (examples/gsm8k_grpo/
+# run_sagemaker_launch.py) uses — pinned <3 because SDK v3 is an API rewrite
+# that dropped sagemaker.estimator.Estimator (F3 §1, verified live 2026-06-09).
 aws = [
     "boto3>=1.34",
+    "sagemaker>=2.200,<3",
 ]
 # Replaysim dataset normalization (per ADR-004)
 #

scripts/build_and_push_ecr.sh

@@ -0,0 +1,43 @@
+#!/usr/bin/env bash
+# Build + push the baked SageMaker RL image to this account's private ECR (F3 §3.2).
+# One-time (Admin). The one-shot smoke does NOT need this — use
+#   python examples/gsm8k_grpo/run_sagemaker_launch.py --image dlc
+# which runs on the stock DLC + source_dir with no local build. Use this script
+# for the repeatable path (no per-job pip-install) and the DiLoCo N-replica
+# executor (which passes ContainerEntrypoint and wants the framework baked in).
+#
+# NOTE: the DLC base is GPU/linux-amd64 (~15 GB). On an Apple-Silicon host you
+# must cross-build: pass --platform linux/amd64 (set below). This is slow under
+# emulation; prefer building on a linux/amd64 host or CodeBuild for real use.
+set -euo pipefail
+
+REGION="${REGION:-us-west-2}"
+ACCOUNT="${ACCOUNT:-386931836011}"
+REPO="${REPO:-composer-rl}"
+TAG="${TAG:-smoke}"
+DLC_ACCOUNT="763104351884"
+
+REGISTRY="${ACCOUNT}.dkr.ecr.${REGION}.amazonaws.com"
+IMAGE="${REGISTRY}/${REPO}:${TAG}"
+REPO_ROOT="$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd)"
+
+echo "[ecr] region=${REGION} image=${IMAGE}"
+
+# 1. Ensure the ECR repo exists (idempotent).
+aws ecr describe-repositories --repository-names "${REPO}" --region "${REGION}" >/dev/null 2>&1 \
+  || aws ecr create-repository --repository-name "${REPO}" --region "${REGION}" >/dev/null
+
+# 2. Log in to BOTH the DLC registry (to pull the base) and our own (to push).
+aws ecr get-login-password --region "${REGION}" \
+  | docker login --username AWS --password-stdin "${DLC_ACCOUNT}.dkr.ecr.${REGION}.amazonaws.com"
+aws ecr get-login-password --region "${REGION}" \
+  | docker login --username AWS --password-stdin "${REGISTRY}"
+
+# 3. Build (cross-arch on Apple Silicon) + push.
+docker build --platform linux/amd64 \
+  -f "${REPO_ROOT}/docker/Dockerfile.sagemaker" \
+  -t "${IMAGE}" "${REPO_ROOT}"
+docker push "${IMAGE}"
+
+echo "[ecr] pushed ${IMAGE}"
+echo "[ecr] launch with: python examples/gsm8k_grpo/run_sagemaker_launch.py --image ${IMAGE}"