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

- 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

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: .

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]`;
+- `runtime.gpu.devices` is `all` (default) or a single-device list such as `[0]`;
-  the list becomes `CUDA_VISIBLE_DEVICES`. Compose passes `count: 1`.
+  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
+rendered Compose configuration. `up` starts the container detached; `run`
-rebuilds the image when the code changed (`--build`); `run` keeps it in the
+keeps it in the foreground. Both reuse the existing image; run
-foreground. The wrapper manages a fixed container name
+`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

@@ -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
+code is baked into `/app`; `start` reuses the existing image, so run
-image when necessary.
+`bash scripts/train.sh build [CONFIG]` after code changes.
 
 ## Operations
 

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"

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):

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):

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

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"] == ""