Implement UP_STOP_ON_FAIL option to manage GPU state on installation failure

- Added a new configuration option `UP_STOP_ON_FAIL` to control whether the GPU should be stopped automatically if the `up` command fails, enhancing user control over resource management. - Updated the CLI to include a `--keep-on-fail` flag, allowing users to prevent GPU shutdown during installation errors. - Enhanced the installation scripts and documentation to reflect these changes, providing clearer guidance on the new behavior and configuration options. - Improved error handling in the CLI to ensure proper cleanup of resources in case of failure, preventing unexpected billing for unused GPU resources.
2026-08-21 08:32:33 +03:00
parent 5081a6bc1e
commit 6871c511c4
9 changed files with 179 additions and 14 deletions
@@ -57,6 +57,21 @@ def _die(exc: BaseException) -> None:
    raise typer.Exit(1)


+def _stop_after_failed_up(cfg, cause: BaseException) -> None:
+    """Delete compute after a failed up so billing does not continue unnoticed."""
+    warn(
+        "up упал — гашу GPU (UP_STOP_ON_FAIL; оставить: --keep-on-fail / UP_STOP_ON_FAIL=false)"
+    )
+    try:
+        cmd_stop(cfg, no_pull=True, log=log)
+        ok("compute остановлен, диски на месте")
+    except Exception as stop_exc:
+        err(
+            f"auto-stop не удался: {stop_exc} — срочно: gpu-rent stop "
+            f"(причина up: {cause})"
+        )
+
+
 def _print_checks(checks, *, quiet: bool = False) -> int:
    failed = blocking_failed(checks)
    if quiet:
@@ -439,6 +454,11 @@ def up(
        "--no-update",
        help="Не делать git pull SwarmUI и установленных extensions",
    ),
+    keep_on_fail: bool = typer.Option(
+        False,
+        "--keep-on-fail",
+        help="Не гасить GPU если up упал (по умолчанию stop; UP_STOP_ON_FAIL=false)",
+    ),
    verbose: bool = typer.Option(
        False,
        "--verbose",
@@ -458,6 +478,8 @@ def up(
    ),
 ) -> None:
    """Create/unshelve GPU; SwarmUI и/или LLM; по умолчанию туннель."""
+    cfg = None
+    up_ok = False
    try:
        from dataclasses import replace

@@ -680,6 +702,7 @@ def up(
            ask=None if yes else ask,
            log=log,
        )
+        up_ok = True
        if no_tunnel:
            from gpu_rent.access_card import print_access_card

@@ -705,7 +728,24 @@ def up(
            open_browser=open_browser,
            log=log,
        )
+    except KeyboardInterrupt as exc:
+        if (
+            not up_ok
+            and cfg is not None
+            and not keep_on_fail
+            and bool(getattr(cfg, "up_stop_on_fail", True))
+        ):
+            _stop_after_failed_up(cfg, exc)
+        err("прервано (Ctrl+C)")
+        raise typer.Exit(130) from exc
    except GpuRentError as exc:
+        if (
+            not up_ok
+            and cfg is not None
+            and not keep_on_fail
+            and bool(getattr(cfg, "up_stop_on_fail", True))
+        ):
+            _stop_after_failed_up(cfg, exc)
        _die(exc)


@@ -121,6 +121,7 @@ class Config:
    idle_grace_minutes: int
    pull_output: bool
    notify_ready: bool
+    up_stop_on_fail: bool
    selectel_api_token: str
    balance_notify_step_rub: float
    balance_notify_low_rub: float
@@ -258,6 +259,8 @@ def load_config(*, require_auth: bool = True) -> Config:
        idle_grace_minutes=_as_int(os.environ.get("IDLE_GRACE_MINUTES"), 90),
        pull_output=_as_bool(os.environ.get("PULL_OUTPUT"), False),
        notify_ready=_as_bool(os.environ.get("NOTIFY_READY"), True),
+        # Default on: failed up must not leave a billing GPU running.
+        up_stop_on_fail=_as_bool(os.environ.get("UP_STOP_ON_FAIL"), True),
        selectel_api_token=(os.environ.get("SELECTEL_API_TOKEN") or "").strip(),
        balance_notify_step_rub=_as_float(os.environ.get("BALANCE_NOTIFY_STEP_RUB"), 200.0),
        balance_notify_low_rub=_as_float(os.environ.get("BALANCE_NOTIFY_LOW_RUB"), 0.0),
