diff --git a/docs/huggingface.md b/docs/huggingface.md index 7cca14a9e..fdbe0fd64 100644 --- a/docs/huggingface.md +++ b/docs/huggingface.md @@ -86,9 +86,18 @@ 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 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: @@ -147,9 +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 `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. 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 7a6136c03..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 @@ -1107,7 +1171,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 +1188,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 ``` @@ -1164,7 +1228,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 +1236,24 @@ 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 -- 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/native/csrc/bindings.cpp b/native/csrc/bindings.cpp index 9fa95f7b4..a41cd3106 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,12 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { py::call_guard()) .def("failures", &DMXHostEngine::failures) .def("raise_if_failed", &DMXHostEngine::raise_if_failed) + // 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", @@ -304,6 +349,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..11a71768c 100644 --- a/native/csrc/dmx_host_engine.h +++ b/native/csrc/dmx_host_engine.h @@ -8,11 +8,20 @@ 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: + // 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{}){} + 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>; 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/__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 89667ce67..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,6 +23,97 @@ def _ring_module() -> Any: return importlib.import_module("dmi.transport.ring") +@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}" + ) + + +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: + 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: + fields["suppressed"] = int(ring_engine.suppressed_submit_failures()) + except Exception: + 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.""" @@ -87,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 @@ -277,8 +371,31 @@ 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) if self._ring_transport is not None: # Best-effort reset of the device-global native null flag. This is @@ -290,27 +407,75 @@ 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 + if host_engine is not None: 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 + + # 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: + 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/src/dmi/storage/internals.py b/src/dmi/storage/internals.py index 0e456d5a3..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 @@ -198,7 +209,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 +221,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 +251,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 +317,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, @@ -416,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, @@ -437,25 +455,52 @@ 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 ): 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_engine_runtime_api.py b/tests/test_engine_runtime_api.py index 108fe7cf1..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 @@ -53,12 +59,65 @@ def init(self): def start(self): self.start_calls += 1 + def stop(self): + self.stop_calls += 1 + if self.fail_stop: + raise RuntimeError("ring stop failed") + + def suppressed_submit_failures(self): + return self.suppressed + + +class _FakeQueueStats: + 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 + + +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 + *, + null_offload=False, + force_eager=False, + fail_null_mode=False, + suppressed=0, + fail_stop=False, ): engine = MonitoringEngine(enable_ring_transport=False) transport = SimpleNamespace( @@ -68,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 @@ -143,6 +204,218 @@ 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(suppressed=7) + 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 "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 + assert engine._host_engine is None + 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 diff --git a/tests/test_storage_internals.py b/tests/test_storage_internals.py index 2e4d4727a..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, @@ -372,6 +375,239 @@ 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_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 = [