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}
|
- BASE_MODEL=${BASE_MODEL:-/models/base}
|
||||||
- CHECKPOINT_ROOT=/checkpoints
|
- CHECKPOINT_ROOT=/checkpoints
|
||||||
- TRAIN_GPU_COUNT=${TRAIN_GPU_COUNT:-all}
|
- TRAIN_GPU_COUNT=${TRAIN_GPU_COUNT:-all}
|
||||||
|
- TRAIN_PARALLEL_MODE=${TRAIN_PARALLEL_MODE:-auto}
|
||||||
- CUDA_VISIBLE_DEVICES
|
- CUDA_VISIBLE_DEVICES
|
||||||
- NCCL_P2P_DISABLE
|
|
||||||
- NCCL_NET_GDR_LEVEL
|
|
||||||
entrypoint: ["bash", "/app/scripts/docker/train-entrypoint.sh"]
|
entrypoint: ["bash", "/app/scripts/docker/train-entrypoint.sh"]
|
||||||
ipc: ${TRAIN_IPC_MODE:-host}
|
ipc: ${TRAIN_IPC_MODE:-host}
|
||||||
stop_grace_period: ${TRAIN_STOP_GRACE_PERIOD:-10m}
|
stop_grace_period: ${TRAIN_STOP_GRACE_PERIOD:-10m}
|
||||||
|
|||||||
@@ -1,58 +1,125 @@
|
|||||||
# Containerized Training Deployment
|
# 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
|
## Architecture
|
||||||
|
|
||||||
```
|
```text
|
||||||
scripts/train.sh host-side CLI: env loading, preflight, compose wrapper, lifecycle
|
train.yaml
|
||||||
└── docker-compose.yml GPU passthrough, mounts, in-container env vars, entrypoint
|
├── runtime parsed on the host before Docker starts
|
||||||
└── train-entrypoint.sh GPU-count resolution, parallel-mode selection, auto-resume
|
└── 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
|
└── train.py --config /run/astrai/train.yaml
|
||||||
```
|
```
|
||||||
|
|
||||||
| Layer | Responsible for | NOT responsible for |
|
The two parsers deliberately own different sections. `scripts/tools/train_runtime.py`
|
||||||
|-------|-----------------|---------------------|
|
reads only `runtime`; `scripts/tools/train.py` reads only
|
||||||
| `train.sh` | host paths, `.env.train` + `infra:` YAML overrides, preflight, lifecycle | training args, GPU selection, parallel mode |
|
`model/data/parallel/training/ckpt/log`. Explicit trainer arguments after `--`
|
||||||
| compose | GPU passthrough, mounts, in-container env (NCCL) | training args (beyond `TRAIN_*` forwarding) |
|
override training YAML values.
|
||||||
| 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 |
|
|
||||||
|
|
||||||
## Path Conventions
|
## Runtime Schema
|
||||||
|
|
||||||
| Host var | Container | Perm | Purpose |
|
```yaml
|
||||||
|---|---|---|---|
|
runtime:
|
||||||
| `TRAIN_DATA_DIR` | `/data` | ro | dataset (`data_root_path` must be `/data`) |
|
job_name: astrai-train
|
||||||
| `TRAIN_MODEL_DIR` | `/models/base` | ro | base model (`config.json` + `model.safetensors`) |
|
paths:
|
||||||
| `TRAIN_CHECKPOINT_DIR` | `/checkpoints` | rw | checkpoint root, per-`TRAIN_JOB_NAME` subdirs |
|
data: ./data
|
||||||
| `TRAIN_CONFIG_FILE` | `/run/astrai/train.yaml` | ro | training YAML (mounted only on `start`) |
|
model: ./params
|
||||||
| code | `/app` | image | **not a mount**; rebuild image for code changes |
|
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`).
|
## Fixed Container Paths
|
||||||
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).
|
| Runtime path | Container path | Access |
|
||||||
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.
|
| `runtime.paths.data` | `/data` | read-only |
|
||||||
6. **Checkpoint complete =** `meta.json + config.json + model.safetensors + optimizer.pt + scheduler.pt`; `start` auto-resumes the latest complete one.
|
| `runtime.paths.model` | `/models/base` | read-only |
|
||||||
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.
|
| `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
|
## Operations
|
||||||
|
|
||||||
|
The config argument defaults to `./train.yaml`:
|
||||||
|
|
||||||
```bash
|
```bash
|
||||||
bash scripts/train.sh init # first run: dirs + .env.train (edit per machine)
|
bash scripts/train.sh init [CONFIG]
|
||||||
bash scripts/train.sh preflight # validate Docker/paths/GPU/model/YAML/compose
|
bash scripts/train.sh preflight [CONFIG]
|
||||||
bash scripts/train.sh start # build + start in background (auto-resume)
|
bash scripts/train.sh start [CONFIG]
|
||||||
bash scripts/train.sh start --foreground -- --dry-run # print plan only
|
bash scripts/train.sh start [CONFIG] --foreground -- --dry-run
|
||||||
bash scripts/train.sh logs | status | stop | restart
|
bash scripts/train.sh logs [CONFIG]
|
||||||
bash scripts/train.sh clean --keep 5 # prune old checkpoints (--force to delete)
|
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
|
## Checkpoint Recovery
|
||||||
- `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`
|
Checkpoints are stored below
|
||||||
- `scripts/docker/train-entrypoint.sh` — GPU-count resolution, parallel mode, resume
|
`runtime.paths.checkpoints/<job_name>/epoch_<N>_step_<N>`. A checkpoint is
|
||||||
- `.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
|
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}"
|
BASE_MODEL="${BASE_MODEL:-/models/base}"
|
||||||
TRAIN_CONFIG="${TRAIN_CONFIG:-}"
|
TRAIN_CONFIG="${TRAIN_CONFIG:-}"
|
||||||
TRAIN_GPU_COUNT="${TRAIN_GPU_COUNT:-all}"
|
TRAIN_GPU_COUNT="${TRAIN_GPU_COUNT:-all}"
|
||||||
|
TRAIN_PARALLEL_MODE="${TRAIN_PARALLEL_MODE:-auto}"
|
||||||
|
|
||||||
validate_job_name "${TRAIN_JOB_NAME}"
|
validate_job_name "${TRAIN_JOB_NAME}"
|
||||||
if [[ "${TRAIN_GPU_COUNT}" == "all" ]]; then
|
if [[ "${TRAIN_GPU_COUNT}" == "all" ]]; then
|
||||||
TRAIN_GPU_COUNT="$(python -c 'import torch; print(torch.cuda.device_count())')"
|
TRAIN_GPU_COUNT="$(python -c 'import torch; print(torch.cuda.device_count())')"
|
||||||
fi
|
fi
|
||||||
[[ "${TRAIN_GPU_COUNT}" =~ ^[1-9][0-9]*$ ]] || die "No visible GPU found"
|
[[ "${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
|
if [[ -n "${TRAIN_CONFIG}" ]]; then
|
||||||
[[ -f "${TRAIN_CONFIG}" ]] || die "Training config not found: ${TRAIN_CONFIG}"
|
[[ -f "${TRAIN_CONFIG}" ]] || die "Training config not found: ${TRAIN_CONFIG}"
|
||||||
fi
|
fi
|
||||||
@@ -36,11 +51,7 @@ if [[ -n "${TRAIN_CONFIG}" ]]; then
|
|||||||
train_args+=(--config "${TRAIN_CONFIG}")
|
train_args+=(--config "${TRAIN_CONFIG}")
|
||||||
fi
|
fi
|
||||||
|
|
||||||
if (( TRAIN_GPU_COUNT > 1 )); then
|
train_args+=(--parallel_mode "${TRAIN_PARALLEL_MODE}")
|
||||||
train_args+=(--parallel_mode ddp)
|
|
||||||
else
|
|
||||||
train_args+=(--parallel_mode none)
|
|
||||||
fi
|
|
||||||
|
|
||||||
if [[ -n "${latest_checkpoint}" ]]; then
|
if [[ -n "${latest_checkpoint}" ]]; then
|
||||||
log_info "Resuming ${TRAIN_JOB_NAME} from ${latest_checkpoint}"
|
log_info "Resuming ${TRAIN_JOB_NAME} from ${latest_checkpoint}"
|
||||||
@@ -52,7 +63,7 @@ else
|
|||||||
train_args+=(--param_path "${BASE_MODEL}")
|
train_args+=(--param_path "${BASE_MODEL}")
|
||||||
fi
|
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.
|
# Replace the shell so the container init forwards SIGTERM to the trainer.
|
||||||
exec "${train_args[@]}" "$@"
|
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)"
|
ROOT_DIR="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")/.." && pwd)"
|
||||||
source "${ROOT_DIR}/scripts/docker/lib/train-common.sh"
|
source "${ROOT_DIR}/scripts/docker/lib/train-common.sh"
|
||||||
|
|
||||||
ENV_FILE="${TRAIN_ENV_FILE:-${ROOT_DIR}/.env.train}"
|
|
||||||
COMPOSE_BASE=(
|
COMPOSE_BASE=(
|
||||||
docker compose
|
docker compose
|
||||||
--project-directory "${ROOT_DIR}"
|
--project-directory "${ROOT_DIR}"
|
||||||
@@ -14,50 +13,28 @@ COMPOSE_BASE=(
|
|||||||
|
|
||||||
usage() {
|
usage() {
|
||||||
cat <<'EOF'
|
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:
|
Commands:
|
||||||
init Create local directories and .env.train
|
init [CONFIG] Create runtime directories
|
||||||
preflight Validate Docker, paths, GPU settings, and Compose
|
preflight [CONFIG] Validate Docker, paths, GPUs, and Compose
|
||||||
build Build the trainer image
|
build [CONFIG] Build the trainer image
|
||||||
start [--foreground] [-- ARGS...] Start or resume training
|
start [CONFIG] [--foreground] [-- ARGS...]
|
||||||
stop Gracefully stop and checkpoint training
|
Start or resume training
|
||||||
restart Stop, then start training
|
stop [CONFIG] Gracefully stop and checkpoint training
|
||||||
logs Follow trainer logs
|
restart [CONFIG] Stop, then start training
|
||||||
status Show container and latest checkpoint status
|
logs [CONFIG] Follow trainer logs
|
||||||
latest Print the latest complete checkpoint path
|
status [CONFIG] Show container and checkpoint status
|
||||||
list List all complete checkpoints
|
latest [CONFIG] Print the latest complete checkpoint
|
||||||
clean [--keep N] Preview old checkpoint removal
|
list [CONFIG] List complete checkpoints
|
||||||
clean --force Remove old checkpoints after previewing
|
clean [CONFIG] [--keep N] [--force]
|
||||||
|
Preview or remove old checkpoints
|
||||||
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.
|
|
||||||
EOF
|
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() {
|
resolve_path() {
|
||||||
if [[ "$1" = /* ]]; then
|
if [[ "$1" = /* ]]; then
|
||||||
printf '%s\n' "$1"
|
printf '%s\n' "$1"
|
||||||
@@ -66,203 +43,147 @@ resolve_path() {
|
|||||||
fi
|
fi
|
||||||
}
|
}
|
||||||
|
|
||||||
# Read the optional top-level `infra:` section from TRAIN_CONFIG_FILE and
|
load_config() {
|
||||||
# export the host-side variables it overrides (job name, mount paths, GPU
|
CONFIG_FILE="$(resolve_path "$1")"
|
||||||
# filter). Compose interpolation prefers the shell environment over the
|
[[ -f "${CONFIG_FILE}" ]] || die "Training config not found: ${CONFIG_FILE}"
|
||||||
# --env-file, so these exports win over .env.train; keys absent from `infra`
|
require_command python3
|
||||||
# fall back to the env file. Requires python3 with PyYAML on the host.
|
python3 -c 'import yaml' >/dev/null 2>&1 ||
|
||||||
load_infra() {
|
die "PyYAML is required on the host (install python3-yaml)"
|
||||||
local infra_file exports
|
|
||||||
|
|
||||||
[[ -n "${TRAIN_CONFIG_FILE:-}" ]] || return 0
|
local exports
|
||||||
infra_file="$(resolve_path "${TRAIN_CONFIG_FILE}")"
|
exports="$(python3 "${ROOT_DIR}/scripts/tools/train_runtime.py" exports "${CONFIG_FILE}")" ||
|
||||||
[[ -f "${infra_file}" ]] || return 0
|
die "Failed to load runtime configuration"
|
||||||
|
|
||||||
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}"
|
eval "${exports}"
|
||||||
log_info "Applied infra overrides from ${infra_file}"
|
|
||||||
fi
|
|
||||||
|
|
||||||
validate_job_name "${TRAIN_JOB_NAME}"
|
validate_job_name "${TRAIN_JOB_NAME}"
|
||||||
}
|
}
|
||||||
|
|
||||||
checkpoint_dir() {
|
compose() {
|
||||||
printf '%s/%s\n' "$(resolve_path "${TRAIN_CHECKPOINT_DIR}")" "${TRAIN_JOB_NAME}"
|
ASTRAI_UID="$(id -u)" ASTRAI_GID="$(id -g)" "${COMPOSE_BASE[@]}" "$@"
|
||||||
}
|
}
|
||||||
|
|
||||||
compose() {
|
checkpoint_dir() {
|
||||||
local -a command=("${COMPOSE_BASE[@]}")
|
printf '%s/%s\n' "${TRAIN_CHECKPOINT_DIR}" "${TRAIN_JOB_NAME}"
|
||||||
|
}
|
||||||
|
|
||||||
if [[ -f "${ENV_FILE}" ]]; then
|
container_name() {
|
||||||
command+=(--env-file "${ENV_FILE}")
|
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
|
fi
|
||||||
|
rm -f -- "${pid_file}"
|
||||||
|
}
|
||||||
|
|
||||||
# Inject the host user into compose so container processes share the
|
schedule_timer() {
|
||||||
# checkpoint directory ownership (bash UID/GID are readonly).
|
(( TRAIN_MAX_DURATION_SECONDS > 0 )) || return 0
|
||||||
ASTRAI_UID="$(id -u)" ASTRAI_GID="$(id -g)" "${command[@]}" "$@"
|
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() {
|
init_environment() {
|
||||||
local data_dir model_dir checkpoints_dir
|
mkdir -p "${TRAIN_DATA_DIR}" "${TRAIN_MODEL_DIR}" "${TRAIN_CHECKPOINT_DIR}"
|
||||||
|
log_info "Data: ${TRAIN_DATA_DIR}"
|
||||||
data_dir="$(resolve_path "${TRAIN_DATA_DIR}")"
|
log_info "Model: ${TRAIN_MODEL_DIR}"
|
||||||
model_dir="$(resolve_path "${TRAIN_MODEL_DIR}")"
|
log_info "Checkpoints: ${TRAIN_CHECKPOINT_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}"
|
|
||||||
}
|
}
|
||||||
|
|
||||||
preflight() {
|
preflight() {
|
||||||
local data_dir model_dir checkpoints_dir config_file latest visible_count
|
local latest visible_count
|
||||||
|
|
||||||
require_command docker
|
require_command docker
|
||||||
docker info >/dev/null 2>&1 || die "Docker daemon is unavailable"
|
docker info >/dev/null 2>&1 || die "Docker daemon is unavailable"
|
||||||
[[ "${TRAIN_GPU_COUNT}" == "all" || "${TRAIN_GPU_COUNT}" =~ ^[1-9][0-9]*$ ]] ||
|
[[ -d "${TRAIN_DATA_DIR}" ]] || die "Training data directory not found: ${TRAIN_DATA_DIR}"
|
||||||
die "TRAIN_GPU_COUNT must be 'all' or a positive integer"
|
mkdir -p "$(checkpoint_dir)"
|
||||||
|
[[ -w "$(checkpoint_dir)" ]] || die "Checkpoint directory is not writable: $(checkpoint_dir)"
|
||||||
|
|
||||||
data_dir="$(resolve_path "${TRAIN_DATA_DIR}")"
|
latest="$(find_latest_checkpoint "$(checkpoint_dir)" || true)"
|
||||||
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)"
|
|
||||||
if [[ -z "${latest}" ]]; then
|
if [[ -z "${latest}" ]]; then
|
||||||
[[ -s "${model_dir}/config.json" ]] || die "Model config not found: ${model_dir}/config.json"
|
[[ -s "${TRAIN_MODEL_DIR}/config.json" ]] ||
|
||||||
[[ -s "${model_dir}/model.safetensors" ]] || die "Model weights not found: ${model_dir}/model.safetensors"
|
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
|
else
|
||||||
log_info "Resume candidate: ${latest}"
|
log_info "Resume candidate: ${latest}"
|
||||||
fi
|
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}"
|
IFS=',' read -r -a visible_gpus <<<"${CUDA_VISIBLE_DEVICES}"
|
||||||
visible_count="${#visible_gpus[@]}"
|
visible_count="${#visible_gpus[@]}"
|
||||||
(( visible_count == TRAIN_GPU_COUNT )) ||
|
(( 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
|
fi
|
||||||
|
|
||||||
compose config --quiet
|
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() {
|
start_training() {
|
||||||
local foreground="$1"
|
local foreground="$1"
|
||||||
local config_file container running
|
|
||||||
local -a run_options=()
|
|
||||||
shift
|
shift
|
||||||
|
local container running
|
||||||
|
local -a run_options
|
||||||
preflight
|
preflight
|
||||||
if [[ -n "${TRAIN_CONFIG_FILE:-}" ]]; then
|
runtime_environment_args
|
||||||
config_file="$(resolve_path "${TRAIN_CONFIG_FILE}")"
|
container="$(container_name)"
|
||||||
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}"
|
|
||||||
running="$(docker inspect --format '{{.State.Running}}' "${container}" 2>/dev/null || true)"
|
running="$(docker inspect --format '{{.State.Running}}' "${container}" 2>/dev/null || true)"
|
||||||
[[ "${running}" != "true" ]] || die "Trainer is already running: ${container}"
|
[[ "${running}" != "true" ]] || die "Trainer is already running: ${container}"
|
||||||
docker rm "${container}" >/dev/null 2>&1 || true
|
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
|
if [[ "${foreground}" == "true" ]]; then
|
||||||
compose run --build --rm "${run_options[@]}" trainer "$@"
|
compose run --build --rm "${run_options[@]}" trainer "$@"
|
||||||
else
|
else
|
||||||
compose run -d --build --name "${container}" \
|
compose run -d --build --name "${container}" "${run_options[@]}" trainer "$@"
|
||||||
"${run_options[@]}" trainer "$@"
|
schedule_timer
|
||||||
log_info "Training started; run scripts/train.sh logs to follow it"
|
log_info "Training started; run scripts/train.sh logs ${CONFIG_FILE} to follow it"
|
||||||
fi
|
fi
|
||||||
}
|
}
|
||||||
|
|
||||||
stop_training() {
|
stop_training() {
|
||||||
|
local from_timer="$1"
|
||||||
|
[[ "${from_timer}" == "true" ]] || cancel_timer
|
||||||
log_info "Stopping trainer with ${TRAIN_STOP_TIMEOUT}s grace period"
|
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"
|
log_warn "Trainer container is not running"
|
||||||
}
|
[[ "${from_timer}" != "true" ]] || rm -f -- "$(timer_pid_file)"
|
||||||
|
|
||||||
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
|
|
||||||
}
|
}
|
||||||
|
|
||||||
show_status() {
|
show_status() {
|
||||||
local latest
|
local latest
|
||||||
|
docker ps -a --filter "name=^/$(container_name)$"
|
||||||
docker ps -a --filter "name=^/astrai-trainer-${TRAIN_JOB_NAME}$"
|
|
||||||
latest="$(find_latest_checkpoint "$(checkpoint_dir)" || true)"
|
latest="$(find_latest_checkpoint "$(checkpoint_dir)" || true)"
|
||||||
if [[ -n "${latest}" ]]; then
|
if [[ -n "${latest}" ]]; then
|
||||||
log_info "Latest checkpoint: ${latest}"
|
log_info "Latest checkpoint: ${latest}"
|
||||||
@@ -274,7 +195,6 @@ show_status() {
|
|||||||
clean_checkpoints() {
|
clean_checkpoints() {
|
||||||
local keep="$1" force="$2" dir count remove_count index path
|
local keep="$1" force="$2" dir count remove_count index path
|
||||||
local -a checkpoints=()
|
local -a checkpoints=()
|
||||||
|
|
||||||
[[ "${keep}" =~ ^[1-9][0-9]*$ ]] || die "--keep must be a positive integer"
|
[[ "${keep}" =~ ^[1-9][0-9]*$ ]] || die "--keep must be a positive integer"
|
||||||
dir="$(checkpoint_dir)"
|
dir="$(checkpoint_dir)"
|
||||||
while IFS= read -r line; do
|
while IFS= read -r line; do
|
||||||
@@ -287,7 +207,6 @@ clean_checkpoints() {
|
|||||||
log_info "Nothing to clean; ${count} complete checkpoint(s), keeping ${keep}"
|
log_info "Nothing to clean; ${count} complete checkpoint(s), keeping ${keep}"
|
||||||
return
|
return
|
||||||
fi
|
fi
|
||||||
|
|
||||||
for ((index = 0; index < remove_count; index++)); do
|
for ((index = 0; index < remove_count; index++)); do
|
||||||
path="${checkpoints[index]}"
|
path="${checkpoints[index]}"
|
||||||
if [[ "${force}" == "true" ]]; then
|
if [[ "${force}" == "true" ]]; then
|
||||||
@@ -301,58 +220,49 @@ clean_checkpoints() {
|
|||||||
}
|
}
|
||||||
|
|
||||||
main() {
|
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=()
|
local -a train_args=()
|
||||||
[[ -n "${command}" ]] || { usage; exit 1; }
|
[[ -n "${command}" ]] || { usage; exit 1; }
|
||||||
shift || true
|
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
|
case "${command}" in
|
||||||
init)
|
init) init_environment ;;
|
||||||
init_environment
|
preflight) preflight ;;
|
||||||
;;
|
build) preflight; compose build trainer ;;
|
||||||
preflight)
|
|
||||||
preflight
|
|
||||||
;;
|
|
||||||
build)
|
|
||||||
preflight
|
|
||||||
compose build trainer
|
|
||||||
;;
|
|
||||||
start)
|
start)
|
||||||
while [[ $# -gt 0 ]]; do
|
while [[ $# -gt 0 ]]; do
|
||||||
case "$1" in
|
case "$1" in
|
||||||
--foreground)
|
--foreground) foreground=true; shift ;;
|
||||||
foreground=true
|
--) shift; train_args=("$@"); break ;;
|
||||||
shift
|
*) die "Unknown start option: $1 (put trainer arguments after --)" ;;
|
||||||
;;
|
|
||||||
--)
|
|
||||||
shift
|
|
||||||
train_args=("$@")
|
|
||||||
break
|
|
||||||
;;
|
|
||||||
*)
|
|
||||||
die "Unknown start option: $1 (put trainer arguments after --)"
|
|
||||||
;;
|
|
||||||
esac
|
esac
|
||||||
done
|
done
|
||||||
start_training "${foreground}" "${train_args[@]}"
|
start_training "${foreground}" "${train_args[@]}"
|
||||||
;;
|
;;
|
||||||
stop)
|
stop)
|
||||||
stop_training
|
[[ "${1:-}" != "--from-timer" ]] || from_timer=true
|
||||||
|
stop_training "${from_timer}"
|
||||||
;;
|
;;
|
||||||
restart)
|
restart)
|
||||||
restart_training
|
stop_training false
|
||||||
;;
|
start_training false
|
||||||
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"
|
|
||||||
;;
|
;;
|
||||||
|
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)
|
||||||
list_complete_checkpoints "$(checkpoint_dir)" | while read -r _epoch _step path; do
|
list_complete_checkpoints "$(checkpoint_dir)" | while read -r _epoch _step path; do
|
||||||
printf '%s\n' "${path}"
|
printf '%s\n' "${path}"
|
||||||
@@ -361,28 +271,14 @@ main() {
|
|||||||
clean)
|
clean)
|
||||||
while [[ $# -gt 0 ]]; do
|
while [[ $# -gt 0 ]]; do
|
||||||
case "$1" in
|
case "$1" in
|
||||||
--keep)
|
--keep) [[ $# -ge 2 ]] || die "--keep requires a value"; keep="$2"; shift 2 ;;
|
||||||
[[ $# -ge 2 ]] || die "--keep requires a value"
|
--force) force=true; shift ;;
|
||||||
keep="$2"
|
*) die "Unknown clean option: $1" ;;
|
||||||
shift 2
|
|
||||||
;;
|
|
||||||
--force)
|
|
||||||
force=true
|
|
||||||
shift
|
|
||||||
;;
|
|
||||||
*)
|
|
||||||
die "Unknown clean option: $1"
|
|
||||||
;;
|
|
||||||
esac
|
esac
|
||||||
done
|
done
|
||||||
clean_checkpoints "${keep}" "${force}"
|
clean_checkpoints "${keep}" "${force}"
|
||||||
;;
|
;;
|
||||||
help|-h|--help)
|
*) die "Unknown command: ${command}" ;;
|
||||||
usage
|
|
||||||
;;
|
|
||||||
*)
|
|
||||||
die "Unknown command: ${command}"
|
|
||||||
;;
|
|
||||||
esac
|
esac
|
||||||
}
|
}
|
||||||
|
|
||||||
|
|||||||
Reference in New Issue
Block a user