diff --git a/docker-compose.yml b/docker-compose.yml index 98ed2c4..021f601 100644 --- a/docker-compose.yml +++ b/docker-compose.yml @@ -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: . diff --git a/docs/developer/docker-serving.md b/docs/developer/docker-serving.md index 698854c..0241a4d 100644 --- a/docs/developer/docker-serving.md +++ b/docs/developer/docker-serving.md @@ -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-`); 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 diff --git a/docs/developer/docker-training.md b/docs/developer/docker-training.md index 3b6266b..74ab421 100644 --- a/docs/developer/docker-training.md +++ b/docs/developer/docker-training.md @@ -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 diff --git a/scripts/serve.sh b/scripts/serve.sh index 0dade2d..b687764 100644 --- a/scripts/serve.sh +++ b/scripts/serve.sh @@ -55,7 +55,12 @@ load_config() { } compose() { - ASTRAI_UID="$(id -u)" ASTRAI_GID="$(id -g)" "${COMPOSE_BASE[@]}" "$@" + 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" diff --git a/scripts/tools/serve_runtime.py b/scripts/tools/serve_runtime.py index 5b6402e..d4deb84 100644 --- a/scripts/tools/serve_runtime.py +++ b/scripts/tools/serve_runtime.py @@ -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): diff --git a/scripts/tools/train_runtime.py b/scripts/tools/train_runtime.py index 325c0c2..e53b3fd 100644 --- a/scripts/tools/train_runtime.py +++ b/scripts/tools/train_runtime.py @@ -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): diff --git a/scripts/train.sh b/scripts/train.sh index 539ee7a..0b8cee5 100755 --- a/scripts/train.sh +++ b/scripts/train.sh @@ -58,7 +58,12 @@ load_config() { } compose() { - ASTRAI_UID="$(id -u)" ASTRAI_GID="$(id -g)" "${COMPOSE_BASE[@]}" "$@" + 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 diff --git a/tests/test_serve_runtime.py b/tests/test_serve_runtime.py index 6ada9f4..196f8a8 100644 --- a/tests/test_serve_runtime.py +++ b/tests/test_serve_runtime.py @@ -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"] == ""