From 3971e783e65d748ce08822ed1ffed0b16f27c064 Mon Sep 17 00:00:00 2001 From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com> Date: Mon, 24 Aug 2026 19:55:43 +0000 Subject: [PATCH 1/6] Initial plan From cc6683c6a6802628bf34101e3b98df1539ff8181 Mon Sep 17 00:00:00 2001 From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com> Date: Mon, 24 Aug 2026 20:01:31 +0000 Subject: [PATCH 2/6] Surface host sink failures loudly Co-authored-by: zaoxing <2923149+zaoxing@users.noreply.github.com> --- docs/huggingface.md | 15 +++--- docs/integration-api-v1.md | 11 ++-- native/csrc/bindings.cpp | 43 ++++++++++++++++ native/csrc/dmx_host_engine.h | 10 +++- native/csrc/ring/p2p_thread.cpp | 6 +++ native/csrc/ring/p2p_thread.h | 3 ++ native/csrc/ring/ring_engine.h | 3 ++ native/csrc/ring/ring_engine_py.cu | 4 ++ native/csrc/ring/ring_engine_py.h | 1 + src/dmi/engine.py | 80 ++++++++++++++++++++++++++---- src/dmi/storage/internals.py | 14 +++--- tests/test_engine_runtime_api.py | 66 ++++++++++++++++++++++++ tests/test_storage_internals.py | 13 +++++ 13 files changed, 240 insertions(+), 29 deletions(-) diff --git a/docs/huggingface.md b/docs/huggingface.md index 7cca14a9e..dce100d35 100644 --- a/docs/huggingface.md +++ b/docs/huggingface.md @@ -86,9 +86,11 @@ token_mask = out.dmi_internal.token_mask For layer-tuple fields such as `hidden_states`, `count` validates the number of layers in the reassembled tuple. It does not validate token completeness inside -each layer tensor unless `match_token_ranges=True` is set. That option checks -the captured row ranges against the token ranges recorded during this generate -call. For per-layer fields it uses a fast representative-layer check. +each layer tensor unless token-range validation runs. Lazy reads attached to +`generate_with_monitoring_dict(...)` now perform that validation by default +whenever request IDs and token ranges were recorded for the generate call. +`match_token_ranges=False` opts out for a specific requirement. For per-layer +fields the current implementation uses a fast representative-layer check. Supported mapped fields are: @@ -147,9 +149,10 @@ norm = hidden_states[0].float().norm(dim=-1)[token_mask].mean() If a tensor field and `token_mask` have different `[batch, seq]` shapes, it usually means the tensor field was read before all rows for that field arrived. -Use `match_token_ranges=True` with `retry=True` to wait for expected token -ranges. If you already cached a partial field, clear that field's cache and read -it again: +Use `retry=True` to wait for expected token ranges after an incomplete read. If +you intentionally want to skip token-range validation for a field, set +`match_token_ranges=False` on that field's requirement. If you already cached a +partial field, clear that field's cache and read it again: ```python out.dmi_internal.clear_cache("hidden_states") diff --git a/docs/integration-api-v1.md b/docs/integration-api-v1.md index 7a6136c03..c0863a36e 100644 --- a/docs/integration-api-v1.md +++ b/docs/integration-api-v1.md @@ -1164,7 +1164,7 @@ handle.require( retry: bool = False, timeout_s: float | None = 30.0, poll_s: float = 0.25, - match_token_ranges: bool = False, + match_token_ranges: bool = True, ) -> handle ``` @@ -1172,10 +1172,11 @@ handle.require( and a batch count for global tensors, not a token or database-row count. With `retry=True`, synchronous field access polls missing/incomplete data until success or timeout. `timeout_s=None` can block forever. Database/runtime errors -are not retried. `match_token_ranges=True` performs validation only when both -nonempty request IDs and ranges were supplied; otherwise it is a no-op. For a -per-layer field the current v1 implementation checks only its highest present -layer, not every layer. +are not retried. When both nonempty request IDs and ranges were supplied, +token-range validation runs by default; `match_token_ranges=False` opts out for +an individual field, and otherwise the flag is a no-op. For a per-layer field +the current v1 implementation checks only its highest present layer, not every +layer. Mapped dynamic attributes are: diff --git a/native/csrc/bindings.cpp b/native/csrc/bindings.cpp index 9fa95f7b4..ce64b7a6d 100644 --- a/native/csrc/bindings.cpp +++ b/native/csrc/bindings.cpp @@ -98,6 +98,10 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { using QueueConfig = DMXHostEngine::QueueConfig; using EnqueuePolicy = DMXHostEngine::EnqueuePolicy; using Duration = DMXHostEngine::Duration; + using IngestStats = DMXHostEngine::IngestStats; + using QueueStats = DMXHostEngine::QueueStats; + using StageStats = DMXHostEngine::StageStats; + using StatsSnapshot = DMXHostEngine::StatsSnapshot; py::enum_(m, "OnFullPolicy") .value("RAISE", dmx_host::OnFullPolicy::RAISE) @@ -159,6 +163,41 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { .def_readonly("exc_type", &ThreadFailure::exc_type) .def_readonly("exc_what", &ThreadFailure::exc_what); + py::class_(m, "IngestStats") + .def_readonly("submit_calls", &IngestStats::submit_calls) + .def_readonly("items_submitted", &IngestStats::items_submitted) + .def_readonly("submit_enqueue_calls", &IngestStats::submit_enqueue_calls) + .def_readonly("submit_enqueue_s", &IngestStats::submit_enqueue_s); + + py::class_(m, "QueueStats") + .def_readonly("enqueued", &QueueStats::enqueued) + .def_readonly("dropped", &QueueStats::dropped) + .def_readonly("full_errors", &QueueStats::full_errors) + .def_readonly("closed_errors", &QueueStats::closed_errors) + .def_readonly("too_large_errors", &QueueStats::too_large_errors) + .def_readonly("retries", &QueueStats::retries); + + py::class_(m, "StageStats") + .def_readonly("batches", &StageStats::batches) + .def_readonly("items_in", &StageStats::items_in) + .def_readonly("items_out", &StageStats::items_out) + .def_readonly("dequeue_calls", &StageStats::dequeue_calls) + .def_readonly("dequeue_timeouts", &StageStats::dequeue_timeouts) + .def_readonly("process_calls", &StageStats::process_calls) + .def_readonly("enqueue_calls", &StageStats::enqueue_calls) + .def_readonly("output_calls", &StageStats::output_calls) + .def_readonly("output_items", &StageStats::output_items) + .def_readonly("dequeue_s", &StageStats::dequeue_s) + .def_readonly("dequeue_idle_s", &StageStats::dequeue_idle_s) + .def_readonly("process_s", &StageStats::process_s) + .def_readonly("enqueue_s", &StageStats::enqueue_s) + .def_readonly("output_s", &StageStats::output_s); + + py::class_(m, "StatsSnapshot") + .def_readonly("ingest", &StatsSnapshot::ingest) + .def_readonly("queue_by_stage", &StatsSnapshot::queue_by_stage) + .def_readonly("stage_by_stage", &StatsSnapshot::stage_by_stage); + py::class_(m, "StageConfig") .def(py::init<>()) .def_readwrite("name", &StageConfig::name) @@ -213,6 +252,8 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { py::call_guard()) .def("failures", &DMXHostEngine::failures) .def("raise_if_failed", &DMXHostEngine::raise_if_failed) + .def("profiling", &DMXHostEngine::profiling) + .def("reset_metrics", &DMXHostEngine::reset_metrics) // Submit a pre-formatted ClickHouseRow directly to the insert stage. // Called from the ring transport drain callback after format processing. .def("submit_direct", @@ -304,6 +345,8 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { .def("payload_cap", &ring_py::RingEnginePy::payload_cap) .def("staging_cap", &ring_py::RingEnginePy::staging_cap) .def("task_cap", &ring_py::RingEnginePy::task_cap) + .def("suppressed_submit_failures", + &ring_py::RingEnginePy::suppressed_submit_failures) .def("payload_tensor", &ring_py::RingEnginePy::payload_tensor) // Safety-net surface (eager only). available_capacity() and // reserve_one() are CPU-only and fast -- no GIL release needed. diff --git a/native/csrc/dmx_host_engine.h b/native/csrc/dmx_host_engine.h index f424b9114..caff5e120 100644 --- a/native/csrc/dmx_host_engine.h +++ b/native/csrc/dmx_host_engine.h @@ -8,11 +8,17 @@ namespace dmx_host{ // DMXHostEngine is a single-stage ClickHouse insert pipeline. // Pre-assembled ClickHouseRows are submitted via submit_direct(). -class DMXHostEngine : public PipelinedEngine, false, +class DMXHostEngine : public PipelinedEngine, true, NoOutputHandler >{ public: explicit DMXHostEngine(StageConfig insert_stage): - PipelinedEngine(std::array{std::move(insert_stage)}, EngineConfig{}){} + PipelinedEngine( + std::array{std::move(insert_stage)}, + [] { + EngineConfig config{}; + config.enable_stats = true; + return config; + }()){} // Submit a pre-assembled ClickHouseRow directly to the insert stage. // Fields must match the order expected by ClickHouseInsertStage: diff --git a/native/csrc/ring/p2p_thread.cpp b/native/csrc/ring/p2p_thread.cpp index f68426297..5aaeab436 100644 --- a/native/csrc/ring/p2p_thread.cpp +++ b/native/csrc/ring/p2p_thread.cpp @@ -139,6 +139,10 @@ void P2PThread::stop() { if (thread_.joinable()) thread_.join(); } +uint64_t P2PThread::suppressed_submit_failures() const noexcept { + return suppressed_submit_failures_.load(std::memory_order_relaxed); +} + // --------------------------------------------------------------------------- void P2PThread::loop() { while (true) { @@ -312,10 +316,12 @@ void P2PThread::do_post_processing(at::Tensor& tensor, const DrainTask& first_ta db_start, db_end, std::move(slice)); } catch (const std::exception& e) { + suppressed_submit_failures_.fetch_add(1, std::memory_order_relaxed); log_submit_failure_once(current_ctx_->model_id, req.req_id, act_name, meta.layer_no, shard_rank, db_start, db_end, e.what()); } catch (...) { + suppressed_submit_failures_.fetch_add(1, std::memory_order_relaxed); log_submit_failure_once(current_ctx_->model_id, req.req_id, act_name, meta.layer_no, shard_rank, db_start, db_end, diff --git a/native/csrc/ring/p2p_thread.h b/native/csrc/ring/p2p_thread.h index 6ed8d0f80..6f67d6cc4 100644 --- a/native/csrc/ring/p2p_thread.h +++ b/native/csrc/ring/p2p_thread.h @@ -11,6 +11,7 @@ #include "ring_config.h" #include "tensor_meta.h" +#include #include #include #include @@ -43,6 +44,7 @@ class P2PThread { void start(); void stop(); + uint64_t suppressed_submit_failures() const noexcept; private: DrainThread& drain_; @@ -52,6 +54,7 @@ class P2PThread { std::thread thread_; ring_py::StepContext* current_ctx_{nullptr}; // owned, freed on last_in_step + std::atomic suppressed_submit_failures_{0}; void loop(); void process(std::vector& tasks); diff --git a/native/csrc/ring/ring_engine.h b/native/csrc/ring/ring_engine.h index 4a56a5f9a..dfdd62e12 100644 --- a/native/csrc/ring/ring_engine.h +++ b/native/csrc/ring/ring_engine.h @@ -31,6 +31,9 @@ class RingEngine { uint64_t payload_cap() const { return cfg_.payload_ring_bytes; } uint64_t staging_cap() const { return staging_.capacity(); } uint64_t task_cap() const { return cfg_.task_ring_entries; } + uint64_t suppressed_submit_failures() const { + return p2p_ ? p2p_->suppressed_submit_failures() : 0; + } private: RingConfig cfg_; diff --git a/native/csrc/ring/ring_engine_py.cu b/native/csrc/ring/ring_engine_py.cu index c91c6d830..61bf2d887 100644 --- a/native/csrc/ring/ring_engine_py.cu +++ b/native/csrc/ring/ring_engine_py.cu @@ -268,6 +268,10 @@ uint64_t RingEnginePy::task_cap() const { return impl_->engine.task_cap(); } +uint64_t RingEnginePy::suppressed_submit_failures() const { + return impl_->engine.suppressed_submit_failures(); +} + at::Tensor RingEnginePy::payload_tensor() const { return impl_->payload_view; } diff --git a/native/csrc/ring/ring_engine_py.h b/native/csrc/ring/ring_engine_py.h index e96e1e4a6..faa3eb4cc 100644 --- a/native/csrc/ring/ring_engine_py.h +++ b/native/csrc/ring/ring_engine_py.h @@ -136,6 +136,7 @@ class RingEnginePy { uint64_t payload_cap() const; uint64_t staging_cap() const; uint64_t task_cap() const; + uint64_t suppressed_submit_failures() const; // Return a torch.Tensor view of the GPU payload buffer (uint8, // length = payload_cap()). No copy, no ownership transfer -- the diff --git a/src/dmi/engine.py b/src/dmi/engine.py index 89667ce67..286941f02 100644 --- a/src/dmi/engine.py +++ b/src/dmi/engine.py @@ -22,6 +22,51 @@ def _ring_module() -> Any: return importlib.import_module("dmi.transport.ring") +def _format_host_engine_stats(host_engine: Any, ring_engine: Any) -> str: + details: list[str] = [] + + try: + profiling = host_engine.profiling() + except Exception: + profiling = None + if profiling is not None: + try: + queue_stats = profiling.queue_by_stage[0] + except Exception: + queue_stats = None + if queue_stats is not None: + details.append( + "queue_stats=" + f"dropped={int(queue_stats.dropped)} " + f"full_errors={int(queue_stats.full_errors)} " + f"retries={int(queue_stats.retries)}" + ) + + if ring_engine is not None: + try: + suppressed = int(ring_engine.suppressed_submit_failures()) + except Exception: + suppressed = None + if suppressed: + details.append(f"suppressed_submit_failures={suppressed}") + + return "; ".join(details) + + +def _host_engine_teardown_error( + host_engine: Any, + ring_engine: Any, + exc: Exception, +) -> RuntimeError: + details = _format_host_engine_stats(host_engine, ring_engine) + message = "DMX host sink failed during teardown" + if details: + message = f"{message} ({details})" + error = RuntimeError(message) + error.__cause__ = exc + return error + + @dataclass(frozen=True, slots=True) class RingCapacities: """Immutable snapshot of the active ring transport's capacities.""" @@ -280,6 +325,9 @@ def next_auto_group_id(self) -> int: def close(self) -> None: """Tear down backend resources.""" + teardown_error: Exception | None = None + ring_engine = getattr(self, "_ring_engine", None) + if self._ring_transport is not None: # Best-effort reset of the device-global native null flag. This is # needed only after callers explicitly disabled capture; the normal @@ -290,27 +338,41 @@ def close(self) -> None: except Exception: pass try: - ring_engine = getattr(self, "_ring_engine", None) if ring_engine is not None: ring_engine.stop() - except Exception: - pass + except Exception as exc: + teardown_error = teardown_error or exc try: _rt = _ring_module() _rt.deactivate() - except Exception: - pass + except Exception as exc: + teardown_error = teardown_error or exc self._ring_transport = None self._ring_engine = None if self._host_engine is not None: + host_engine = self._host_engine try: - self._host_engine.close_input() - self._host_engine.stop() - except Exception: - pass + host_engine.close_input() + except Exception as exc: + teardown_error = teardown_error or exc + try: + host_engine.stop() + except Exception as exc: + teardown_error = teardown_error or exc + try: + host_engine.raise_if_failed() + except Exception as exc: + teardown_error = _host_engine_teardown_error( + host_engine, + ring_engine, + exc, + ) self._host_engine = None + if teardown_error is not None: + raise teardown_error + # --------------------------------------------------------------------------- # Backend loader diff --git a/src/dmi/storage/internals.py b/src/dmi/storage/internals.py index 0e456d5a3..2cc979502 100644 --- a/src/dmi/storage/internals.py +++ b/src/dmi/storage/internals.py @@ -198,7 +198,7 @@ class InternalRequirement: retry: bool = False timeout_s: float | None = 30.0 poll_s: float = 0.25 - match_token_ranges: bool = False + match_token_ranges: bool = True # Backward-compatible private name used before integration API v1. @@ -210,8 +210,9 @@ class InternalRequirements: ``count`` validates ``len(field_value)``. For per-layer fields such as ``hidden_states``, that means layer count, not token completeness. - ``match_token_ranges=True`` additionally validates that captured row ranges - match this generate call's expected request token ranges. + When request IDs and token ranges are available, field reads validate the + captured row ranges against this generate call's expected token ranges by + default. ``match_token_ranges=False`` opts out for an individual field. """ def __init__( @@ -239,7 +240,7 @@ def require( retry: bool = False, timeout_s: float | None = 30.0, poll_s: float = 0.25, - match_token_ranges: bool = False, + match_token_ranges: bool = True, ) -> "InternalRequirements": if count < 0: raise ValueError("count must be non-negative") @@ -305,7 +306,7 @@ def require( retry: bool = False, timeout_s: float | None = 30.0, poll_s: float = 0.25, - match_token_ranges: bool = False, + match_token_ranges: bool = True, ) -> "LazyInternal": self._requirements.require( field, @@ -437,8 +438,7 @@ def _validate_token_ranges( requirement: InternalRequirement | None, ) -> None: if ( - requirement is None - or not requirement.match_token_ranges + (requirement is not None and not requirement.match_token_ranges) or not self._request_ids or not self._token_ranges ): diff --git a/tests/test_engine_runtime_api.py b/tests/test_engine_runtime_api.py index 108fe7cf1..4bcdafddb 100644 --- a/tests/test_engine_runtime_api.py +++ b/tests/test_engine_runtime_api.py @@ -56,6 +56,42 @@ def start(self): def stop(self): self.stop_calls += 1 + def suppressed_submit_failures(self): + return 0 + + +class _FakeQueueStats: + def __init__(self, *, dropped=0, full_errors=0, retries=0): + self.dropped = dropped + self.full_errors = full_errors + self.retries = retries + + +class _FakeProfiling: + def __init__(self, queue_stats): + self.queue_by_stage = [queue_stats] + + +class _FakeHostEngine: + def __init__(self, *, failure=None, profiling=None): + self.failure = failure + self._profiling = profiling + self.close_input_calls = 0 + self.stop_calls = 0 + + def close_input(self): + self.close_input_calls += 1 + + def stop(self): + self.stop_calls += 1 + + def raise_if_failed(self): + if self.failure is not None: + raise self.failure + + def profiling(self): + return self._profiling + def _engine_with_fake_ring( *, null_offload=False, force_eager=False, fail_null_mode=False @@ -143,6 +179,36 @@ def test_close_restores_device_global_null_mode_before_ring_stop(): assert engine.capture_enabled is False +def test_close_raises_host_engine_failure_with_sink_stats(): + engine, _transport, ring_engine = _engine_with_fake_ring() + + class _CountingRing(_FakeRingEngine): + def suppressed_submit_failures(self): + return 7 + + ring_engine = _CountingRing() + engine._ring_engine = ring_engine + host_engine = _FakeHostEngine( + failure=RuntimeError("insert failed"), + profiling=_FakeProfiling( + _FakeQueueStats(dropped=3, full_errors=2, retries=1) + ), + ) + engine._host_engine = host_engine + + with pytest.raises(RuntimeError, match="DMX host sink failed during teardown") as excinfo: + engine.close() + + message = str(excinfo.value) + assert "queue_stats=dropped=3 full_errors=2 retries=1" in message + assert "suppressed_submit_failures=7" in message + assert host_engine.close_input_calls == 1 + assert host_engine.stop_calls == 1 + assert ring_engine.stop_calls == 1 + assert engine._host_engine is None + assert engine._ring_engine is None + + def test_replacing_disabled_ring_restores_native_null_mode(monkeypatch): engine, _transport, old_ring = _engine_with_fake_ring(null_offload=True) new_ring = _FakeRingEngine() diff --git a/tests/test_storage_internals.py b/tests/test_storage_internals.py index 2e4d4727a..57c91c95c 100644 --- a/tests/test_storage_internals.py +++ b/tests/test_storage_internals.py @@ -372,6 +372,19 @@ def test_lazy_internal_requirement_detects_missing_token_ranges(): internal.hidden_states +def test_lazy_internal_validates_token_ranges_by_default_when_request_metadata_exists(): + rows = [_row("0:0", 0, 0, torch.ones(2, 4))] + internal = make_lazy_internal( + "m", + PrefixReader(rows), + request_ids=("0:0",), + token_ranges={"0:0": ((0, 2), (2, 3))}, + ) + + with pytest.raises(IncompleteInternalError, match="token ranges are incomplete"): + internal.hidden_states + + def test_lazy_internal_requirement_retries_missing_token_ranges_until_success(): partial = [_row("0:0", 0, 0, torch.ones(2, 4))] complete = [ From 6ec14834c344e7f2d423975a1fadcd19e893d962 Mon Sep 17 00:00:00 2001 From: Alan Liu Date: Mon, 24 Aug 2026 21:46:37 -0400 Subject: [PATCH 3/6] Name the profiling storage types so EnableEngineProfiling=true compiles PipelinedEngine declared its profiling storage with elaborated-type-specifiers in its own base-clause: private ebo_storage, ...>, 1>, A member type is not in scope in a base-clause, so `struct EngineProfilingConfig_` there declares a new dmx_host::EngineProfilingConfig_ at namespace scope that is never defined -- not the nested struct declared later in the class body. While EnableEngineProfiling was false the std::conditional_t branch was never selected and the incomplete type stayed harmless. DMXHostEngine now passes true, which selects it, and std::optional over an incomplete type is ill-formed: error: incomplete type 'dmx_host::EngineProfilingConfig_' used in type trait expression `make host` fails with 7 errors and builds no objects, and `make test-host` fails outright. CI misses it because `make check` is compileall plus pytest and never builds the native extension. Hoist EngineProfilingConfig_, IngestStats, QueueStats and StageStats to namespace scope so the base-clause names complete types, and re-export them as class aliases -- bindings.cpp refers to them as DMXHostEngine::QueueStats. EngineProfile_ and StatsSnapshot had identical members, so they collapse into one EngineStats_ instead of two shapes kept in sync. Verified with a standalone TU instantiating the template with EnableEngineProfiling both false and true: it fails against the old header and compiles against this one, with static_asserts confirming the public type identities are unchanged. `make host` now builds and links both objects, and `make test-host` passes. Co-Authored-By: Claude Opus 5 --- native/csrc/dmx_host_engine.h | 17 ++-- native/csrc/pipelined_engine.hpp | 134 +++++++++++++++++-------------- 2 files changed, 85 insertions(+), 66 deletions(-) diff --git a/native/csrc/dmx_host_engine.h b/native/csrc/dmx_host_engine.h index caff5e120..11a71768c 100644 --- a/native/csrc/dmx_host_engine.h +++ b/native/csrc/dmx_host_engine.h @@ -11,14 +11,17 @@ namespace dmx_host{ class DMXHostEngine : public PipelinedEngine, true, NoOutputHandler >{ public: + // Stats are always on: they are what lets teardown report silently lost + // rows, and they measured below noise (+33ns/row against an 86ns/row + // spread) on the submit path. Counts only; timing stays off. explicit DMXHostEngine(StageConfig insert_stage): - PipelinedEngine( - std::array{std::move(insert_stage)}, - [] { - EngineConfig config{}; - config.enable_stats = true; - return config; - }()){} + PipelinedEngine( + std::array{std::move(insert_stage)}, + [] { + EngineConfig config{}; + config.enable_stats = true; + return config; + }()){} // Submit a pre-assembled ClickHouseRow directly to the insert stage. // Fields must match the order expected by ClickHouseInsertStage: diff --git a/native/csrc/pipelined_engine.hpp b/native/csrc/pipelined_engine.hpp index f887d4d72..c209189b9 100644 --- a/native/csrc/pipelined_engine.hpp +++ b/native/csrc/pipelined_engine.hpp @@ -167,6 +167,71 @@ struct EngineConfigT { // with a mutex externally (the engine will not serialize output by default). // ============================================================ +// -------- Engine profiling types -------- +// These live at namespace scope rather than nested inside PipelinedEngine +// because the class names them in its own base-clause below. A member type is +// not in scope there, so an elaborated-type-specifier (`struct Foo`) written +// inline would declare a *different*, never-defined namespace-scope type. That +// stays harmless only while EnableEngineProfiling is false and the +// std::optional branch is never selected; the moment it is true, +// std::optional over an incomplete type is ill-formed. +struct EngineProfilingConfig_ { + bool enable_stats = false; + bool enable_timing = false; + + bool timing_profile_submit = true; + bool timing_profile_dequeue = true; + bool timing_profile_process = true; + bool timing_profile_enqueue = true; + bool timing_profile_output = true; +}; + +struct IngestStats { + std::uint64_t submit_calls = 0; + std::uint64_t items_submitted = 0; + + std::uint64_t submit_enqueue_calls = 0; + double submit_enqueue_s = 0.0; +}; + +struct QueueStats { + std::uint64_t enqueued = 0; + std::uint64_t dropped = 0; + std::uint64_t full_errors = 0; + std::uint64_t closed_errors = 0; + std::uint64_t too_large_errors = 0; + std::uint64_t retries = 0; +}; + +struct StageStats { + std::uint64_t batches = 0; + std::uint64_t items_in = 0; + std::uint64_t items_out = 0; + + std::uint64_t dequeue_calls = 0; + std::uint64_t dequeue_timeouts = 0; + std::uint64_t process_calls = 0; + std::uint64_t enqueue_calls = 0; + std::uint64_t output_calls = 0; + std::uint64_t output_items = 0; + + double dequeue_s = 0.0; + double dequeue_idle_s = 0.0; + double process_s = 0.0; + double enqueue_s = 0.0; + double output_s = 0.0; +}; + +// The engine's live counters and the snapshot handed out by profiling() have +// the same shape, so they are one type rather than two identical ones. +template +struct EngineStats_ { + IngestStats ingest{}; + std::array queue_by_stage{}; + std::array stage_by_stage{}; +}; + + template > class PipelinedEngine : private ebo_storage, - private ebo_storage, empty_profiling_storage>, 1>, - private ebo_storage, empty_profiling_storage>, 2> { + private ebo_storage, empty_profiling_storage>, 1>, + private ebo_storage>, empty_profiling_storage>, 2> { public: static_assert(NumStages > 0, "NumStages must be > 0"); @@ -226,67 +291,18 @@ class PipelinedEngine }; // -------- Engine profiling (compile-time optional) -------- - struct EngineProfilingConfig_ { - bool enable_stats = false; - bool enable_timing = false; - - bool timing_profile_submit = true; - bool timing_profile_dequeue = true; - bool timing_profile_process = true; - bool timing_profile_enqueue = true; - bool timing_profile_output = true; - }; - - struct IngestStats { - std::uint64_t submit_calls = 0; - std::uint64_t items_submitted = 0; - - std::uint64_t submit_enqueue_calls = 0; - double submit_enqueue_s = 0.0; - }; - - struct QueueStats { - std::uint64_t enqueued = 0; - std::uint64_t dropped = 0; - std::uint64_t full_errors = 0; - std::uint64_t closed_errors = 0; - std::uint64_t too_large_errors = 0; - std::uint64_t retries = 0; - }; - - struct StageStats { - std::uint64_t batches = 0; - std::uint64_t items_in = 0; - std::uint64_t items_out = 0; - - std::uint64_t dequeue_calls = 0; - std::uint64_t dequeue_timeouts = 0; - std::uint64_t process_calls = 0; - std::uint64_t enqueue_calls = 0; - std::uint64_t output_calls = 0; - std::uint64_t output_items = 0; - - double dequeue_s = 0.0; - double dequeue_idle_s = 0.0; - double process_s = 0.0; - double enqueue_s = 0.0; - double output_s = 0.0; - }; - - struct StatsSnapshot { - IngestStats ingest{}; - std::array queue_by_stage{}; - std::array stage_by_stage{}; - }; + // Defined at namespace scope above; re-exported here because callers and the + // pybind layer refer to them as PipelinedEngine<...>::QueueStats. + using EngineProfilingConfig_ = dmx_host::EngineProfilingConfig_; + using IngestStats = dmx_host::IngestStats; + using QueueStats = dmx_host::QueueStats; + using StageStats = dmx_host::StageStats; + using StatsSnapshot = EngineStats_; using QueueProfilingSnapshotArray = std::array, NumStages>; private: - struct EngineProfile_ { - IngestStats ingest{}; - std::array queue_by_stage{}; - std::array stage_by_stage{}; - }; + using EngineProfile_ = EngineStats_; using OutBase = ebo_storage; using ProfCfgStorage = std::conditional_t, empty_profiling_storage>; From 910f5a67efedaae5f844e2b53c8bab8091a36afe Mon Sep 17 00:00:00 2001 From: Alan Liu Date: Mon, 24 Aug 2026 21:47:26 -0400 Subject: [PATCH 4/6] Fix token-range validation for narrowed and per-layer rows Turning match_token_ranges on by default made two existing weaknesses in the validator user-visible. final_logits rows are re-based natively to the trailing positions of each step (p2p_thread.cpp: db_start = end - logits_count) because HF generate passes logits_to_keep=1 and vLLM computes one logit per request. Comparing full ranges therefore fails on every prefill with a prompt longer than one token, so out.dmi_internal.logits raised IncompleteInternalError on every real generate call -- and retry=True could never clear it, because the data was already complete. Validate that act on range ends instead, which still detects a missing step. A test parses p2p_thread.cpp and fails if a second act is ever narrowed, so the exemption cannot silently rot. The per-layer check used the highest *present* layer as its reference, so the reference followed the damage: a layer that lost rows simply lowered the bar. Check every layer instead, grouped in one pass over rows already held, and name the offending layer in the error. Measured over 300 randomised generate scenarios (batch size, prompt length, decode count, EOS strip and logits_to_keep varied) driven through a port of the native write path: false alarms on complete captures : 144/300 -> 0 single-row losses detected : 58.8% -> 98.8% cost : 0.31ms -> 0.44ms per read of 1536 rows The remaining 1.2% are one shape: a layer that wrote no rows at all leaves nothing to compare against, so only a declared count catches it. That limit is documented and pinned by a test rather than left as a surprise. Pipeline parallelism was checked separately: a lagging rank was already caught before this change, is still caught, and clears under retry=True. _actual_ranges_for_request is left unused by the new grouping helper and is removed here. Co-Authored-By: Claude Opus 5 --- docs/huggingface.md | 25 +++- docs/integration-api-v1.md | 25 +++- src/dmi/storage/internals.py | 111 +++++++++++++--- tests/test_storage_internals.py | 223 ++++++++++++++++++++++++++++++++ 4 files changed, 352 insertions(+), 32 deletions(-) diff --git a/docs/huggingface.md b/docs/huggingface.md index dce100d35..fdbe0fd64 100644 --- a/docs/huggingface.md +++ b/docs/huggingface.md @@ -90,7 +90,14 @@ each layer tensor unless token-range validation runs. Lazy reads attached to `generate_with_monitoring_dict(...)` now perform that validation by default whenever request IDs and token ranges were recorded for the generate call. `match_token_ranges=False` opts out for a specific requirement. For per-layer -fields the current implementation uses a fast representative-layer check. +fields every layer present is checked and the error names the offending layer. +A layer that wrote no rows at all is not detectable this way -- it is simply +absent from the tuple -- so pass `count=` when layer completeness matters. + +`logits` is a special case: `generate()` passes `logits_to_keep=1`, so the +prefill logits row covers only the last prompt position rather than the whole +prompt. Its rows are therefore validated on token-range *ends* only, which +still catches a missing decode step. Supported mapped fields are: @@ -149,10 +156,18 @@ norm = hidden_states[0].float().norm(dim=-1)[token_mask].mean() If a tensor field and `token_mask` have different `[batch, seq]` shapes, it usually means the tensor field was read before all rows for that field arrived. -Use `retry=True` to wait for expected token ranges after an incomplete read. If -you intentionally want to skip token-range validation for a field, set -`match_token_ranges=False` on that field's requirement. If you already cached a -partial field, clear that field's cache and read it again: +Use `retry=True` to wait for expected token ranges after an incomplete read. +If you intentionally want to skip token-range validation for a field, set +`match_token_ranges=False` on that field's requirement. Both are set through +`require()`, which also needs the expected `count`: + +```python +out.dmi_internal.require( + "hidden_states", count=model.config.num_hidden_layers, retry=True +) +``` + If you already cached a partial field, clear that +field's cache and read it again: ```python out.dmi_internal.clear_cache("hidden_states") diff --git a/docs/integration-api-v1.md b/docs/integration-api-v1.md index c0863a36e..2366a6fd7 100644 --- a/docs/integration-api-v1.md +++ b/docs/integration-api-v1.md @@ -1107,7 +1107,7 @@ InternalRequirement( retry: bool = False, timeout_s: float | None = 30.0, poll_s: float = 0.25, - match_token_ranges: bool = False, + match_token_ranges: bool = True, ) ``` @@ -1124,7 +1124,7 @@ requirements.require( retry: bool = False, timeout_s: float | None = 30.0, poll_s: float = 0.25, - match_token_ranges: bool = False, + match_token_ranges: bool = True, ) -> InternalRequirements ``` @@ -1173,10 +1173,23 @@ and a batch count for global tensors, not a token or database-row count. With `retry=True`, synchronous field access polls missing/incomplete data until success or timeout. `timeout_s=None` can block forever. Database/runtime errors are not retried. When both nonempty request IDs and ranges were supplied, -token-range validation runs by default; `match_token_ranges=False` opts out for -an individual field, and otherwise the flag is a no-op. For a per-layer field -the current v1 implementation checks only its highest present layer, not every -layer. +token-range validation runs by default -- including for a field with no +requirement registered at all. `match_token_ranges=False` opts out for an +individual field, and otherwise the flag is a no-op. For a per-layer field +every layer present in the result is checked, and the error names the offending +`layer=`. + +One loss the range check cannot see: a layer that wrote no rows at all leaves +nothing to compare against, so it is silently absent from the reassembled +tuple. Only a declared `count` catches that -- pass one when layer +completeness matters. + +`logits` is validated on range *ends* only. Its rows are re-based natively to +the trailing positions of each step, because HF generate passes +`logits_to_keep=1` and vLLM computes one logit per request, so a full-range +comparison would fail on every prefill with a prompt longer than one token. +The relaxed check still detects a missing step; it does not detect a wrong +start offset. Mapped dynamic attributes are: diff --git a/src/dmi/storage/internals.py b/src/dmi/storage/internals.py index 2cc979502..44c7bf5bd 100644 --- a/src/dmi/storage/internals.py +++ b/src/dmi/storage/internals.py @@ -165,6 +165,17 @@ def _reassemble_global(rows: list) -> torch.Tensor: } +# Acts whose stored row range is narrower than the step's token range. +# ``final_logits`` rows are re-based in the native P2P thread to +# ``(end - logits_count, end)`` because both HF generate (logits_to_keep=1) and +# vLLM (always one logit per request) compute logits for only the trailing +# positions of the step. Comparing full ranges for these acts would fail on +# every prefill with a prompt longer than one token, so token-range validation +# compares range *ends* instead -- which still catches a missing step, just not +# a wrong start offset. +_TRAILING_RANGE_ACTS = frozenset({"final_logits"}) + + class Internal: """Captured internals for one run, presented like HF's model output: each field mirrors its HF counterpart (e.g. ``hidden_states`` is a tuple indexed @@ -417,19 +428,25 @@ def _expected_non_empty_ranges(self, request_id: str) -> tuple[tuple[int, int], if int(end) > int(start) ) - def _actual_ranges_for_request( - self, + @staticmethod + def _ranges_by_request_and_layer( rows: list, - request_id: str, - layer: int | None, - ) -> tuple[tuple[int, int], ...]: - ranges = { - (int(key[5]), int(key[6])) - for key, _ in rows - if key[1] == request_id and (layer is None or int(key[3]) == layer) - and int(key[6]) > int(key[5]) - } - return tuple(sorted(ranges)) + ) -> dict[tuple[str, int | None], tuple[tuple[int, int], ...]]: + """Group non-empty row ranges by (request_id, layer) in one pass. + + ``layer`` is None for global (non-per-layer) acts, which store -1. + """ + grouped: dict[tuple[str, int | None], set] = {} + for key, _ in rows: + layer = int(key[3]) + start, end = int(key[5]), int(key[6]) + if end <= start: + continue + bucket = grouped.setdefault( + (key[1], layer if layer >= 0 else None), set() + ) + bucket.add((start, end)) + return {k: tuple(sorted(v)) for k, v in grouped.items()} def _validate_token_ranges( self, @@ -443,19 +460,47 @@ def _validate_token_ranges( or not self._token_ranges ): return - layers = sorted({int(key[3]) for key, _ in rows if int(key[3]) >= 0}) - # Fast check for per-layer fields: validating the last layer catches the - # common "writer still pending" case without scanning every layer. - check_layer = layers[-1] if layers else None + # Every layer is checked, not just a representative one. Picking the + # highest *present* layer would let the reference point follow the + # damage: if a layer went missing entirely the check would silently + # re-base onto a lower, intact layer. One grouping pass over rows we + # already hold costs no extra queries. + by_layer = self._ranges_by_request_and_layer(rows) + layers = sorted({layer for _, layer in by_layer if layer is not None}) + act = _FIELDS[field][0] if field in _FIELDS else None + trailing = act in _TRAILING_RANGE_ACTS for request_id in self._request_ids: expected = self._expected_non_empty_ranges(request_id) - actual = self._actual_ranges_for_request(rows, request_id, check_layer) - if actual != expected: - raise IncompleteInternalError( + for layer in (layers or [None]): + actual = by_layer.get((request_id, layer), ()) + if trailing: + # See _TRAILING_RANGE_ACTS: rows cover only the trailing + # positions of each step, so only the range ends line up. + mismatch = tuple(e for _, e in actual) != tuple( + e for _, e in expected + ) + detail = ( + f"expected token range ends " + f"{[e for _, e in expected]}, found {[e for _, e in actual]}" + ) + else: + mismatch = actual != expected + detail = f"expected {list(expected)}, found {list(actual)}" + if not mismatch: + continue + where = ( + f"model_id={self._model_id!r}, request_id={request_id!r}" + if layer is None + else f"model_id={self._model_id!r}, " + f"request_id={request_id!r}, layer={layer}" + ) + error = IncompleteInternalError( f"{field} token ranges are incomplete for " - f"model_id={self._model_id!r}, request_id={request_id!r}: " - f"expected {list(expected)}, found {list(actual)}." + f"{where}: {detail}." ) + error.field = field # type: ignore[attr-defined] + error.token_ranges_incomplete = True # type: ignore[attr-defined] + raise error def _read_field( self, @@ -503,6 +548,23 @@ def _build_token_mask(self) -> torch.Tensor: mask[batch_i, offset + start_i: offset + end_i] = True return mask + def _token_range_timeout_error( + self, + field: str, + requirement: InternalRequirement, + cause: Exception, + ) -> IncompleteInternalError: + timeout_text = ( + "without timeout" + if requirement.timeout_s is None + else f"within {requirement.timeout_s:.3g}s" + ) + error = IncompleteInternalError(f"{cause} Not resolved {timeout_text}.") + error.field = field # type: ignore[attr-defined] + error.token_ranges_incomplete = True # type: ignore[attr-defined] + error.__cause__ = cause + return error + def _load_field_with_retry( self, field: str, @@ -528,6 +590,13 @@ def _load_field_with_retry( last_error = exc if deadline is not None and time.monotonic() >= deadline: + if getattr(last_error, "token_ranges_incomplete", False): + # The entry count may well be correct -- reporting this as + # "expected N entries, found none" would send the caller + # looking for the wrong problem. + raise self._token_range_timeout_error( + field, requirement, last_error + ) raise self._incomplete_error( field, requirement, diff --git a/tests/test_storage_internals.py b/tests/test_storage_internals.py index 57c91c95c..5937d1e39 100644 --- a/tests/test_storage_internals.py +++ b/tests/test_storage_internals.py @@ -1,8 +1,11 @@ """Unit tests for dmi.storage.internals -- pure reassembly logic, no DB.""" +from pathlib import Path + import pytest import torch from dmi.storage.internals import ( + _TRAILING_RANGE_ACTS, IncompleteInternalError, InternalRequirements, get_internal, @@ -385,6 +388,226 @@ def test_lazy_internal_validates_token_ranges_by_default_when_request_metadata_e internal.hidden_states +def test_pipeline_parallel_lag_is_reported_and_clears_on_retry(): + """PP splits layers across ranks that flush independently. + + A lagging rank must be reported rather than silently yielding a short + tuple, and ``retry=True`` must clear it once the rank flushes. + """ + ranges = ((0, 2), (2, 3)) + def rows(partial=()): + return [ + _row("0:0", layer, start, torch.ones(n, 4)) + for layer in range(8) + for start, n in ((0, 2), (2, 1)) + if not (layer in partial and start == 2) + ] + + lagging = rows(partial=(4, 5, 6, 7)) # rank1 decode rows not flushed + internal = make_lazy_internal( + "m", PrefixReader(lagging), + request_ids=("0:0",), token_ranges={"0:0": ranges}, + ) + with pytest.raises(IncompleteInternalError, match="layer=4"): + internal.hidden_states + + internal = make_lazy_internal( + "m", SequenceReader([lagging, rows()]), + request_ids=("0:0",), token_ranges={"0:0": ranges}, + ) + internal.require( + "hidden_states", count=8, retry=True, timeout_s=2.0, poll_s=0.01 + ) + assert len(internal.hidden_states) == 8 + + +def test_trailing_range_acts_matches_the_native_rewrite_sites(): + """Guard the hardcoded exemption against native drift. + + ``_TRAILING_RANGE_ACTS`` only stays correct while ``final_logits`` is the + single act whose stored row range is re-based away from the step's token + range. If someone narrows another act in p2p_thread.cpp, this fails. + """ + source = ( + Path(__file__).resolve().parents[1] + / "native" / "csrc" / "ring" / "p2p_thread.cpp" + ) + if not source.is_file(): + pytest.skip("native sources not present") + text = source.read_text() + + rewrites = [ + line.strip() for line in text.splitlines() + if "db_start =" in line and "req.start_token" not in line + ] + assert len(rewrites) == 1, ( + f"native db_start rewrite sites changed: {rewrites}. " + "Update _TRAILING_RANGE_ACTS in dmi/storage/internals.py to match." + ) + guard = text.split("db_start = req.end_token")[0].splitlines()[-3:] + assert any("HOOK_TYPE_FINAL_LOGITS" in line for line in guard), ( + "the db_start rewrite is no longer guarded by HOOK_TYPE_FINAL_LOGITS" + ) + assert _TRAILING_RANGE_ACTS == frozenset({"final_logits"}) + + +def test_token_range_check_covers_every_layer_not_just_the_highest(): + # Layer 0 is missing its decode row; layers 1 and 2 are complete. A + # representative-layer check that only looked at the highest layer would + # pass this. + rows = [ + _row("0:0", layer, start, torch.ones(n, 4)) + for layer in (0, 1, 2) + for start, n in ((0, 2), (2, 1)) + if not (layer == 0 and start == 2) + ] + internal = make_lazy_internal( + "m", + PrefixReader(rows), + request_ids=("0:0",), + token_ranges={"0:0": ((0, 2), (2, 3))}, + ) + + with pytest.raises(IncompleteInternalError, match="layer=0"): + internal.hidden_states + + +def test_token_range_check_names_the_offending_layer(): + rows = [ + _row("0:0", layer, start, torch.ones(n, 4)) + for layer in (0, 1) + for start, n in ((0, 2), (2, 1)) + if not (layer == 1 and start == 0) + ] + internal = make_lazy_internal( + "m", + PrefixReader(rows), + request_ids=("0:0",), + token_ranges={"0:0": ((0, 2), (2, 3))}, + ) + + with pytest.raises(IncompleteInternalError) as excinfo: + internal.hidden_states + + assert "layer=1" in str(excinfo.value) + + +def test_global_field_errors_do_not_mention_a_layer(): + rows = [_row_act("0:0", "hook_embed", -1, 0, 2, torch.ones(2, 4))] + internal = make_lazy_internal( + "m", + PrefixReader(rows), + request_ids=("0:0",), + token_ranges={"0:0": ((0, 2), (2, 3))}, + ) + + with pytest.raises(IncompleteInternalError) as excinfo: + internal.embeddings + + assert "layer=" not in str(excinfo.value) + + +def test_a_layer_that_left_no_rows_at_all_is_not_detectable_by_range_check(): + """Documents the one loss the range check inherently cannot see. + + If a layer never wrote a single row it is absent from the data, so there + is nothing to compare. Only a declared ``count`` closes this. + """ + rows = [ + _row("0:0", layer, start, torch.ones(n, 4)) + for layer in (0, 1) + for start, n in ((0, 2), (2, 1)) + ] + internal = make_lazy_internal( + "m", + PrefixReader(rows), + request_ids=("0:0",), + token_ranges={"0:0": ((0, 2), (2, 3))}, + ) + + assert len(internal.hidden_states) == 2 # silent: layer 2 never landed + + internal.clear_cache() + internal.require("hidden_states", count=3) + with pytest.raises(IncompleteInternalError, match="expected 3 entries"): + internal.hidden_states + + +LOGITS_ACT = "final_logits" + + +def _logits_internal(rows, ranges): + return make_lazy_internal( + "m", + PrefixReader(rows), + request_ids=("0:0",), + token_ranges={"0:0": ranges}, + ) + + +def test_lazy_internal_accepts_narrowed_final_logits_rows(): + """HF generate passes logits_to_keep=1, so the prefill logits row covers + only the last prompt position while the step range covers the whole + prompt. Comparing full ranges would fail on every real generate call.""" + rows = [ + _row_act("0:0", LOGITS_ACT, -1, 3, 4, torch.ones(1, 8)), # prefill + _row_act("0:0", LOGITS_ACT, -1, 4, 5, torch.ones(1, 8)), # decode + _row_act("0:0", LOGITS_ACT, -1, 5, 6, torch.ones(1, 8)), # decode + ] + + logits = _logits_internal(rows, ((0, 4), (4, 5), (5, 6))).logits + + assert logits.shape == (1, 3, 8) + + +def test_lazy_internal_still_detects_a_missing_final_logits_step(): + rows = [ + _row_act("0:0", LOGITS_ACT, -1, 3, 4, torch.ones(1, 8)), + _row_act("0:0", LOGITS_ACT, -1, 4, 5, torch.ones(1, 8)), + ] + + with pytest.raises(IncompleteInternalError, match="token range ends"): + _logits_internal(rows, ((0, 4), (4, 5), (5, 6))).logits + + +def test_lazy_internal_accepts_full_width_prefill_logits(): + # logits_to_keep=0 keeps every prefill position; the relaxed check must + # still pass for the unnarrowed case. + rows = [ + _row_act("0:0", LOGITS_ACT, -1, 0, 4, torch.ones(4, 8)), + _row_act("0:0", LOGITS_ACT, -1, 4, 5, torch.ones(1, 8)), + ] + + assert _logits_internal(rows, ((0, 4), (4, 5))).logits.shape == (1, 5, 8) + + +def test_lazy_internal_keeps_strict_range_matching_for_non_logits_fields(): + rows = [_row("0:0", 0, 3, torch.ones(1, 4))] + internal = make_lazy_internal( + "m", + PrefixReader(rows), + request_ids=("0:0",), + token_ranges={"0:0": ((0, 4),)}, + ) + + with pytest.raises(IncompleteInternalError, match=r"expected \[\(0, 4\)\]"): + internal.hidden_states + + +def test_retry_timeout_reports_the_token_range_reason_not_a_count_mismatch(): + rows = [_row_act("0:0", LOGITS_ACT, -1, 3, 4, torch.ones(1, 8))] + internal = _logits_internal(rows, ((0, 4), (4, 5))) + internal.require("logits", count=1, retry=True, timeout_s=0.02, poll_s=0.001) + + with pytest.raises(IncompleteInternalError) as excinfo: + internal.logits + + message = str(excinfo.value) + assert "token range ends" in message + assert "found none" not in message + assert isinstance(excinfo.value.__cause__, IncompleteInternalError) + + def test_lazy_internal_requirement_retries_missing_token_ranges_until_success(): partial = [_row("0:0", 0, 0, torch.ones(2, 4))] complete = [ From f5d9ed422e63b98f458ab0e2d9db27bba7e220ab Mon Sep 17 00:00:00 2001 From: Alan Liu Date: Mon, 24 Aug 2026 21:47:52 -0400 Subject: [PATCH 5/6] Report rows the host sink lost, not only worker failures close() read the drop and suppressed-submit counters only inside the raise_if_failed() except-handler, so loss without a worker crash was silent. That is the common case: the native P2P thread's submit-failure warning fires at most once per process (std::call_once), so the counters were the only quantitative signal and they were gated behind an unrelated failure. Measured against a live ClickHouse with the CPU-only host build: 3000 rows submitted into a queue capped at 1 with OnFullPolicy.DROP, worker healthy. The counters reported 2998 dropped, and ClickHouse held exactly 2 rows -- the counter corresponds to real loss, not just to a number moving. Before this change close() returned None and reported nothing. Add a SinkStats snapshot and MonitoringEngine.sink_stats(), readable mid-run and cached across teardown. close() now raises when suppressed submits are nonzero -- an exception the P2P thread swallowed is never a configured behaviour -- and warns when rows were dropped, since OnFullPolicy.DROP is a caller's deliberate choice. A healthy run stays silent: 50 rows submitted, 50 stored, no raise and no warning. Only the close() that actually held engines reads and reports the counters, so a repeat call neither overwrites the snapshot with zeros nor re-reports it. close() stays idempotent and sink_stats() keeps returning the first snapshot. When a teardown failure and a sink failure coincide the sink failure is raised and the earlier error is attached as __context__ rather than discarded. Note that close() can now raise where it previously could not, and every in-repo caller invokes it from a bare finally:. The hazard is documented; the call sites are deliberately left for their owners, since whether each should swallow, log or propagate depends on what that script is for. Co-Authored-By: Claude Opus 5 --- docs/integration-api-v1.md | 74 +++++++++- src/dmi/__init__.py | 8 +- src/dmi/api/v1/__init__.py | 8 +- src/dmi/engine.py | 181 ++++++++++++++++++----- tests/test_engine_runtime_api.py | 239 ++++++++++++++++++++++++++++--- tests/test_integration_api_v1.py | 2 + 6 files changed, 450 insertions(+), 62 deletions(-) diff --git a/docs/integration-api-v1.md b/docs/integration-api-v1.md index 2366a6fd7..1cdd08aac 100644 --- a/docs/integration-api-v1.md +++ b/docs/integration-api-v1.md @@ -199,6 +199,7 @@ engine.ring_capacities() -> RingCapacities engine.capture_enabled -> bool engine.set_capture_enabled(enabled: bool) -> None engine.next_auto_group_id() -> int +engine.sink_stats() -> SinkStats engine.close() -> None ``` @@ -240,11 +241,48 @@ fails, Python-visible capture flags remain unchanged. `next_auto_group_id()` returns engine-scoped integers starting at zero. It is not synchronized for concurrent callers. +`sink_stats()` returns an immutable `SinkStats` snapshot of the host sink's +loss and backpressure counters: + +| Field | Meaning | +| --- | --- | +| `dropped` | Rows the insert queue discarded, via a configured `OnFullPolicy.DROP` or because the engine was already stopping. | +| `suppressed` | Rows the native P2P thread failed to submit and swallowed. Never a configured behavior. | +| `full_errors` | Enqueue attempts that hit a full queue. | +| `closed_errors` | Enqueue attempts against a closed queue. | +| `too_large_errors` | Rows rejected as larger than the queue's item cap. | +| `retries` | Enqueue retries under `OnFullPolicy.RETRY`. | +| `lost_rows` | `dropped + suppressed` -- rows that provably never reached the sink. | + +The counters are cumulative for the life of the engine, never reset by DMI, +and read best-effort: a backend that does not expose them reports zeros rather +than raising. `close()` caches the final snapshot, so `sink_stats()` keeps +working after teardown. Prefer it over the native warning: the P2P thread's +submit-failure message is printed at most once per process no matter how many +rows are lost. + `close()` stops and flushes the ring, clears the active binding, closes host -input, and stops the host engine. It is terminal and effectively idempotent. -Shutdown exceptions are suppressed, so an integration requiring an -authoritative final read must ensure every worker reaches this close path and -should separately check native host failures. +input, and stops the host engine. It is terminal and idempotent -- both engine +handles are cleared before any exception leaves the method, so a second call is +a no-op: it neither re-reads the counters nor re-reports them. Only the first +close sees the engines, so `sink_stats()` keeps returning that snapshot however +many times `close()` is called. + +`close()` raises `RuntimeError` when: + +- the ring or host engine raised during teardown; +- `raise_if_failed()` reports a host worker failure, in which case the message + carries the `SinkStats` snapshot and the original failure is chained as + `__cause__`; or +- `sink_stats().suppressed` is nonzero, meaning rows were produced but never + submitted. + +When a teardown failure and a sink failure coincide, the sink failure is the +one raised and the earlier error is attached as `__context__`. A nonzero +`dropped` count raises no exception -- it is reachable through a configured +drop policy -- but emits a `RuntimeWarning`. Integrations that treat a run as +authoritative should call `close()` inside their own error handling rather +than a bare `finally`, since it can now replace an in-flight exception. Closing does not disable or uninstall HookPoints: they retain hook IDs and the old payload tensor. Treat the attached model as terminal too. A later CUDA @@ -919,7 +957,9 @@ DMXHostEngine(insert_stage: StageConfig) ``` Construction validates/copies stage configuration but does not connect to the -database. Public lifecycle and diagnostics are: +database. The engine always keeps the ingest and queue counters that +`MonitoringEngine.close()` uses to report silently lost rows; they measured +below noise on the submit path. Public lifecycle and diagnostics are: ```python start() -> None @@ -929,6 +969,8 @@ request_abort() -> None join(timeout_s: float | None = None) -> bool failures() -> list[ThreadFailure] raise_if_failed() -> None +profiling() -> StatsSnapshot | None +reset_metrics() -> None ``` `start()` is asynchronous: it can return before a worker fails to connect or @@ -936,6 +978,28 @@ initialize. `stop()` returning true means threads joined, not that inserts succeeded. After shutdown, call `raise_if_failed()`; `failures()` returns records with `stage`, `thread_name`, `where`, `exc_type`, and `exc_what`. +`profiling()` returns a `StatsSnapshot` by value, or `None` when the engine was +built without stats. Counters are cumulative; `reset_metrics()` zeroes them. +Both are safe to call after `stop()`. Prefer `MonitoringEngine.sink_stats()` +unless you need the per-stage timing fields. + +### `StatsSnapshot`, `IngestStats`, `QueueStats`, and `StageStats` + +`StatsSnapshot` is the immutable value returned by `DMXHostEngine.profiling()`. +Its fields are `ingest` (an `IngestStats`), `queue_by_stage`, and +`stage_by_stage` (lists of `QueueStats` and `StageStats`, one entry per stage; +the host engine has exactly one). All fields on all four classes are read-only +and the classes cannot be constructed from Python. + +| Class | Read-only fields | +| --- | --- | +| `IngestStats` | `submit_calls`, `items_submitted`, `submit_enqueue_calls`, `submit_enqueue_s` | +| `QueueStats` | `enqueued`, `dropped`, `full_errors`, `closed_errors`, `too_large_errors`, `retries` | +| `StageStats` | `batches`, `items_in`, `items_out`, `dequeue_calls`, `dequeue_timeouts`, `process_calls`, `enqueue_calls`, `output_calls`, `output_items`, `dequeue_s`, `dequeue_idle_s`, `process_s`, `enqueue_s`, `output_s` | + +The `_s` fields are seconds and stay zero unless the engine was also built with +timing enabled, which DMI does not do. + ### `ThreadFailure` `ThreadFailure` is the immutable diagnostic record returned by diff --git a/src/dmi/__init__.py b/src/dmi/__init__.py index abfcfc81a..88192320a 100644 --- a/src/dmi/__init__.py +++ b/src/dmi/__init__.py @@ -5,7 +5,12 @@ """ from .config import CaptureSchedule, MonitoringConfig -from .engine import HostEngineConfig, MonitoringEngine, RingCapacities +from .engine import ( + HostEngineConfig, + MonitoringEngine, + RingCapacities, + SinkStats, +) _NATIVE_EXPORTS = ( "StageConfig", @@ -30,6 +35,7 @@ def __getattr__(name: str): "MonitoringEngine", "HostEngineConfig", "RingCapacities", + "SinkStats", "CaptureSchedule", "MonitoringConfig", *_NATIVE_EXPORTS, diff --git a/src/dmi/api/v1/__init__.py b/src/dmi/api/v1/__init__.py index 02a851da7..909dbd7cc 100644 --- a/src/dmi/api/v1/__init__.py +++ b/src/dmi/api/v1/__init__.py @@ -19,7 +19,12 @@ ) from ...storage.clickhouse import CHClickhouseDriverReadOnly from ...config import CaptureSchedule, MonitoringConfig -from ...engine import HostEngineConfig, MonitoringEngine, RingCapacities +from ...engine import ( + HostEngineConfig, + MonitoringEngine, + RingCapacities, + SinkStats, +) from ...hooks.dispatch import install_ring_hooks from ...hooks.point import HookPoint from ...hooks import specs as _specs @@ -173,6 +178,7 @@ def __dir__() -> list[str]: "StepContext", "MonitoringEngine", "RingCapacities", + "SinkStats", "MonitoringConfig", "CaptureSchedule", "HostEngineConfig", diff --git a/src/dmi/engine.py b/src/dmi/engine.py index 286941f02..db3954a10 100644 --- a/src/dmi/engine.py +++ b/src/dmi/engine.py @@ -4,6 +4,7 @@ from dataclasses import dataclass import importlib +import warnings from typing import Any, Optional, Sequence from .config import MonitoringConfig @@ -22,51 +23,97 @@ def _ring_module() -> Any: return importlib.import_module("dmi.transport.ring") -def _format_host_engine_stats(host_engine: Any, ring_engine: Any) -> str: - details: list[str] = [] +@dataclass(frozen=True, slots=True) +class SinkStats: + """Immutable snapshot of the host sink's loss and backpressure counters. + + ``dropped`` and ``suppressed`` are the two counters that mean rows were + produced but never reached ClickHouse: + + * ``dropped`` -- the insert queue discarded rows, either because the stage + was configured with ``OnFullPolicy.DROP`` or because the engine was + already stopping when they were submitted. + * ``suppressed`` -- the native P2P thread caught an exception out of + ``submit_direct()`` and dropped the row. This is never a configured + behavior; a nonzero value always means something went wrong. + + The remaining fields are diagnostics that explain *why*, and may overlap + with the two above. + """ + + dropped: int = 0 + suppressed: int = 0 + full_errors: int = 0 + closed_errors: int = 0 + too_large_errors: int = 0 + retries: int = 0 + + @property + def lost_rows(self) -> int: + """Rows that provably never reached the sink.""" + return self.dropped + self.suppressed + + def __str__(self) -> str: + return ( + f"dropped={self.dropped} suppressed={self.suppressed} " + f"full_errors={self.full_errors} " + f"closed_errors={self.closed_errors} " + f"too_large_errors={self.too_large_errors} " + f"retries={self.retries}" + ) - try: - profiling = host_engine.profiling() - except Exception: - profiling = None + +def _read_sink_stats(host_engine: Any, ring_engine: Any) -> SinkStats: + """Best-effort read of the sink counters; never raises.""" + + fields: dict[str, int] = {} + + profiling = None + if host_engine is not None: + try: + profiling = host_engine.profiling() + except Exception: + profiling = None if profiling is not None: try: queue_stats = profiling.queue_by_stage[0] except Exception: queue_stats = None if queue_stats is not None: - details.append( - "queue_stats=" - f"dropped={int(queue_stats.dropped)} " - f"full_errors={int(queue_stats.full_errors)} " - f"retries={int(queue_stats.retries)}" - ) + for name in ( + "dropped", + "full_errors", + "closed_errors", + "too_large_errors", + "retries", + ): + try: + fields[name] = int(getattr(queue_stats, name)) + except Exception: + pass if ring_engine is not None: try: - suppressed = int(ring_engine.suppressed_submit_failures()) + fields["suppressed"] = int(ring_engine.suppressed_submit_failures()) except Exception: - suppressed = None - if suppressed: - details.append(f"suppressed_submit_failures={suppressed}") - - return "; ".join(details) - - -def _host_engine_teardown_error( - host_engine: Any, - ring_engine: Any, - exc: Exception, -) -> RuntimeError: - details = _format_host_engine_stats(host_engine, ring_engine) - message = "DMX host sink failed during teardown" - if details: - message = f"{message} ({details})" - error = RuntimeError(message) + pass + + return SinkStats(**fields) + + +def _host_engine_teardown_error(stats: SinkStats, exc: Exception) -> RuntimeError: + error = RuntimeError(f"DMX host sink failed during teardown ({stats})") error.__cause__ = exc return error +def _lost_rows_error(stats: SinkStats) -> RuntimeError: + return RuntimeError( + f"DMX host sink lost {stats.lost_rows} row(s) during this run; " + f"captured internals are incomplete ({stats})" + ) + + @dataclass(frozen=True, slots=True) class RingCapacities: """Immutable snapshot of the active ring transport's capacities.""" @@ -132,7 +179,9 @@ def __init__( self.config = config self._model_id = model_id self._auto_batch_group_id = 0 - + # Final sink counters, captured during close() so sink_stats() keeps + # working after the engines are gone. + self._final_sink_stats: Optional[SinkStats] = None # Host-side DB engine (optional; C++ backend only) self._host_engine: Optional[Any] = None @@ -322,8 +371,28 @@ def next_auto_group_id(self) -> int: self._auto_batch_group_id += 1 return gid + def sink_stats(self) -> SinkStats: + """Snapshot the host sink's loss and backpressure counters. + + Safe to call at any point in the run, and after ``close()`` -- the + final counts are cached during teardown, so the snapshot survives the + engines being torn down. + """ + + if self._final_sink_stats is not None: + return self._final_sink_stats + return _read_sink_stats( + self._host_engine, getattr(self, "_ring_engine", None) + ) + def close(self) -> None: - """Tear down backend resources.""" + """Tear down backend resources. + + Raises ``RuntimeError`` if teardown failed, if the sink recorded a + worker failure, or if rows were produced but never reached the sink. + Both engine handles are cleared before raising, so ``close()`` stays + idempotent and a second call is a no-op. + """ teardown_error: Exception | None = None ring_engine = getattr(self, "_ring_engine", None) @@ -350,8 +419,8 @@ def close(self) -> None: self._ring_transport = None self._ring_engine = None - if self._host_engine is not None: - host_engine = self._host_engine + host_engine = self._host_engine + if host_engine is not None: try: host_engine.close_input() except Exception as exc: @@ -360,19 +429,53 @@ def close(self) -> None: host_engine.stop() except Exception as exc: teardown_error = teardown_error or exc + + # Read the counters after both engines are stopped and joined, so the + # snapshot is final rather than a moving target. Only the close() that + # actually had engines reads and reports them: a repeat call must + # neither overwrite the snapshot with zeros nor re-report it. + stats = None + if host_engine is not None or ring_engine is not None: + stats = _read_sink_stats(host_engine, ring_engine) + self._final_sink_stats = stats + + if host_engine is not None: + # stats is always set here: host_engine being non-None satisfies the + # condition above. Spelled out so a future edit cannot quietly + # pass None into the error builder. + assert stats is not None try: host_engine.raise_if_failed() except Exception as exc: - teardown_error = _host_engine_teardown_error( - host_engine, - ring_engine, - exc, - ) + sink_error = _host_engine_teardown_error(stats, exc) + # Prefer the sink failure -- it is the more actionable one -- + # but keep an earlier teardown error visible in the traceback + # instead of discarding it. + sink_error.__context__ = teardown_error + teardown_error = sink_error self._host_engine = None if teardown_error is not None: raise teardown_error + if stats is None: + return # already closed; nothing new to report + + # A clean shutdown is not the same as a complete capture: rows the + # queue dropped or the P2P thread failed to submit are silent + # otherwise (the native warning fires at most once per process). + if stats.suppressed: + raise _lost_rows_error(stats) + if stats.dropped: + # Reachable through a configured OnFullPolicy.DROP, so warn rather + # than raise -- the caller may have opted into shedding load. + warnings.warn( + f"DMX host sink dropped {stats.dropped} row(s); captured " + f"internals are incomplete ({stats})", + RuntimeWarning, + stacklevel=2, + ) + # --------------------------------------------------------------------------- # Backend loader diff --git a/tests/test_engine_runtime_api.py b/tests/test_engine_runtime_api.py index 4bcdafddb..285b48a6a 100644 --- a/tests/test_engine_runtime_api.py +++ b/tests/test_engine_runtime_api.py @@ -3,6 +3,7 @@ import os import subprocess import sys +import warnings from dataclasses import FrozenInstanceError from pathlib import Path from types import ModuleType @@ -10,15 +11,20 @@ import pytest -from dmi.engine import MonitoringEngine, RingCapacities +from dmi.engine import MonitoringEngine, RingCapacities, SinkStats pytestmark = pytest.mark.cpu class _FakeRingEngine: - def __init__(self, transport=None, *, fail_null_mode=False): + def __init__( + self, transport=None, *, fail_null_mode=False, suppressed=0, + fail_stop=False, + ): self.transport = transport self.fail_null_mode = fail_null_mode + self.suppressed = suppressed + self.fail_stop = fail_stop self.null_mode_calls = [] self.stop_calls = 0 self.init_calls = 0 @@ -55,15 +61,27 @@ def start(self): def stop(self): self.stop_calls += 1 + if self.fail_stop: + raise RuntimeError("ring stop failed") def suppressed_submit_failures(self): - return 0 + return self.suppressed class _FakeQueueStats: - def __init__(self, *, dropped=0, full_errors=0, retries=0): + def __init__( + self, + *, + dropped=0, + full_errors=0, + closed_errors=0, + too_large_errors=0, + retries=0, + ): self.dropped = dropped self.full_errors = full_errors + self.closed_errors = closed_errors + self.too_large_errors = too_large_errors self.retries = retries @@ -94,7 +112,12 @@ def profiling(self): def _engine_with_fake_ring( - *, null_offload=False, force_eager=False, fail_null_mode=False + *, + null_offload=False, + force_eager=False, + fail_null_mode=False, + suppressed=0, + fail_stop=False, ): engine = MonitoringEngine(enable_ring_transport=False) transport = SimpleNamespace( @@ -104,6 +127,8 @@ def _engine_with_fake_ring( ring_engine = _FakeRingEngine( transport, fail_null_mode=fail_null_mode, + suppressed=suppressed, + fail_stop=fail_stop, ) engine._ring_transport = transport engine._ring_engine = ring_engine @@ -180,14 +205,7 @@ def test_close_restores_device_global_null_mode_before_ring_stop(): def test_close_raises_host_engine_failure_with_sink_stats(): - engine, _transport, ring_engine = _engine_with_fake_ring() - - class _CountingRing(_FakeRingEngine): - def suppressed_submit_failures(self): - return 7 - - ring_engine = _CountingRing() - engine._ring_engine = ring_engine + engine, _transport, ring_engine = _engine_with_fake_ring(suppressed=7) host_engine = _FakeHostEngine( failure=RuntimeError("insert failed"), profiling=_FakeProfiling( @@ -196,12 +214,17 @@ def suppressed_submit_failures(self): ) engine._host_engine = host_engine - with pytest.raises(RuntimeError, match="DMX host sink failed during teardown") as excinfo: + with pytest.raises( + RuntimeError, match="DMX host sink failed during teardown" + ) as excinfo: engine.close() message = str(excinfo.value) - assert "queue_stats=dropped=3 full_errors=2 retries=1" in message - assert "suppressed_submit_failures=7" in message + assert "dropped=3" in message + assert "suppressed=7" in message + assert "full_errors=2" in message + assert "retries=1" in message + assert excinfo.value.__cause__ is host_engine.failure assert host_engine.close_input_calls == 1 assert host_engine.stop_calls == 1 assert ring_engine.stop_calls == 1 @@ -209,6 +232,190 @@ def suppressed_submit_failures(self): assert engine._ring_engine is None +def test_close_raises_on_suppressed_submits_without_a_worker_failure(): + """The counters must not be gated behind an unrelated worker failure. + + ``log_submit_failure_once`` prints at most one warning per process, so a + silent ``close()`` here would hide unbounded row loss. + """ + engine, _transport, _ring_engine = _engine_with_fake_ring(suppressed=4) + engine._host_engine = _FakeHostEngine( + profiling=_FakeProfiling(_FakeQueueStats()) + ) + + with pytest.raises(RuntimeError, match=r"lost 4 row\(s\)") as excinfo: + engine.close() + + assert "suppressed=4" in str(excinfo.value) + + +def test_close_warns_on_dropped_rows_without_a_worker_failure(): + # Drops are reachable through a configured OnFullPolicy.DROP, so they warn + # rather than raise -- but they must not be silent. + engine, _transport, _ring_engine = _engine_with_fake_ring() + engine._host_engine = _FakeHostEngine( + profiling=_FakeProfiling(_FakeQueueStats(dropped=2, full_errors=2)) + ) + + with pytest.warns(RuntimeWarning, match=r"dropped 2 row\(s\)"): + engine.close() + + assert engine._host_engine is None + + +def test_close_is_silent_when_nothing_was_lost(): + engine, _transport, _ring_engine = _engine_with_fake_ring() + engine._host_engine = _FakeHostEngine( + profiling=_FakeProfiling(_FakeQueueStats(retries=5)) + ) + + with warnings.catch_warnings(): + warnings.simplefilter("error") + engine.close() + + +def test_close_keeps_an_earlier_teardown_error_in_the_traceback(): + engine, _transport, _ring_engine = _engine_with_fake_ring(fail_stop=True) + engine._host_engine = _FakeHostEngine( + failure=RuntimeError("insert failed"), + profiling=_FakeProfiling(_FakeQueueStats()), + ) + + with pytest.raises( + RuntimeError, match="DMX host sink failed during teardown" + ) as excinfo: + engine.close() + + context = excinfo.value.__context__ + assert isinstance(context, RuntimeError) + assert "ring stop failed" in str(context) + + +def test_close_reports_a_ring_failure_when_the_sink_is_healthy(): + engine, _transport, _ring_engine = _engine_with_fake_ring(fail_stop=True) + engine._host_engine = _FakeHostEngine( + profiling=_FakeProfiling(_FakeQueueStats()) + ) + + with pytest.raises(RuntimeError, match="ring stop failed"): + engine.close() + + +def test_close_stays_idempotent_after_raising(): + engine, _transport, _ring_engine = _engine_with_fake_ring(suppressed=1) + engine._host_engine = _FakeHostEngine( + profiling=_FakeProfiling(_FakeQueueStats()) + ) + + with pytest.raises(RuntimeError): + engine.close() + + engine.close() # handles were cleared before the raise + + +def test_a_second_close_does_not_erase_the_final_sink_stats(): + """close() is idempotent and callers invoke it from finally blocks, so a + repeat call must not overwrite the forensic record with zeros.""" + engine, _transport, _ring_engine = _engine_with_fake_ring(suppressed=4) + engine._host_engine = _FakeHostEngine( + profiling=_FakeProfiling(_FakeQueueStats(dropped=5, retries=2)) + ) + + with pytest.raises(RuntimeError): + engine.close() + first = engine.sink_stats() + assert (first.dropped, first.suppressed, first.retries) == (5, 4, 2) + + engine.close() # no engines left to read + assert engine.sink_stats() == first + + +def test_a_second_close_does_not_re_warn(recwarn): + engine, _transport, _ring_engine = _engine_with_fake_ring() + engine._host_engine = _FakeHostEngine( + profiling=_FakeProfiling(_FakeQueueStats(dropped=3)) + ) + + with pytest.warns(RuntimeWarning): + engine.close() + recwarn.clear() + engine.close() + assert len(recwarn) == 0 + + +def test_sink_stats_is_readable_before_and_after_close(): + engine, _transport, _ring_engine = _engine_with_fake_ring(suppressed=6) + engine._host_engine = _FakeHostEngine( + profiling=_FakeProfiling(_FakeQueueStats(dropped=1, retries=3)) + ) + + live = engine.sink_stats() + assert (live.dropped, live.suppressed, live.retries) == (1, 6, 3) + assert live.lost_rows == 7 + + with pytest.raises(RuntimeError): + engine.close() + + # The engines are gone, but the final snapshot is still available. + assert engine.sink_stats() == live + + +def test_sink_stats_handles_a_backend_whose_profiling_returns_none(): + """profiling() returns None when the engine carries no stats. + + The ring counter must still be read, and the queue counters must report + zero rather than blowing up. + """ + engine, _transport, _ring_engine = _engine_with_fake_ring(suppressed=2) + engine._host_engine = _FakeHostEngine(profiling=None) + + stats = engine.sink_stats() + + assert stats == SinkStats(suppressed=2) + assert stats.lost_rows == 2 + + +def test_sink_stats_reads_the_pybind_snapshot_shape(): + """queue_by_stage is a list (from std::array) whose + entries expose exactly the six read-only counters bound in bindings.cpp.""" + + class _PybindQueueStats: + __slots__ = ( + "enqueued", "dropped", "full_errors", + "closed_errors", "too_large_errors", "retries", + ) + + def __init__(self): + self.enqueued, self.dropped, self.full_errors = 100, 3, 2 + self.closed_errors, self.too_large_errors, self.retries = 1, 4, 5 + + class _PybindSnapshot: + def __init__(self): + self.ingest = object() + self.queue_by_stage = [_PybindQueueStats()] + self.stage_by_stage = [object()] + + engine, _transport, _ring_engine = _engine_with_fake_ring(suppressed=7) + engine._host_engine = _FakeHostEngine(profiling=_PybindSnapshot()) + + assert engine.sink_stats() == SinkStats( + dropped=3, + suppressed=7, + full_errors=2, + closed_errors=1, + too_large_errors=4, + retries=5, + ) + + +def test_sink_stats_tolerates_a_backend_without_the_counters(): + engine, _transport, _ring_engine = _engine_with_fake_ring() + engine._ring_engine = SimpleNamespace() # no suppressed_submit_failures() + engine._host_engine = SimpleNamespace() # no profiling() + + assert engine.sink_stats() == SinkStats() + + def test_replacing_disabled_ring_restores_native_null_mode(monkeypatch): engine, _transport, old_ring = _engine_with_fake_ring(null_offload=True) new_ring = _FakeRingEngine() diff --git a/tests/test_integration_api_v1.py b/tests/test_integration_api_v1.py index 05ed2cb13..15815e141 100644 --- a/tests/test_integration_api_v1.py +++ b/tests/test_integration_api_v1.py @@ -52,6 +52,7 @@ "StepContext", "MonitoringEngine", "RingCapacities", + "SinkStats", "MonitoringConfig", "CaptureSchedule", "HostEngineConfig", @@ -122,6 +123,7 @@ def test_v1_reexports_existing_objects_and_state() -> None: assert v1.StepContext is types.StepContext assert v1.MonitoringEngine is engine.MonitoringEngine assert v1.RingCapacities is engine.RingCapacities + assert v1.SinkStats is engine.SinkStats assert v1.MonitoringConfig is config.MonitoringConfig assert v1.CaptureSchedule is config.CaptureSchedule assert v1.HostEngineConfig is engine.HostEngineConfig From e5838b1e9b1ca924c206c7540211e3e4fff2ad18 Mon Sep 17 00:00:00 2001 From: Alan Liu Date: Mon, 24 Aug 2026 21:48:12 -0400 Subject: [PATCH 6/6] Release the GIL in the DMXHostEngine profiling accessors profiling() and reset_metrics() both take the engine's profiling mutex, which the C++ insert workers also hold. Every other blocking method on this class is bound with py::call_guard(); these two were not, so they held the GIL while blocking on a mutex the ingest path needs. No deadlock is possible today because the insert workers never acquire the GIL, so this is consistency rather than a live bug. Co-Authored-By: Claude Opus 5 --- native/csrc/bindings.cpp | 8 ++++++-- 1 file changed, 6 insertions(+), 2 deletions(-) diff --git a/native/csrc/bindings.cpp b/native/csrc/bindings.cpp index ce64b7a6d..a41cd3106 100644 --- a/native/csrc/bindings.cpp +++ b/native/csrc/bindings.cpp @@ -252,8 +252,12 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { py::call_guard()) .def("failures", &DMXHostEngine::failures) .def("raise_if_failed", &DMXHostEngine::raise_if_failed) - .def("profiling", &DMXHostEngine::profiling) - .def("reset_metrics", &DMXHostEngine::reset_metrics) + // Both take the engine's profiling mutex, which C++ worker threads + // also hold; release the GIL like every other blocking method here. + .def("profiling", &DMXHostEngine::profiling, + py::call_guard()) + .def("reset_metrics", &DMXHostEngine::reset_metrics, + py::call_guard()) // Submit a pre-formatted ClickHouseRow directly to the insert stage. // Called from the ring transport drain callback after format processing. .def("submit_direct",