Codeseys/composer-replication-framework

1# Baked SageMaker training image for the Composer-replication RL stack (F3 §3.2).2#3# The repeatable path: bake trl+vllm+the framework into an image so jobs don't4# pip-install at startup (saves ~5-10 min/job and removes a flaky failure5# surface). The one-shot smoke can instead use the stock DLC + source_dir6# (run_sagemaker_launch.py --image dlc), which needs no local build.7#8# Base: AWS PyTorch DLC, tag RESOLVED LIVE against the us-west-2 registry.9# MUST be torch-2.7: trl 1.5 → transformers>=4.56.2 → torch.float8_e8m0fnu10# (torch>=2.7). The torch-2.6 DLC fails AutoModel.from_pretrained on that dtype.11# cu128, -v1.26 build suffix required (no bare floating tag exists).12FROM 763104351884.dkr.ecr.us-west-2.amazonaws.com/pytorch-training:2.7.1-gpu-py312-cu128-ubuntu22.04-sagemaker-v1.2613 14# RL stack baked in. torch 2.7 + CUDA 12.8 already in the DLC — do NOT reinstall15# torch. vllm>=0.9 is the torch-2.7 line (0.8.x hard-pins torch 2.6 and would16# fight this base); pin to a 2.7-matched vllm to avoid a wheel/CUDA mismatch.17RUN pip install --no-cache-dir \18      "trl>=1.5,<2" "peft>=0.13" "accelerate>=1.0" "datasets>=3.0" \19      "vllm>=0.9" "fsspec>=2024.6" "s3fs>=2024.6" "hf_transfer>=0.1.6"20 21# The framework itself (train + serverless extras → trainer, loss, executors,22# replica_entrypoint, s3fs all present).23COPY . /opt/composer_replication24RUN pip install --no-cache-dir -e "/opt/composer_replication[train,serverless]"25 26ENV HF_HOME=/opt/ml/input/hf_cache \27    HF_HUB_ENABLE_HF_TRANSFER=128