From f8b68f2452c424d22de4a69527a427207dcfbca5 Mon Sep 17 00:00:00 2001 From: Gahow Wang 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