ViperEkura/AstrAI
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

fix: serve and train reuse the built image and expose GPUs correctly

Browse Source 
- serve.sh/train.sh no longer pass --build on up/run; the build subcommand is the only path that rebuilds
- compose services pin image: astrai:latest so run reuses the existing image instead of triggering a rebuild
- runtime parsers leave CUDA_VISIBLE_DEVICES unset for gpu.devices: all; an empty string hid every GPU inside the container
- server service reserves count: all GPUs so CUDA_VISIBLE_DEVICES performs the only filtering, matching the trainer
- wrapper compose() strips an empty host CUDA_VISIBLE_DEVICES before invoking docker compose
This commit is contained in:
ViperEkura
2026-08-27 12:22:02 +08:00
parent 7cfcc6c86a
commit ba8beb81be
8 changed files with 37 additions and 21 deletions
Download Patch File Download Diff File Expand all files Collapse all files
docker-compose.yml
+4 -1
View File
@@ -1,5 +1,6 @@
services:
server:
image: astrai:latest
build:
context: .
dockerfile: Dockerfile
@@ -20,7 +21,7 @@ services:
reservations:
devices:
- driver: nvidia
count: 1
count: all
capabilities: [gpu]
healthcheck:
test: ["CMD", "curl", "-f", "http://localhost:8000/health"]
@@ -31,6 +32,7 @@ services:
restart: unless-stopped
server-cpu:
image: astrai:latest
profiles: [cpu]
build:
context: .
@@ -54,6 +56,7 @@ services:
restart: unless-stopped
trainer:
image: astrai:latest
profiles: [train]
build:
context: .
docs/developer/docker-serving.md
+7 -6
View File
@@ -54,8 +54,9 @@ server:
- `runtime.gpu.enabled: true` (default) selects the `server` service with an
NVIDIA device reservation; `false` selects `server-cpu` (no GPU passthrough).
When disabled, `server.device` must be `cpu`.
- `runtime.gpu.devices` is either `all` or a single-device list such as `[0]`;
the list becomes `CUDA_VISIBLE_DEVICES`. Compose passes `count: 1`.
- `runtime.gpu.devices` is `all` (default) or a single-device list such as `[0]`;
the list becomes `CUDA_VISIBLE_DEVICES`. Compose passes `count: all`; the
env var performs the only filtering.
- `environment` values are explicitly passed to the serving container. Keep
host-specific settings here; they are not universal defaults.
- `server.device` must agree with `runtime.gpu.enabled`; `preflight` enforces it.
@@ -89,9 +90,9 @@ bash scripts/serve.sh status [CONFIG]
`preflight` validates Docker, the model directory
(`config.json` + `model.safetensors`), GPU/device consistency, and the
rendered Compose configuration. `up` starts the container detached and
rebuilds the image when the code changed (`--build`); `run` keeps it in the
foreground. The wrapper manages a fixed container name
rendered Compose configuration. `up` starts the container detached; `run`
keeps it in the foreground. Both reuse the existing image; run
`bash scripts/serve.sh build [CONFIG]` after code changes. The wrapper manages a fixed container name
(`astrai-server` or `astrai-server-<job_name>`); the plain
`docker compose up -d` / `docker compose --profile cpu up -d` path keeps
working with defaults (port 8000, `./params`).
@@ -99,7 +100,7 @@ working with defaults (port 8000, `./params`).
## Hard Rules
1. Keep Docker settings in `runtime` and server settings in `server`.
2. Filter GPUs once: the `server` service reserves one device; a `devices`
2. Filter GPUs once: Compose passes `count: all`; a `devices`
list becomes `CUDA_VISIBLE_DEVICES`.
3. `runtime.gpu.enabled: false` requires `server.device: cpu`.
4. In Docker, `server.port` must match the published container port (default
docs/developer/docker-training.md
+2 -2
View File
@@ -72,8 +72,8 @@ runtime:
| 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.
code is baked into `/app`; `start` reuses the existing image, so run
`bash scripts/train.sh build [CONFIG]` after code changes.
## Operations
scripts/serve.sh
+7 -2
View File
@@ -55,7 +55,12 @@ load_config() {
}
compose() {
if [[ -n "${CUDA_VISIBLE_DEVICES:-}" ]]; then
ASTRAI_UID="$(id -u)" ASTRAI_GID="$(id -g)" "${COMPOSE_BASE[@]}" "$@"
else
ASTRAI_UID="$(id -u)" ASTRAI_GID="$(id -g)" \
env -u CUDA_VISIBLE_DEVICES "${COMPOSE_BASE[@]}" "$@"
fi
}
container_name() {
@@ -129,11 +134,11 @@ start_server() {
"${RUNTIME_ENV_ARGS[@]}"
)
if [[ "${foreground}" == "true" ]]; then
compose "${PROFILE_ARGS[@]}" run --build --rm --service-ports \
compose "${PROFILE_ARGS[@]}" run --rm --service-ports \
"${run_options[@]}" "$(service_name)" \
python -m scripts.tools.server --config /run/astrai/serve.yaml "$@"
else
compose "${PROFILE_ARGS[@]}" run -d --build --service-ports \
compose "${PROFILE_ARGS[@]}" run -d --service-ports \
--name "${container}" "${run_options[@]}" "$(service_name)" \
python -m scripts.tools.server --config /run/astrai/serve.yaml "$@"
log_info "Server started; run scripts/serve.sh logs ${CONFIG_FILE} to follow it"
scripts/tools/serve_runtime.py
+4 -3
View File
@@ -78,9 +78,10 @@ def load_runtime(config_path: str) -> dict[str, str]:
raise ValueError("runtime.gpu.enabled must be a boolean")
devices = gpu.get("devices", "all")
visible_devices = None
if gpu_enabled:
if devices == "all":
visible_devices = ""
pass
elif isinstance(devices, list) and len(devices) == 1:
text = str(devices[0])
if not text.isdigit():
@@ -93,7 +94,6 @@ def load_runtime(config_path: str) -> dict[str, str]:
"runtime.gpu.devices must be 'all' or a single-device list such as [0]"
)
else:
visible_devices = ""
if devices != "all":
raise ValueError(
"runtime.gpu.devices is ignored when runtime.gpu.enabled is false"
@@ -110,9 +110,10 @@ def load_runtime(config_path: str) -> dict[str, str]:
"SERVE_PARAM_DIR": _path(paths.get("param", "./params"), "param", path.parent),
"SERVE_GPU_ENABLED": "true" if gpu_enabled else "false",
"SERVE_DEVICE": device,
"CUDA_VISIBLE_DEVICES": visible_devices,
"CUDA_TAG": str(container.get("cuda_tag", "cu128")),
}
if visible_devices is not None:
values["CUDA_VISIBLE_DEVICES"] = visible_devices
for name, value in environment.items():
if not isinstance(name, str) or not ENV_NAME.fullmatch(name):
scripts/tools/train_runtime.py
+3 -2
View File
@@ -53,9 +53,9 @@ def load_runtime(config_path: str) -> dict[str, str]:
)
devices = gpu.get("devices", "all")
visible_devices = None
if devices == "all":
gpu_count = "all"
visible_devices = ""
elif isinstance(devices, list) and devices:
normalized = []
for device in devices:
@@ -102,7 +102,6 @@ def load_runtime(config_path: str) -> dict[str, str]:
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")),
@@ -111,6 +110,8 @@ def load_runtime(config_path: str) -> dict[str, str]:
"CHECKPOINT_KEEP_LAST": str(container.get("checkpoint_keep_last", 5)),
"TRAIN_MAX_DURATION_SECONDS": str(max_seconds),
}
if visible_devices is not None:
values["CUDA_VISIBLE_DEVICES"] = visible_devices
for name, value in environment.items():
if not isinstance(name, str) or not ENV_NAME.fullmatch(name):
scripts/train.sh
+7 -2
View File
@@ -58,7 +58,12 @@ load_config() {
}
compose() {
if [[ -n "${CUDA_VISIBLE_DEVICES:-}" ]]; then
ASTRAI_UID="$(id -u)" ASTRAI_GID="$(id -g)" "${COMPOSE_BASE[@]}" "$@"
else
ASTRAI_UID="$(id -u)" ASTRAI_GID="$(id -g)" \
env -u CUDA_VISIBLE_DEVICES "${COMPOSE_BASE[@]}" "$@"
fi
}
checkpoint_dir() {
@@ -164,9 +169,9 @@ start_training() {
"${RUNTIME_ENV_ARGS[@]}"
)
if [[ "${foreground}" == "true" ]]; then
compose run --build --rm "${run_options[@]}" trainer "$@"
compose run --rm "${run_options[@]}" trainer "$@"
else
compose run -d --build --name "${container}" "${run_options[@]}" trainer "$@"
compose run -d --name "${container}" "${run_options[@]}" trainer "$@"
schedule_timer
log_info "Training started; run scripts/train.sh logs ${CONFIG_FILE} to follow it"
fi
tests/test_serve_runtime.py
+1 -1
View File
@@ -26,7 +26,7 @@ def test_runtime_exports_defaults(tmp_path):
assert runtime["SERVE_CONTAINER_PORT"] == "8000"
assert runtime["SERVE_PARAM_DIR"] == str((tmp_path / "params").resolve())
assert runtime["SERVE_GPU_ENABLED"] == "true"
assert runtime["CUDA_VISIBLE_DEVICES"] == ""
assert "CUDA_VISIBLE_DEVICES" not in runtime
assert runtime["SERVE_DEVICE"] == "cuda"
assert runtime["CUDA_TAG"] == "cu128"
assert runtime["SERVE_JOB_NAME"] == ""