#!/usr/bin/env python3
"""Make SGLang's post-load monitored_barrier timeout configurable.

The launcher boots this lane with a per-node load time that varies wildly
(warm NFS cache ~40 min, cold up to 3.5 h on GB10 UMA where page-cache
pressure slows the tail ranks). The stock post-load barrier is 480 s — a
slow-but-fine rank gets murdered right after rank 0 finishes
("Rank N failed to pass monitoredBarrier in 480000 ms").

Patch (idempotent, anchor-based — refuses to touch a file whose anchor does
not match):
  UNBALANCED_MODEL_LOADING_TIMEOUT_S = 480
  →
  UNBALANCED_MODEL_LOADING_TIMEOUT_S = int(
      os.environ.get("SGLANG_UNBALANCED_LOAD_TIMEOUT_S", "480")
  )

`os` is already imported at module top. Runtime knob:
  SGLANG_UNBALANCED_LOAD_TIMEOUT_S=<seconds>  (recipe [env]).
"""

import hashlib
import re
import sys
from pathlib import Path

ANCHOR = re.compile(
    r"^UNBALANCED_MODEL_LOADING_TIMEOUT_S = 480"
    r"(  # leave more time for post data processing)?$",
    re.M,
)
NEW = (
    "UNBALANCED_MODEL_LOADING_TIMEOUT_S = int("
    "\n    os.environ.get(\"SGLANG_UNBALANCED_LOAD_TIMEOUT_S\", \"480\")"
    ")  # leave more time for post data processing; launcher-tunable"
)

TARGET = Path(
    "/sgl-workspace/sglang/python/sglang/srt/model_executor"
    "/model_runner_components/load_model_utils.py"
)


def main() -> None:
    args = sys.argv[1:]
    check = "--check" in args
    src = TARGET.read_text()
    matches = ANCHOR.findall(src)
    patched = NEW in src
    if patched:
        print(f"already patched: {TARGET}")
        sys.exit(0 if check else 0)
    if len(matches) != 1:
        sys.exit(
            f"refusing: anchor matched {len(matches)} times in {TARGET} "
            "(upstream changed; the patch must be re-anchored)"
        )
    if check:
        print(f"anchor present, not applied: {TARGET}")
        return
    out = ANCHOR.sub(lambda m: NEW, src, count=1)
    before = hashlib.md5(src.encode()).hexdigest()
    TARGET.write_text(out)
    after_hash = hashlib.md5(TARGET.read_text().encode()).hexdigest()
    print(f"patched {TARGET} (md5 {before} -> {after_hash})")


if __name__ == "__main__":
    main()