@@ -36,6 +36,7 @@ _LLAMACPP_INSTALL_ENV = (
    "LLAMACPP_ASSET_URL",
    "LLAMACPP_SHA256",
    "LLAMACPP_BUILD_CUDA",
+    "LLAMACPP_BACKEND",
    "LLAMACPP_FORCE_REINSTALL",
    "LLAMACPP_NGL",
    "LLAMACPP_CTX",
@@ -544,13 +545,13 @@ def provision_llm(cfg: Config, host: str, log: Log) -> None:
        log("LLM: ставим/запускаем llama.cpp server")
        import os

-        build_cuda = (os.environ.get("LLAMACPP_BUILD_CUDA") or "").strip() == "1"
+        # Vulkan finishes in seconds; CUDA compile needs up to ~15–20 min.
        run_script_sudo(
            cfg,
            host,
            _pkg_text("install_llamacpp.sh"),
            remote_path="/tmp/gpu-rent-install_llamacpp.sh",
-            timeout=3600 if build_cuda else 1200,
+            timeout=3600,
            env=_remote_llm_env(cfg, *_LLAMACPP_INSTALL_ENV),
            log=log,
        )
@@ -31,8 +31,10 @@ SERVER_BIN="${BIN_DIR}/llama-server"
 LLAMACPP_TAG="${LLAMACPP_TAG:-}"
 LLAMACPP_ASSET_URL="${LLAMACPP_ASSET_URL:-}"
 LLAMACPP_SHA256="${LLAMACPP_SHA256:-}"
-# 1 = force CUDA compile; 0/empty = prebuilt first, build only if prebuilt fails
+# 1 = force CUDA compile; 0 = never compile (Vulkan/CPU prebuilt only)
 LLAMACPP_BUILD_CUDA="${LLAMACPP_BUILD_CUDA:-}"
+# auto | cuda | vulkan — default auto: CUDA if nvcc already on VM, else Vulkan prebuilt
+LLAMACPP_BACKEND="${LLAMACPP_BACKEND:-auto}"
 LLAMACPP_FORCE_REINSTALL="${LLAMACPP_FORCE_REINSTALL:-}"
 LLAMACPP_NGL="${LLAMACPP_NGL:-}"
 LLAMACPP_CTX="${LLAMACPP_CTX:-}"
@@ -40,11 +42,48 @@ LLAMACPP_HOST="${LLAMACPP_HOST:-127.0.0.1}"
 LLAMACPP_PORT="${LLAMACPP_PORT:-8080}"
 LLAMACPP_EXTRA_ARGS="${LLAMACPP_EXTRA_ARGS:-}"

+have_nvcc() {
+  if command -v nvcc >/dev/null 2>&1; then
+    return 0
+  fi
+  if [[ -x /usr/local/cuda/bin/nvcc ]]; then
+    export PATH="/usr/local/cuda/bin:${PATH}"
+    return 0
+  fi
+  return 1
+}
+
+# Prefer CUDA when toolkit already present (GPU images / prior up). Vulkan = fast no-compile.
+want_cuda_build() {
+  case "${LLAMACPP_BUILD_CUDA}" in
+    1|yes|true) return 0 ;;
+    0|no|false) return 1 ;;
+  esac
+  case "${LLAMACPP_BACKEND}" in
+    cuda) return 0 ;;
+    vulkan) return 1 ;;
+    *)
+      if have_nvcc; then
+        return 0
+      fi
+      return 1
+      ;;
+  esac
+}
+
 if [[ "${LLAMACPP_FORCE_REINSTALL}" == "1" ]]; then
  log "LLAMACPP_FORCE_REINSTALL=1 — удаляю старый бинарь"
  rm -f "$SERVER_BIN" "$STAMP"
 fi

+# Upgrade path: previous default was Vulkan prebuilt; if nvcc is here, prefer CUDA.
+if [[ -x "$SERVER_BIN" && -f "$STAMP" && "${LLAMACPP_BACKEND}" != "vulkan" && "${LLAMACPP_BUILD_CUDA}" != "0" ]]; then
+  if grep -q '^asset:' "$STAMP" 2>/dev/null && want_cuda_build; then
+    log "был Vulkan/CPU prebuilt, nvcc есть — пересобираю CUDA (лучше на NVIDIA)"
+    rm -f "$SERVER_BIN" "$STAMP"
+  fi
+fi
+
 # Prefer ubuntu CUDA (rare) → vulkan → cpu. Never Windows/macOS/cudart-only.
 pick_linux_asset_url() {
  python3 -c '
@@ -321,12 +360,18 @@ else
    fi

    installed=0
-    if [[ "${LLAMACPP_BUILD_CUDA}" == "1" ]]; then
-      log "LLAMACPP_BUILD_CUDA=1 — сразу CUDA-сборка"
+    if want_cuda_build; then
+      log "backend: CUDA (nvcc есть или LLAMACPP_BACKEND/BUILD_CUDA) — сборка, Vulkan только если упадёт"
      if build_cuda_from_source "$tag"; then
        installed=1
+      else
+        log "CUDA-сборка не вышла — fallback на Linux prebuilt (Vulkan/CPU)"
+        if install_linux_release "$tag"; then
+          installed=1
+        fi
      fi
    else
+      log "backend: Linux prebuilt (нет nvcc / LLAMACPP_BACKEND=vulkan) — без compile"
      if install_linux_release "$tag"; then
        installed=1
      else