diff --git a/src/tether/runtime/record.py b/src/tether/runtime/record.py index 6fba229f..c94b2f0d 100644 --- a/src/tether/runtime/record.py +++ b/src/tether/runtime/record.py @@ -10,6 +10,11 @@ performs file I/O so the /act event loop is not blocked by large JSONL lines. The worker still flushes per record, so a writer crash leaves at most one partial line (reader skips it per D.1.11). +- Every live writer is tracked in a module-level weak registry and drained by + an `atexit` hook, so queued records reach disk even when the ASGI lifespan + shutdown never runs (startup crash, bare `sys.exit`, non-uvicorn embedding). + The worker is a daemon thread — `atexit` handlers run *before* daemon threads + are killed, which is what makes the drain reachable at all. - Disk-full → degrade silently. Catches OSError, sets `self.degraded`, stops writing, but lets `tether serve` continue. /health surfaces this via `getattr(server, '_recorder', None).degraded` if a consumer wants. @@ -44,6 +49,7 @@ from __future__ import annotations +import atexit import copy import gzip import hashlib @@ -53,6 +59,7 @@ import queue import threading import uuid +import weakref from datetime import datetime, timezone from pathlib import Path from typing import Any, Literal @@ -63,8 +70,39 @@ RECORD_QUEUE_MAXSIZE = 1000 _QUEUE_STOP = object() +# Upper bound on how long the atexit drain waits for one writer's worker to +# finish. Unbounded would be "correct" but turns a stalled disk (NFS hang, +# full device retrying) into a process that never exits; a lost tail of records +# beats an unkillable `tether serve`. +RECORD_EXIT_DRAIN_TIMEOUT = 5.0 + ImageRedaction = Literal["full", "hash_only", "none"] +# Live writers, drained by _close_all_writers on interpreter exit. Weak so a +# recorder that is dropped without close() can still be collected — note the +# worker thread holds a strong ref to its writer, so entries stay alive for +# exactly as long as there is a thread that might still have queued work. +_active_writers: weakref.WeakSet[RecordWriter] = weakref.WeakSet() +_active_writers_lock = threading.Lock() + + +def _close_all_writers(timeout: float = RECORD_EXIT_DRAIN_TIMEOUT) -> None: + """Flush and close every live RecordWriter. Registered with `atexit`. + + Idempotent per writer (`close()` short-circuits once closed), so this is a + no-op when the server's lifespan shutdown already closed the recorder. + """ + with _active_writers_lock: + writers = list(_active_writers) + for writer in writers: + try: + writer.close(timeout=timeout) + except Exception as exc: # noqa: BLE001 — exit must not raise + logger.warning("RecordWriter atexit close failed: %s", exc) + + +atexit.register(_close_all_writers) + def _utc_now_iso() -> str: """UTC timestamp with ms precision, ISO-8601, trailing 'Z'.""" @@ -358,6 +396,10 @@ def __init__( daemon=True, ) self._worker.start() + # Track for the atexit drain only once the worker is actually running — + # a writer with no worker has nothing to flush. + with _active_writers_lock: + _active_writers.add(self) # Curate dual-write: when a FreeContributorCollector is attached, # write_request emits to BOTH the JSONL trace (audit) AND the # curate queue (training corpus). Independent failure modes; if @@ -705,10 +747,20 @@ def write_footer(self, totals: dict[str, int]) -> None: } self._enqueue_records((record,)) - def close(self) -> None: + def close(self, timeout: float | None = None) -> None: + """Drain the queue, close the file, and stop the worker. + + Blocks until every record queued so far is on disk. `timeout` bounds + that wait in seconds (None = wait indefinitely); on expiry the + still-queued tail is abandoned rather than hanging the caller. Safe to + call twice — the second call is a no-op — and safe to call from the + worker thread itself, where joining would otherwise deadlock. + """ if self._closed: return self._closed = True + with _active_writers_lock: + _active_writers.discard(self) # Stop the curate collector first so its drain has a chance to # flush queued events before the process exits. if self._curate_collector is not None: @@ -716,13 +768,32 @@ def close(self) -> None: self._curate_collector.stop() except Exception as exc: # noqa: BLE001 logger.warning("curate_collector.stop failed: %s", exc) - self.flush_sync() - self._record_queue.put(_QUEUE_STOP) - self._record_queue.join() - if self._worker is not threading.current_thread(): - self._worker.join(timeout=1.0) - if self._worker.is_alive(): - logger.warning("RecordWriter worker did not stop cleanly") + if self._worker is threading.current_thread(): + # Reached from inside the worker (e.g. a close() in a record hook). + # Nothing left to hand off — close the file in place. + self._close_file() + return + # The queue is FIFO and the worker only returns on _QUEUE_STOP, so the + # sentinel landing on the queue is enough to guarantee everything ahead + # of it is written; joining the thread is the drain. + try: + self._record_queue.put(_QUEUE_STOP, timeout=timeout) + except queue.Full: + logger.warning( + "RecordWriter close: queue still full after %ss — %d record(s) not flushed (%s)", + timeout, + self._record_queue.qsize(), + self.filepath, + ) + return + self._worker.join(timeout=timeout) + if self._worker.is_alive(): + logger.warning( + "RecordWriter worker did not stop cleanly within %ss — " + "trailing records may be missing (%s)", + timeout, + self.filepath, + ) # --------------------------------------------------------------- # Convenience diff --git a/tests/test_record.py b/tests/test_record.py index beb7e29d..d8f6e6ec 100644 --- a/tests/test_record.py +++ b/tests/test_record.py @@ -11,12 +11,16 @@ import gzip import hashlib import json +import subprocess +import sys import threading import time from pathlib import Path +from typing import Any import pytest +import tether.runtime.record as record_mod from tether.runtime.record import ( SCHEMA_VERSION, RecordWriter, @@ -30,22 +34,25 @@ # --------------------------------------------------------------------------- -def _make_writer(tmp_path: Path, **kwargs) -> RecordWriter: +def _make_writer(tmp_path: Path, **kwargs: Any) -> RecordWriter: """Factory with sensible defaults for tests.""" - defaults = dict( - model_hash="abc123def4567890", - config_hash="0123456789abcdef", - export_dir=str(tmp_path / "fake_export"), - model_type="pi0.5", - export_kind="monolithic", - providers=["CUDAExecutionProvider"], - gpu="test-gpu", - cuda_version="12.6", - ort_version="1.20.1", - embodiment="franka", - image_redaction="hash_only", - tether_version="0.0.0-test", - ) + # Annotated dict[str, Any]: inference would otherwise narrow this to the + # join of the value types (Sequence[str]) and reject the ** unpack against + # RecordWriter's precisely-typed keywords. + defaults: dict[str, Any] = { + "model_hash": "abc123def4567890", + "config_hash": "0123456789abcdef", + "export_dir": str(tmp_path / "fake_export"), + "model_type": "pi0.5", + "export_kind": "monolithic", + "providers": ["CUDAExecutionProvider"], + "gpu": "test-gpu", + "cuda_version": "12.6", + "ort_version": "1.20.1", + "embodiment": "franka", + "image_redaction": "hash_only", + "tether_version": "0.0.0-test", + } defaults.update(kwargs) return RecordWriter(record_dir=tmp_path, **defaults) @@ -57,17 +64,17 @@ def _read_all(path: Path) -> list[dict]: return [json.loads(line) for line in f if line.strip()] -def _dummy_request(rec: RecordWriter, i: int = 0, **overrides) -> int: - kw = dict( - chunk_id=i, - image_b64="aGVsbG8gd29ybGQ=", - instruction=f"test instruction {i}", - state=[0.1, 0.2, 0.3], - actions=[[0.0] * 7] * 50, - action_dim=7, - latency_total_ms=100.0 + i, - mode="onnx_gpu", - ) +def _dummy_request(rec: RecordWriter, i: int = 0, **overrides: Any) -> int: + kw: dict[str, Any] = { + "chunk_id": i, + "image_b64": "aGVsbG8gd29ybGQ=", + "instruction": f"test instruction {i}", + "state": [0.1, 0.2, 0.3], + "actions": [[0.0] * 7] * 50, + "action_dim": 7, + "latency_total_ms": 100.0 + i, + "mode": "onnx_gpu", + } kw.update(overrides) return rec.write_request(**kw) @@ -490,6 +497,138 @@ def test_close_drains_pending_records(self, tmp_path): assert records[-1]["total_requests"] == 10 +# --------------------------------------------------------------------------- +# Graceful shutdown / process exit +# --------------------------------------------------------------------------- + + +_EXIT_SCRIPT = """ +import sys +from tether.runtime.record import RecordWriter + +rec = RecordWriter( + record_dir=sys.argv[1], + model_hash="abc123def4567890", + config_hash="0123456789abcdef", + export_dir=sys.argv[1] + "/fake_export", + model_type="pi0.5", + export_kind="monolithic", + providers=["CPUExecutionProvider"], + gzip_output={gzip_output}, +) +for i in range(25): + rec.write_request( + chunk_id=i, + image_b64="aGVsbG8gd29ybGQ=", + instruction="exit test %d" % i, + state=[0.1, 0.2, 0.3], + actions=[[0.0] * 7] * 50, + action_dim=7, + latency_total_ms=1.0, + ) +print(rec.filepath) +# Deliberately no close() and no lifespan shutdown — the atexit hook is the +# only thing that can get these records onto disk before the daemon worker +# is killed at interpreter exit. +""" + + +def _run_exit_script(tmp_path: Path, *, gzip_output: bool) -> Path: + """Run a child interpreter that records and exits without close().""" + proc = subprocess.run( + [ + sys.executable, + "-c", + _EXIT_SCRIPT.format(gzip_output=gzip_output), + str(tmp_path), + ], + capture_output=True, + text=True, + timeout=60, + check=False, # assert on returncode below so stderr lands in the failure + ) + assert proc.returncode == 0, f"child failed: {proc.stderr}" + return Path(proc.stdout.strip().splitlines()[-1]) + + +class TestGracefulShutdown: + def test_atexit_flushes_records_on_process_exit(self, tmp_path): + """A process that exits without calling close() still lands every + queued record on disk, via the atexit drain.""" + filepath = _run_exit_script(tmp_path, gzip_output=False) + + records = _read_all(filepath) + requests = [r for r in records if r["kind"] == "request"] + assert records[0]["kind"] == "header" + assert len(requests) == 25 + assert [r["seq"] for r in requests] == list(range(25)) + + def test_atexit_closes_gzip_stream_cleanly(self, tmp_path): + """Gzip is the strict case: a stream whose writer was killed mid-flight + has no trailer and fails to decompress. Reading it back proves close() + actually ran rather than the file merely happening to have bytes.""" + filepath = _run_exit_script(tmp_path, gzip_output=True) + + assert filepath.suffix == ".gz" + records = _read_all(filepath) # raises/truncates if the trailer is missing + assert len([r for r in records if r["kind"] == "request"]) == 25 + + def test_writer_deregisters_itself_on_close(self, tmp_path): + """Closed writers drop out of the exit registry so the atexit hook + doesn't touch them (and doesn't pin them in memory).""" + rec = _make_writer(tmp_path, gzip_output=False) + assert rec in record_mod._active_writers + + rec.close() + assert rec not in record_mod._active_writers + + def test_close_all_writers_drains_a_live_writer(self, tmp_path): + rec = _make_writer(tmp_path, gzip_output=False) + for i in range(5): + _dummy_request(rec, i=i) + + record_mod._close_all_writers() + + assert rec._closed is True + assert len([r for r in _read_all(rec.filepath) if r["kind"] == "request"]) == 5 + + def test_close_is_idempotent(self, tmp_path): + """The lifespan shutdown and the atexit hook both call close(); the + second call must be a harmless no-op.""" + rec = _make_writer(tmp_path, gzip_output=False) + _dummy_request(rec) + rec.close() + size_after_first = rec.filepath.stat().st_size + + rec.close() # lifespan already closed it; atexit runs anyway + record_mod._close_all_writers() + + assert rec.filepath.stat().st_size == size_after_first + assert len(_read_all(rec.filepath)) == 2 # header + request + + def test_close_timeout_does_not_hang_on_stuck_worker(self, tmp_path, monkeypatch): + """A wedged worker (stalled disk) must not make close() block forever — + exit gives up on the tail instead of hanging the process.""" + rec = _make_writer(tmp_path, gzip_output=False) + release_emit = threading.Event() + original_emit = rec._emit + + def blocking_emit(record): + release_emit.wait(30.0) + original_emit(record) + + monkeypatch.setattr(rec, "_emit", blocking_emit) + _dummy_request(rec) + + started_at = time.perf_counter() + rec.close(timeout=0.2) + elapsed = time.perf_counter() - started_at + + assert elapsed < 5.0 + assert rec._closed is True + release_emit.set() + + # --------------------------------------------------------------------------- # Disk-full degraded path (D.1.11) # ---------------------------------------------------------------------------