From 09b6bf6c3937fd6f42dcfaf30d34237394a6dd5d Mon Sep 17 00:00:00 2001 From: Souravrajvi0 <144546710+Souravrajvi0@users.noreply.github.com> Date: Fri, 24 Jul 2026 07:43:54 +0530 Subject: [PATCH] fix(studio): opt-in source-build GPU smoke validation (#7322) * fix(studio): opt-in source-build GPU smoke validation (#5854) Gap 1 (empty CUDA arch -> CPU) already landed in #6481. Wire gap 2: after a GPU source build, optionally run the same staged llama-server smoke test as the prebuilt path, then CPU-fallback on failure. Gated by UNSLOTH_LLAMA_STAGED_VALIDATION (default off) to avoid Blackwell JIT stalls. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * fix(install): normalize staged validation env in setup.sh (#7322) Strip and lowercase UNSLOTH_LLAMA_STAGED_VALIDATION before the shell gate so values like True and surrounding whitespace match the Python staged_validation_enabled() helper. * Rebuild visual server after staged-validation CPU fallback (#5854) Mirror the primary source-build path by best-effort building llama-diffusion-gemma-visual-server after smoke-failure CPU fallback. --------- Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> --- .github/workflows/studio-backend-ci.yml | 1 + studio/install_llama_prebuilt.py | 97 ++++++++++++++- studio/setup.sh | 63 ++++++++++ tests/run_all.sh | 1 + tests/sh/test_staged_validation_enabled.sh | 116 ++++++++++++++++++ .../test_install_llama_prebuilt_logic.py | 60 +++++++++ 6 files changed, 334 insertions(+), 4 deletions(-) create mode 100755 tests/sh/test_staged_validation_enabled.sh diff --git a/.github/workflows/studio-backend-ci.yml b/.github/workflows/studio-backend-ci.yml index b8f587b63e..d926d5c3e4 100644 --- a/.github/workflows/studio-backend-ci.yml +++ b/.github/workflows/studio-backend-ci.yml @@ -228,6 +228,7 @@ jobs: tests/sh/test_system_node_readonly.sh \ tests/sh/test_nvcc_meets_llama_minimum.sh \ tests/sh/test_resolve_cuda_archs.sh \ + tests/sh/test_staged_validation_enabled.sh \ tests/sh/test_tauri_install_exit_order.sh \ tests/sh/test_torch_constraint.sh \ tests/sh/test_torch_flavor.sh \ diff --git a/studio/install_llama_prebuilt.py b/studio/install_llama_prebuilt.py index 6026962478..6ea850139e 100644 --- a/studio/install_llama_prebuilt.py +++ b/studio/install_llama_prebuilt.py @@ -186,8 +186,24 @@ VALIDATION_MODEL_CACHE_FILENAME = "stories260K.gguf" # in validate_prebuilt_choice. Disabled for now: the llama-server GPU forward pass # JIT-compiles CUDA kernels on first load and stalls every install and update by # minutes on Blackwell (sm_100). The check and the source-build fallback it triggers -# are kept intact -- set this to True to re-enable them. +# are kept intact -- set this to True, or set UNSLOTH_LLAMA_STAGED_VALIDATION=1, to +# re-enable them (#5854 gap 2). _RUN_STAGED_PREBUILT_VALIDATION = False + + +def staged_validation_enabled() -> bool: + """True when the expensive llama-server GPU smoke test should run. + + Default off (Blackwell CUDA JIT stalls installs). Opt in via the module + constant or ``UNSLOTH_LLAMA_STAGED_VALIDATION`` (1/true/yes/on). Used by both + the prebuilt path and setup.sh's source-build post-check (#5854). + """ + if _RUN_STAGED_PREBUILT_VALIDATION: + return True + raw = os.environ.get("UNSLOTH_LLAMA_STAGED_VALIDATION", "").strip().lower() + return raw in ("1", "true", "yes", "on") + + INSTALL_LOCK_TIMEOUT_SECONDS = 300 INSTALL_STAGING_ROOT_NAME = ".staging" GITHUB_AUTH_HOSTS = {"api.github.com", "github.com"} @@ -5868,9 +5884,10 @@ def validate_prebuilt_choice( # so they are always validated. For an approved bundle the sha256 manifest # already proves integrity, so its runtime smoke test -- a cold CUDA-JIT pass # costing minutes on Blackwell sm_100 -- is gated behind - # _RUN_STAGED_PREBUILT_VALIDATION, disabled for now. The check and the - # source-build fallback it triggers are kept intact; flip the flag to restore it. - if choice.expected_sha256 is None or _RUN_STAGED_PREBUILT_VALIDATION: + # staged_validation_enabled() (constant or UNSLOTH_LLAMA_STAGED_VALIDATION), + # disabled for now. The check and the source-build fallback it triggers are + # kept intact; flip the flag / env to restore it (#5854). + if choice.expected_sha256 is None or staged_validation_enabled(): validate_quantize( quantize_path, probe_path, @@ -5891,6 +5908,49 @@ def validate_prebuilt_choice( return server_path, quantize_path +def validate_existing_install( + install_dir: Path, + *, + install_kind: str | None = None, + host: HostInfo | None = None, +) -> None: + """Run the staged smoke test against an already-built llama.cpp tree (#5854). + + Used by setup.sh after a GPU source build when ``UNSLOTH_LLAMA_STAGED_VALIDATION`` + is set. Raises ``PrebuiltFallback`` on failure so the caller can retry CPU. + """ + host = host or detect_host() + bin_dir = install_dir / "build" / "bin" + server_name = "llama-server.exe" if host.is_windows else "llama-server" + quantize_name = "llama-quantize.exe" if host.is_windows else "llama-quantize" + server_path = bin_dir / server_name + quantize_path = bin_dir / quantize_name + if not server_path.is_file(): + raise PrebuiltFallback(f"llama-server not found at {server_path}") + + with tempfile.TemporaryDirectory(prefix = "unsloth-llama-source-validate-") as tmp: + work_dir = Path(tmp) + probe_path = work_dir / "stories260K.gguf" + quantized_path = work_dir / "stories260K-q4.gguf" + download_validation_model(probe_path, validation_model_cache_path(install_dir)) + if quantize_path.is_file(): + validate_quantize( + quantize_path, + probe_path, + quantized_path, + install_dir, + host, + ) + validate_server( + server_path, + probe_path, + host, + install_dir, + install_kind = install_kind, + ) + log(f"staged source-build validation succeeded for {install_dir}") + + def validate_prebuilt_attempts( attempts: Iterable[AssetChoice], host: HostInfo, @@ -6345,6 +6405,24 @@ def parse_args() -> argparse.Namespace: "fork). Use --output-format json." ), ) + resolve_group.add_argument( + "--validate-install", + metavar = "DIR", + help = ( + "Run the staged llama-server smoke test against an existing build " + "tree (setup.sh source-build post-check, #5854). Exit 2 on failure. " + "Normally gated by UNSLOTH_LLAMA_STAGED_VALIDATION; this flag always " + "runs the check." + ), + ) + parser.add_argument( + "--install-kind", + default = None, + help = ( + "Install kind for --validate-install GPU offload (e.g. linux-cuda, " + "linux-rocm, macos-arm64). When omitted, host detection decides." + ), + ) parser.add_argument( "--output-format", choices = ("plain", "json"), @@ -6381,6 +6459,17 @@ def emit_resolver_output(payload: dict[str, Any], *, output_format: str) -> None def main() -> int: args = parse_args() + if args.validate_install is not None: + try: + validate_existing_install( + Path(args.validate_install), + install_kind = args.install_kind, + ) + except PrebuiltFallback as exc: + print(str(exc), file = sys.stderr) + raise SystemExit(EXIT_FALLBACK) from exc + return EXIT_SUCCESS + if args.resolve_llama_tag is not None: resolved = resolve_requested_llama_tag( args.resolve_llama_tag, diff --git a/studio/setup.sh b/studio/setup.sh index 3c97f77065..f6a6bc346b 100755 --- a/studio/setup.sh +++ b/studio/setup.sh @@ -251,6 +251,38 @@ _resolve_cuda_archs() { printf '%s' "$_archs" } +# Opt-in staged GPU smoke test after a source build (#5854 gap 2). Default off: +# llama-server's first GPU forward pass JIT-compiles CUDA kernels and stalls +# installs for minutes on Blackwell. Same env as install_llama_prebuilt.py. +_staged_validation_enabled() { + local _raw="${UNSLOTH_LLAMA_STAGED_VALIDATION:-}" + # Match install_llama_prebuilt.py staged_validation_enabled(): strip + lowercase. + _raw="$(printf '%s' "$_raw" | sed -e 's/^[[:space:]]*//' -e 's/[[:space:]]*$//' | tr '[:upper:]' '[:lower:]')" + case "$_raw" in + 1|true|yes|on) return 0 ;; + *) return 1 ;; + esac +} + +# Map the source-build GPU backend to install_llama_prebuilt --install-kind so +# validate_server enables --n-gpu-layers for the right backends. +_source_smoke_install_kind() { + if [ "${_TRY_METAL_CPU_FALLBACK:-false}" = true ]; then + printf '%s' "macos-arm64" + return 0 + fi + case "${GPU_BACKEND:-}" in + cuda) + case "$(uname -m 2>/dev/null || true)" in + aarch64|arm64) printf '%s' "linux-arm64-cuda" ;; + *) printf '%s' "linux-cuda" ;; + esac + ;; + rocm) printf '%s' "linux-rocm" ;; + *) printf '%s' "" ;; + esac +} + # Run a GPU probe under a 10s timeout when `timeout` is available so a wedged # NVIDIA driver cannot hang setup; fall back to a bare call where it is not. _setup_run_smi() { @@ -1900,6 +1932,37 @@ else run_quiet_no_exit "build diffusion visual server" cmake --build "$_BUILD_TMP/build" --config Release --target llama-diffusion-gemma-visual-server -j"$NCPU" || true fi + # Opt-in post-build GPU smoke test (#5854 gap 2). Default off (Blackwell + # CUDA JIT stalls). On failure, reuse the CPU fallback path so the user + # still gets a working llama-server. Runs before the install swap. + if [ "$BUILD_OK" = true ] && _staged_validation_enabled; then + _FB_LABEL="$(_gpu_fallback_label)" + _SMOKE_KIND="$(_source_smoke_install_kind)" + if [ -n "$_FB_LABEL" ]; then + _SMOKE_CMD=( + python "$SCRIPT_DIR/install_llama_prebuilt.py" + --validate-install "$_BUILD_TMP" + ) + [ -n "$_SMOKE_KIND" ] && _SMOKE_CMD+=(--install-kind "$_SMOKE_KIND") + if ! run_quiet_no_exit "validate source llama.cpp" "${_SMOKE_CMD[@]}"; then + substep "$_FB_LABEL source build failed smoke test; retrying CPU build..." "$C_WARN" + _TRY_METAL_CPU_FALLBACK=false + rm -rf "$_BUILD_TMP/build" + if run_quiet_no_exit "cmake llama.cpp (cpu fallback)" cmake $CMAKE_GENERATOR_ARGS -S "$_BUILD_TMP" -B "$_BUILD_TMP/build" $CPU_FALLBACK_CMAKE_ARGS; then + _BUILD_DESC="building (CPU fallback after $_FB_LABEL smoke failed)" + GPU_BACKEND="" + run_quiet_no_exit "build llama-server (cpu fallback)" cmake --build "$_BUILD_TMP/build" --config Release --target llama-server -j"$NCPU" || BUILD_OK=false + if [ "$BUILD_OK" = true ]; then + run_quiet_no_exit "build llama-quantize (cpu fallback)" cmake --build "$_BUILD_TMP/build" --config Release --target llama-quantize -j"$NCPU" || true + run_quiet_no_exit "build diffusion visual server (cpu fallback)" cmake --build "$_BUILD_TMP/build" --config Release --target llama-diffusion-gemma-visual-server -j"$NCPU" || true + fi + else + BUILD_OK=false + fi + fi + fi + fi + # Swap only after build succeeds -- preserves existing install on failure if [ "$BUILD_OK" = true ]; then _assert_studio_owned_or_absent "$LLAMA_CPP_DIR" "llama.cpp install" diff --git a/tests/run_all.sh b/tests/run_all.sh index eaa726f73c..6eccffc75f 100755 --- a/tests/run_all.sh +++ b/tests/run_all.sh @@ -12,6 +12,7 @@ sh "$TESTS_DIR/sh/test_mac_intel_compat.sh" sh "$TESTS_DIR/sh/test_torch_constraint.sh" sh "$TESTS_DIR/sh/test_nvcc_meets_llama_minimum.sh" sh "$TESTS_DIR/sh/test_resolve_cuda_archs.sh" +sh "$TESTS_DIR/sh/test_staged_validation_enabled.sh" sh "$TESTS_DIR/sh/test_strixhalo_wsl_reroute.sh" sh "$TESTS_DIR/sh/test_uninstall_shared_icon.sh" sh "$TESTS_DIR/sh/test_torch_flavor.sh" diff --git a/tests/sh/test_staged_validation_enabled.sh b/tests/sh/test_staged_validation_enabled.sh new file mode 100755 index 0000000000..da6a0bdd27 --- /dev/null +++ b/tests/sh/test_staged_validation_enabled.sh @@ -0,0 +1,116 @@ +#!/bin/bash +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 +# Unit tests for setup.sh staged-validation helpers (#5854 gap 2). +# Opt-in GPU smoke after a source build; default off (Blackwell JIT stall). +set -e + +SCRIPT_DIR="$(cd "$(dirname "$0")" && pwd)" +SETUP_SH="$SCRIPT_DIR/../../studio/setup.sh" +PASS=0 +FAIL=0 + +_FUNC_FILE=$(mktemp) +{ + sed -n '/^_staged_validation_enabled()/,/^}/p' "$SETUP_SH" + sed -n '/^_source_smoke_install_kind()/,/^}/p' "$SETUP_SH" +} > "$_FUNC_FILE" +# shellcheck disable=SC1090 +. "$_FUNC_FILE" +rm -f "$_FUNC_FILE" + +assert_eq() { + _label="$1"; _expected="$2"; _actual="$3" + if [ "$_actual" = "$_expected" ]; then + echo " PASS: $_label"; PASS=$((PASS + 1)) + else + echo " FAIL: $_label (expected '$_expected', got '$_actual')"; FAIL=$((FAIL + 1)) + fi +} + +assert_rc() { + _label="$1"; _expected="$2" + shift 2 + set +e + "$@" >/dev/null 2>&1 + _rc=$? + set -e + if [ "$_rc" -eq "$_expected" ]; then + echo " PASS: $_label"; PASS=$((PASS + 1)) + else + echo " FAIL: $_label (expected rc $_expected, got $_rc)"; FAIL=$((FAIL + 1)) + fi +} + +echo "=== _staged_validation_enabled ===" +unset UNSLOTH_LLAMA_STAGED_VALIDATION +assert_rc "default off" 1 _staged_validation_enabled + +UNSLOTH_LLAMA_STAGED_VALIDATION=0 +assert_rc "0 is off" 1 _staged_validation_enabled + +UNSLOTH_LLAMA_STAGED_VALIDATION=1 +assert_rc "1 is on" 0 _staged_validation_enabled + +UNSLOTH_LLAMA_STAGED_VALIDATION=true +assert_rc "true is on" 0 _staged_validation_enabled + +UNSLOTH_LLAMA_STAGED_VALIDATION=yes +assert_rc "yes is on" 0 _staged_validation_enabled + +UNSLOTH_LLAMA_STAGED_VALIDATION=on +assert_rc "on is on" 0 _staged_validation_enabled + +UNSLOTH_LLAMA_STAGED_VALIDATION=True +assert_rc "True is on" 0 _staged_validation_enabled + +UNSLOTH_LLAMA_STAGED_VALIDATION=' yes ' +assert_rc "whitespace yes is on" 0 _staged_validation_enabled + +UNSLOTH_LLAMA_STAGED_VALIDATION=maybe +assert_rc "maybe is off" 1 _staged_validation_enabled +unset UNSLOTH_LLAMA_STAGED_VALIDATION + +echo "=== _source_smoke_install_kind ===" +_TRY_METAL_CPU_FALLBACK=true +GPU_BACKEND="" +assert_eq "metal" "macos-arm64" "$(_source_smoke_install_kind)" + +_TRY_METAL_CPU_FALLBACK=false +GPU_BACKEND=cuda +_kind="$(_source_smoke_install_kind)" +case "$(uname -m)" in + aarch64|arm64) assert_eq "cuda arm" "linux-arm64-cuda" "$_kind" ;; + *) assert_eq "cuda x86" "linux-cuda" "$_kind" ;; +esac + +GPU_BACKEND=rocm +assert_eq "rocm" "linux-rocm" "$(_source_smoke_install_kind)" + +GPU_BACKEND="" +assert_eq "cpu empty" "" "$(_source_smoke_install_kind)" + +echo "=== setup.sh source smoke contract ===" +assert_contains() { + _label="$1"; _hay="$2"; _needle="$3" + case "$_hay" in + *"$_needle"*) echo " PASS: $_label"; PASS=$((PASS + 1)) ;; + *) echo " FAIL: $_label (missing '$_needle')"; FAIL=$((FAIL + 1)) ;; + esac +} +_src=$(cat "$SETUP_SH") +assert_contains "env gate present" "$_src" "UNSLOTH_LLAMA_STAGED_VALIDATION" +assert_contains "calls validate-install" "$_src" "--validate-install" +assert_contains "smoke fail retries CPU" "$_src" "source build failed smoke test; retrying CPU build" +# Smoke must run before the install swap. +_smoke_pos=$(printf '%s' "$_src" | awk '/validate source llama.cpp/{print NR; exit}') +_swap_pos=$(printf '%s' "$_src" | awk '/mv "\$_BUILD_TMP" "\$LLAMA_CPP_DIR"/{print NR; exit}') +if [ -n "$_smoke_pos" ] && [ -n "$_swap_pos" ] && [ "$_smoke_pos" -lt "$_swap_pos" ]; then + echo " PASS: smoke before install swap"; PASS=$((PASS + 1)) +else + echo " FAIL: smoke before install swap (smoke=$_smoke_pos swap=$_swap_pos)"; FAIL=$((FAIL + 1)) +fi + +echo "" +echo "Results: $PASS passed, $FAIL failed" +[ "$FAIL" -eq 0 ] diff --git a/tests/studio/install/test_install_llama_prebuilt_logic.py b/tests/studio/install/test_install_llama_prebuilt_logic.py index 0964c047d5..3eaf56d15c 100644 --- a/tests/studio/install/test_install_llama_prebuilt_logic.py +++ b/tests/studio/install/test_install_llama_prebuilt_logic.py @@ -3324,6 +3324,66 @@ def test_validate_prebuilt_choice_approved_validation_runs_when_flag_enabled(tmp assert calls == {"quantize": 1, "server": 1} +def test_staged_validation_enabled_default_off(monkeypatch): + monkeypatch.setattr(INSTALL_LLAMA_PREBUILT, "_RUN_STAGED_PREBUILT_VALIDATION", False) + monkeypatch.delenv("UNSLOTH_LLAMA_STAGED_VALIDATION", raising = False) + assert INSTALL_LLAMA_PREBUILT.staged_validation_enabled() is False + + +@pytest.mark.parametrize("value", ["1", "true", "YES", "on"]) +def test_staged_validation_enabled_env_opt_in(monkeypatch, value): + monkeypatch.setattr(INSTALL_LLAMA_PREBUILT, "_RUN_STAGED_PREBUILT_VALIDATION", False) + monkeypatch.setenv("UNSLOTH_LLAMA_STAGED_VALIDATION", value) + assert INSTALL_LLAMA_PREBUILT.staged_validation_enabled() is True + + +def test_validate_prebuilt_choice_approved_validation_runs_when_env_enabled(tmp_path, monkeypatch): + monkeypatch.setattr(INSTALL_LLAMA_PREBUILT, "_RUN_STAGED_PREBUILT_VALIDATION", False) + monkeypatch.setenv("UNSLOTH_LLAMA_STAGED_VALIDATION", "1") + calls = _run_validate_prebuilt_choice(monkeypatch, tmp_path, expected_sha256 = "ab" * 32) + assert calls == {"quantize": 1, "server": 1} + + +def test_validate_existing_install_runs_server_smoke(tmp_path, monkeypatch): + # setup.sh --validate-install path: exercise smoke helpers without a real GPU. + install_dir = tmp_path / "llama.cpp" + bin_dir = install_dir / "build" / "bin" + bin_dir.mkdir(parents = True) + (bin_dir / "llama-server").write_text("#!/bin/sh\n", encoding = "utf-8") + (bin_dir / "llama-quantize").write_text("#!/bin/sh\n", encoding = "utf-8") + calls: dict[str, int] = {"quantize": 0, "server": 0, "download": 0} + + monkeypatch.setattr( + INSTALL_LLAMA_PREBUILT, + "download_validation_model", + lambda path, cache = None: calls.__setitem__("download", calls["download"] + 1), + ) + monkeypatch.setattr( + INSTALL_LLAMA_PREBUILT, + "validate_quantize", + lambda *a, **k: calls.__setitem__("quantize", calls["quantize"] + 1), + ) + monkeypatch.setattr( + INSTALL_LLAMA_PREBUILT, + "validate_server", + lambda *a, **k: calls.__setitem__("server", calls["server"] + 1), + ) + monkeypatch.setattr( + INSTALL_LLAMA_PREBUILT, + "detect_host", + lambda: linux_host(), + ) + + INSTALL_LLAMA_PREBUILT.validate_existing_install(install_dir, install_kind = "linux-cuda") + assert calls == {"quantize": 1, "server": 1, "download": 1} + + +def test_validate_existing_install_missing_server_raises(tmp_path, monkeypatch): + monkeypatch.setattr(INSTALL_LLAMA_PREBUILT, "detect_host", lambda: linux_host()) + with pytest.raises(INSTALL_LLAMA_PREBUILT.PrebuiltFallback, match = "llama-server not found"): + INSTALL_LLAMA_PREBUILT.validate_existing_install(tmp_path / "missing") + + def test_diffusion_visual_server_uses_approved_checksum_download(monkeypatch, tmp_path: Path): asset_name = "llama-diffusion-gemma-visual-server-linux-x64" expected_sha = "a" * 64