构建自定义镜像

FROM nvidia/cuda:12.4.1-devel-ubuntu22.04
RUN pip install torch transformers trl accelerate
COPY . /workspace

# syntax=docker/dockerfile:1
ARG BASE_IMAGE=nvidia/cuda:12.8.0-devel-ubuntu24.04
FROM ${BASE_IMAGE}

ARG PYTHON_VERSION=3.12
ENV DEBIAN_FRONTEND=noninteractive \
    HF_HUB_ENABLE_HF_TRANSFER=1

RUN apt-get update && apt-get install -y --no-install-recommends \
        git curl ca-certificates \
    && rm -rf /var/lib/apt/lists/*

COPY --from=ghcr.io/astral-sh/uv:0.8 /uv /usr/local/bin/uv

RUN uv python install ${PYTHON_VERSION} \
 && uv venv --python ${PYTHON_VERSION} /opt/train-venv
ENV VIRTUAL_ENV=/opt/train-venv PATH=/opt/train-venv/bin:$PATH

# 生产请钉版本，不要留未 pin 的 extra。
RUN uv pip install --no-cache \
        "torch" \
        "transformers" \
        "accelerate" \
        "trl" \
        "hf-transfer" \
        "starforge-core"

# Job Capsule 用镜像里的 Python。保持这个 venv 在 PATH 最前。
CMD ["python", "-c", "import torch; print(torch.__version__)"]

docker build -t myregistry.io/my-train:v1 .
docker push myregistry.io/my-train:v1

# 记下 digest 再提交：
docker buildx imagetools inspect myregistry.io/my-train:v1

sf submit my-custom --profile h200:8 \
  --image myregistry.io/my-train:v1 \
  --observability-url https://wandb.example/my-proj

sf submit my-custom --profile h200:8 \
  --image myregistry.io/my-train@sha256:abcd… \
  --observability-url https://wandb.example/my-proj