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

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