diff --git a/.github/workflows/pr-test.yml b/.github/workflows/pr-test.yml index f0f8f45e3..b647d8af4 100644 --- a/.github/workflows/pr-test.yml +++ b/.github/workflows/pr-test.yml @@ -780,23 +780,20 @@ jobs: CUSTOM_BUILD_SGL_KERNEL=${{needs.check-changes.outputs.sgl_kernel}} bash scripts/ci/cuda/ci_install_dependency.sh - name: Warmup DeepGEMM JIT Compilation - timeout-minutes: 25 + # Per-model TP must match the test's launch in test/registered/ — see + # FALLBACK_ARGS in scripts/ci/cuda/warmup_deep_gemm.py for extra dp/ep + # flags. Only models that actually invoke DeepGEMM kernels at runtime + # are listed. Cold-cache ~13 min; warm-cache ≤30 s via marker file. + timeout-minutes: 60 run: | - # Activate venv if available (GITHUB_ENV may have failed to propagate) [ -f "${SGLANG_CI_VENV_PATH:-/dev/null}/bin/activate" ] && source "${SGLANG_CI_VENV_PATH}/bin/activate" [ -f "${SGLANG_CI_VENV_PATH:-/dev/null}/env.sh" ] && source "${SGLANG_CI_VENV_PATH}/env.sh" python3 scripts/ci/cuda/warmup_deep_gemm.py \ deepseek-ai/DeepSeek-V3-0324:8 \ - deepseek-ai/DeepSeek-V3.2-Exp:8 - - - name: Warmup Server CUDA Graphs - timeout-minutes: 25 - run: | - [ -f "${SGLANG_CI_VENV_PATH:-/dev/null}/bin/activate" ] && source "${SGLANG_CI_VENV_PATH}/bin/activate" - [ -f "${SGLANG_CI_VENV_PATH:-/dev/null}/env.sh" ] && source "${SGLANG_CI_VENV_PATH}/env.sh" - python3 scripts/ci/cuda/warmup_server.py \ - deepseek-ai/DeepSeek-V3-0324:8 \ - inclusionAI/Ring-2.5-1T:8 + deepseek-ai/DeepSeek-V3.2:8 \ + zai-org/GLM-5-FP8:8 \ + XiaomiMiMo/MiMo-V2-Flash:4 \ + XiaomiMiMo/MiMo-V2.5:8 - name: Run test timeout-minutes: 30 @@ -998,7 +995,7 @@ jobs: [ -f "${SGLANG_CI_VENV_PATH:-/dev/null}/env.sh" ] && source "${SGLANG_CI_VENV_PATH}/env.sh" python3 scripts/ci/cuda/warmup_deep_gemm.py \ deepseek-ai/DeepSeek-V3-0324:8 \ - deepseek-ai/DeepSeek-V3.2-Exp:8 + deepseek-ai/DeepSeek-V3.2:8 - name: Warmup Server CUDA Graphs timeout-minutes: 25 diff --git a/scripts/ci/cuda/warmup_deep_gemm.py b/scripts/ci/cuda/warmup_deep_gemm.py index 0e8a0442d..58b7c752c 100644 --- a/scripts/ci/cuda/warmup_deep_gemm.py +++ b/scripts/ci/cuda/warmup_deep_gemm.py @@ -11,16 +11,61 @@ for unsupported architectures. Usage: python3 scripts/ci/cuda/warmup_deep_gemm.py \ deepseek-ai/DeepSeek-V3-0324:8 \ - deepseek-ai/DeepSeek-V3.2-Exp:8 + deepseek-ai/DeepSeek-V3.2:8 """ +import hashlib import json import os +import signal import subprocess import sys +import threading import time from math import ceil from pathlib import Path +from typing import Dict, List + +# Shared with warmup_server.py. Wipe alongside /root/.cache/deep_gemm if you +# clear the DeepGEMM JIT cache — a stale marker → in-test JIT compile. +MARKER_DIR = os.path.join(os.path.expanduser("~"), ".cache", "sglang", "warmup_markers") + +# Outer cap for stuck fallback subprocesses; CRASH_MARKERS abort sooner. +FALLBACK_TIMEOUT_SEC = 600 + +# Per-model launch flags forwarded to `sglang.compile_deep_gemm`. DeepGEMM +# cache key includes per-rank N/K (depends on tp/dp/ep) — must match each +# model's `other_args` in test/registered/ or warmed shapes won't be hit. +FALLBACK_ARGS: Dict[str, List[str]] = { + "deepseek-ai/DeepSeek-V3.2": ["--dp", "8", "--enable-dp-attention"], + "zai-org/GLM-5-FP8": ["--dp", "8", "--enable-dp-attention"], + "XiaomiMiMo/MiMo-V2-Flash": [ + "--dp", + "2", + "--enable-dp-attention", + "--attention-backend", + "fa3", + ], + # --mm-enable-dp-encoder is required: without it DP0 runs the vision + # encoder alone and DP1 deadlocks at the next collective. + "XiaomiMiMo/MiMo-V2.5": [ + "--dp", + "2", + "--enable-dp-attention", + "--mm-enable-dp-encoder", + "--attention-backend", + "fa3", + "--mm-attention-backend", + "fa3", + ], +} + +# compile_deep_gemm polls /v1/models for the full timeout even after a TP rank +# dies; the watcher uses these to kill the group within seconds instead. +CRASH_MARKERS = ( + "Scheduler hit an exception", + "Received sigquit from a child", +) # Configure DeepGEMM cache before importing deep_gemm os.environ["DG_JIT_CACHE_DIR"] = os.getenv( @@ -294,9 +339,89 @@ def compile_shapes_lightweight(shapes, m_list): print(f" Done in {elapsed:.1f}s") +def _kill_pg_and_wait(proc): + """SIGTERM the subprocess's process group, escalate to SIGKILL if needed.""" + try: + os.killpg(os.getpgid(proc.pid), signal.SIGTERM) + except (ProcessLookupError, OSError): + pass + try: + return proc.wait(timeout=10) + except subprocess.TimeoutExpired: + try: + os.killpg(os.getpgid(proc.pid), signal.SIGKILL) + except (ProcessLookupError, OSError): + pass + try: + return proc.wait(timeout=5) + except subprocess.TimeoutExpired: + return -1 + + +def get_version_key(): + """Hash of Python + Triton + PyTorch versions; invalidates markers on upgrade.""" + parts = [sys.version] + try: + import triton # noqa: WPS433 + + parts.append(f"triton={triton.__version__}") + except ImportError: + parts.append("triton=none") + try: + import torch # noqa: WPS433 + + parts.append(f"torch={torch.__version__}") + except ImportError: + parts.append("torch=none") + return hashlib.sha256("|".join(parts).encode()).hexdigest()[:12] + + +def get_fallback_marker_path(model, tp, extra_args): + """Marker path for one (model, tp, extra_args) fallback invocation.""" + args_blob = json.dumps(list(extra_args)) + args_hash = hashlib.md5(args_blob.encode()).hexdigest()[:8] + safe_model = model.replace("/", "--") + return os.path.join( + MARKER_DIR, + f"deepgemm_fallback_{safe_model}_tp{tp}_{args_hash}_{get_version_key()}.done", + ) + + +def check_fallback_marker(model, tp, extra_args): + return os.path.exists(get_fallback_marker_path(model, tp, extra_args)) + + +def write_fallback_marker(model, tp, extra_args): + marker = get_fallback_marker_path(model, tp, extra_args) + os.makedirs(os.path.dirname(marker), exist_ok=True) + Path(marker).write_text( + json.dumps( + { + "model": model, + "tp": tp, + "extra_args": list(extra_args), + "version_key": get_version_key(), + "timestamp": time.time(), + } + ) + ) + print(f" Wrote marker: {marker}") + + def fallback_compile_deep_gemm(model, tp): - """Fall back to full sglang.compile_deep_gemm (loads model weights).""" - print(f"Falling back to full compile_deep_gemm for {model} (tp={tp})...") + """Fall back to full sglang.compile_deep_gemm (loads model weights). + + Runs in its own process group so a hung subprocess (e.g. one TP rank + crashes and the rest deadlock on NCCL collectives) can be killed + cleanly without leaking children. Watches subprocess output for crash + markers so a deterministic failure aborts in seconds rather than burning + the full FALLBACK_TIMEOUT_SEC. + """ + extra_args = FALLBACK_ARGS.get(model, []) + print( + f"Falling back to full compile_deep_gemm for {model} " + f"(tp={tp}, extra_args={extra_args})..." + ) cmd = [ sys.executable, "-m", @@ -308,11 +433,57 @@ def fallback_compile_deep_gemm(model, tp): "--trust-remote-code", "--model-loader-extra-config", '{"enable_multithread_load": true, "num_threads": 64}', + # Cap compile_deep_gemm's own /v1/models polling loop so it gives up + # before our outer timeout has to SIGTERM it. + "--timeout", + str(FALLBACK_TIMEOUT_SEC), + *extra_args, ] - result = subprocess.run(cmd) - if result.returncode != 0: - print(f"Warning: fallback failed for {model} (exit code {result.returncode})") - return result.returncode == 0 + + crashed = threading.Event() + proc = subprocess.Popen( + cmd, + preexec_fn=os.setsid, + stdout=subprocess.PIPE, + stderr=subprocess.STDOUT, + bufsize=1, + text=True, + ) + + def _watch(): + # Stream child output to our stdout while scanning for crash markers. + for line in proc.stdout: + sys.stdout.write(line) + sys.stdout.flush() + if not crashed.is_set() and any(m in line for m in CRASH_MARKERS): + crashed.set() + + watcher = threading.Thread(target=_watch, daemon=True) + watcher.start() + + deadline = time.monotonic() + FALLBACK_TIMEOUT_SEC + while True: + rc = proc.poll() + if rc is not None: + watcher.join(timeout=2) + if rc != 0: + print(f"Warning: fallback failed for {model} (exit code {rc})") + return rc == 0 + if crashed.is_set(): + print( + f"Warning: detected crash marker in {model} (tp={tp}) subprocess; " + "killing process group and continuing." + ) + _kill_pg_and_wait(proc) + return False + if time.monotonic() >= deadline: + print( + f"Warning: fallback timed out after {FALLBACK_TIMEOUT_SEC}s for " + f"{model} (tp={tp}); killing process group and continuing." + ) + _kill_pg_and_wait(proc) + return False + time.sleep(2) def main(): @@ -353,22 +524,39 @@ def main(): print(f" SKIP {model} (tp={tp}): config.json not in HF cache") continue + # Models with FALLBACK_ARGS launch with extra dp/ep/dp-attention flags + # that change per-rank N/K. The lightweight path doesn't model those — + # it computes attention shapes assuming TP-only sharding — so dedup'ing + # such a model to a no-override DeepSeek V2/V3 lookalike (e.g. V3.2 → + # V3-0324) silently picks the wrong attention shapes for the test's + # actual launch config. Force these through fallback so the populated + # cache matches the real test. + has_fallback_override = model in FALLBACK_ARGS + key = get_architecture_key(config, tp) - if key in seen_keys: + if key in seen_keys and not has_fallback_override: print(f" DEDUP {model} (tp={tp}): same shapes as {seen_keys[key]}") continue - if is_deepseek_v2v3(config): + if is_deepseek_v2v3(config) and not has_fallback_override: shapes = compute_deepseek_v2v3_shapes(config, tp) seen_keys[key] = model to_process.append((model, tp, config, shapes)) print(f" FOUND {model} (tp={tp}): {len(shapes)} DeepGEMM shape(s)") else: - # Unknown architecture: will use fallback seen_keys[key] = model to_process.append((model, tp, config, None)) - arch = config.get("architectures", ["unknown"]) - print(f" FOUND {model} (tp={tp}): unknown arch {arch}, will use fallback") + if has_fallback_override: + print( + f" FOUND {model} (tp={tp}): forced fallback (extra args " + f"{FALLBACK_ARGS[model]})" + ) + else: + arch = config.get("architectures", ["unknown"]) + print( + f" FOUND {model} (tp={tp}): unknown arch {arch}, " + "will use fallback" + ) if not to_process: print("\nNo models to process. Done.") @@ -383,8 +571,20 @@ def main(): print(f"{'=' * 60}") if shapes is None: - # Unknown architecture: fall back to full compile_deep_gemm - fallback_compile_deep_gemm(model, tp) + # Fallback path: full sglang.compile_deep_gemm with the test's launch + # flags. Loading model weights inside that subprocess is the dominant + # cost (45-170s/model) and dwarfs the actual DeepGEMM compile, so we + # skip the whole fallback when the marker says we already populated + # this cache. Cache is keyed on (model, tp, extra_args, version_key). + extra_args = FALLBACK_ARGS.get(model, []) + if check_fallback_marker(model, tp, extra_args): + print( + f" SKIP fallback (warm marker found): {model} (tp={tp}, " + f"extra_args={extra_args})" + ) + continue + if fallback_compile_deep_gemm(model, tp): + write_fallback_marker(model, tp, extra_args) continue # Print shape summary diff --git a/scripts/ci/cuda/warmup_server.py b/scripts/ci/cuda/warmup_server.py index d93541b05..c4d1abad5 100644 --- a/scripts/ci/cuda/warmup_server.py +++ b/scripts/ci/cuda/warmup_server.py @@ -11,8 +11,7 @@ invalidated when Python, Triton, or PyTorch versions change. Usage: python3 scripts/ci/cuda/warmup_server.py \ - deepseek-ai/DeepSeek-V3-0324:8 \ - inclusionAI/Ring-2.5-1T:8 + deepseek-ai/DeepSeek-V3-0324:8 """ import hashlib @@ -87,23 +86,41 @@ def write_marker(model, tp): def kill_server(proc): """Kill server process tree.""" - if proc.poll() is not None: - return - try: - os.killpg(os.getpgid(proc.pid), signal.SIGTERM) - except (ProcessLookupError, OSError): - pass - try: - proc.wait(timeout=15) - except subprocess.TimeoutExpired: + if proc.poll() is None: try: - os.killpg(os.getpgid(proc.pid), signal.SIGKILL) + os.killpg(os.getpgid(proc.pid), signal.SIGTERM) except (ProcessLookupError, OSError): pass try: - proc.wait(timeout=5) + proc.wait(timeout=15) except subprocess.TimeoutExpired: + try: + os.killpg(os.getpgid(proc.pid), signal.SIGKILL) + except (ProcessLookupError, OSError): + pass + try: + proc.wait(timeout=5) + except subprocess.TimeoutExpired: + pass + + # sglang's scheduler_TP* and detokenizer workers spawn through + # multiprocessing with their own session/process group, so they escape + # killpg on launch_server and stay alive holding GPU memory after a + # readiness-timeout or unclean exit. Kill any survivors by name so the + # next model (or the next CI step) starts with empty GPUs. + for pattern in ("sglang::scheduler", "sglang::detokenizer"): + try: + subprocess.run( + ["pkill", "-9", "-f", pattern], + timeout=5, + check=False, + stdout=subprocess.DEVNULL, + stderr=subprocess.DEVNULL, + ) + except (FileNotFoundError, subprocess.TimeoutExpired): pass + # Let the driver release device memory before the caller measures it. + time.sleep(2) def wait_for_server(base_url, proc, timeout): @@ -192,15 +209,6 @@ def warmup_one_model(model, tp, port): ok, err = wait_for_server(base_url, proc, SERVER_STARTUP_TIMEOUT) if not ok: print(f" Warning: server not ready: {err}") - # Dump last lines of server log for debugging - try: - log_file.flush() - with open(log_path) as f: - lines = f.readlines() - for line in lines[-20:]: - print(f" | {line.rstrip()}") - except Exception: - pass return False print(" Server ready, sending generate request...") @@ -208,6 +216,20 @@ def warmup_one_model(model, tp, port): return True finally: + # Surface the tail of the server log so CI captures validation + # messages, exceptions, and warmup progress (the launch_server + # subprocess writes stdout/stderr to the tempfile, not our stdout). + try: + log_file.flush() + with open(log_path) as f: + lines = f.readlines() + print(f" --- server log tail ({len(lines)} lines, last 30) ---") + for line in lines[-30:]: + print(f" | {line.rstrip()}") + print(" --- end server log ---") + except Exception: + pass + print(" Killing server...") kill_server(proc) log_file.close()