feat: unify Docker training configuration in YAML
This commit is contained in:
+1
-2
@@ -72,9 +72,8 @@ services:
|
||||
- BASE_MODEL=${BASE_MODEL:-/models/base}
|
||||
- CHECKPOINT_ROOT=/checkpoints
|
||||
- TRAIN_GPU_COUNT=${TRAIN_GPU_COUNT:-all}
|
||||
- TRAIN_PARALLEL_MODE=${TRAIN_PARALLEL_MODE:-auto}
|
||||
- CUDA_VISIBLE_DEVICES
|
||||
- NCCL_P2P_DISABLE
|
||||
- NCCL_NET_GDR_LEVEL
|
||||
entrypoint: ["bash", "/app/scripts/docker/train-entrypoint.sh"]
|
||||
ipc: ${TRAIN_IPC_MODE:-host}
|
||||
stop_grace_period: ${TRAIN_STOP_GRACE_PERIOD:-10m}
|
||||
|
||||
@@ -1,58 +1,125 @@
|
||||
# Containerized Training Deployment
|
||||
|
||||
Rules for running AstrAI distributed training in containers, distilled from real deployment failures. Read before touching `Dockerfile`, `docker-compose.yml`, `scripts/train.sh`, `train-entrypoint.sh`. AGENTS.md mirrors this locally; this file is the committed version.
|
||||
AstrAI uses one training YAML as the declaration for both host-side container
|
||||
runtime settings and in-container training settings. Do not invoke the trainer
|
||||
with raw `docker compose up`; use `scripts/train.sh` so preflight validation,
|
||||
checkpoint recovery, and graceful shutdown remain active.
|
||||
|
||||
## Architecture
|
||||
|
||||
```
|
||||
scripts/train.sh host-side CLI: env loading, preflight, compose wrapper, lifecycle
|
||||
└── docker-compose.yml GPU passthrough, mounts, in-container env vars, entrypoint
|
||||
└── train-entrypoint.sh GPU-count resolution, parallel-mode selection, auto-resume
|
||||
```text
|
||||
train.yaml
|
||||
├── runtime parsed on the host before Docker starts
|
||||
└── model/data/... parsed by train.py inside the container
|
||||
│
|
||||
scripts/train.sh preflight, Compose wrapper, lifecycle, timer
|
||||
└── docker-compose.yml GPU passthrough, mounts, image, container limits
|
||||
└── train-entrypoint.sh process count, parallel mode, auto-resume
|
||||
└── train.py --config /run/astrai/train.yaml
|
||||
```
|
||||
|
||||
| Layer | Responsible for | NOT responsible for |
|
||||
|-------|-----------------|---------------------|
|
||||
| `train.sh` | host paths, `.env.train` + `infra:` YAML overrides, preflight, lifecycle | training args, GPU selection, parallel mode |
|
||||
| compose | GPU passthrough, mounts, in-container env (NCCL) | training args (beyond `TRAIN_*` forwarding) |
|
||||
| entrypoint | `--ckpt_dir/--nprocs/--parallel_mode/--param_path`, resume | hyperparameters (YAML/CLI) |
|
||||
| `train.yaml` | hyperparameters (`_merge_yaml_into_kwargs`, CLI wins); top-level `infra:` host vars | container paths, process count |
|
||||
The two parsers deliberately own different sections. `scripts/tools/train_runtime.py`
|
||||
reads only `runtime`; `scripts/tools/train.py` reads only
|
||||
`model/data/parallel/training/ckpt/log`. Explicit trainer arguments after `--`
|
||||
override training YAML values.
|
||||
|
||||
## Path Conventions
|
||||
## Runtime Schema
|
||||
|
||||
| Host var | Container | Perm | Purpose |
|
||||
|---|---|---|---|
|
||||
| `TRAIN_DATA_DIR` | `/data` | ro | dataset (`data_root_path` must be `/data`) |
|
||||
| `TRAIN_MODEL_DIR` | `/models/base` | ro | base model (`config.json` + `model.safetensors`) |
|
||||
| `TRAIN_CHECKPOINT_DIR` | `/checkpoints` | rw | checkpoint root, per-`TRAIN_JOB_NAME` subdirs |
|
||||
| `TRAIN_CONFIG_FILE` | `/run/astrai/train.yaml` | ro | training YAML (mounted only on `start`) |
|
||||
| code | `/app` | image | **not a mount**; rebuild image for code changes |
|
||||
```yaml
|
||||
runtime:
|
||||
job_name: astrai-train
|
||||
paths:
|
||||
data: ./data
|
||||
model: ./params
|
||||
checkpoints: ./checkpoints
|
||||
gpu:
|
||||
devices: all
|
||||
parallel_mode: auto # one GPU: none; multiple GPUs: ddp
|
||||
container:
|
||||
cuda_tag: cu128
|
||||
ipc: host
|
||||
stop_grace_period: 10m
|
||||
stop_timeout_seconds: 600
|
||||
checkpoint_keep_last: 5
|
||||
# max_duration_hours: 12
|
||||
# Add host-specific workarounds only when required:
|
||||
# environment:
|
||||
# NCCL_P2P_DISABLE: "1"
|
||||
# NCCL_NET_GDR_LEVEL: "0"
|
||||
```
|
||||
|
||||
## Hard Rules
|
||||
- Relative paths resolve from the YAML file's directory, not the current shell.
|
||||
- `devices` is either `all` or a non-empty physical GPU index list. Compose
|
||||
passes all GPUs once; `CUDA_VISIBLE_DEVICES` performs the only filtering.
|
||||
- The process count is derived from `devices`. With `all`, the entrypoint uses
|
||||
`torch.cuda.device_count()` after Docker starts.
|
||||
- `parallel_mode: auto` selects `none` for one GPU and `ddp` for multiple GPUs.
|
||||
Use `fsdp` explicitly when model sharding is required.
|
||||
- To select specific physical GPUs, replace `all` with a list such as
|
||||
`devices: [0, 1]`.
|
||||
- `environment` values are explicitly passed to the training container. Keep
|
||||
host-specific NCCL workarounds here; they are not universal defaults.
|
||||
- `max_duration_hours` starts a detached host timer that calls the same graceful
|
||||
`stop` command. A manual stop cancels the timer.
|
||||
|
||||
1. **Filter GPUs once**: compose passes the full physical set (`count: all`); `CUDA_VISIBLE_DEVICES` filters inside by physical index. Never `count: N` + physical indices (double filter leaves 1 card → `device_id out of range`).
|
||||
2. **In-container UID = host UID**: Dockerfile builds the user via `USER_UID/USER_GID` args; `train.sh` injects `ASTRAI_UID/GID` (bash `UID` is readonly). compose `user:` alone does not create the /etc/passwd entry — torch's `getpass.getuser()` then dies with `uid not found`.
|
||||
3. **In-container env vars are explicit**: `.env.train` (`--env-file`) is only compose's interpolation dictionary — never reaches the container. A var arrives only via a value-less `environment` entry (`- VAR`, read from the calling process env).
|
||||
4. **Host vars can come from `train.yaml` instead**: `scripts/train.sh load_infra()` reads the top-level `infra:` section of `TRAIN_CONFIG_FILE` (host side, before the container exists) and exports `TRAIN_JOB_NAME`, `TRAIN_DATA_DIR`, `TRAIN_MODEL_DIR`, `TRAIN_CHECKPOINT_DIR`, `TRAIN_GPU_COUNT`, `CUDA_VISIBLE_DEVICES`. Compose interpolation prefers the shell environment over `--env-file`, so `infra:` wins; keys absent from it fall back to `.env.train`. Requires python3 + PyYAML on the host. train.py only merges the `model/data/parallel/training/ckpt/log` sections, so the `infra` section is invisible to the trainer.
|
||||
5. **NCCL hang workaround** (this host): `NCCL_P2P_DISABLE=1` + `NCCL_NET_GDR_LEVEL=0` must be in-container.
|
||||
6. **Checkpoint complete =** `meta.json + config.json + model.safetensors + optimizer.pt + scheduler.pt`; `start` auto-resumes the latest complete one.
|
||||
7. **tqdm is silent without a TTY**: add `disable=False` in `astrai/trainer/train_callback.py`; `metric.jsonl` (per step) works as progress evidence regardless.
|
||||
## Fixed Container Paths
|
||||
|
||||
| Runtime path | Container path | Access |
|
||||
|---|---|---|
|
||||
| `runtime.paths.data` | `/data` | read-only |
|
||||
| `runtime.paths.model` | `/models/base` | read-only |
|
||||
| `runtime.paths.checkpoints` | `/checkpoints` | read-write |
|
||||
| the selected YAML | `/run/astrai/train.yaml` | read-only |
|
||||
|
||||
Training configuration must therefore use `data_root_path: /data`. The source
|
||||
code is baked into `/app`; `start` uses `--build`, so code changes rebuild the
|
||||
image when necessary.
|
||||
|
||||
## Operations
|
||||
|
||||
The config argument defaults to `./train.yaml`:
|
||||
|
||||
```bash
|
||||
bash scripts/train.sh init # first run: dirs + .env.train (edit per machine)
|
||||
bash scripts/train.sh preflight # validate Docker/paths/GPU/model/YAML/compose
|
||||
bash scripts/train.sh start # build + start in background (auto-resume)
|
||||
bash scripts/train.sh start --foreground -- --dry-run # print plan only
|
||||
bash scripts/train.sh logs | status | stop | restart
|
||||
bash scripts/train.sh clean --keep 5 # prune old checkpoints (--force to delete)
|
||||
bash scripts/train.sh init [CONFIG]
|
||||
bash scripts/train.sh preflight [CONFIG]
|
||||
bash scripts/train.sh start [CONFIG]
|
||||
bash scripts/train.sh start [CONFIG] --foreground -- --dry-run
|
||||
bash scripts/train.sh logs [CONFIG]
|
||||
bash scripts/train.sh status [CONFIG]
|
||||
bash scripts/train.sh stop [CONFIG]
|
||||
bash scripts/train.sh restart [CONFIG]
|
||||
bash scripts/train.sh clean [CONFIG] --keep 5
|
||||
bash scripts/train.sh clean [CONFIG] --keep 5 --force
|
||||
```
|
||||
|
||||
## Files
|
||||
`init` creates the declared runtime directories but does not generate or mutate
|
||||
the YAML. `preflight` validates Docker, paths, base model files, checkpoint
|
||||
writability, GPU configuration, and rendered Compose configuration.
|
||||
|
||||
- `docker-compose.yml` — trainer service: `count: all`, `ASTRAI_UID/GID` build args + `user:`, env whitelist, mounts
|
||||
- `Dockerfile` — production stage builds user from `USER_UID/USER_GID`; `ENV HOME=/home/astrai`; `USER astrai`
|
||||
- `scripts/train.sh` — `load_env` filters `UID=` lines (readonly var); `load_infra` reads the `infra:` section from `TRAIN_CONFIG_FILE`; `compose()` injects `ASTRAI_UID/GID`
|
||||
- `scripts/docker/train-entrypoint.sh` — GPU-count resolution, parallel mode, resume
|
||||
- `.env.train`, `train.yaml` — host-specific; templates from `scripts/train.sh init`; scientific-notation floats (`2e-5`) parse correctly since train.py uses the YAML 1.2 float schema; `.env.train` is the fallback for host vars not present in the `infra:` section
|
||||
## Checkpoint Recovery
|
||||
|
||||
Checkpoints are stored below
|
||||
`runtime.paths.checkpoints/<job_name>/epoch_<N>_step_<N>`. A checkpoint is
|
||||
complete only when it contains:
|
||||
|
||||
```text
|
||||
meta.json
|
||||
config.json
|
||||
model.safetensors
|
||||
optimizer.pt
|
||||
scheduler.pt
|
||||
```
|
||||
|
||||
`start` resumes the latest complete checkpoint and ignores partial writes. If no
|
||||
complete checkpoint exists, `/models/base/config.json` and
|
||||
`/models/base/model.safetensors` are required. `stop` sends `SIGTERM`; the
|
||||
trainer finishes at a batch boundary and saves an emergency checkpoint before
|
||||
the Docker timeout expires.
|
||||
|
||||
## Hard Rules
|
||||
|
||||
1. Keep Docker settings in `runtime` and trainer settings in the remaining YAML sections.
|
||||
2. Filter GPUs once: Compose passes `count: all`; `devices` becomes `CUDA_VISIBLE_DEVICES`.
|
||||
3. Do not force DDP for a model that requires FSDP; declare the mode explicitly.
|
||||
4. Do not use `kill -9` for routine shutdown; use `scripts/train.sh stop CONFIG`.
|
||||
5. The image user is built with the host UID/GID so mounted checkpoints retain usable ownership.
|
||||
|
||||
@@ -10,12 +10,27 @@ CHECKPOINT_DIR="${CHECKPOINT_ROOT}/${TRAIN_JOB_NAME}"
|
||||
BASE_MODEL="${BASE_MODEL:-/models/base}"
|
||||
TRAIN_CONFIG="${TRAIN_CONFIG:-}"
|
||||
TRAIN_GPU_COUNT="${TRAIN_GPU_COUNT:-all}"
|
||||
TRAIN_PARALLEL_MODE="${TRAIN_PARALLEL_MODE:-auto}"
|
||||
|
||||
validate_job_name "${TRAIN_JOB_NAME}"
|
||||
if [[ "${TRAIN_GPU_COUNT}" == "all" ]]; then
|
||||
TRAIN_GPU_COUNT="$(python -c 'import torch; print(torch.cuda.device_count())')"
|
||||
fi
|
||||
[[ "${TRAIN_GPU_COUNT}" =~ ^[1-9][0-9]*$ ]] || die "No visible GPU found"
|
||||
if [[ "${TRAIN_PARALLEL_MODE}" == "auto" ]]; then
|
||||
if (( TRAIN_GPU_COUNT > 1 )); then
|
||||
TRAIN_PARALLEL_MODE=ddp
|
||||
else
|
||||
TRAIN_PARALLEL_MODE=none
|
||||
fi
|
||||
fi
|
||||
[[ "${TRAIN_PARALLEL_MODE}" =~ ^(none|ddp|fsdp)$ ]] || die "Invalid parallel mode: ${TRAIN_PARALLEL_MODE}"
|
||||
if [[ "${TRAIN_PARALLEL_MODE}" == "none" ]] && (( TRAIN_GPU_COUNT != 1 )); then
|
||||
die "Parallel mode none requires exactly one GPU"
|
||||
fi
|
||||
if [[ "${TRAIN_PARALLEL_MODE}" != "none" ]] && (( TRAIN_GPU_COUNT < 2 )); then
|
||||
die "Parallel mode ${TRAIN_PARALLEL_MODE} requires at least two GPUs"
|
||||
fi
|
||||
if [[ -n "${TRAIN_CONFIG}" ]]; then
|
||||
[[ -f "${TRAIN_CONFIG}" ]] || die "Training config not found: ${TRAIN_CONFIG}"
|
||||
fi
|
||||
@@ -36,11 +51,7 @@ if [[ -n "${TRAIN_CONFIG}" ]]; then
|
||||
train_args+=(--config "${TRAIN_CONFIG}")
|
||||
fi
|
||||
|
||||
if (( TRAIN_GPU_COUNT > 1 )); then
|
||||
train_args+=(--parallel_mode ddp)
|
||||
else
|
||||
train_args+=(--parallel_mode none)
|
||||
fi
|
||||
train_args+=(--parallel_mode "${TRAIN_PARALLEL_MODE}")
|
||||
|
||||
if [[ -n "${latest_checkpoint}" ]]; then
|
||||
log_info "Resuming ${TRAIN_JOB_NAME} from ${latest_checkpoint}"
|
||||
@@ -52,7 +63,7 @@ else
|
||||
train_args+=(--param_path "${BASE_MODEL}")
|
||||
fi
|
||||
|
||||
log_info "GPUs=${TRAIN_GPU_COUNT}, checkpoints=${CHECKPOINT_DIR}"
|
||||
log_info "GPUs=${TRAIN_GPU_COUNT}, parallel=${TRAIN_PARALLEL_MODE}, checkpoints=${CHECKPOINT_DIR}"
|
||||
|
||||
# Replace the shell so the container init forwards SIGTERM to the trainer.
|
||||
exec "${train_args[@]}" "$@"
|
||||
|
||||
@@ -0,0 +1,152 @@
|
||||
"""Parse the host-side runtime section of a training configuration."""
|
||||
|
||||
import argparse
|
||||
import math
|
||||
import re
|
||||
import shlex
|
||||
from pathlib import Path
|
||||
|
||||
import yaml
|
||||
|
||||
ENV_NAME = re.compile(r"^[A-Za-z_][A-Za-z0-9_]*$")
|
||||
PARALLEL_MODES = {"auto", "none", "ddp", "fsdp"}
|
||||
|
||||
|
||||
def _mapping(value, name: str) -> dict:
|
||||
if value is None:
|
||||
return {}
|
||||
if not isinstance(value, dict):
|
||||
raise ValueError(f"runtime.{name} must be a mapping")
|
||||
return value
|
||||
|
||||
|
||||
def _path(value, name: str, config_dir: Path) -> str:
|
||||
if not isinstance(value, str) or not value.strip():
|
||||
raise ValueError(f"runtime.paths.{name} is required")
|
||||
path = Path(value).expanduser()
|
||||
if not path.is_absolute():
|
||||
path = config_dir / path
|
||||
return str(path.resolve())
|
||||
|
||||
|
||||
def load_runtime(config_path: str) -> dict[str, str]:
|
||||
path = Path(config_path).resolve()
|
||||
with path.open(encoding="utf-8") as file:
|
||||
config = yaml.safe_load(file) or {}
|
||||
if not isinstance(config, dict):
|
||||
raise ValueError("training configuration must be a mapping")
|
||||
|
||||
runtime = _mapping(config.get("runtime"), "runtime")
|
||||
if not runtime:
|
||||
raise ValueError("top-level runtime section is required")
|
||||
paths = _mapping(runtime.get("paths"), "paths")
|
||||
gpu = _mapping(runtime.get("gpu"), "gpu")
|
||||
container = _mapping(runtime.get("container"), "container")
|
||||
environment = _mapping(runtime.get("environment"), "environment")
|
||||
|
||||
job_name = runtime.get("job_name")
|
||||
if not isinstance(job_name, str) or not re.fullmatch(
|
||||
r"[A-Za-z0-9][A-Za-z0-9._-]*", job_name
|
||||
):
|
||||
raise ValueError(
|
||||
"runtime.job_name must use letters, numbers, dot, underscore, or dash"
|
||||
)
|
||||
|
||||
devices = gpu.get("devices", "all")
|
||||
if devices == "all":
|
||||
gpu_count = "all"
|
||||
visible_devices = ""
|
||||
elif isinstance(devices, list) and devices:
|
||||
normalized = []
|
||||
for device in devices:
|
||||
text = str(device)
|
||||
if not text.isdigit():
|
||||
raise ValueError(
|
||||
"runtime.gpu.devices entries must be non-negative integers"
|
||||
)
|
||||
normalized.append(text)
|
||||
if len(set(normalized)) != len(normalized):
|
||||
raise ValueError("runtime.gpu.devices must not contain duplicates")
|
||||
gpu_count = str(len(normalized))
|
||||
visible_devices = ",".join(normalized)
|
||||
else:
|
||||
raise ValueError("runtime.gpu.devices must be 'all' or a non-empty list")
|
||||
|
||||
parallel_mode = str(gpu.get("parallel_mode", "auto"))
|
||||
if parallel_mode not in PARALLEL_MODES:
|
||||
raise ValueError("runtime.gpu.parallel_mode must be auto, none, ddp, or fsdp")
|
||||
if gpu_count != "all":
|
||||
count = int(gpu_count)
|
||||
if parallel_mode == "none" and count != 1:
|
||||
raise ValueError("parallel_mode none requires exactly one GPU")
|
||||
if parallel_mode in {"ddp", "fsdp"} and count < 2:
|
||||
raise ValueError(
|
||||
f"parallel_mode {parallel_mode} requires at least two GPUs"
|
||||
)
|
||||
|
||||
max_hours = container.get("max_duration_hours", 0)
|
||||
try:
|
||||
max_seconds = math.ceil(float(max_hours) * 3600) if max_hours else 0
|
||||
except (TypeError, ValueError) as exc:
|
||||
raise ValueError(
|
||||
"runtime.container.max_duration_hours must be a number"
|
||||
) from exc
|
||||
if max_seconds < 0:
|
||||
raise ValueError("runtime.container.max_duration_hours must not be negative")
|
||||
|
||||
values = {
|
||||
"TRAIN_JOB_NAME": job_name,
|
||||
"TRAIN_DATA_DIR": _path(paths.get("data"), "data", path.parent),
|
||||
"TRAIN_MODEL_DIR": _path(paths.get("model"), "model", path.parent),
|
||||
"TRAIN_CHECKPOINT_DIR": _path(
|
||||
paths.get("checkpoints"), "checkpoints", path.parent
|
||||
),
|
||||
"TRAIN_GPU_COUNT": gpu_count,
|
||||
"CUDA_VISIBLE_DEVICES": visible_devices,
|
||||
"TRAIN_PARALLEL_MODE": parallel_mode,
|
||||
"CUDA_TAG": str(container.get("cuda_tag", "cu128")),
|
||||
"TRAIN_IPC_MODE": str(container.get("ipc", "host")),
|
||||
"TRAIN_STOP_GRACE_PERIOD": str(container.get("stop_grace_period", "10m")),
|
||||
"TRAIN_STOP_TIMEOUT": str(container.get("stop_timeout_seconds", 600)),
|
||||
"CHECKPOINT_KEEP_LAST": str(container.get("checkpoint_keep_last", 5)),
|
||||
"TRAIN_MAX_DURATION_SECONDS": str(max_seconds),
|
||||
}
|
||||
|
||||
for name, value in environment.items():
|
||||
if not isinstance(name, str) or not ENV_NAME.fullmatch(name):
|
||||
raise ValueError(f"invalid runtime.environment name: {name!r}")
|
||||
if value is not None and not isinstance(value, (str, int, float, bool)):
|
||||
raise ValueError(f"runtime.environment.{name} must be a scalar")
|
||||
values["environment"] = environment
|
||||
return values
|
||||
|
||||
|
||||
def shell_exports(runtime: dict[str, str]) -> str:
|
||||
return "\n".join(
|
||||
f"export {name}={shlex.quote(value)}"
|
||||
for name, value in runtime.items()
|
||||
if name != "environment"
|
||||
)
|
||||
|
||||
|
||||
def main() -> None:
|
||||
parser = argparse.ArgumentParser()
|
||||
parser.add_argument("command", choices=("exports", "environment"))
|
||||
parser.add_argument("config")
|
||||
args = parser.parse_args()
|
||||
|
||||
try:
|
||||
runtime = load_runtime(args.config)
|
||||
except (OSError, ValueError, yaml.YAMLError) as exc:
|
||||
parser.error(str(exc))
|
||||
|
||||
if args.command == "exports":
|
||||
print(shell_exports(runtime))
|
||||
return
|
||||
for name, value in runtime["environment"].items():
|
||||
rendered = "" if value is None else str(value)
|
||||
print(f"{name}={rendered}", end="\0")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
+137
-241
@@ -4,7 +4,6 @@ set -euo pipefail
|
||||
ROOT_DIR="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")/.." && pwd)"
|
||||
source "${ROOT_DIR}/scripts/docker/lib/train-common.sh"
|
||||
|
||||
ENV_FILE="${TRAIN_ENV_FILE:-${ROOT_DIR}/.env.train}"
|
||||
COMPOSE_BASE=(
|
||||
docker compose
|
||||
--project-directory "${ROOT_DIR}"
|
||||
@@ -14,50 +13,28 @@ COMPOSE_BASE=(
|
||||
|
||||
usage() {
|
||||
cat <<'EOF'
|
||||
Usage: scripts/train.sh <command> [options]
|
||||
Usage: scripts/train.sh <command> [CONFIG] [options]
|
||||
|
||||
CONFIG defaults to ./train.yaml. The same file declares host runtime settings
|
||||
under `runtime:` and trainer settings under model/data/parallel/training/ckpt/log.
|
||||
|
||||
Commands:
|
||||
init Create local directories and .env.train
|
||||
preflight Validate Docker, paths, GPU settings, and Compose
|
||||
build Build the trainer image
|
||||
start [--foreground] [-- ARGS...] Start or resume training
|
||||
stop Gracefully stop and checkpoint training
|
||||
restart Stop, then start training
|
||||
logs Follow trainer logs
|
||||
status Show container and latest checkpoint status
|
||||
latest Print the latest complete checkpoint path
|
||||
list List all complete checkpoints
|
||||
clean [--keep N] Preview old checkpoint removal
|
||||
clean --force Remove old checkpoints after previewing
|
||||
|
||||
Environment:
|
||||
TRAIN_ENV_FILE Env file path (default: .env.train)
|
||||
TRAIN_CONFIG_FILE Optional host YAML mounted only when the job starts
|
||||
|
||||
Training arguments come from an externally mounted TRAIN_CONFIG or ARGS passed
|
||||
after --. The image does not contain experiment configuration.
|
||||
init [CONFIG] Create runtime directories
|
||||
preflight [CONFIG] Validate Docker, paths, GPUs, and Compose
|
||||
build [CONFIG] Build the trainer image
|
||||
start [CONFIG] [--foreground] [-- ARGS...]
|
||||
Start or resume training
|
||||
stop [CONFIG] Gracefully stop and checkpoint training
|
||||
restart [CONFIG] Stop, then start training
|
||||
logs [CONFIG] Follow trainer logs
|
||||
status [CONFIG] Show container and checkpoint status
|
||||
latest [CONFIG] Print the latest complete checkpoint
|
||||
list [CONFIG] List complete checkpoints
|
||||
clean [CONFIG] [--keep N] [--force]
|
||||
Preview or remove old checkpoints
|
||||
EOF
|
||||
}
|
||||
|
||||
load_env() {
|
||||
if [[ -f "${ENV_FILE}" ]]; then
|
||||
set -a
|
||||
# UID/GID are readonly in bash; compose gets them via ASTRAI_UID/GID in compose()
|
||||
# shellcheck disable=SC1090
|
||||
source <(grep -v -E '^[[:space:]]*(UID|GID)=' "${ENV_FILE}")
|
||||
set +a
|
||||
fi
|
||||
|
||||
TRAIN_JOB_NAME="${TRAIN_JOB_NAME:-astrai-train}"
|
||||
TRAIN_DATA_DIR="${TRAIN_DATA_DIR:-./data}"
|
||||
TRAIN_MODEL_DIR="${TRAIN_MODEL_DIR:-./params}"
|
||||
TRAIN_CHECKPOINT_DIR="${TRAIN_CHECKPOINT_DIR:-./checkpoints}"
|
||||
TRAIN_GPU_COUNT="${TRAIN_GPU_COUNT:-all}"
|
||||
TRAIN_STOP_TIMEOUT="${TRAIN_STOP_TIMEOUT:-600}"
|
||||
|
||||
validate_job_name "${TRAIN_JOB_NAME}"
|
||||
}
|
||||
|
||||
resolve_path() {
|
||||
if [[ "$1" = /* ]]; then
|
||||
printf '%s\n' "$1"
|
||||
@@ -66,203 +43,147 @@ resolve_path() {
|
||||
fi
|
||||
}
|
||||
|
||||
# Read the optional top-level `infra:` section from TRAIN_CONFIG_FILE and
|
||||
# export the host-side variables it overrides (job name, mount paths, GPU
|
||||
# filter). Compose interpolation prefers the shell environment over the
|
||||
# --env-file, so these exports win over .env.train; keys absent from `infra`
|
||||
# fall back to the env file. Requires python3 with PyYAML on the host.
|
||||
load_infra() {
|
||||
local infra_file exports
|
||||
load_config() {
|
||||
CONFIG_FILE="$(resolve_path "$1")"
|
||||
[[ -f "${CONFIG_FILE}" ]] || die "Training config not found: ${CONFIG_FILE}"
|
||||
require_command python3
|
||||
python3 -c 'import yaml' >/dev/null 2>&1 ||
|
||||
die "PyYAML is required on the host (install python3-yaml)"
|
||||
|
||||
[[ -n "${TRAIN_CONFIG_FILE:-}" ]] || return 0
|
||||
infra_file="$(resolve_path "${TRAIN_CONFIG_FILE}")"
|
||||
[[ -f "${infra_file}" ]] || return 0
|
||||
|
||||
if ! command -v python3 >/dev/null 2>&1; then
|
||||
die "TRAIN_CONFIG_FILE is set but python3 is missing; it is needed to read the 'infra' section"
|
||||
fi
|
||||
if ! python3 -c 'import yaml' >/dev/null 2>&1; then
|
||||
die "TRAIN_CONFIG_FILE is set but PyYAML is missing on the host (install python3-yaml)"
|
||||
fi
|
||||
|
||||
exports="$(TRAIN_INFRA_FILE="${infra_file}" python3 - <<'PYEOF'
|
||||
import os
|
||||
import shlex
|
||||
import sys
|
||||
import yaml
|
||||
|
||||
path = os.environ["TRAIN_INFRA_FILE"]
|
||||
try:
|
||||
with open(path) as f:
|
||||
cfg = yaml.safe_load(f) or {}
|
||||
except Exception as exc:
|
||||
print(f"failed to parse {path}: {exc}", file=sys.stderr)
|
||||
sys.exit(1)
|
||||
|
||||
infra = cfg.get("infra") or {}
|
||||
if not isinstance(infra, dict):
|
||||
print(f"the 'infra' section in {path} must be a mapping", file=sys.stderr)
|
||||
sys.exit(1)
|
||||
|
||||
mapping = {
|
||||
"job_name": "TRAIN_JOB_NAME",
|
||||
"data_dir": "TRAIN_DATA_DIR",
|
||||
"model_dir": "TRAIN_MODEL_DIR",
|
||||
"checkpoint_dir": "TRAIN_CHECKPOINT_DIR",
|
||||
"gpu_count": "TRAIN_GPU_COUNT",
|
||||
"cuda_visible_devices": "CUDA_VISIBLE_DEVICES",
|
||||
}
|
||||
for key, env_name in mapping.items():
|
||||
if key in infra:
|
||||
print(f"export {env_name}={shlex.quote(str(infra[key]))}")
|
||||
PYEOF
|
||||
)"
|
||||
if [[ -n "${exports}" ]]; then
|
||||
local exports
|
||||
exports="$(python3 "${ROOT_DIR}/scripts/tools/train_runtime.py" exports "${CONFIG_FILE}")" ||
|
||||
die "Failed to load runtime configuration"
|
||||
eval "${exports}"
|
||||
log_info "Applied infra overrides from ${infra_file}"
|
||||
fi
|
||||
|
||||
validate_job_name "${TRAIN_JOB_NAME}"
|
||||
}
|
||||
|
||||
checkpoint_dir() {
|
||||
printf '%s/%s\n' "$(resolve_path "${TRAIN_CHECKPOINT_DIR}")" "${TRAIN_JOB_NAME}"
|
||||
compose() {
|
||||
ASTRAI_UID="$(id -u)" ASTRAI_GID="$(id -g)" "${COMPOSE_BASE[@]}" "$@"
|
||||
}
|
||||
|
||||
compose() {
|
||||
local -a command=("${COMPOSE_BASE[@]}")
|
||||
checkpoint_dir() {
|
||||
printf '%s/%s\n' "${TRAIN_CHECKPOINT_DIR}" "${TRAIN_JOB_NAME}"
|
||||
}
|
||||
|
||||
if [[ -f "${ENV_FILE}" ]]; then
|
||||
command+=(--env-file "${ENV_FILE}")
|
||||
container_name() {
|
||||
printf 'astrai-trainer-%s\n' "${TRAIN_JOB_NAME}"
|
||||
}
|
||||
|
||||
timer_pid_file() {
|
||||
printf '/tmp/astrai-timer-%s.pid\n' "${TRAIN_JOB_NAME}"
|
||||
}
|
||||
|
||||
timer_log_file() {
|
||||
printf '/tmp/astrai-timer-%s.log\n' "${TRAIN_JOB_NAME}"
|
||||
}
|
||||
|
||||
cancel_timer() {
|
||||
local pid_file pid
|
||||
pid_file="$(timer_pid_file)"
|
||||
[[ -f "${pid_file}" ]] || return 0
|
||||
pid="$(<"${pid_file}")"
|
||||
if [[ "${pid}" =~ ^[1-9][0-9]*$ ]] && kill -0 "${pid}" 2>/dev/null; then
|
||||
kill "${pid}" 2>/dev/null || true
|
||||
fi
|
||||
rm -f -- "${pid_file}"
|
||||
}
|
||||
|
||||
# Inject the host user into compose so container processes share the
|
||||
# checkpoint directory ownership (bash UID/GID are readonly).
|
||||
ASTRAI_UID="$(id -u)" ASTRAI_GID="$(id -g)" "${command[@]}" "$@"
|
||||
schedule_timer() {
|
||||
(( TRAIN_MAX_DURATION_SECONDS > 0 )) || return 0
|
||||
cancel_timer
|
||||
local pid_file log_file
|
||||
pid_file="$(timer_pid_file)"
|
||||
log_file="$(timer_log_file)"
|
||||
(
|
||||
sleep "${TRAIN_MAX_DURATION_SECONDS}"
|
||||
"${ROOT_DIR}/scripts/train.sh" stop "${CONFIG_FILE}" --from-timer
|
||||
) >"${log_file}" 2>&1 &
|
||||
printf '%s\n' "$!" >"${pid_file}"
|
||||
log_info "Automatic stop scheduled in ${TRAIN_MAX_DURATION_SECONDS}s"
|
||||
}
|
||||
|
||||
init_environment() {
|
||||
local data_dir model_dir checkpoints_dir
|
||||
|
||||
data_dir="$(resolve_path "${TRAIN_DATA_DIR}")"
|
||||
model_dir="$(resolve_path "${TRAIN_MODEL_DIR}")"
|
||||
checkpoints_dir="$(resolve_path "${TRAIN_CHECKPOINT_DIR}")"
|
||||
mkdir -p "${data_dir}" "${model_dir}" "${checkpoints_dir}"
|
||||
|
||||
if [[ ! -f "${ENV_FILE}" ]]; then
|
||||
cat >"${ENV_FILE}" <<'EOF'
|
||||
TRAIN_JOB_NAME=astrai-train
|
||||
TRAIN_DATA_DIR=./data
|
||||
TRAIN_MODEL_DIR=./params
|
||||
TRAIN_CHECKPOINT_DIR=./checkpoints
|
||||
TRAIN_CONFIG_FILE=
|
||||
TRAIN_GPU_COUNT=all
|
||||
# CUDA_VISIBLE_DEVICES=0,1
|
||||
# TRAIN_* vars above can be overridden per-job via the top-level `infra:`
|
||||
# section in TRAIN_CONFIG_FILE (see docs/developer/docker-training.md).
|
||||
CUDA_TAG=cu128
|
||||
TRAIN_IPC_MODE=host
|
||||
TRAIN_STOP_GRACE_PERIOD=10m
|
||||
TRAIN_STOP_TIMEOUT=600
|
||||
CHECKPOINT_KEEP_LAST=5
|
||||
EOF
|
||||
log_info "Created ${ENV_FILE}"
|
||||
else
|
||||
log_info "Keeping existing ${ENV_FILE}"
|
||||
fi
|
||||
log_info "Data: ${data_dir}"
|
||||
log_info "Model: ${model_dir}"
|
||||
log_info "Checkpoints: ${checkpoints_dir}"
|
||||
mkdir -p "${TRAIN_DATA_DIR}" "${TRAIN_MODEL_DIR}" "${TRAIN_CHECKPOINT_DIR}"
|
||||
log_info "Data: ${TRAIN_DATA_DIR}"
|
||||
log_info "Model: ${TRAIN_MODEL_DIR}"
|
||||
log_info "Checkpoints: ${TRAIN_CHECKPOINT_DIR}"
|
||||
}
|
||||
|
||||
preflight() {
|
||||
local data_dir model_dir checkpoints_dir config_file latest visible_count
|
||||
|
||||
local latest visible_count
|
||||
require_command docker
|
||||
docker info >/dev/null 2>&1 || die "Docker daemon is unavailable"
|
||||
[[ "${TRAIN_GPU_COUNT}" == "all" || "${TRAIN_GPU_COUNT}" =~ ^[1-9][0-9]*$ ]] ||
|
||||
die "TRAIN_GPU_COUNT must be 'all' or a positive integer"
|
||||
[[ -d "${TRAIN_DATA_DIR}" ]] || die "Training data directory not found: ${TRAIN_DATA_DIR}"
|
||||
mkdir -p "$(checkpoint_dir)"
|
||||
[[ -w "$(checkpoint_dir)" ]] || die "Checkpoint directory is not writable: $(checkpoint_dir)"
|
||||
|
||||
data_dir="$(resolve_path "${TRAIN_DATA_DIR}")"
|
||||
model_dir="$(resolve_path "${TRAIN_MODEL_DIR}")"
|
||||
checkpoints_dir="$(resolve_path "${TRAIN_CHECKPOINT_DIR}")"
|
||||
[[ -d "${data_dir}" ]] || die "Training data directory not found: ${data_dir}"
|
||||
mkdir -p "${checkpoints_dir}/${TRAIN_JOB_NAME}"
|
||||
[[ -w "${checkpoints_dir}/${TRAIN_JOB_NAME}" ]] || die "Checkpoint directory is not writable"
|
||||
|
||||
if [[ -n "${TRAIN_CONFIG_FILE:-}" ]]; then
|
||||
config_file="$(resolve_path "${TRAIN_CONFIG_FILE}")"
|
||||
[[ -f "${config_file}" ]] || die "Training config not found: ${config_file}"
|
||||
fi
|
||||
|
||||
latest="$(find_latest_checkpoint "${checkpoints_dir}/${TRAIN_JOB_NAME}" || true)"
|
||||
latest="$(find_latest_checkpoint "$(checkpoint_dir)" || true)"
|
||||
if [[ -z "${latest}" ]]; then
|
||||
[[ -s "${model_dir}/config.json" ]] || die "Model config not found: ${model_dir}/config.json"
|
||||
[[ -s "${model_dir}/model.safetensors" ]] || die "Model weights not found: ${model_dir}/model.safetensors"
|
||||
[[ -s "${TRAIN_MODEL_DIR}/config.json" ]] ||
|
||||
die "Model config not found: ${TRAIN_MODEL_DIR}/config.json"
|
||||
[[ -s "${TRAIN_MODEL_DIR}/model.safetensors" ]] ||
|
||||
die "Model weights not found: ${TRAIN_MODEL_DIR}/model.safetensors"
|
||||
else
|
||||
log_info "Resume candidate: ${latest}"
|
||||
fi
|
||||
|
||||
if [[ -n "${CUDA_VISIBLE_DEVICES:-}" && "${TRAIN_GPU_COUNT}" != "all" ]]; then
|
||||
if [[ -n "${CUDA_VISIBLE_DEVICES}" ]]; then
|
||||
IFS=',' read -r -a visible_gpus <<<"${CUDA_VISIBLE_DEVICES}"
|
||||
visible_count="${#visible_gpus[@]}"
|
||||
(( visible_count == TRAIN_GPU_COUNT )) ||
|
||||
die "TRAIN_GPU_COUNT=${TRAIN_GPU_COUNT}, but CUDA_VISIBLE_DEVICES exposes ${visible_count} GPU(s)"
|
||||
die "Configured GPU count and visible device list disagree"
|
||||
fi
|
||||
|
||||
compose config --quiet
|
||||
log_info "Preflight passed for ${TRAIN_JOB_NAME} (GPU request: ${TRAIN_GPU_COUNT})"
|
||||
log_info "Preflight passed for ${TRAIN_JOB_NAME} (GPU request: ${TRAIN_GPU_COUNT}, parallel: ${TRAIN_PARALLEL_MODE})"
|
||||
}
|
||||
|
||||
runtime_environment_args() {
|
||||
RUNTIME_ENV_ARGS=()
|
||||
local pair
|
||||
while IFS= read -r -d '' pair; do
|
||||
RUNTIME_ENV_ARGS+=(--env "${pair}")
|
||||
done < <(python3 "${ROOT_DIR}/scripts/tools/train_runtime.py" environment "${CONFIG_FILE}")
|
||||
}
|
||||
|
||||
start_training() {
|
||||
local foreground="$1"
|
||||
local config_file container running
|
||||
local -a run_options=()
|
||||
shift
|
||||
|
||||
local container running
|
||||
local -a run_options
|
||||
preflight
|
||||
if [[ -n "${TRAIN_CONFIG_FILE:-}" ]]; then
|
||||
config_file="$(resolve_path "${TRAIN_CONFIG_FILE}")"
|
||||
run_options+=(
|
||||
--volume "${config_file}:/run/astrai/train.yaml:ro"
|
||||
--env TRAIN_CONFIG=/run/astrai/train.yaml
|
||||
)
|
||||
elif [[ -z "${TRAIN_CONFIG:-}" && $# -eq 0 ]]; then
|
||||
die "Set TRAIN_CONFIG_FILE or pass complete trainer arguments after --"
|
||||
fi
|
||||
|
||||
container="astrai-trainer-${TRAIN_JOB_NAME}"
|
||||
runtime_environment_args
|
||||
container="$(container_name)"
|
||||
running="$(docker inspect --format '{{.State.Running}}' "${container}" 2>/dev/null || true)"
|
||||
[[ "${running}" != "true" ]] || die "Trainer is already running: ${container}"
|
||||
docker rm "${container}" >/dev/null 2>&1 || true
|
||||
|
||||
run_options=(
|
||||
--volume "${CONFIG_FILE}:/run/astrai/train.yaml:ro"
|
||||
--env TRAIN_CONFIG=/run/astrai/train.yaml
|
||||
"${RUNTIME_ENV_ARGS[@]}"
|
||||
)
|
||||
if [[ "${foreground}" == "true" ]]; then
|
||||
compose run --build --rm "${run_options[@]}" trainer "$@"
|
||||
else
|
||||
compose run -d --build --name "${container}" \
|
||||
"${run_options[@]}" trainer "$@"
|
||||
log_info "Training started; run scripts/train.sh logs to follow it"
|
||||
compose run -d --build --name "${container}" "${run_options[@]}" trainer "$@"
|
||||
schedule_timer
|
||||
log_info "Training started; run scripts/train.sh logs ${CONFIG_FILE} to follow it"
|
||||
fi
|
||||
}
|
||||
|
||||
stop_training() {
|
||||
local from_timer="$1"
|
||||
[[ "${from_timer}" == "true" ]] || cancel_timer
|
||||
log_info "Stopping trainer with ${TRAIN_STOP_TIMEOUT}s grace period"
|
||||
docker stop --timeout "${TRAIN_STOP_TIMEOUT}" "astrai-trainer-${TRAIN_JOB_NAME}" >/dev/null 2>&1 ||
|
||||
docker stop --timeout "${TRAIN_STOP_TIMEOUT}" "$(container_name)" >/dev/null 2>&1 ||
|
||||
log_warn "Trainer container is not running"
|
||||
}
|
||||
|
||||
restart_training() {
|
||||
local container="astrai-trainer-${TRAIN_JOB_NAME}"
|
||||
|
||||
docker inspect "${container}" >/dev/null 2>&1 ||
|
||||
die "Trainer container not found; use start with a config or CLI arguments first"
|
||||
log_info "Restarting trainer with ${TRAIN_STOP_TIMEOUT}s grace period"
|
||||
docker restart --timeout "${TRAIN_STOP_TIMEOUT}" "${container}" >/dev/null
|
||||
[[ "${from_timer}" != "true" ]] || rm -f -- "$(timer_pid_file)"
|
||||
}
|
||||
|
||||
show_status() {
|
||||
local latest
|
||||
|
||||
docker ps -a --filter "name=^/astrai-trainer-${TRAIN_JOB_NAME}$"
|
||||
docker ps -a --filter "name=^/$(container_name)$"
|
||||
latest="$(find_latest_checkpoint "$(checkpoint_dir)" || true)"
|
||||
if [[ -n "${latest}" ]]; then
|
||||
log_info "Latest checkpoint: ${latest}"
|
||||
@@ -274,7 +195,6 @@ show_status() {
|
||||
clean_checkpoints() {
|
||||
local keep="$1" force="$2" dir count remove_count index path
|
||||
local -a checkpoints=()
|
||||
|
||||
[[ "${keep}" =~ ^[1-9][0-9]*$ ]] || die "--keep must be a positive integer"
|
||||
dir="$(checkpoint_dir)"
|
||||
while IFS= read -r line; do
|
||||
@@ -287,7 +207,6 @@ clean_checkpoints() {
|
||||
log_info "Nothing to clean; ${count} complete checkpoint(s), keeping ${keep}"
|
||||
return
|
||||
fi
|
||||
|
||||
for ((index = 0; index < remove_count; index++)); do
|
||||
path="${checkpoints[index]}"
|
||||
if [[ "${force}" == "true" ]]; then
|
||||
@@ -301,58 +220,49 @@ clean_checkpoints() {
|
||||
}
|
||||
|
||||
main() {
|
||||
local command="${1:-}" foreground=false keep="${CHECKPOINT_KEEP_LAST:-5}" force=false
|
||||
local command="${1:-}" config="${TRAIN_CONFIG_FILE:-${ROOT_DIR}/train.yaml}"
|
||||
local foreground=false keep force=false from_timer=false
|
||||
local -a train_args=()
|
||||
[[ -n "${command}" ]] || { usage; exit 1; }
|
||||
shift || true
|
||||
load_env
|
||||
load_infra
|
||||
|
||||
if [[ "${command}" =~ ^(help|-h|--help)$ ]]; then
|
||||
usage
|
||||
return
|
||||
fi
|
||||
|
||||
if [[ $# -gt 0 && "$1" != --* ]]; then
|
||||
config="$1"
|
||||
shift
|
||||
fi
|
||||
load_config "${config}"
|
||||
keep="${CHECKPOINT_KEEP_LAST}"
|
||||
|
||||
case "${command}" in
|
||||
init)
|
||||
init_environment
|
||||
;;
|
||||
preflight)
|
||||
preflight
|
||||
;;
|
||||
build)
|
||||
preflight
|
||||
compose build trainer
|
||||
;;
|
||||
init) init_environment ;;
|
||||
preflight) preflight ;;
|
||||
build) preflight; compose build trainer ;;
|
||||
start)
|
||||
while [[ $# -gt 0 ]]; do
|
||||
case "$1" in
|
||||
--foreground)
|
||||
foreground=true
|
||||
shift
|
||||
;;
|
||||
--)
|
||||
shift
|
||||
train_args=("$@")
|
||||
break
|
||||
;;
|
||||
*)
|
||||
die "Unknown start option: $1 (put trainer arguments after --)"
|
||||
;;
|
||||
--foreground) foreground=true; shift ;;
|
||||
--) shift; train_args=("$@"); break ;;
|
||||
*) die "Unknown start option: $1 (put trainer arguments after --)" ;;
|
||||
esac
|
||||
done
|
||||
start_training "${foreground}" "${train_args[@]}"
|
||||
;;
|
||||
stop)
|
||||
stop_training
|
||||
[[ "${1:-}" != "--from-timer" ]] || from_timer=true
|
||||
stop_training "${from_timer}"
|
||||
;;
|
||||
restart)
|
||||
restart_training
|
||||
;;
|
||||
logs)
|
||||
docker logs -f --tail "${TRAIN_LOG_TAIL:-200}" "astrai-trainer-${TRAIN_JOB_NAME}"
|
||||
;;
|
||||
status)
|
||||
show_status
|
||||
;;
|
||||
latest)
|
||||
find_latest_checkpoint "$(checkpoint_dir)" || die "No complete checkpoint found"
|
||||
stop_training false
|
||||
start_training false
|
||||
;;
|
||||
logs) docker logs -f --tail "${TRAIN_LOG_TAIL:-200}" "$(container_name)" ;;
|
||||
status) show_status ;;
|
||||
latest) find_latest_checkpoint "$(checkpoint_dir)" || die "No complete checkpoint found" ;;
|
||||
list)
|
||||
list_complete_checkpoints "$(checkpoint_dir)" | while read -r _epoch _step path; do
|
||||
printf '%s\n' "${path}"
|
||||
@@ -361,28 +271,14 @@ main() {
|
||||
clean)
|
||||
while [[ $# -gt 0 ]]; do
|
||||
case "$1" in
|
||||
--keep)
|
||||
[[ $# -ge 2 ]] || die "--keep requires a value"
|
||||
keep="$2"
|
||||
shift 2
|
||||
;;
|
||||
--force)
|
||||
force=true
|
||||
shift
|
||||
;;
|
||||
*)
|
||||
die "Unknown clean option: $1"
|
||||
;;
|
||||
--keep) [[ $# -ge 2 ]] || die "--keep requires a value"; keep="$2"; shift 2 ;;
|
||||
--force) force=true; shift ;;
|
||||
*) die "Unknown clean option: $1" ;;
|
||||
esac
|
||||
done
|
||||
clean_checkpoints "${keep}" "${force}"
|
||||
;;
|
||||
help|-h|--help)
|
||||
usage
|
||||
;;
|
||||
*)
|
||||
die "Unknown command: ${command}"
|
||||
;;
|
||||
*) die "Unknown command: ${command}" ;;
|
||||
esac
|
||||
}
|
||||
|
||||
|
||||
Reference in New Issue
Block a user