Просмотр исходного кода

readiness: watch() must not call an API that never opened 'dead'

--no-wait follow declared 'containers died mid-follow' while sglang was
still loading weights (17+ min on the 510 GB snapshot) — all containers
were Up. Liveness now counts misses only after the head's API port has
answered once (or listens but misbehaves); connection-refused before
first readiness is 'still loading'.
Plus sparse-copy header repair: stage_sparse_layer left the safetensors
header region hole-zeroed (aiden overlay reads it first -> 'header is not
JSON'); header is now copied/repaired without recopying ranges.
Alexander Korolev 3 недель назад
Родитель
Сommit
6ae3c5159c
2 измененных файлов с 56 добавлено и 0 удалено
  1. 21 0
      src/readiness.rs
  2. 35 0
      src/staging/scripts/sparse_rows.py

+ 21 - 0
src/readiness.rs

@@ -7,6 +7,7 @@
 use crate::config::Cluster;
 use anyhow::Result;
 use std::time::{Duration, Instant};
+use tokio::net::TcpStream;
 
 /// `sm121 serve`: big checkpoints load for many minutes.
 pub(crate) const SERVE_WAIT_SECS: u64 = 1800;
@@ -16,6 +17,15 @@ pub(crate) const BENCH_WAIT_SECS: u64 = 900;
 const WATCH_POLL_SECS: u64 = 30;
 const WATCH_MISSES: u32 = 2;
 
+fn head_host_port(cluster: &Cluster, port: u16) -> (String, u16) {
+    let head = cluster
+        .nodes
+        .iter()
+        .find(|n| n.rank == 0)
+        .or_else(|| cluster.nodes.first());
+    (head.map_or("", |n| n.ssh_ip.as_str()).to_string(), port)
+}
+
 fn head_url(cluster: &Cluster, port: u16) -> String {
     let head = cluster
         .nodes
@@ -77,13 +87,24 @@ pub(crate) async fn alive(cluster: &Cluster, port: u16) -> bool {
 
 /// Follow-mode liveness watch: resolves `Err` once the engine misses
 /// `WATCH_MISSES` consecutive probes — the containers died mid-follow.
+/// The API port opens only when the engine is ready (weight load can take
+/// 15+ min), so a connection-level failure before first readiness is "still
+/// loading", not "died"; misses count only after the engine answered once
+/// (or the port is listening but /v1/models misbehaves).
 pub(crate) async fn watch(cluster: &Cluster, port: u16) -> Result<()> {
+    let (host, p) = head_host_port(cluster, port);
     let mut misses = 0;
+    let mut ever_ready = false;
     loop {
         tokio::time::sleep(Duration::from_secs(WATCH_POLL_SECS)).await;
         if alive(cluster, port).await {
+            ever_ready = true;
             misses = 0;
         } else {
+            let listening = TcpStream::connect((host.as_str(), p)).await.is_ok();
+            if !ever_ready && !listening {
+                continue; // API not open yet — normal while loading
+            }
             misses += 1;
             if misses >= WATCH_MISSES {
                 anyhow::bail!(

+ 35 - 0
src/staging/scripts/sparse_rows.py

@@ -183,6 +183,33 @@ def copy_range(src_f, dst_f, off: int, length: int, started: float, tag: str) ->
                   flush=True)
 
 
+def headers_ok(src: Path, dst: Path) -> bool:
+    """The safetensors header (8-byte length + JSON) must be byte-identical —
+    the overlay loaders parse it before touching rows; holes read as len 0."""
+    try:
+        n = struct.unpack("<Q", _head_read(src, 0, 8))[0]
+        return _head_read(src, 0, 8 + n) == _head_read(dst, 0, 8 + n)
+    except Exception:  # noqa: BLE001
+        return False
+
+
+def _head_read(path: Path, off: int, n: int) -> bytes:
+    with open(path, "rb", buffering=0) as f:
+        f.seek(off)
+        b = f.read(n)
+    if len(b) != n:
+        raise OSError(f"{path}: short read {len(b)} < {n}")
+    return b
+
+
+def copy_header(src: Path, dst: Path) -> None:
+    n = struct.unpack("<Q", _head_read(src, 0, 8))[0]
+    head = _head_read(src, 0, 8 + n)
+    with open(dst, "r+b", buffering=0) as f:
+        f.seek(0)
+        f.write(head)
+
+
 def stage_sparse_layer(out: Path, layer: int, recs, lo: int, hi: int, ledger: Ledger,
                        samples_n: int, force: bool, verify: str) -> None:
     """One rank slice into node-local sparse shard copies at the ORIGINAL
@@ -214,6 +241,11 @@ def stage_sparse_layer(out: Path, layer: int, recs, lo: int, hi: int, ledger: Le
     t0 = time.monotonic()
     verified_only = False
     if not force and files_ok():
+        # header repair fast path: builds before 2026-09-27 left the header
+        # region hole-zeroed, which the loaders read as "not JSON"
+        for srcp, dstp in ((wpath, dst_w), (spath, dst_s)):
+            if not headers_ok(srcp, dstp):
+                copy_header(srcp, dstp)
         if rec is not None and rec.get("lo") == lo and rec.get("hi") == hi:
             if verify == "presence" or check([(int(k), v) for k, v in rec["sample"].items()]):
                 verified_only = True
@@ -239,6 +271,9 @@ def stage_sparse_layer(out: Path, layer: int, recs, lo: int, hi: int, ledger: Le
                 os.truncate(dst, size)
         elif dst.stat().st_size != size:
             os.truncate(dst, size)
+    for srcp, dstp in ((wpath, dst_w), (spath, dst_s)):
+        if not headers_ok(srcp, dstp):
+            copy_header(srcp, dstp)
     with open(wpath, "rb", buffering=0) as sf, open(dst_w, "r+b", buffering=0) as df:
         copy_range(sf, df, woff + lo * WEIGHT_BYTES, (hi - lo) * WEIGHT_BYTES,
                    t0, f"l{layer}-weight")