From 00c2c80c8f7c3b8c40f3e030d9a2fcc750792d75 Mon Sep 17 00:00:00 2001 From: ViperEkura <3081035982@qq.com> Date: Wed, 19 Aug 2026 14:04:35 +0800 Subject: [PATCH] feat: unify Docker training configuration in YAML --- docker-compose.yml | 3 +- docs/developer/docker-training.md | 145 ++++++++--- scripts/docker/train-entrypoint.sh | 23 +- scripts/tools/train_runtime.py | 152 ++++++++++++ scripts/train.sh | 380 +++++++++++------------------ 5 files changed, 414 insertions(+), 289 deletions(-) create mode 100644 scripts/tools/train_runtime.py diff --git a/docker-compose.yml b/docker-compose.yml index b7809c4..f61d9bd 100644 --- a/docker-compose.yml +++ b/docker-compose.yml @@ -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} diff --git a/docs/developer/docker-training.md b/docs/developer/docker-training.md index 1bdc7c7..b5e4f93 100644 --- a/docs/developer/docker-training.md +++ b/docs/developer/docker-training.md @@ -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//epoch__step_`. 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. diff --git a/scripts/docker/train-entrypoint.sh b/scripts/docker/train-entrypoint.sh index af590d0..8d65a08 100755 --- a/scripts/docker/train-entrypoint.sh +++ b/scripts/docker/train-entrypoint.sh @@ -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[@]}" "$@" diff --git a/scripts/tools/train_runtime.py b/scripts/tools/train_runtime.py new file mode 100644 index 0000000..325c0c2 --- /dev/null +++ b/scripts/tools/train_runtime.py @@ -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() diff --git a/scripts/train.sh b/scripts/train.sh index 467d0e3..539ee7a 100755 --- a/scripts/train.sh +++ b/scripts/train.sh @@ -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 [options] +Usage: scripts/train.sh [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 - - [[ -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 - eval "${exports}" - log_info "Applied infra overrides from ${infra_file}" - fi +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)" + local exports + exports="$(python3 "${ROOT_DIR}/scripts/tools/train_runtime.py" exports "${CONFIG_FILE}")" || + die "Failed to load runtime configuration" + eval "${exports}" 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 }