From f8b68f2452c424d22de4a69527a427207dcfbca5 Mon Sep 17 00:00:00 2001
From: Gahow Wang <gahow.wang@gmail.com>
Date: Sun, 12 Jul 2026 12:56:39 +0800
Subject: [PATCH] Checkpoint OpProf accounting across hard kills

---
 tests/v1/core/test_opprof.py | 118 +++++++++++++++++++++++++++++++++++
 vllm/v1/opprof.py            |  84 ++++++++++++++++++++++---
 2 files changed, 195 insertions(+), 7 deletions(-)

diff --git a/tests/v1/core/test_opprof.py b/tests/v1/core/test_opprof.py
index 0b8a8a1..007c5bb 100644
--- a/tests/v1/core/test_opprof.py
+++ b/tests/v1/core/test_opprof.py
@@ -5,6 +5,8 @@
 import errno
 import importlib.util
 import logging
+import os
+import subprocess
 import sys
 import threading
 from pathlib import Path
@@ -121,6 +123,12 @@ def read_jsonl(path):
     return [msgspec.json.decode(line) for line in path.read_bytes().splitlines()]
 
 
+def read_sidecar(path):
+    return msgspec.json.decode(
+        path.with_name(f"{path.name}.footer.json").read_bytes()
+    )
+
+
 def test_import_light_and_approved_constants():
     assert "torch" not in sys.modules
     assert "vllm" not in sys.modules
@@ -366,6 +374,116 @@ def test_shutdown_flush_is_idempotent(tmp_path):
     assert path.stat().st_size > 0
 
 
+def test_sidecar_updates_are_atomic(tmp_path, monkeypatch):
+    writer = opprof.JSONLWriter(tmp_path / "atomic.jsonl", start=False)
+    writer._write_sidecar(
+        encoded_records=1,
+        written_records=1,
+        dropped_records=0,
+        last_step_index=0,
+        final=False,
+    )
+    original = writer.sidecar_path.read_bytes()
+    real_replace = os.replace
+    replacements = []
+
+    def inspect_replace(source, destination):
+        assert Path(destination) == writer.sidecar_path
+        assert writer.sidecar_path.read_bytes() == original
+        candidate = msgspec.json.decode(Path(source).read_bytes())
+        assert candidate["written_records"] == 2
+        replacements.append(candidate)
+        real_replace(source, destination)
+
+    monkeypatch.setattr(opprof.os, "replace", inspect_replace)
+    writer._write_sidecar(
+        encoded_records=3,
+        written_records=2,
+        dropped_records=1,
+        last_step_index=2,
+        final=False,
+    )
+    monkeypatch.setattr(opprof.os, "replace", real_replace)
+
+    assert len(replacements) == 1
+    assert read_sidecar(writer.path)["encoded_records"] == 3
+    assert not list(tmp_path.glob(".*.tmp-*"))
+    writer.close()
+
+
+def test_hard_kill_sidecar_balances_last_flush(tmp_path):
+    path = tmp_path / "hard-kill.jsonl"
+    child = """
+import importlib.util
+import sys
+import time
+from pathlib import Path
+
+source, output = sys.argv[1:]
+spec = importlib.util.spec_from_file_location("opprof_hard_kill", source)
+module = importlib.util.module_from_spec(spec)
+sys.modules[spec.name] = module
+spec.loader.exec_module(module)
+writer = module.JSONLWriter(Path(output))
+for step in range(3):
+    assert writer.submit({"schema": 1, "step_index": step})
+writer._queue.join()
+while not writer.sidecar_path.exists():
+    time.sleep(0.01)
+print("READY", flush=True)
+time.sleep(60)
+"""
+    process = subprocess.Popen(
+        [
+            sys.executable,
+            "-c",
+            child,
+            str(_ROOT / "vllm/v1/opprof.py"),
+            str(path),
+        ],
+        stdout=subprocess.PIPE,
+        text=True,
+    )
+    try:
+        assert process.stdout is not None
+        assert process.stdout.readline().strip() == "READY"
+        process.kill()
+        assert process.wait(timeout=5) < 0
+    finally:
+        if process.poll() is None:
+            process.kill()
+            process.wait(timeout=5)
+
+    records = read_jsonl(path)
+    sidecar = read_sidecar(path)
+    assert len(records) == 3
+    assert all(record.get("record_type") != "footer" for record in records)
+    assert sidecar["final"] is False
+    assert sidecar["written_records"] == len(records)
+    assert sidecar["encoded_records"] == (
+        sidecar["written_records"] + sidecar["dropped_records"]
+    )
+    assert sidecar["last_step_index"] == records[-1]["step_index"] == 2
+
+
+def test_clean_footer_and_final_sidecar_agree(tmp_path):
+    sched = scheduler({"decode": request(8, 9)})
+    rec, path = recorder(tmp_path)
+    assert emit(rec, sched, schedule_output({"decode": 1}))
+    rec.close()
+
+    record, footer = read_jsonl(path)
+    sidecar = read_sidecar(path)
+    assert sidecar["record_type"] == "footer_checkpoint"
+    assert sidecar["stream"] == path.name
+    assert sidecar["final"] is True
+    assert sidecar["last_step_index"] == record["step_index"] == 0
+    assert sidecar["checkpoint_wall_ns"] > 0
+    assert sidecar["flush_interval_seconds"] == opprof.FLUSH_INTERVAL_SECONDS
+    for counter in ("encoded_records", "written_records", "dropped_records"):
+        assert sidecar[counter] == footer[counter]
+
+
 def test_piecewise_cudagraph_record_preserved(tmp_path):
     sched = scheduler({"decode": request(4096, 4097)})
     output = schedule_output({"decode": 1})
