126 lines
4.6 KiB
Markdown
126 lines
4.6 KiB
Markdown
# Containerized Training Deployment
|
|
|
|
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
|
|
|
|
```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
|
|
```
|
|
|
|
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.
|
|
|
|
## Runtime Schema
|
|
|
|
```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"
|
|
```
|
|
|
|
- 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.
|
|
|
|
## 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 [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
|
|
```
|
|
|
|
`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.
|
|
|
|
## 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.
|