diff --git a/vllm/v1/opprof.py b/vllm/v1/opprof.py
index 75f63de..28d9635 100644
--- a/vllm/v1/opprof.py
+++ b/vllm/v1/opprof.py
@@ -7,6 +7,7 @@ import queue
 import threading
 import time
 from bisect import bisect_left
+from contextlib import suppress
 from pathlib import Path
 from typing import Any
 
@@ -18,6 +19,7 @@ SCHEMA_VERSION = 1
 CONTEXT_LENGTH_EDGES = tuple(1 << exponent for exponent in range(7, 18))
 CHUNK_SIZE_EDGES = tuple(1 << exponent for exponent in range(4, 12))
 DEFAULT_QUEUE_CAPACITY = 8192
+FLUSH_INTERVAL_SECONDS = 1.0
 _CLOSE_TIMEOUT_SECONDS = 1.0
 _STOP = object()
 _PREFIX_FIELDS = (  # noqa: SIM905
@@ -69,6 +71,7 @@ class JSONLWriter:
     ) -> None:
         self._queue: queue.Queue[Any] = queue.Queue(capacity)
         self.path = path
+        self.sidecar_path = path.with_name(f"{path.name}.footer.json")
         self._encoder = msgspec.json.Encoder()
         self._file = path.open("xb", buffering=1 << 20)
         self._thread = threading.Thread(target=self._run, daemon=True)
@@ -78,6 +81,9 @@ class JSONLWriter:
         self.failure: Exception | None = None
         self.encoded_records = self.written_records = 0
         self.dropped_records = self._unreported_drops = 0
+        self._checkpoint_encoded_records = 0
+        self._checkpoint_dropped_records = 0
+        self._last_written_step_index: int | None = None
         if start:
             self.start()
 
@@ -119,33 +125,90 @@ class JSONLWriter:
         if self._writer_unavailable():
             return self._drop(pending)
         try:
-            self._queue.put_nowait(payload)
+            self._queue.put_nowait(
+                (
+                    payload,
+                    self.encoded_records,
+                    self.dropped_records,
+                    int(record["step_index"]),
+                )
+            )
         except queue.Full:
             return self._drop(pending)
         self._unreported_drops = 0
         return True
 
+    def _write_sidecar(
+        self,
+        *,
+        encoded_records: int,
+        written_records: int,
+        dropped_records: int,
+        last_step_index: int | None,
+        final: bool,
+    ) -> None:
+        sidecar = dict(
+            schema=SCHEMA_VERSION,
+            record_type="footer_checkpoint",
+            stream=self.path.name,
+            encoded_records=encoded_records,
+            written_records=written_records,
+            dropped_records=dropped_records,
+            last_step_index=last_step_index,
+            checkpoint_wall_ns=time.time_ns(),
+            flush_interval_seconds=FLUSH_INTERVAL_SECONDS,
+            final=final,
+        )
+        temporary_path = self.sidecar_path.with_name(
+            f".{self.sidecar_path.name}.tmp-{os.getpid()}-{threading.get_ident()}"
+        )
+        try:
+            with temporary_path.open("wb") as temporary_file:
+                temporary_file.write(self._encoder.encode(sidecar) + b"\n")
+                temporary_file.flush()
+            os.replace(temporary_path, self.sidecar_path)
+        finally:
+            with suppress(FileNotFoundError):
+                temporary_path.unlink()
+
+    def _flush_checkpoint(self) -> None:
+        self._file.flush()
+        self._write_sidecar(
+            encoded_records=self._checkpoint_encoded_records,
+            written_records=self.written_records,
+            dropped_records=self._checkpoint_dropped_records,
+            last_step_index=self._last_written_step_index,
+            final=False,
+        )
+
     def _run(self) -> None:
         buffered = 0
         last_flush = time.monotonic()
         try:
             while True:
                 try:
-                    item = self._queue.get(timeout=1.0)
+                    item = self._queue.get(timeout=FLUSH_INTERVAL_SECONDS)
                 except queue.Empty:
-                    self._file.flush()
+                    self._flush_checkpoint()
                     buffered = 0
                     last_flush = time.monotonic()
                     continue
                 try:
                     if item is _STOP:
                         break
-                    self._file.write(item)
-                    buffered += len(item)
+                    payload, encoded, dropped, step_index = item
+                    self._file.write(payload)
+                    buffered += len(payload)
                     self.written_records += 1
+                    self._checkpoint_encoded_records = encoded
+                    self._checkpoint_dropped_records = dropped
+                    self._last_written_step_index = step_index
                     now = time.monotonic()
-                    if buffered >= 1 << 20 or now - last_flush >= 1.0:
-                        self._file.flush()
+                    if (
+                        buffered >= 1 << 20
+                        or now - last_flush >= FLUSH_INTERVAL_SECONDS
+                    ):
+                        self._flush_checkpoint()
                         buffered = 0
                         last_flush = now
                 finally:
@@ -163,6 +226,13 @@ class JSONLWriter:
             try:
                 self._file.write(self._encoder.encode(footer) + b"\n")
                 self._file.flush()
+                self._write_sidecar(
+                    encoded_records=footer["encoded_records"],
+                    written_records=footer["written_records"],
+                    dropped_records=footer["dropped_records"],
+                    last_step_index=self._last_written_step_index,
+                    final=True,
+                )
             except Exception as error:
                 self._record_failure(error)
             finally:
-- 
2.43.0

