diff --git a/docs/integration-api-v1.md b/docs/integration-api-v1.md index 1e9d857de..1d60f24ab 100644 --- a/docs/integration-api-v1.md +++ b/docs/integration-api-v1.md @@ -136,8 +136,12 @@ The predicates apply warmup, then offset, then stride. Step selection also honors `capture_prefill`/`capture_decode`; an unknown phase raises `ValueError`. `MonitoringConfig` carries this schedule plus three storage fields: `storage_backend`, `capture_sink_config` (a `NativeSinkConfig`, or `None`) and -`capture_storage_config` (a `NativeCaptureStorageConfig` from -`dmi.storage.native_capture`, or `None`). +`capture_storage_config` (a `NativeCaptureStorageConfig`, or `None`), both +from `dmi.storage.native_capture`. `NativeSinkConfig` bounds the native pack +sink's queue and packs, and sets its admission policy: `overload="block"` +(the default) waits up to `admission_timeout_s` (2 s; `None` waits without +bound) for queue room before refusing a record, and `"drop_newest"` refuses +it at once. `storage_backend` is the user's storage choice, one of `dmi.config.USER_STORAGE_CHOICES`: @@ -408,6 +412,13 @@ runtime = engine.create_record_runtime( record_sink=explicit_native_sink, ) +# Keep the forward running when the sink refuses or stalls. +runtime = engine.create_record_runtime( + record_format, + failure_policy="disable_capture", + step_stall_budget_ms=2000, +) + runtime.bind_hook( hook, hook_runtime=hook_runtime, @@ -415,14 +426,17 @@ runtime.bind_hook( gate_value=0, ) +runtime.begin_step() runtime.emit_output(entry, metadata, output) runtime.prepare_replay(plan, metadata) engine.flush_and_wait(timeout_s=600.0) +engine.capture_status() ``` The public operations are `create_record_runtime()`, `bind_hook()`, -`emit_output()`, `prepare_replay()`, and `flush_and_wait()`. +`begin_step()`, `emit_output()`, `prepare_replay()`, `flush_and_wait()`, +`capture_status()`, and `validate_capture_bounds()`. `RecordRuntime.bind_hook()` assigns stable output IDs and binds the hook to the record ring. `RecordRuntime.emit_output()` reserves and publishes one eager @@ -456,6 +470,89 @@ sink reached that sink's durability boundary within one timeout. The current ClickHouse adapter waits for acknowledged inserts. A timeout raises `TimeoutError`; asynchronous failures propagate unchanged. +A sink can refuse a record (dropped, timed out, too large), lose one it +admitted (a record that fits no empty pack, a duplicate id), or fail, and it +does so on the record worker, after the forward that produced the record has +moved on. The refusal latches the runtime, and `failure_policy` decides what +the forward sees next: + +| `failure_policy` | After a latch | +| --- | --- | +| `"raise"` (default) | The next `RecordRuntime.begin_step()`, record reservation or descriptor push raises it. | +| `"disable_capture"` | Capture stops; the forward keeps running. Descriptors and payloads that follow are discarded and counted. | + +Under both, `flush_and_wait()` raises the failure, `capture_status()` reports +it, and `close()` logs it at WARNING. A `NativePackSink` reports a record lost +after admission at its next submit and at the checked flush, as the reference +adapter does, so such a loss also latches the runtime. Ring failures that are +not sink refusals, such as a CUDA error in the drain, raise under both +policies. + +A slow sink reaches the forward as a stall instead: a reservation that does +not fit waits for the drain, which waits for the sink. +`step_stall_budget_ms` caps the time the reservations of one step may wait. +A step is what lies between two `RecordRuntime.begin_step()` calls. The +runtime cannot find step boundaries itself (eager hooks reserve once per +output, output ids repeat within a step, and a CUDA-graph step replays one +plan, or several with piecewise graphs), so with a budget the integration +calls `begin_step()` once per model step, before the step's first +reservation, and a reservation before the first call is refused. Past the +budget the rest of that step is skipped: the step's remaining records and +every record still queued for the sink, from any step, are discarded, and +capture resumes at the next step. The stall bound needs that whole queue out +of the drain's way, so records of earlier, already finished steps that the +sink had not taken yet are lost too. The discards are counted in +`discarded_*`; `skipped_steps` counts the skips (always equal to +`stall_budget_exhaustions`), and `steps_with_discards` the distinct steps +that lost at least one record to them, which can be more. +Under `"disable_capture"` a skipped step is not a failure: capture stays +active and `flush_and_wait()` succeeds. Under `"raise"` the exhaustion is +raised at the next `begin_step()` and at `flush_and_wait()`, never inside +the forward that ran out of budget, and the runtime is failed from then on. +`None`, the default, waits without bound. + +The skip still waits for the one envelope the sink is admitting, so a budget +needs a sink whose admission is bounded, and `create_record_runtime()` +refuses a budget otherwise, under either policy: + +- the ClickHouse host path (`record_sink=None` without the persistent + backend), which enqueues row by row under the host engine's ingress + policy, with no bound per envelope; +- a `record_sink` whose `admission_bound_s` is None: a `NativePackSink` with + `overload="block"` and no `admission_timeout_s`, or the reference bridge; +- a persistent `capture_sink_config` with `overload="block"` and + `admission_timeout_s=None`. + +A `NativePackSink` admits all rows of one envelope against one +`admission_timeout_s`, so its `admission_bound_s` is that timeout under +`"block"` and zero under `"drop_newest"`. With such a sink, a step stalls the +forward for at most `step_stall_budget_ms` plus `admission_bound_s`, plus the +time to move what the ring holds out to the record worker, which discards it +without copying it to pageable memory. A ring that has not drained by the +bound plus a drain grace fails from the reservation under either policy, +rather than hanging the forward: the error cannot tell a sink that held its +admission past its bound from a slow drain, and names both. The grace is +2 s plus the bytes of the payload ring and its pinned staging at 1 GB/s, +about 10.6 s with the engine's default 4 GiB of each. + +`MonitoringEngine.capture_status()` returns a plain dict: `record_mode`, +`capture_active` (False once a failure latched, and under `"raise"` once a +spent budget is waiting to be raised), `failure_policy`, `failure`, +`discarded_descriptors`, `discarded_payloads`, `step_stall_budget_ms`, +`stall_budget_exhaustions`, `skipped_steps`, `steps_with_discards`, +`reserve_wait_s` and `max_step_wait_s` (time reservations waited for the +sink, in total and in the worst step), and the `sink` and `storage` +snapshots when the engine holds a native pack sink or storage service. +Without a record runtime every field is empty. + +`MonitoringEngine.validate_capture_bounds(max_record_bytes)` refuses, with +`ConfigurationError`, capture bounds under which the largest record an +integration emits could not be stored: a `max_queue_bytes` below it, a +`max_pack_bytes` below it plus 64 KiB of pack framing, or a `max_pack_bytes` +above the storage service's `uploader_max_in_flight_bytes`. Integrations call +it at attach, so these refusals never reach the forward. Without a +`capture_sink_config` there is nothing to check. + ### `deactivate_ring_transport` ```python diff --git a/native/csrc/bindings.cpp b/native/csrc/bindings.cpp index 0cbe6e8b1..feac37904 100644 --- a/native/csrc/bindings.cpp +++ b/native/csrc/bindings.cpp @@ -141,6 +141,14 @@ std::shared_ptr MakeClickHouseRecordSink( [host] { host->raise_if_failed(); }); } +// RecordSink.admission_bound_s: the sink's admission bound in seconds, or +// None when it has none. +std::optional AdmissionBoundSeconds(const ring::RecordSink& sink) { + const auto bound = sink.admission_bound(); + if (!bound) return std::nullopt; + return std::chrono::duration(*bound).count(); +} + template std::shared_ptr MakeRingEngine(Args&&... args) { // Ring destruction may join a worker that is completing a callback. Never @@ -157,6 +165,22 @@ std::shared_ptr MakeRingEngine(Args&&... args) { }); } +ring::RecordFailurePolicy ParseRecordFailurePolicy(const std::string& name) { + if (name == "raise") return ring::RecordFailurePolicy::kRaiseAtProducer; + if (name == "disable_capture") { + return ring::RecordFailurePolicy::kDisableCapture; + } + throw py::value_error( + "failure_policy must be 'raise' or 'disable_capture', got '" + name + + "'"); +} + +const char* RecordFailurePolicyName(ring::RecordFailurePolicy policy) { + return policy == ring::RecordFailurePolicy::kDisableCapture + ? "disable_capture" + : "raise"; +} + ring::RecordDescriptor CopyRecordDescriptor( const py::handle& descriptor_py, const dmx_host::RecordSchema& schema) { @@ -630,7 +654,8 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { .def("_acquire_engine", [](std::shared_ptr sink) { return ring::RecordSinkLease::acquire(std::move(sink)); - }); + }) + .def_property_readonly("admission_bound_s", &AdmissionBoundSeconds); py::class_>( m, "ClickHouseRecordSink") @@ -641,12 +666,25 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { py::class_>( m, "ReferencePythonCaptureSink") - .def(py::init([](py::object target, std::string layout) { + .def(py::init([](py::object target, std::string layout, + std::optional admission_bound_s) { + std::optional bound; + if (admission_bound_s) { + if (!std::isfinite(*admission_bound_s) || + *admission_bound_s < 0) { + throw py::value_error( + "admission_bound_s must be None or a finite, " + "non-negative number"); + } + bound = std::chrono::ceil( + std::chrono::duration(*admission_bound_s)); + } return std::make_shared< dmi_capture::ReferencePythonCaptureSink>( - target.ptr(), std::move(layout)); + target.ptr(), std::move(layout), bound); }), - py::arg("target"), py::arg("layout")) + py::arg("target"), py::arg("layout"), + py::arg("admission_bound_s") = py::none()) .def_property_readonly( "attached", &dmi_capture::ReferencePythonCaptureSink::engine_owned); @@ -678,7 +716,12 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { py::arg("config"), py::arg("host_engine") = py::none()) .def_static( "create_record", - [](ring_py::RingConfig cfg, py::object sink_or_host) { + [](ring_py::RingConfig cfg, py::object sink_or_host, + const std::string& failure_policy, + uint64_t step_stall_budget_ms) { + ring_py::RecordRuntimeOptions options; + options.failure_policy = ParseRecordFailurePolicy(failure_policy); + options.step_stall_budget_ms = step_stall_budget_ms; std::shared_ptr lease; if (!sink_or_host.is_none()) { if (py::isinstance(sink_or_host)) { @@ -694,9 +737,11 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { MakeClickHouseRecordSink(std::move(host))); } } - return MakeRingEngine(std::move(cfg), std::move(lease)); + return MakeRingEngine(std::move(cfg), std::move(lease), options); }, - py::arg("config"), py::arg("sink_or_host") = py::none()) + py::arg("config"), py::arg("sink_or_host") = py::none(), + py::arg("failure_policy") = "raise", + py::arg("step_stall_budget_ms") = uint64_t{0}) .def("init", &ring_py::RingEnginePy::init, py::arg("stream_handle") = uint64_t{0}) .def("start", &ring_py::RingEnginePy::start) @@ -711,6 +756,26 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { &ring_py::RingEnginePy::reserve_record, py::arg("reservation_items"), py::call_guard()) + .def("begin_record_step", &ring_py::RingEnginePy::begin_record_step) + .def("capture_status", + [](const ring_py::RingEnginePy& self) { + const ring_py::RecordCaptureStatus status = + self.record_capture_status(); + py::dict out; + out["failure_policy"] = + RecordFailurePolicyName(status.failure_policy); + out["failed"] = status.failed; + out["failure"] = status.failure; + out["discarded_descriptors"] = status.discarded_descriptors; + out["discarded_payloads"] = status.discarded_payloads; + out["step_stall_budget_ms"] = status.step_stall_budget_ms; + out["stall_budget_exhaustions"] = status.stall_budget_exhaustions; + out["skipped_steps"] = status.skipped_steps; + out["steps_with_discards"] = status.steps_with_discards; + out["reserve_wait_ns"] = status.reserve_wait_ns; + out["max_step_wait_ns"] = status.max_step_wait_ns; + return out; + }) .def("push_record_descriptors", [](ring_py::RingEnginePy& self, py::sequence descriptors, py::object schema) { diff --git a/native/csrc/reference_python_capture_sink.cpp b/native/csrc/reference_python_capture_sink.cpp index 47cd7f4d5..482dc2e98 100644 --- a/native/csrc/reference_python_capture_sink.cpp +++ b/native/csrc/reference_python_capture_sink.cpp @@ -99,8 +99,10 @@ py::object target_method(PyObject* target, const char* name) { } // namespace ReferencePythonCaptureSink::ReferencePythonCaptureSink( - PyObject* target, std::string layout) - : target_(target), layout_(std::move(layout)) { + PyObject* target, std::string layout, + std::optional admission_bound) + : target_(target), layout_(std::move(layout)), + admission_bound_(admission_bound) { if (target == nullptr || target == Py_None) { throw std::invalid_argument( "ReferencePythonCaptureSink requires a Python target"); diff --git a/native/csrc/reference_python_capture_sink.h b/native/csrc/reference_python_capture_sink.h index bd399c69a..33553e22f 100644 --- a/native/csrc/reference_python_capture_sink.h +++ b/native/csrc/reference_python_capture_sink.h @@ -6,13 +6,19 @@ #include +#include #include namespace dmi_capture { class ReferencePythonCaptureSink final : public ring::RecordSink { public: - ReferencePythonCaptureSink(PyObject* target, std::string layout); + // `admission_bound` is what the target promises about its submit: the + // longest one call can take. The bridge cannot know it, so it has none + // unless the caller says so (tests that stand in for a bounded sink). + ReferencePythonCaptureSink( + PyObject* target, std::string layout, + std::optional admission_bound = std::nullopt); ~ReferencePythonCaptureSink() override; ReferencePythonCaptureSink(const ReferencePythonCaptureSink&) = delete; @@ -22,6 +28,9 @@ class ReferencePythonCaptureSink final : public ring::RecordSink { void submit(ring::RecordEnvelope envelope) override; bool flush_and_wait(Duration timeout) override; void rethrow_if_failed() const override; + std::optional admission_bound() const override { + return admission_bound_; + } protected: void on_engine_acquire() override; @@ -30,6 +39,7 @@ class ReferencePythonCaptureSink final : public ring::RecordSink { private: PyObject* target_; const std::string layout_; + const std::optional admission_bound_; }; } // namespace dmi_capture diff --git a/native/csrc/ring/p2p_thread.cpp b/native/csrc/ring/p2p_thread.cpp index 9c0001d9a..fb73ca8d1 100644 --- a/native/csrc/ring/p2p_thread.cpp +++ b/native/csrc/ring/p2p_thread.cpp @@ -333,8 +333,9 @@ void P2PThread::do_post_processing(at::Tensor& tensor, const DrainTask& first_ta // RecordP2PThread -- fixed schema-driven consumer path. // --------------------------------------------------------------------------- RecordP2PThread::RecordP2PThread( - DrainThread& drain, std::shared_ptr sink) - : drain_(drain), consumer_(std::move(sink)) {} + DrainThread& drain, std::shared_ptr sink, + RecordFailurePolicy failure_policy) + : drain_(drain), consumer_(std::move(sink), failure_policy) {} RecordP2PThread::~RecordP2PThread() noexcept { stop(); @@ -362,6 +363,17 @@ void RecordP2PThread::loop() { void RecordP2PThread::process(std::vector& tasks) { for (DrainTask& task : tasks) { try { + // A skipped step's record, or capture is off: nobody stores + // these bytes, so free the staging range without copying them. + // The consumer has accounted the discard exactly as + // consume_payload would, so the pairing is unchanged. + if (consumer_.discard_next_payload_if_unwanted()) { + if (task.alloc_bytes > 0) { + drain_.notify_staging_freed_bytes(task.alloc_bytes); + task.alloc_bytes = 0; + } + continue; + } at::Tensor payload; if (task.cpu_paged_tensor.defined()) { payload = std::move(task.cpu_paged_tensor); diff --git a/native/csrc/ring/p2p_thread.h b/native/csrc/ring/p2p_thread.h index 651d8b005..a7e87dc10 100644 --- a/native/csrc/ring/p2p_thread.h +++ b/native/csrc/ring/p2p_thread.h @@ -64,7 +64,9 @@ class P2PThread { // payload. Construction fixes the consumer for the lifetime of the ring. class RecordP2PThread { public: - RecordP2PThread(DrainThread& drain, std::shared_ptr sink); + RecordP2PThread(DrainThread& drain, std::shared_ptr sink, + RecordFailurePolicy failure_policy = + RecordFailurePolicy::kRaiseAtProducer); ~RecordP2PThread() noexcept; RecordP2PThread(const RecordP2PThread&) = delete; diff --git a/native/csrc/ring/record_consumer.cpp b/native/csrc/ring/record_consumer.cpp index 96702a30c..255a3faf0 100644 --- a/native/csrc/ring/record_consumer.cpp +++ b/native/csrc/ring/record_consumer.cpp @@ -3,6 +3,7 @@ #include #include +#include #include namespace ring { @@ -13,37 +14,124 @@ namespace { throw std::runtime_error("record consumer: " + message); } +std::string describe_failure(const std::exception_ptr& failure) { + try { + std::rethrow_exception(failure); + } catch (const std::exception& error) { + return error.what(); + } catch (...) { + return "unknown record failure"; + } +} + } // namespace -RecordConsumer::RecordConsumer(std::shared_ptr sink) - : sink_(std::move(sink)) {} +RecordConsumer::RecordConsumer( + std::shared_ptr sink, RecordFailurePolicy policy) + : sink_(std::move(sink)), policy_(policy) {} + +void RecordConsumer::latch_locked(std::exception_ptr failure) { + if (!failure_) failure_ = std::move(failure); + if (policy_ != RecordFailurePolicy::kDisableCapture) return; + // Their payloads will be discarded on arrival, so no descriptor queued + // at the latch can ever be paired again. + discarded_descriptors_ += descriptors_.size(); + descriptors_.clear(); + // Every later payload is discarded by the latch itself. Inert -- each + // read of the count comes after a failure_ check -- but reset so it never + // describes payloads the latch's own discard already accounts for. + payloads_to_discard_ = 0; +} + +void RecordConsumer::note_discarded_step_locked(uint64_t step) { + if (last_discarded_step_ && *last_discarded_step_ == step) return; + last_discarded_step_ = step; + ++steps_with_discards_; +} + +bool RecordConsumer::take_unwanted_payload_locked() { + if (failure_) { + // Under kRaiseAtProducer consume_payload raises instead. + if (policy_ != RecordFailurePolicy::kDisableCapture) return false; + // Capture is off: the drain still delivers what the forward + // produced, and the payload is dropped here, never submitted. + ++discarded_payloads_; + return true; + } + if (payloads_to_discard_ == 0) return false; + // A skipped step's record: its descriptor was dropped by a discard + // window. + --payloads_to_discard_; + ++discarded_payloads_; + if (payloads_to_discard_ == 0) idle_cv_.notify_all(); + return true; +} + +bool RecordConsumer::discard_next_payload_if_unwanted() { + std::lock_guard lock(mu_); + return take_unwanted_payload_locked(); +} + +void RecordConsumer::begin_step() { + std::lock_guard lock(mu_); + ++step_; +} void RecordConsumer::push_descriptor(RecordDescriptor descriptor) { std::lock_guard lock(mu_); - if (failure_) std::rethrow_exception(failure_); - descriptors_.push_back(std::move(descriptor)); + if (failure_) { + if (policy_ != RecordFailurePolicy::kDisableCapture) { + std::rethrow_exception(failure_); + } + ++discarded_descriptors_; + return; + } + if (discarding_) { + ++discarded_descriptors_; + ++payloads_to_discard_; + note_discarded_step_locked(step_); + return; + } + descriptors_.push_back({std::move(descriptor), step_}); } void RecordConsumer::push_descriptors( std::vector descriptors) { std::lock_guard lock(mu_); - if (failure_) std::rethrow_exception(failure_); + if (failure_) { + if (policy_ != RecordFailurePolicy::kDisableCapture) { + std::rethrow_exception(failure_); + } + discarded_descriptors_ += descriptors.size(); + return; + } + if (discarding_) { + discarded_descriptors_ += descriptors.size(); + payloads_to_discard_ += descriptors.size(); + if (!descriptors.empty()) note_discarded_step_locked(step_); + return; + } for (auto& descriptor : descriptors) { - descriptors_.push_back(std::move(descriptor)); + descriptors_.push_back({std::move(descriptor), step_}); } } void RecordConsumer::consume_payload(at::Tensor payload) { + const bool disable = policy_ == RecordFailurePolicy::kDisableCapture; RecordDescriptor descriptor; { std::lock_guard lock(mu_); - if (failure_) std::rethrow_exception(failure_); + if (failure_ && !disable) std::rethrow_exception(failure_); + if (take_unwanted_payload_locked()) return; if (descriptors_.empty()) { - failure_ = std::make_exception_ptr(std::runtime_error( - "record consumer: physical payload arrived without an encoded descriptor")); - std::rethrow_exception(failure_); + latch_locked(std::make_exception_ptr(std::runtime_error( + "record consumer: physical payload arrived without an encoded descriptor"))); + if (!disable) std::rethrow_exception(failure_); + ++discarded_payloads_; + idle_cv_.notify_all(); + return; } - descriptor = std::move(descriptors_.front()); + descriptor = std::move(descriptors_.front().descriptor); descriptors_.pop_front(); ++active_payloads_; } @@ -73,11 +161,14 @@ void RecordConsumer::consume_payload(at::Tensor payload) { // waiter must never observe an idle consumer before the submit // failure becomes visible. std::lock_guard lock(mu_); - if (!failure_) failure_ = failure; + latch_locked(failure); --active_payloads_; + if (disable) ++discarded_payloads_; } idle_cv_.notify_all(); - throw; + // Under kDisableCapture the refusal stops capture, not the worker: + // it is reported at flush, in the snapshot and at close. + if (!disable) throw; } } @@ -85,7 +176,7 @@ void RecordConsumer::record_failure(std::exception_ptr failure) noexcept { if (!failure) return; try { std::lock_guard lock(mu_); - if (!failure_) failure_ = std::move(failure); + latch_locked(std::move(failure)); idle_cv_.notify_all(); } catch (...) { // Failure reporting must not terminate a worker while unwinding. @@ -100,7 +191,8 @@ bool RecordConsumer::wait_until_idle( } std::unique_lock lock(mu_); const bool ready = idle_cv_.wait_for(lock, timeout, [this] { - return failure_ || (descriptors_.empty() && active_payloads_ == 0); + return failure_ || (descriptors_.empty() && active_payloads_ == 0 && + payloads_to_discard_ == 0); }); if (failure_) std::rethrow_exception(failure_); return ready; @@ -118,8 +210,12 @@ void RecordConsumer::rethrow_if_failed() const { void RecordConsumer::finish() const { std::lock_guard lock(mu_); if (failure_) std::rethrow_exception(failure_); - if (!descriptors_.empty()) { - invalid("durable completion found leftover encoded descriptors"); + if (!descriptors_.empty() || payloads_to_discard_ != 0) { + invalid("durable completion found " + + std::to_string(descriptors_.size()) + + " leftover encoded descriptors and " + + std::to_string(payloads_to_discard_) + + " payloads still owed to a discard window"); } if (active_payloads_ != 0) { invalid("durable completion found active sink submission"); @@ -131,4 +227,37 @@ size_t RecordConsumer::pending_descriptors() const { return descriptors_.size(); } +bool RecordConsumer::failed() const { + std::lock_guard lock(mu_); + return static_cast(failure_); +} + +RecordConsumerSnapshot RecordConsumer::snapshot() const { + std::lock_guard lock(mu_); + RecordConsumerSnapshot snapshot; + snapshot.policy = policy_; + snapshot.failed = static_cast(failure_); + if (failure_) snapshot.failure = describe_failure(failure_); + snapshot.discarded_descriptors = discarded_descriptors_; + snapshot.discarded_payloads = discarded_payloads_; + snapshot.steps_with_discards = steps_with_discards_; + return snapshot; +} + +void RecordConsumer::begin_discard_window() { + std::lock_guard lock(mu_); + discarding_ = true; + for (const QueuedDescriptor& queued : descriptors_) { + note_discarded_step_locked(queued.step); + } + discarded_descriptors_ += descriptors_.size(); + payloads_to_discard_ += descriptors_.size(); + descriptors_.clear(); +} + +void RecordConsumer::end_discard_window() { + std::lock_guard lock(mu_); + discarding_ = false; +} + } // namespace ring diff --git a/native/csrc/ring/record_consumer.h b/native/csrc/ring/record_consumer.h index 721892dcb..4f65ddb3b 100644 --- a/native/csrc/ring/record_consumer.h +++ b/native/csrc/ring/record_consumer.h @@ -3,6 +3,7 @@ #pragma once #include "record_descriptor.h" +#include "record_failure_policy.h" #include "record_sink.h" #include @@ -15,14 +16,36 @@ #include #include #include +#include #include #include namespace ring { +struct RecordConsumerSnapshot { + RecordFailurePolicy policy{RecordFailurePolicy::kRaiseAtProducer}; + bool failed{false}; + std::string failure; + // Descriptors dropped without being stored: by a kDisableCapture latch + // (those still queued when it happened, plus every push after it) and, + // under either policy, by a discard window (a skipped step). + uint64_t discarded_descriptors{0}; + // Payloads not stored: under kDisableCapture the one whose submission + // failed and every payload delivered after the latch; under either + // policy the payloads of the descriptors a discard window dropped. + uint64_t discarded_payloads{0}; + // Distinct steps (see begin_step) that lost at least one record to a + // discard window. A window drops every record still queued, from any + // step, so this can exceed the number of windows. A latch's discards + // are not counted here. + uint64_t steps_with_discards{0}; +}; + class RecordConsumer { public: - explicit RecordConsumer(std::shared_ptr sink); + explicit RecordConsumer( + std::shared_ptr sink, + RecordFailurePolicy policy = RecordFailurePolicy::kRaiseAtProducer); RecordConsumer(const RecordConsumer&) = delete; RecordConsumer& operator=(const RecordConsumer&) = delete; @@ -33,7 +56,8 @@ class RecordConsumer { // Consume exactly one descriptor for one physical payload. The payload // must be a contiguous CPU byte tensor containing the actual produced - // bytes from its ready publication. + // bytes from its ready publication. Throws on failure only under + // kRaiseAtProducer; under kDisableCapture it latches and returns. void consume_payload(at::Tensor payload); // Latch an asynchronous worker failure. The first failure is retained. @@ -44,15 +68,71 @@ class RecordConsumer { bool wait_until_idle(std::chrono::milliseconds timeout) const; void finish() const; size_t pending_descriptors() const; + bool failed() const; + RecordConsumerSnapshot snapshot() const; + + // Under kDisableCapture after a latch, or while payloads owed to a + // discard window remain: account the next payload as discarded, exactly + // as consume_payload would, and return true, so the caller can skip + // copying bytes nobody will store. Otherwise change nothing and return + // false; the caller then delivers the payload to consume_payload. + bool discard_next_payload_if_unwanted(); + + // Start a new step: descriptors pushed from now on are tagged with it, + // so steps_with_discards can count the steps a discard window reaches. + void begin_step(); + + // Skip the rest of a step without latching (a spent stall budget). + // begin_discard_window() drops every descriptor still queued -- the + // step's own and those of earlier steps whose records the sink has not + // taken yet, which the stall bound needs out of the drain's way -- and + // every descriptor pushed until end_discard_window() is dropped on + // arrival. The payloads of all of them are discarded as the drain + // delivers them, so the ring's descriptor/payload pairing is kept. + // Counted in discarded_* and steps_with_discards. A payload already in + // the sink is not recalled. + void begin_discard_window(); + void end_discard_window(); private: + // A queued descriptor and the step it was pushed in. + struct QueuedDescriptor { + RecordDescriptor descriptor; + uint64_t step{0}; + }; + + // Retain the first failure; under kDisableCapture also drop the queued + // descriptors. Caller holds mu_. + void latch_locked(std::exception_ptr failure); + // Count `step` in steps_with_discards_ unless it already is. Windows + // see steps in non-decreasing order (FIFO, and a window drops all of + // its own step's pushes), so the last counted step is enough. Caller + // holds mu_. + void note_discarded_step_locked(uint64_t step); + // The discard branches shared by consume_payload and + // discard_next_payload_if_unwanted; false under kRaiseAtProducer after + // a latch. Caller holds mu_. + bool take_unwanted_payload_locked(); + std::shared_ptr sink_; + const RecordFailurePolicy policy_; mutable std::mutex mu_; mutable std::condition_variable idle_cv_; - std::deque descriptors_; + std::deque descriptors_; + uint64_t step_{0}; std::exception_ptr failure_; size_t active_payloads_{0}; + // Discard window: while open, pushes are dropped on arrival. Every + // dropped descriptor still owes the consumer one payload, and those + // payloads arrive before any descriptor queued after the window + // (FIFO), so a count is enough to keep the pairing. + bool discarding_{false}; + uint64_t payloads_to_discard_{0}; + uint64_t discarded_descriptors_{0}; + uint64_t discarded_payloads_{0}; + uint64_t steps_with_discards_{0}; + std::optional last_discarded_step_; }; } // namespace ring diff --git a/native/csrc/ring/record_failure_policy.h b/native/csrc/ring/record_failure_policy.h new file mode 100644 index 000000000..741391945 --- /dev/null +++ b/native/csrc/ring/record_failure_policy.h @@ -0,0 +1,19 @@ +// What a latched record failure does to the rest of a record runtime. +// +// Plain C++ (no ATen/CUDA) so ring_engine_py.h can name it. + +#pragma once + +namespace ring { + +// kRaiseAtProducer: the failure reaches the forward that publishes the next +// record. kDisableCapture: capture stops and the forward keeps running -- +// descriptors pushed and payloads delivered after the latch are discarded +// and counted. Under both, the failure still surfaces at every checked +// completion (flush_records_and_wait) and in the capture status. +enum class RecordFailurePolicy { + kRaiseAtProducer, + kDisableCapture, +}; + +} // namespace ring diff --git a/native/csrc/ring/record_sink.h b/native/csrc/ring/record_sink.h index bb084d4e5..b3907a1c5 100644 --- a/native/csrc/ring/record_sink.h +++ b/native/csrc/ring/record_sink.h @@ -10,6 +10,7 @@ #include #include #include +#include #include namespace ring { @@ -39,6 +40,14 @@ class RecordSink { virtual bool flush_and_wait(Duration timeout) = 0; virtual void rethrow_if_failed() const = 0; + // The longest one submit() can block waiting for the sink to admit an + // envelope, or nullopt when that wait has no bound (or the sink cannot + // say). A record ring with a step stall budget waits for one in-flight + // submit after the budget is spent, so it refuses a sink without one. + virtual std::optional admission_bound() const { + return std::nullopt; + } + bool engine_owned() const { std::lock_guard lock(engine_mu_); return engine_state_ != EngineState::AVAILABLE; diff --git a/native/csrc/ring/record_stall_budget.h b/native/csrc/ring/record_stall_budget.h new file mode 100644 index 000000000..229ae2f99 --- /dev/null +++ b/native/csrc/ring/record_stall_budget.h @@ -0,0 +1,42 @@ +// How long a record reservation past a spent step stall budget may wait for +// the drain before the ring is failed. +// +// Plain C++ (no ATen/CUDA) so the CPU tests can check it. + +#pragma once + +#include +#include +#include + +namespace ring { + +// Once a step's stall budget is spent the sink is out of the path, but the +// drain still has to move what the ring holds -- device ring to pinned +// staging, then the worker takes each payload off the staging and discards +// it -- before the reservation fits. The reservation waits for that at most +// the sink's admission bound plus this grace: a fixed 2 s, plus the bytes the +// payload ring and its pinned staging can hold, at a conservative 1 GB/s, +// which is one byte per nanosecond. With the engine's default 4 GiB ring +// and 4 GiB staging that is about 10.6 s. +inline constexpr std::chrono::milliseconds kRecordDrainGraceBase{2000}; +inline constexpr uint64_t kRecordDrainGraceBytesPerSecond = 1'000'000'000; +static_assert(kRecordDrainGraceBytesPerSecond == 1'000'000'000, + "record_drain_grace converts one byte to one nanosecond"); + +inline std::chrono::nanoseconds record_drain_grace( + uint64_t payload_ring_bytes, uint64_t staging_bytes) { + // Saturate far below the clock's range, so that a deadline adding this + // grace to steady_clock::now() cannot overflow. + constexpr uint64_t kCap = + static_cast(std::numeric_limits::max() / 4); + const uint64_t base = static_cast( + std::chrono::nanoseconds(kRecordDrainGraceBase).count()); + uint64_t total = base; + for (const uint64_t bytes : {payload_ring_bytes, staging_bytes}) { + total = bytes > kCap - total ? kCap : total + bytes; + } + return std::chrono::nanoseconds(static_cast(total)); +} + +} // namespace ring diff --git a/native/csrc/ring/ring_engine.cu b/native/csrc/ring/ring_engine.cu index 0d75eaf21..ea43b7500 100644 --- a/native/csrc/ring/ring_engine.cu +++ b/native/csrc/ring/ring_engine.cu @@ -36,7 +36,8 @@ RingEngine::RingEngine(const RingConfig& cfg, ring_py::TensorMetaFifo& fifo, } RingEngine::RingEngine(const RingConfig& cfg, - std::shared_ptr lease) + std::shared_ptr lease, + RecordFailurePolicy failure_policy) : cfg_(cfg), ring_(cfg), record_sink_lease_(std::move(lease)) { if (cfg_.payload_ring_bytes % PAYLOAD_ALIGN != 0) { @@ -64,7 +65,7 @@ RingEngine::RingEngine(const RingConfig& cfg, ? record_sink_lease_->claim() : nullptr; try { record_p2p_ = std::make_unique( - *drain_, record_sink_); + *drain_, record_sink_, failure_policy); } catch (...) { release_record_sink(); throw; diff --git a/native/csrc/ring/ring_engine.h b/native/csrc/ring/ring_engine.h index 0053ca3e2..a793b1fc1 100644 --- a/native/csrc/ring/ring_engine.h +++ b/native/csrc/ring/ring_engine.h @@ -17,7 +17,9 @@ class RingEngine { public: RingEngine(const RingConfig& cfg, ring_py::TensorMetaFifo& fifo, SubmitFn submit_fn); - RingEngine(const RingConfig& cfg, std::shared_ptr lease); + RingEngine(const RingConfig& cfg, std::shared_ptr lease, + RecordFailurePolicy failure_policy = + RecordFailurePolicy::kRaiseAtProducer); ~RingEngine() noexcept; RingEngine(const RingEngine&) = delete; diff --git a/native/csrc/ring/ring_engine_py.cu b/native/csrc/ring/ring_engine_py.cu index 78fe08187..541486193 100644 --- a/native/csrc/ring/ring_engine_py.cu +++ b/native/csrc/ring/ring_engine_py.cu @@ -10,10 +10,14 @@ #include "ring/ring_torch_op.h" #include "ring/producer.cuh" #include "ring/ring_debug.h" +#include "ring/record_stall_budget.h" #include // at::cuda::getCurrentCUDAStream #include +#include #include #include +#include +#include #include #include #include @@ -29,6 +33,16 @@ namespace { using FlushClock = std::chrono::steady_clock; +std::string describe(const std::exception_ptr& failure) { + try { + std::rethrow_exception(failure); + } catch (const std::exception& error) { + return error.what(); + } catch (...) { + return "unknown record failure"; + } +} + void check_flush_cuda(cudaError_t error, const char* operation) { if (error == cudaSuccess) return; throw std::runtime_error( @@ -138,6 +152,28 @@ struct RingEnginePy::Impl { at::Tensor payload_view; bool record_mode{false}; + // Record rings only. The step state is written by the reserving + // thread; the atomics let record_capture_status read from any thread. + RecordRuntimeOptions record_options; + // The sink's admission bound, when a budget needs one (see the ctor), + // and the time the drain may take past it to move what the ring holds + // once a budget is spent (ring::record_drain_grace). + FlushClock::duration sink_admission_bound{}; + FlushClock::duration drain_grace{}; + std::atomic step_wait_ns{0}; + std::atomic reserve_wait_ns{0}; + std::atomic max_step_wait_ns{0}; + std::atomic stall_budget_exhaustions{0}; + std::atomic skipped_steps{0}; + // begin_record_step() has run at least once. + std::atomic step_started{false}; + // This step's budget ran out: its remaining records are being skipped. + std::atomic skipping_step{false}; + // kRaiseAtProducer: the spent budget, raised at the next step boundary + // and at flush rather than inside the forward that ran out of it. + mutable std::mutex step_failure_mu; + std::exception_ptr step_failure; + Impl(ring::RingConfig cfg, SubmitFn sf) : engine(std::move(cfg), fifo, std::move(sf)) { @@ -150,9 +186,31 @@ struct RingEnginePy::Impl { at::TensorOptions().dtype(at::kByte).device(at::kCUDA, dev_idx)); } - Impl(ring::RingConfig cfg, std::shared_ptr lease) - : engine(std::move(cfg), std::move(lease)), record_mode(true) + Impl(ring::RingConfig cfg, std::shared_ptr lease, + RecordRuntimeOptions options) + : engine(std::move(cfg), std::move(lease), options.failure_policy), + record_mode(true), record_options(options) { + if (options.step_stall_budget_ms != 0) { + // Past the budget the forward still waits for the one envelope + // the sink is admitting; only a sink that bounds its admission + // makes that wait, and so the step's stall, bounded. + const auto sink = engine.record_sink(); + const auto bound = sink ? sink->admission_bound() : std::nullopt; + if (!bound) { + throw std::invalid_argument( + "step_stall_budget_ms needs a record sink with an " + "admission bound (a NativePackSink with overload " + "'drop_newest', or 'block' with an admission_timeout_s); " + "this sink's admission has no bound, so a stall past the " + "budget would not be bounded either"); + } + sink_admission_bound = + std::chrono::duration_cast(*bound); + drain_grace = std::chrono::duration_cast( + ring::record_drain_grace( + engine.payload_cap(), engine.staging_cap())); + } const auto& state = engine.ring_state(); int dev_idx = 0; cudaGetDevice(&dev_idx); @@ -161,6 +219,116 @@ struct RingEnginePy::Impl { {static_cast(state.payload_cap)}, at::TensorOptions().dtype(at::kByte).device(at::kCUDA, dev_idx)); } + + void account_record_wait(FlushClock::time_point started) { + const uint64_t waited = static_cast( + std::chrono::duration_cast( + FlushClock::now() - started).count()); + reserve_wait_ns.fetch_add(waited, std::memory_order_relaxed); + const uint64_t step = step_wait_ns.fetch_add( + waited, std::memory_order_relaxed) + waited; + uint64_t worst = max_step_wait_ns.load(std::memory_order_relaxed); + while (step > worst && !max_step_wait_ns.compare_exchange_weak( + worst, step, std::memory_order_relaxed)) { + } + } + + // Make room for a record reservation that did not fit: finish the + // producers already queued, then drain the ring. The drain is held back + // by the record worker, and the worker by the sink, so this is where a + // slow or stuck sink reaches the forward. The wait is bounded by what + // is left of this step's stall budget. Past it the rest of the step is + // skipped: every record still queued for the sink, from any step, and + // every record the step still reserves are discarded, which leaves the + // drain waiting only for the envelope the sink is admitting. + void wait_for_record_space() { + cudaStream_t stream = at::cuda::getCurrentCUDAStream().stream(); + cudaStreamSynchronize(stream); + auto& drain = engine.drain_thread(); + auto& consumer = engine.record_consumer(); + const bool disable = record_options.failure_policy == + ring::RecordFailurePolicy::kDisableCapture; + const auto started = FlushClock::now(); + if (record_options.step_stall_budget_ms == 0) { + drain.force_flush_and_wait_until(FlushClock::time_point::max()); + account_record_wait(started); + return; + } + // Already skipping this step, or capture is off: the worker discards + // instead of submitting, so the sink is out of the path. + if (skipping_step.load(std::memory_order_relaxed) || + (disable && consumer.failed())) { + wait_past_the_budget(started); + return; + } + const auto budget = std::chrono::duration_cast( + std::chrono::milliseconds(record_options.step_stall_budget_ms)); + const auto spent = std::chrono::duration_cast( + std::chrono::nanoseconds( + step_wait_ns.load(std::memory_order_relaxed))); + const auto deadline = + started + (spent < budget ? budget - spent : FlushClock::duration{}); + const bool drained = drain.force_flush_and_wait_until(deadline); + account_record_wait(started); + if (drained) return; + + stall_budget_exhaustions.fetch_add(1, std::memory_order_relaxed); + skipped_steps.fetch_add(1, std::memory_order_relaxed); + skipping_step.store(true, std::memory_order_relaxed); + consumer.begin_discard_window(); + if (!disable) { + const auto step_ms = + step_wait_ns.load(std::memory_order_relaxed) / 1'000'000ull; + std::lock_guard lock(step_failure_mu); + if (!step_failure) { + step_failure = std::make_exception_ptr(std::runtime_error( + "record capture stall budget exhausted: reservations " + "waited " + std::to_string(step_ms) + " ms in one step " + "for the sink to free ring space (step_stall_budget_ms=" + + std::to_string(record_options.step_stall_budget_ms) + + "); the step's remaining records and every record still " + "queued for the sink, from any step, were discarded")); + } + } + // The reservation must still complete: under CUDA-graph replay the + // producers launch whatever the host decides. + wait_past_the_budget(FlushClock::now()); + } + + // With the sink out of the path, what is left is the envelope it is + // admitting (at most its admission bound) and moving what the ring + // holds out to the worker, which discards it (at most drain_grace). + // Past both the ring is failed under either policy: the alternative is + // a forward that never returns. Either the sink held its admission past + // its bound or the drain was slower than the grace assumes; the wait + // cannot tell which, so the error names both. + void wait_past_the_budget(FlushClock::time_point started) { + auto& drain = engine.drain_thread(); + const auto deadline = started + sink_admission_bound + drain_grace; + const bool drained = drain.force_flush_and_wait_until(deadline); + account_record_wait(started); + if (drained) return; + const auto ms = [](FlushClock::duration duration) { + return std::to_string(std::chrono::duration_cast< + std::chrono::milliseconds>(duration).count()); + }; + const std::exception_ptr failure = + std::make_exception_ptr(std::runtime_error( + "record ring did not drain, with the sink out of the path, " + "within the sink's admission bound (" + + ms(sink_admission_bound) + " ms) plus the drain grace (" + + ms(drain_grace) + " ms: 2 s plus the ring and staging bytes " + "at 1 GB/s); either the sink held an admission past its bound " + "or the drain was slower than the grace allows")); + engine.record_consumer().record_failure(failure); + std::rethrow_exception(failure); + } + + // kRaiseAtProducer: the spent budget of an earlier step, if any. + std::exception_ptr pending_step_failure() const { + std::lock_guard lock(step_failure_mu); + return step_failure; + } }; // --------------------------------------------------------------------------- @@ -187,13 +355,16 @@ RingEnginePy::RingEnginePy(RingConfig cfg, SubmitFn submit_fn) { } RingEnginePy::RingEnginePy( - RingConfig cfg, std::shared_ptr sink) + RingConfig cfg, std::shared_ptr sink, + RecordRuntimeOptions options) : RingEnginePy( - std::move(cfg), ring::RecordSinkLease::acquire(std::move(sink))) {} + std::move(cfg), ring::RecordSinkLease::acquire(std::move(sink)), + options) {} RingEnginePy::RingEnginePy( - RingConfig cfg, std::shared_ptr lease) { - impl_ = std::make_unique(convert(cfg), std::move(lease)); + RingConfig cfg, std::shared_ptr lease, + RecordRuntimeOptions options) { + impl_ = std::make_unique(convert(cfg), std::move(lease), options); } RingEnginePy::~RingEnginePy() = default; @@ -499,11 +670,26 @@ int RingEnginePy::reserve_record( drain.rethrow_drain_failure(); drain.rethrow_record_reclaim_failure(); drain.apply_pending_record_reclaims(); + // A latched failure raises here, before anything is reserved, rather + // than from the descriptor push after a reservation no producer uses. + // Under kDisableCapture it never raises: the ring keeps its protocol + // and the consumer discards. + if (impl_->record_options.failure_policy == + ring::RecordFailurePolicy::kRaiseAtProducer) { + impl_->engine.record_consumer().rethrow_if_failed(); + } + if (impl_->record_options.step_stall_budget_ms != 0 && + !impl_->step_started.load(std::memory_order_relaxed)) { + // Without step boundaries the budget would span the runtime's whole + // life, and once spent every later wait would skip at once. + throw std::logic_error( + "step_stall_budget_ms is set, but no record step has begun: call " + "RecordRuntime.begin_step() once per model step, before its " + "first reservation"); + } if (reservation_bytes > effective_cap || num_tasks > task_cap) { - cudaStream_t stream = at::cuda::getCurrentCUDAStream().stream(); - cudaStreamSynchronize(stream); - drain.force_flush_and_wait(); + impl_->wait_for_record_space(); drain.rethrow_drain_failure(); drain.rethrow_record_reclaim_failure(); drain.apply_pending_record_reclaims(); @@ -524,9 +710,7 @@ int RingEnginePy::reserve_record( return STEP_RING_OK; } - cudaStream_t stream = at::cuda::getCurrentCUDAStream().stream(); - cudaStreamSynchronize(stream); - drain.force_flush_and_wait(); + impl_->wait_for_record_space(); drain.rethrow_drain_failure(); drain.rethrow_record_reclaim_failure(); drain.apply_pending_record_reclaims(); @@ -538,6 +722,59 @@ int RingEnginePy::reserve_record( return STEP_RING_FLUSHED; } +void RingEnginePy::begin_record_step() { + if (!impl_->record_mode) { + throw std::logic_error("record steps require a record ring"); + } + auto& consumer = impl_->engine.record_consumer(); + if (impl_->record_options.failure_policy == + ring::RecordFailurePolicy::kRaiseAtProducer) { + // A spent budget latches here, at the step boundary: from now on + // the runtime is failed like after any other refusal. + if (const std::exception_ptr failure = impl_->pending_step_failure()) { + consumer.record_failure(failure); + } + consumer.rethrow_if_failed(); + } + impl_->step_wait_ns.store(0, std::memory_order_relaxed); + impl_->step_started.store(true, std::memory_order_relaxed); + if (impl_->skipping_step.exchange(false, std::memory_order_relaxed)) { + consumer.end_discard_window(); + } + consumer.begin_step(); +} + +RecordCaptureStatus RingEnginePy::record_capture_status() const { + if (!impl_->record_mode) { + throw std::logic_error("record capture status requires a record ring"); + } + const ring::RecordConsumerSnapshot consumer = + impl_->engine.record_consumer().snapshot(); + RecordCaptureStatus status; + status.failure_policy = consumer.policy; + status.failed = consumer.failed; + status.failure = consumer.failure; + status.discarded_descriptors = consumer.discarded_descriptors; + status.discarded_payloads = consumer.discarded_payloads; + status.step_stall_budget_ms = impl_->record_options.step_stall_budget_ms; + status.stall_budget_exhaustions = + impl_->stall_budget_exhaustions.load(std::memory_order_relaxed); + status.skipped_steps = + impl_->skipped_steps.load(std::memory_order_relaxed); + status.steps_with_discards = consumer.steps_with_discards; + if (!status.failed) { + if (const std::exception_ptr failure = impl_->pending_step_failure()) { + status.failed = true; + status.failure = describe(failure); + } + } + status.reserve_wait_ns = + impl_->reserve_wait_ns.load(std::memory_order_relaxed); + status.max_step_wait_ns = + impl_->max_step_wait_ns.load(std::memory_order_relaxed); + return status; +} + void RingEnginePy::push_record_descriptors( std::vector descriptors) { if (!impl_->record_mode) { @@ -564,7 +801,22 @@ bool RingEnginePy::flush_records_and_wait(uint64_t timeout_ms) { auto& consumer = impl_->engine.record_consumer(); drain.rethrow_drain_failure(); drain.rethrow_record_reclaim_failure(); + // Under kDisableCapture a latched failure is reported only once the + // drain has delivered what the forward already emitted: the consumer + // discards it, so ring and staging space come back and the discard + // counters are complete when the caller hears why capture stopped. + if (consumer.failed() && + impl_->record_options.failure_policy == + ring::RecordFailurePolicy::kDisableCapture) { + cudaStream_t stream = at::cuda::getCurrentCUDAStream().stream(); + if (wait_for_stream_prefix_until(stream, deadline)) { + drain.force_flush_and_wait_until(deadline); + } + } consumer.rethrow_if_failed(); + if (const std::exception_ptr failure = impl_->pending_step_failure()) { + std::rethrow_exception(failure); + } cudaStream_t stream = at::cuda::getCurrentCUDAStream().stream(); if (!wait_for_stream_prefix_until(stream, deadline)) return false; @@ -586,14 +838,22 @@ bool RingEnginePy::flush_records_and_wait(uint64_t timeout_ms) { const auto sink = impl_->engine.record_sink(); if (sink) { - sink->rethrow_if_failed(); - const auto before_sink = FlushClock::now(); - if (before_sink >= deadline) return false; - const auto sink_timeout = - std::chrono::duration_cast( - deadline - before_sink); - if (!sink->flush_and_wait(sink_timeout)) return false; - sink->rethrow_if_failed(); + try { + sink->rethrow_if_failed(); + const auto before_sink = FlushClock::now(); + if (before_sink >= deadline) return false; + const auto sink_timeout = + std::chrono::duration_cast( + deadline - before_sink); + if (!sink->flush_and_wait(sink_timeout)) return false; + sink->rethrow_if_failed(); + } catch (...) { + // A sink failure found at the barrier (a pipeline failure, a + // record lost after admission) is as much a refusal as one at + // submit: it latches, so capture stops and says why. + consumer.record_failure(std::current_exception()); + throw; + } } if (FlushClock::now() > deadline) return false; return true; diff --git a/native/csrc/ring/ring_engine_py.h b/native/csrc/ring/ring_engine_py.h index 76fc307a5..a465f7c33 100644 --- a/native/csrc/ring/ring_engine_py.h +++ b/native/csrc/ring/ring_engine_py.h @@ -15,6 +15,7 @@ #include "ring/tensor_meta.h" // TensorMeta, TensorMetaFifo #include "ring/record_descriptor.h" +#include "ring/record_failure_policy.h" // Forward-declare ATen and the generic sink so this plain interface does not // expose their implementation headers. @@ -45,6 +46,53 @@ struct RingConfig { uint64_t insert_queue_max_items = 65536; }; +// How a record ring treats a sink that refuses or stalls. Fixed for the +// ring's lifetime. +struct RecordRuntimeOptions { + ring::RecordFailurePolicy failure_policy = + ring::RecordFailurePolicy::kRaiseAtProducer; + // Cap on the time the record reservations of one step may wait for the + // drain to free ring space. A step starts at begin_record_step(), which + // the integration calls once per model step: the record ring sees no + // step boundary of its own (see begin_record_step), so with a budget a + // reservation before the first begin_record_step() is refused. Past + // the budget the step's remaining records and every record still queued + // for the sink, from any step, are discarded (counted in discarded_*, + // skipped_steps and steps_with_discards), and capture resumes at the + // next step; under kRaiseAtProducer the exhaustion is then raised at the + // next begin_record_step() and at flush. The skip still waits for the + // one envelope the sink is admitting, so a budget needs a sink with an + // admission bound (RecordSink::admission_bound) and construction + // refuses one without; past that bound plus ring::record_drain_grace + // the reservation fails the ring. 0 waits without bound, as before the + // budget existed. + uint64_t step_stall_budget_ms = 0; +}; + +// The record runtime's failure state and forward-side stall counters. +struct RecordCaptureStatus { + ring::RecordFailurePolicy failure_policy = + ring::RecordFailurePolicy::kRaiseAtProducer; + bool failed = false; + std::string failure; + uint64_t discarded_descriptors = 0; + uint64_t discarded_payloads = 0; + uint64_t step_stall_budget_ms = 0; + uint64_t stall_budget_exhaustions = 0; + // Steps whose budget ran out; always equal to stall_budget_exhaustions. + // Each such skip also discards every record still queued for the sink, + // earlier steps' too, so steps_with_discards can exceed it. + uint64_t skipped_steps = 0; + // Distinct steps that lost at least one record to a spent budget: the + // skipped steps plus the earlier steps whose records were still queued. + uint64_t steps_with_discards = 0; + // Time record reservations spent waiting for the drain, in total and + // for the worst step. The producer-stream synchronisation before each + // wait is not included: that is the forward's own GPU work. + uint64_t reserve_wait_ns = 0; + uint64_t max_step_wait_ns = 0; +}; + // Called by the p2p thread for each per-request tensor slice. using SubmitFn = std::function sink); + std::shared_ptr sink, + RecordRuntimeOptions options = {}); explicit RingEnginePy(RingConfig cfg, - std::shared_ptr lease); + std::shared_ptr lease, + RecordRuntimeOptions options = {}); ~RingEnginePy(); RingEnginePy(const RingEnginePy&) = delete; @@ -171,9 +221,23 @@ class RingEnginePy { int prepare_step(uint64_t step_total_bytes, uint32_t num_hooks); // Each item is (aligned upper bound, needs actual-byte reconciliation). + // When the ring has no room, waits for the drain within the step's + // remaining stall budget; see RecordRuntimeOptions. int reserve_record( const std::vector>& reservation_items); + // Start a new step for the stall budget: the reservations after this + // call share one fresh budget, and a step skipped for a spent budget + // ends. Under kRaiseAtProducer it first raises a failure latched + // during the previous step, including a spent budget, so the error + // surfaces here, outside the forward. Explicit because the ring has no + // reliable step boundary of its own: eager hooks reserve per output + // (and output ids repeat within a step), and a CUDA-graph step replays + // one plan, or several with piecewise graphs. + void begin_record_step(); + + RecordCaptureStatus record_capture_status() const; + void push_record_descriptors(std::vector descriptors); void submit_record_cpu_direct(at::Tensor cpu_tensor, uint64_t tensor_bytes); diff --git a/native/csrc/sink/bindings_sink.cpp b/native/csrc/sink/bindings_sink.cpp index 427edd63d..696256494 100644 --- a/native/csrc/sink/bindings_sink.cpp +++ b/native/csrc/sink/bindings_sink.cpp @@ -12,7 +12,9 @@ #include #include +#include #include +#include #include #include "../ring/record_sink.h" @@ -45,6 +47,27 @@ ring::PayloadSlice ParseSlice(const py::dict& row) { return slice; } +dmi_sink::Overload ParseOverload(const std::string& name) { + if (name == "block") return dmi_sink::Overload::kBlock; + if (name == "drop_newest") return dmi_sink::Overload::kDropNewest; + throw py::value_error("overload must be 'block' or 'drop_newest', got '" + + name + "'"); +} + +const char* OverloadName(dmi_sink::Overload overload) { + return overload == dmi_sink::Overload::kBlock ? "block" : "drop_newest"; +} + +// None waits without bound under block, as SinkConfig's -1 does. +double ParseAdmissionTimeout(const std::optional& timeout_s) { + if (!timeout_s.has_value()) return -1.0; + if (!std::isfinite(*timeout_s) || *timeout_s < 0.0) { + throw py::value_error( + "admission_timeout_s must be None or a finite, non-negative number"); + } + return *timeout_s; +} + } // namespace namespace { @@ -77,6 +100,14 @@ bool RingTypesRegistered() { py::detail::get_type_info(typeid(ring::RecordSinkLease)) != nullptr; } +// RecordSink.admission_bound_s: the sink's admission bound in seconds, or +// None when it has none. +std::optional AdmissionBoundSeconds(const ring::RecordSink& sink) { + const auto bound = sink.admission_bound(); + if (!bound) return std::nullopt; + return std::chrono::duration(*bound).count(); +} + void EnsureRingTypes(py::module_& m) { if (RingTypesRegistered()) return; // The main backend lives beside the dmi package and is loaded from its @@ -100,7 +131,8 @@ void EnsureRingTypes(py::module_& m) { .def("_acquire_engine", [](std::shared_ptr sink) { return ring::RecordSinkLease::acquire(std::move(sink)); - }); + }) + .def_property_readonly("admission_bound_s", &AdmissionBoundSeconds); m.attr("RING_TYPES_ARE_STANDINS") = true; } @@ -123,8 +155,13 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { const std::string& layout, int num_workers, uint64_t max_queue_records, uint64_t max_queue_bytes, uint64_t max_pack_bytes, uint64_t max_pack_records, - uint64_t max_linger_ns, uint64_t spool_max_bytes) { + uint64_t max_linger_ns, uint64_t spool_max_bytes, + const std::string& overload, + std::optional admission_timeout_s) { dmi_sink::SinkConfig config; + config.overload = ParseOverload(overload); + config.admission_timeout_s = + ParseAdmissionTimeout(admission_timeout_s); config.spool_root = spool_root; config.spool_max_bytes = spool_max_bytes; config.num_workers = num_workers; @@ -144,7 +181,11 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { py::arg("max_pack_bytes") = 128ull * 1024 * 1024, py::arg("max_pack_records") = 10'000, py::arg("max_linger_ns") = 1'000'000'000, - py::arg("spool_max_bytes") = 1ull << 40) + py::arg("spool_max_bytes") = 1ull << 40, + // SinkConfig's own defaults: the Python NativeSinkConfig, which + // the ring-fed sink is built from, picks block with 2 s. + py::arg("overload") = "drop_newest", + py::arg("admission_timeout_s") = py::none()) .def("attach", [](std::shared_ptr self) { // Simulates engine ownership for tests (the real engine takes @@ -175,6 +216,18 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { }) .def("rethrow_if_failed", &dmi_sink::NativePackSink::rethrow_if_failed) .def_property_readonly("layout", &dmi_sink::NativePackSink::layout) + .def_property_readonly( + "overload", + [](const dmi_sink::NativePackSink& self) { + return std::string(OverloadName(self.sink().config().overload)); + }) + .def_property_readonly( + "admission_timeout_s", + [](const dmi_sink::NativePackSink& self) -> std::optional { + const double timeout_s = self.sink().config().admission_timeout_s; + if (timeout_s < 0) return std::nullopt; + return timeout_s; + }) .def("snapshot", [](dmi_sink::NativePackSink& self) { const dmi_sink::SinkSnapshot snapshot = self.sink().Snapshot(); py::dict out; @@ -183,6 +236,8 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { out["persisted_records"] = snapshot.persisted_records; out["packs_persisted"] = snapshot.packs_persisted; out["dropped_records"] = snapshot.dropped_records; + out["timed_out_records"] = snapshot.timed_out_records; + out["rejected_closed_records"] = snapshot.rejected_closed_records; out["duplicate_records"] = snapshot.duplicate_records; out["oversized_records"] = snapshot.oversized_records; out["failures"] = snapshot.failures; diff --git a/native/csrc/sink/native_pack_sink.cpp b/native/csrc/sink/native_pack_sink.cpp index ad6c5e54a..9687c73e7 100644 --- a/native/csrc/sink/native_pack_sink.cpp +++ b/native/csrc/sink/native_pack_sink.cpp @@ -111,12 +111,52 @@ NativePackSink::NativePackSink(std::unique_ptr sink, std::string error; const std::string start_error = sink_->Start(&error); if (!start_error.empty()) invalid("sink start failed: " + start_error); + baseline_ = sink_->Snapshot(); +} + +std::string NativePackSink::LossError() const { + const SinkSnapshot now = sink_->Snapshot(); + std::string losses; + const auto add = [&](const char* name, uint64_t current, uint64_t base) { + if (current == base) return; + if (!losses.empty()) losses += ", "; + losses += std::string(name) + "=" + std::to_string(current - base); + }; + // record_adapter.py _LOSS_COUNTERS, in its order. + add("dropped_records", now.dropped_records, baseline_.dropped_records); + add("timed_out_records", now.timed_out_records, + baseline_.timed_out_records); + add("oversized_records", now.oversized_records, + baseline_.oversized_records); + add("duplicate_records", now.duplicate_records, + baseline_.duplicate_records); + add("rejected_closed_records", now.rejected_closed_records, + baseline_.rejected_closed_records); + add("failures", now.failures, baseline_.failures); + if (losses.empty()) return ""; + return "NativePackSink: pipeline reported lost records (" + losses + ")"; +} + +std::optional +NativePackSink::admission_bound() const { + const SinkConfig& config = sink_->config(); + if (config.overload == Overload::kDropNewest) return Duration::zero(); + if (config.admission_timeout_s < 0) return std::nullopt; + // Rounded up: a bound the ring waits on must not be shorter than the + // sink's own deadline. + return std::chrono::ceil( + std::chrono::duration(config.admission_timeout_s)); } NativePackSink::~NativePackSink() = default; void NativePackSink::submit(ring::RecordEnvelope envelope) { if (!engine_owned()) invalid("sink is not attached to a RingEngine"); + // A record lost after admission (oversized framing, a duplicate id) is + // counted on the pack worker, after its submit returned. Refusing the + // next envelope is how the ring hears of it: its record runtime latches + // and reports the loss instead of storing around a hole. + rethrow_if_failed(); const ring::RecordDescriptor& descriptor = envelope.descriptor; if (descriptor.layout != layout_) invalid("unexpected record layout"); if (descriptor.rows.empty()) { @@ -134,6 +174,8 @@ void NativePackSink::submit(ring::RecordEnvelope envelope) { static_cast(payload.numel()) * static_cast(payload.element_size()); + // One admission deadline for the whole envelope, started here. + EnvelopeAdmission admission(*sink_); for (const ring::EncodedRecordRow& row : descriptor.rows) { // Cheap structural check first (mirrors the reference sink). if (row.cells.size() != 2) invalid("descriptor row must contain two cells"); @@ -175,7 +217,7 @@ void NativePackSink::submit(ring::RecordEnvelope envelope) { input.dtype_name = dtype_name; input.shape = std::move(shape); std::string detail; - const RowStatus status = SubmitRow(*sink_, input, &detail); + const RowStatus status = admission.SubmitRow(input, &detail); if (status != RowStatus::kOk) { invalid(std::string(RowStatusName(status)) + (detail.empty() ? "" : ": " + detail)); @@ -186,12 +228,26 @@ void NativePackSink::submit(ring::RecordEnvelope envelope) { bool NativePackSink::flush_and_wait(Duration timeout) { if (!engine_owned()) invalid("sink is not attached to a RingEngine"); const double timeout_s = std::chrono::duration(timeout).count(); + // Everything admitted before the barrier must be persisted by it. + const uint64_t target = sink_->Snapshot().admitted_records; std::string error; const bool ok = sink_->Flush(timeout_s < 0 ? -1.0 : timeout_s, &error); if (!error.empty()) { throw std::runtime_error("NativePackSink: flush failed: " + error); } - return ok; + if (!ok) return false; + const std::string losses = LossError(); + if (!losses.empty()) throw std::runtime_error(losses); + const uint64_t persisted = sink_->Snapshot().persisted_records; + if (persisted - baseline_.persisted_records < + target - baseline_.admitted_records) { + throw std::runtime_error( + "NativePackSink: durability mismatch: admitted=" + + std::to_string(target - baseline_.admitted_records) + + ", persisted=" + + std::to_string(persisted - baseline_.persisted_records)); + } + return true; } void NativePackSink::rethrow_if_failed() const { @@ -200,6 +256,8 @@ void NativePackSink::rethrow_if_failed() const { if (!error.empty()) { throw std::runtime_error("NativePackSink: pipeline failed: " + error); } + const std::string losses = LossError(); + if (!losses.empty()) throw std::runtime_error(losses); } } // namespace dmi_sink diff --git a/native/csrc/sink/native_pack_sink.h b/native/csrc/sink/native_pack_sink.h index d1f288b22..141450848 100644 --- a/native/csrc/sink/native_pack_sink.h +++ b/native/csrc/sink/native_pack_sink.h @@ -12,6 +12,7 @@ #include #include +#include #include #include "../pack/pack_builder.h" @@ -34,9 +35,22 @@ class NativePackSink final : public ring::RecordSink { NativePackSink(const NativePackSink&) = delete; NativePackSink& operator=(const NativePackSink&) = delete; + // Admits the envelope's rows against one admission deadline + // (EnvelopeAdmission). Refuses it first if the pipeline has latched a + // failure or lost an admitted record (see rethrow_if_failed). void submit(ring::RecordEnvelope envelope) override; + // Also fails, after a completed flush, when a record admitted since + // construction was lost instead of persisted -- the reference adapter's + // loss check (record_adapter.py _LOSS_COUNTERS) -- or when fewer records + // were persisted than admitted. bool flush_and_wait(Duration timeout) override; + // Throws on a latched pipeline failure, and on any loss counter + // (dropped, timed out, oversized, duplicate, rejected-closed, failures) + // that has moved since construction. void rethrow_if_failed() const override; + // kDropNewest: zero (a full queue refuses at once). kBlock: the + // admission timeout, now one per envelope; none when it waits forever. + std::optional admission_bound() const override; const PackSink& sink() const { return *sink_; } const std::string& layout() const { return layout_; } @@ -46,8 +60,15 @@ class NativePackSink final : public ring::RecordSink { void on_engine_release() noexcept override {} private: + // "NativePackSink: pipeline reported ..." for the loss counters that + // moved since baseline_; empty when none did. + std::string LossError() const; + std::unique_ptr sink_; const std::string layout_; + // Counters at construction. Losses are judged against it, as the + // reference adapter judges them against its pipeline's baseline. + SinkSnapshot baseline_; }; } // namespace dmi_sink diff --git a/native/csrc/sink/pack_sink.cpp b/native/csrc/sink/pack_sink.cpp index b015a2f6f..adee3592e 100644 --- a/native/csrc/sink/pack_sink.cpp +++ b/native/csrc/sink/pack_sink.cpp @@ -160,8 +160,20 @@ size_t PackSink::RouteWorker(const std::string& tenant, return static_cast(z % static_cast(queues_.size())); } +double PackSink::AdmissionDeadline() const { + return config_.admission_timeout_s < 0 + ? -1.0 + : NowS() + config_.admission_timeout_s; +} + Admission PackSink::Submit(dmi_pack::RecordMetadata metadata, const uint8_t* payload, size_t n) { + return SubmitBy(std::move(metadata), payload, n, AdmissionDeadline()); +} + +Admission PackSink::SubmitBy(dmi_pack::RecordMetadata metadata, + const uint8_t* payload, size_t n, + double deadline) { size_t worker = 0; { std::lock_guard lock(mutex_); @@ -193,10 +205,6 @@ Admission PackSink::Submit(dmi_pack::RecordMetadata metadata, return Admission::kTooLarge; } const bool block = config_.overload == Overload::kBlock; - const double deadline = - config_.admission_timeout_s < 0 - ? -1.0 - : NowS() + config_.admission_timeout_s; SinkRecord record; record.metadata = std::move(metadata); record.payload.assign(payload, payload + n); @@ -642,6 +650,20 @@ void PackSink::Run(size_t w) { ++counters_.duplicate_records; continue; } + if (st == dmi_pack::Status::kCapacity && + assembler.builder->record_count() == 0) { + // The pack was opened for this record and it fits no empty pack: + // admission screened the payload, not the header, footer row and + // trailer around it. The reference's OversizedRecordError: drop the + // record, keep the pipeline. Sealing the empty pack instead failed + // the whole sink ("cannot seal an empty pack"). + assembler.builder.reset(); + assembler.has_first = false; + assembler.opened_ns = -1; + std::lock_guard lock(mutex_); + ++counters_.oversized_records; + continue; + } if (st == dmi_pack::Status::kCapacity) { // Full pack: seal with the reason the reference uses (records bound // hit → RECORDS, else SIZE), then retry into a fresh builder. A diff --git a/native/csrc/sink/pack_sink.h b/native/csrc/sink/pack_sink.h index f9c48bb88..daed49088 100644 --- a/native/csrc/sink/pack_sink.h +++ b/native/csrc/sink/pack_sink.h @@ -142,6 +142,14 @@ class PackSink { Admission Submit(dmi_pack::RecordMetadata metadata, const uint8_t* payload, size_t n); + // Submit against a deadline the caller took earlier from + // AdmissionDeadline(), so several records share one admission bound + // (EnvelopeAdmission). Submit() is SubmitBy(..., AdmissionDeadline()). + Admission SubmitBy(dmi_pack::RecordMetadata metadata, + const uint8_t* payload, size_t n, double deadline_s); + // The kBlock admission deadline for a wait starting now, on the sink's + // steady clock in seconds; -1 when admission_timeout_s waits forever. + double AdmissionDeadline() const; // Persist everything admitted before this call. False on timeout; the // in-flight barrier is kept for the next call to reuse. timeout_s < 0 @@ -152,6 +160,8 @@ class PackSink { SinkSnapshot Snapshot() const; // Latched async failure, if any (for RecordSink::rethrow_if_failed). std::string LastError() const; + // The configuration the sink was built with (num_workers normalised). + const SinkConfig& config() const { return config_; } // Test seam: the sink's spool, so a test can park a stager inside // Spool::Stage through Spool::SetStageHookForTesting and wedge the diff --git a/native/csrc/sink/record_row.cpp b/native/csrc/sink/record_row.cpp index 932e0332c..8d8686792 100644 --- a/native/csrc/sink/record_row.cpp +++ b/native/csrc/sink/record_row.cpp @@ -202,7 +202,12 @@ bool ParseMetadataJson(const std::string& text, return true; } -RowStatus SubmitRow(PackSink& sink, const RowInput& row, std::string* detail) { +namespace { + +// Validate the row against its metadata, then admit it by `deadline_s` (the +// sink's AdmissionDeadline() clock). +RowStatus SubmitRowBy(PackSink& sink, const RowInput& row, double deadline_s, + std::string* detail) { dmi_pack::RecordMetadata metadata; std::string error; if (!ParseMetadataJson(row.metadata_json, &metadata, &error)) { @@ -239,8 +244,8 @@ RowStatus SubmitRow(PackSink& sink, const RowInput& row, std::string* detail) { } return RowStatus::kSizeMismatch; } - const Admission admission = - sink.Submit(std::move(metadata), row.payload, row.payload_bytes); + const Admission admission = sink.SubmitBy( + std::move(metadata), row.payload, row.payload_bytes, deadline_s); if (admission != Admission::kAccepted) { if (detail) *detail = AdmissionName(admission); return RowStatus::kNotAccepted; @@ -248,4 +253,21 @@ RowStatus SubmitRow(PackSink& sink, const RowInput& row, std::string* detail) { return RowStatus::kOk; } +} // namespace + +RowStatus SubmitRow(PackSink& sink, const RowInput& row, std::string* detail) { + // One deadline per call, as Submit takes it. It is taken before the row + // is validated rather than after; validation never blocks, so nothing + // observable changes. + return SubmitRowBy(sink, row, sink.AdmissionDeadline(), detail); +} + +EnvelopeAdmission::EnvelopeAdmission(PackSink& sink) + : sink_(sink), deadline_s_(sink.AdmissionDeadline()) {} + +RowStatus EnvelopeAdmission::SubmitRow(const RowInput& row, + std::string* detail) { + return SubmitRowBy(sink_, row, deadline_s_, detail); +} + } // namespace dmi_sink diff --git a/native/csrc/sink/record_row.h b/native/csrc/sink/record_row.h index 5f2d39246..f0b991dbe 100644 --- a/native/csrc/sink/record_row.h +++ b/native/csrc/sink/record_row.h @@ -62,6 +62,22 @@ int DtypeWidth(const std::string& dtype_name); // On kNotAccepted, `detail` carries the sink's admission name. RowStatus SubmitRow(PackSink& sink, const RowInput& row, std::string* detail); +// The rows of one envelope, admitted against ONE deadline. The ring's record +// worker submits an envelope at a time, and the record ring's stall bound +// counts one admission timeout per envelope: with a deadline per row, a +// steadily slow sink held an N-row envelope for up to N timeouts. The +// deadline starts at construction; PackSink::Submit keeps its per-call +// deadline for its other callers. +class EnvelopeAdmission { + public: + explicit EnvelopeAdmission(PackSink& sink); + RowStatus SubmitRow(const RowInput& row, std::string* detail); + + private: + PackSink& sink_; + const double deadline_s_; // PackSink::AdmissionDeadline(); -1 = none +}; + } // namespace dmi_sink #endif // DMI_SINK_RECORD_ROW_H_ diff --git a/src/dmi/config.py b/src/dmi/config.py index e68c7ec3e..5e9dca10b 100644 --- a/src/dmi/config.py +++ b/src/dmi/config.py @@ -11,8 +11,8 @@ # For the string annotations on MonitoringConfig only. Importing them at # runtime would make this dependency-free module load the storage # packages; the engine imports them where it validates the values. - from .storage.capture.native_sink import NativeSinkConfig - from .storage.native_capture import NativeCaptureStorageConfig + from .storage.native_capture import ( + NativeCaptureStorageConfig, NativeSinkConfig) StorageBackend = Literal[ diff --git a/src/dmi/engine.py b/src/dmi/engine.py index efd1fbc3b..211ff3ebf 100644 --- a/src/dmi/engine.py +++ b/src/dmi/engine.py @@ -20,6 +20,13 @@ DEFAULT_DRAIN_FLUSH_TIMEOUT_US = 0 +# What a record runtime does when its sink refuses a record or stalls the +# forward past the step budget. "raise" (the default) raises in the forward +# that publishes the next record; "disable_capture" stops capture and keeps +# the forward running. Both report the failure at flush_and_wait, in +# capture_status() and at close. +RECORD_FAILURE_POLICIES = ("raise", "disable_capture") + def _native_module() -> Any: """Load the native-extension facade only when the engine needs it.""" @@ -115,6 +122,9 @@ def __init__( self._ring_transport: Optional[Any] = None self._ring_config: Optional[Any] = None self._record_mode = False + # The native sink the record runtime writes to, when the engine + # holds one (the persistent path's pack sink, or an explicit one). + self._record_sink: Optional[Any] = None if host_engine is not None and db_config is not None: raise ValueError("Provide either host_engine or db_config, not both") @@ -135,7 +145,7 @@ def __init__( self._capture_sink_config = ( config.capture_sink_config if config is not None else None) if self._capture_sink_config is not None: - from .storage.capture.native_sink import NativeSinkConfig + from .storage.native_capture import NativeSinkConfig if not isinstance(self._capture_sink_config, NativeSinkConfig): raise TypeError( @@ -266,6 +276,30 @@ def set_capture_enabled(self, enabled: bool) -> None: # lifecycle toggle; the next committed step recomputes it. transport.force_eager = False + def validate_capture_bounds(self, max_record_bytes: int) -> None: + """Refuse this engine's capture bounds if a record cannot be stored. + + Adapters call it at attach with the largest record (one captured + row's payload) they will emit, so a bound too small for it is a + ``ConfigurationError`` there instead of a refusal on the record + worker, which latches the runtime, in the middle of a forward. See + :func:`dmi.storage.native_capture.validate_capture_bounds`. + + Checks the persistent path's ``capture_sink_config`` (and + ``capture_storage_config``); with no sink config there is nothing + to check. + """ + sink_config = self._capture_sink_config + if sink_config is None: + if type(max_record_bytes) is not int or max_record_bytes <= 0: + raise ValueError("max_record_bytes must be a positive int") + return + from .storage.native_capture import validate_capture_bounds + + validate_capture_bounds( + sink_config, max_record_bytes, + storage_config=self._capture_storage_config) + def _backend_label(self) -> str: """The storage backend as the caller wrote it, with its current name when that was a deprecated one: ``'native' (now 'in-memory')``.""" @@ -315,12 +349,50 @@ def create_record_runtime( record_format: "RecordFormat[MetadataT]", *, record_sink: Optional[Any] = None, + failure_policy: str = "raise", + step_stall_budget_ms: Optional[int] = None, ) -> "RecordRuntime[MetadataT]": """Create an opt-in, non-owning runtime for encoded records. ``record_sink=None`` preserves the native ClickHouse host path. An explicit native ``RecordSink`` selects a separate backend for this runtime; the two paths are never active at the same time. + + ``failure_policy`` decides what a sink refusal (a dropped, timed-out + or oversized record, a record lost after admission, or a sink + failure) does to the forward: ``"raise"`` raises it from the next + ``RecordRuntime.begin_step``, record reservation or descriptor + push; ``"disable_capture"`` stops capture, discards what follows + and keeps the forward running. Either way ``flush_and_wait`` raises + the failure, ``capture_status()`` reports it and ``close`` logs it. + + ``step_stall_budget_ms`` caps how long the record reservations of + one step may wait for the sink to free ring space. A step is what + lies between two ``RecordRuntime.begin_step()`` calls; with a budget + the integration must call it once per model step, and a reservation + before the first call is refused. Past the budget the rest of that + step is skipped: the step's remaining records and every record still + queued for the sink, from any step, are discarded (counted in + ``discarded_*``, ``skipped_steps`` and ``steps_with_discards``), and + capture resumes at the next step. Under + ``"raise"`` the exhaustion is then raised at the next ``begin_step`` + and at ``flush_and_wait``, outside the forward. ``None`` waits + without bound. + + The skip still waits for the one envelope the sink is admitting, + so a budget requires a sink that bounds its admission, and is + refused otherwise: the ClickHouse host path (``record_sink=None`` + without the persistent backend), a sink whose ``admission_bound_s`` + is None (a ``NativePackSink`` with ``overload="block"`` and no + ``admission_timeout_s``, the reference bridge), or a persistent + ``capture_sink_config`` with block and no timeout. With a bounded + sink a step stalls the forward for at most the budget plus the + sink's admission bound (``admission_timeout_s`` under block, zero + under drop_newest), plus the time to move what the ring holds out + to the record worker. A ring that has not drained by the bound plus + a drain grace -- 2 s plus the payload ring and pinned staging bytes + at 1 GB/s, about 10.6 s with the default 4 GiB of each -- fails + from the reservation under either policy. """ if getattr(self, "_storage_backend", "auto") == "none": @@ -329,6 +401,8 @@ def create_record_runtime( "has no ring and creates no record runtime. Pick 'in-memory' " "or 'persistent' to capture" ) + self._validate_record_failure_options( + failure_policy, step_stall_budget_ms, record_sink) transport = self._ring_transport ring_config = self._ring_config if transport is None or ring_config is None: @@ -351,7 +425,9 @@ def create_record_runtime( storage = self._start_capture_storage(sweep_spool=record_sink is None) try: runtime = self._attach_record_runtime( - record_format, record_schema, record_sink) + record_format, record_schema, record_sink, + failure_policy=failure_policy, + step_stall_budget_ms=step_stall_budget_ms or 0) except BaseException: if storage is not None: self._capture_storage = None @@ -359,6 +435,57 @@ def create_record_runtime( raise return runtime + def _validate_record_failure_options( + self, failure_policy: Any, step_stall_budget_ms: Any, + record_sink: Optional[Any], + ) -> None: + if failure_policy not in RECORD_FAILURE_POLICIES: + raise ValueError( + f"failure_policy must be one of {RECORD_FAILURE_POLICIES}, " + f"got {failure_policy!r}") + if step_stall_budget_ms is not None and ( + type(step_stall_budget_ms) is not int or step_stall_budget_ms <= 0 + ): + raise ValueError( + "step_stall_budget_ms must be a positive int, or None to wait " + "without bound") + if step_stall_budget_ms is not None: + # Past the budget the forward still waits for the envelope the + # sink is admitting; only a bounded admission bounds that wait. + # The native ring refuses the same, later; this refuses before + # the live ring is torn down. + self._refuse_an_unbounded_sink_admission(record_sink) + + def _refuse_an_unbounded_sink_admission( + self, record_sink: Optional[Any] + ) -> None: + needs = ("step_stall_budget_ms needs a record sink whose admission " + "is bounded, because past the budget the forward still " + "waits for the envelope the sink is admitting") + if record_sink is not None: + if getattr(record_sink, "admission_bound_s", None) is None: + raise ValueError( + f"{needs}; this {type(record_sink).__name__} has no " + "admission bound (admission_bound_s is None). Use a " + "NativePackSink with overload='drop_newest', or 'block' " + "with an admission_timeout_s") + return + sink_config = self._capture_sink_config + if (getattr(self, "_storage_backend", "auto") == "persistent" + and sink_config is not None): + if (sink_config.overload == "block" + and sink_config.admission_timeout_s is None): + raise ValueError( + f"{needs}; set capture_sink_config.admission_timeout_s, " + "or overload='drop_newest'") + return + raise ValueError( + f"{needs}; the ClickHouse host path (record_sink=None) has " + "none: it enqueues row by row under the host engine's ingress " + "policy, with no bound per envelope. Capture through a " + "NativePackSink (storage_backend='persistent', or record_sink=) " + "to use a budget") + def _start_capture_storage(self, *, sweep_spool: bool) -> Optional[Any]: config = self._capture_storage_config if config is None or self._storage_backend != "persistent": @@ -381,6 +508,9 @@ def _attach_record_runtime( record_format: "RecordFormat[MetadataT]", record_schema: Any, record_sink: Optional[Any], + *, + failure_policy: str, + step_stall_budget_ms: int, ) -> "RecordRuntime[MetadataT]": from .records import RecordRuntime @@ -434,6 +564,8 @@ def _attach_record_runtime( record_engine = _native_engine.RingEngine.create_record( ring_config, sink_or_host, + failure_policy=failure_policy, + step_stall_budget_ms=step_stall_budget_ms, ) record_engine.init() record_engine.start() @@ -442,6 +574,7 @@ def _attach_record_runtime( self._ring_engine = record_engine self._ring_transport = record_transport self._record_mode = True + self._record_sink = record_sink _rt.activate(record_transport) return runtime except BaseException: @@ -459,6 +592,7 @@ def _attach_record_runtime( self._ring_transport = None self._ring_engine = None self._record_mode = False + self._record_sink = None # Drop every Python owner of the new native engine. Its native # lease joins the record worker before releasing the sink. runtime = None @@ -484,6 +618,86 @@ def flush_and_wait(self, timeout_s: float = 600.0) -> None: if storage is not None: storage.flush(max(0.0, deadline - time.monotonic())) + def capture_status(self) -> dict[str, Any]: + """The record runtime's capture state, for monitoring and RPCs. + + ``capture_active`` is False once a failure latched (under + ``"raise"``, also once a spent stall budget is waiting to be raised + at the next step): under ``"disable_capture"`` the forward keeps + running and ``failure`` says why capture stopped. ``skipped_steps`` + counts steps whose rest was skipped because the stall budget ran + out (always equal to ``stall_budget_exhaustions``); under + ``"disable_capture"`` capture stays active through them. A skip also + discards every record still queued for the sink, earlier steps' + too, so ``steps_with_discards`` -- the distinct steps that lost a + record to a skip -- can exceed ``skipped_steps``. ``discarded_*`` + count the records dropped by a skip or after a latch. ``reserve_wait_s`` and ``max_step_wait_s`` are the time + record reservations waited for the sink (in total, and in the + worst step). ``sink`` and ``storage`` are the native pack sink's + and storage service's snapshots, when the engine holds them. + Without a record runtime every field is empty. + """ + status: dict[str, Any] = { + "record_mode": False, "capture_active": False, + "failure_policy": None, "failure": None, + "discarded_descriptors": 0, "discarded_payloads": 0, + "step_stall_budget_ms": None, "stall_budget_exhaustions": 0, + "skipped_steps": 0, "steps_with_discards": 0, + "reserve_wait_s": 0.0, "max_step_wait_s": 0.0, + "sink": None, "storage": None, + } + ring_engine = getattr(self, "_ring_engine", None) + if not self._record_mode or ring_engine is None: + return status + native = dict(ring_engine.capture_status()) + status.update( + record_mode=True, + capture_active=not native["failed"], + failure_policy=native["failure_policy"], + failure=native["failure"] or None, + discarded_descriptors=int(native["discarded_descriptors"]), + discarded_payloads=int(native["discarded_payloads"]), + step_stall_budget_ms=int(native["step_stall_budget_ms"]) or None, + stall_budget_exhaustions=int(native["stall_budget_exhaustions"]), + skipped_steps=int(native["skipped_steps"]), + steps_with_discards=int(native["steps_with_discards"]), + reserve_wait_s=int(native["reserve_wait_ns"]) / 1e9, + max_step_wait_s=int(native["max_step_wait_ns"]) / 1e9, + ) + snapshot = getattr(self._record_sink, "snapshot", None) + if callable(snapshot): + status["sink"] = dict(snapshot()) + if self._capture_storage is not None: + status["storage"] = dict(self._capture_storage.snapshot()) + return status + + def _report_capture_failure(self) -> None: + """Log, once, why capture stopped or skipped steps, as a record + ring is retired.""" + try: + status = self.capture_status() + except Exception as exc: + _LOG.warning("capture status unavailable at close: %s", exc) + return + if not status["record_mode"]: + return + if status["skipped_steps"]: + _LOG.warning( + "record capture skipped the rest of %d steps whose stall " + "budget (%s ms) ran out, discarding with them the records " + "still queued for the sink: %d steps lost records; %d " + "descriptors and %d payloads discarded", + status["skipped_steps"], status["step_stall_budget_ms"], + status["steps_with_discards"], + status["discarded_descriptors"], status["discarded_payloads"]) + if status["capture_active"]: + return + _LOG.warning( + "record capture stopped before close (failure_policy=%s): %s; " + "%d descriptors and %d payloads discarded after it", + status["failure_policy"], status["failure"], + status["discarded_descriptors"], status["discarded_payloads"]) + @staticmethod def _make_default_ring_config( *, @@ -549,6 +763,8 @@ def enable_ring_transport( + self._capture_storage_config.close_flush_timeout_s) if storage is not None: self._seal_capture_sink(drain_deadline) + if old_record_mode: + self._report_capture_failure() try: ring_engine = getattr(self, "_ring_engine", None) if ring_engine is not None: @@ -565,6 +781,7 @@ def enable_ring_transport( self._ring_transport = None self._ring_engine = None self._record_mode = False + self._record_sink = None if storage is not None: self._retire_capture_storage(storage, drain_deadline) @@ -656,6 +873,8 @@ def close(self) -> None: pass if record_mode and storage is not None: self._seal_capture_sink(drain_deadline) + if record_mode: + self._report_capture_failure() try: ring_engine = getattr(self, "_ring_engine", None) if ring_engine is not None: @@ -675,6 +894,7 @@ def close(self) -> None: self._ring_transport = None self._ring_engine = None self._record_mode = False + self._record_sink = None if storage is not None: self._retire_capture_storage(storage, drain_deadline) @@ -692,4 +912,4 @@ def close(self) -> None: # Backend loader -__all__ = ["MonitoringEngine", "RingCapacities"] +__all__ = ["MonitoringEngine", "RECORD_FAILURE_POLICIES", "RingCapacities"] diff --git a/src/dmi/records.py b/src/dmi/records.py index 27116ba7d..748aaca13 100644 --- a/src/dmi/records.py +++ b/src/dmi/records.py @@ -309,6 +309,26 @@ def bind_hook( self._device_gated_output_ids.update(output_ids) self._bound_hooks.add(id(hook)) + def begin_step(self) -> None: + """Mark the start of one model step. + + The record reservations after this call share one fresh + ``step_stall_budget_ms`` (see ``create_record_runtime``), and a + step skipped because the previous budget ran out ends here. Call it + once per model step, before the step's first ``emit_output`` or + ``prepare_replay``: with a budget, a reservation before the first + call is refused. The runtime cannot find step boundaries itself -- + eager hooks reserve once per output, output ids repeat within a + step, and a CUDA-graph step replays one plan or, with piecewise + graphs, several. + + Under ``failure_policy="raise"`` it raises a failure latched during + the previous step, including a spent stall budget, so the error + surfaces here, outside the forward. + """ + + self._transport.begin_record_step() + def emit_output( self, entry: ProducerPlanEntry, diff --git a/src/dmi/storage/capture/native_sink.py b/src/dmi/storage/capture/native_sink.py index ade74833a..bc34b942c 100644 --- a/src/dmi/storage/capture/native_sink.py +++ b/src/dmi/storage/capture/native_sink.py @@ -18,42 +18,14 @@ from __future__ import annotations -from dataclasses import dataclass from typing import Any - -LAYOUT_NAME = "capture_pack_reference_v1" +# The value type is live code, so the public config does not depend on this +# backup package; it is re-exported here for this module's callers. +from ..native_capture import NativeSinkConfig -@dataclass(frozen=True, slots=True) -class NativeSinkConfig: - """Pack-sink bounds for the native writer. Field-by-field the same - contract as the pipeline config the reference sink takes.""" - - spool_root: str - spool_max_bytes: int = 1 << 40 - num_workers: int = 1 - max_queue_records: int = 256 - max_queue_bytes: int = 16 * 1024 * 1024 - max_pack_bytes: int = 128 * 1024 * 1024 - max_pack_records: int = 10_000 - max_linger_ns: int = 1_000_000_000 - - def __post_init__(self) -> None: - if not self.spool_root: - raise ValueError("spool_root is required") - for name in ( - "spool_max_bytes", - "num_workers", - "max_queue_records", - "max_queue_bytes", - "max_pack_bytes", - "max_pack_records", - "max_linger_ns", - ): - value = getattr(self, name) - if type(value) is not int or value <= 0: - raise ValueError(f"{name} must be positive") +LAYOUT_NAME = "capture_pack_reference_v1" def _load_native_sink_extension() -> Any: @@ -128,6 +100,8 @@ def __init__(self, config: NativeSinkConfig) -> None: max_pack_records=config.max_pack_records, max_linger_ns=config.max_linger_ns, spool_max_bytes=config.spool_max_bytes, + overload=config.overload, + admission_timeout_s=config.admission_timeout_s, ) # Engine ownership is taken by create_record_runtime; holding no # lease here keeps the handle closable without an engine. diff --git a/src/dmi/storage/native_capture.py b/src/dmi/storage/native_capture.py index 209ac2f56..d8e8022e1 100644 --- a/src/dmi/storage/native_capture.py +++ b/src/dmi/storage/native_capture.py @@ -31,7 +31,7 @@ import re import socket import uuid -from dataclasses import dataclass, field +from dataclasses import MISSING, dataclass, field, fields from typing import Any, Optional, Sequence @@ -61,6 +61,7 @@ def _load_native_store_extension() -> Any: # else would change what the URL the native client builds means. _BARE_HOST = re.compile(r"[A-Za-z0-9._-]+|\[[0-9A-Fa-f:.]+\]") + def _text(name: str, value: Any) -> None: if type(value) is not str: raise TypeError(f"{name} must be str") @@ -85,6 +86,93 @@ def _ns(seconds: float) -> int: _FENCE_MARGIN_NS = 100_000_000 +# The sink's admission policies (native/csrc/sink/pack_sink.h Overload). +SINK_OVERLOAD_POLICIES = ("block", "drop_newest") + + +@dataclass(frozen=True, slots=True) +class NativeSinkConfig: + """Bounds and admission policy for the native pack sink. + + The sink runs on the ring's record worker: every record the forward + captures is admitted here, into a bounded queue ahead of pack assembly. + ``overload`` decides what a full queue does. ``"block"`` waits for room + for up to ``admission_timeout_s`` (``None`` waits without bound) and + then refuses the record as timed out; the rows of one ring record (one + envelope) share that one deadline. ``"drop_newest"`` refuses it at + once, and ``admission_timeout_s`` is not used. Either refusal latches + the record runtime (see ``create_record_runtime``'s failure policy), + and so does a record lost after admission. A ``step_stall_budget_ms`` + needs a bounded admission, so it is refused with block and no timeout. + + The default, block with 2 s, absorbs a burst larger than the queue at + the cost of stalling the record worker, and so the ring, while the sink + catches up. The C++ ``SinkConfig`` keeps drop_newest with no timeout, + the reference pipeline's default; this config is what the ring-fed sink + is built from. + + A record larger than ``max_queue_bytes`` or ``max_pack_bytes`` can + never be admitted; ``validate_capture_bounds`` refuses such a bound at + attach, before any forward runs. + """ + + spool_root: str + spool_max_bytes: int = 1 << 40 + num_workers: int = 1 + max_queue_records: int = 256 + max_queue_bytes: int = 16 * 1024 * 1024 + max_pack_bytes: int = 128 * 1024 * 1024 + max_pack_records: int = 10_000 + max_linger_ns: int = 1_000_000_000 + overload: str = "block" + admission_timeout_s: Optional[float] = 2.0 + + def __post_init__(self) -> None: + if not self.spool_root: + raise ValueError("spool_root is required") + for name in ( + "spool_max_bytes", + "num_workers", + "max_queue_records", + "max_queue_bytes", + "max_pack_bytes", + "max_pack_records", + "max_linger_ns", + ): + value = getattr(self, name) + if type(value) is not int or value <= 0: + raise ValueError(f"{name} must be positive") + if self.overload not in SINK_OVERLOAD_POLICIES: + raise ValueError( + f"overload must be one of {SINK_OVERLOAD_POLICIES}, " + f"got {self.overload!r}") + if self.admission_timeout_s is not None: + _positive("admission_timeout_s", self.admission_timeout_s, float) + + +def _native_sink_config_setstate(self: NativeSinkConfig, state: Any) -> None: + # A frozen slots dataclass pickles its field values as a list, and the + # generated __setstate__ zips them with the current fields, leaving any + # the pickle predates unset. A config pickled before the admission + # fields existed (the old class, at the path that now re-exports this + # one) gets their defaults instead. + names = fields(NativeSinkConfig) + if len(state) > len(names): + raise TypeError( + f"NativeSinkConfig state has {len(state)} values, expected at " + f"most {len(names)}") + for item, value in zip(names, state): + object.__setattr__(self, item.name, value) + for item in names[len(state):]: + if item.default is MISSING: + raise TypeError( + f"NativeSinkConfig state is missing {item.name!r}") + object.__setattr__(self, item.name, item.default) + + +NativeSinkConfig.__setstate__ = _native_sink_config_setstate # type: ignore[method-assign] + + def _positive(name: str, value: Any, kind: type) -> None: if type(value) is not kind and not (kind is float and type(value) is int): raise TypeError(f"{name} must be {kind.__name__}") @@ -164,6 +252,10 @@ class NativeCaptureStorageConfig: # engine.close()'s total budget for draining capture: sealing the sink's # open pack, then getting every staged pack into the catalog. close_flush_timeout_s: float = 60.0 + # Bytes of packs the uploader holds in flight at once. A staged pack + # larger than this is never uploaded, so the sink's max_pack_bytes must + # not exceed it (validate_capture_bounds). The native default. + uploader_max_in_flight_bytes: int = 256 * 1024 * 1024 # The publisher lease. A crashed process holds the catalog for up to # lease_ttl_s; a ClickHouse error of unknown outcome sets the lease aside @@ -228,6 +320,9 @@ def __post_init__(self) -> None: self.clickhouse_request_timeout_s, float) self._validate_clickhouse_connection() _positive("close_flush_timeout_s", self.close_flush_timeout_s, float) + if (type(self.uploader_max_in_flight_bytes) is not int + or self.uploader_max_in_flight_bytes <= 0): + raise ValueError("uploader_max_in_flight_bytes must be positive") if type(self.reconcile_interval_s) not in (int, float): raise TypeError("reconcile_interval_s must be float") if not math.isfinite(self.reconcile_interval_s): @@ -352,6 +447,7 @@ def _native_dict(self) -> dict[str, Any]: "clickhouse_request_timeout_s": float(self.clickhouse_request_timeout_s), "database": self.database, "table_prefix": self.table_prefix, + "uploader_max_in_flight_bytes": self.uploader_max_in_flight_bytes, } def _native_reader_dict(self) -> dict[str, Any]: @@ -363,6 +459,62 @@ def _native_reader_dict(self) -> dict[str, Any]: return native +# What an empty pack needs besides one record's payload: the 64-byte header +# and trailer, up to 63 bytes of payload alignment, and that record's footer +# row (its metadata as JSON, well under a kilobyte for real identifiers). +# Generous on purpose: undershooting it lets a record through that the pack +# worker then counts as oversized and drops. +PACK_FRAMING_RESERVE_BYTES = 64 * 1024 + + +def validate_capture_bounds( + sink_config: NativeSinkConfig, + max_record_bytes: int, + *, + storage_config: Optional[NativeCaptureStorageConfig] = None, +) -> None: + """Refuse capture bounds under which a record or pack cannot be stored. + + ``max_record_bytes`` is the largest single record (one captured row's + payload) the caller will emit. Checked here, at attach, these refusals + never reach the forward: + + - ``max_queue_bytes`` below it: the sink refuses the record outright. + - ``max_pack_bytes`` below it plus ``PACK_FRAMING_RESERVE_BYTES``: the + sink admits the record and the pack worker then drops it as oversized. + - ``max_pack_bytes`` above ``uploader_max_in_flight_bytes``: a full pack + is staged and never uploaded. + + Raises ``ConfigurationError`` naming the bound to raise. + """ + if not isinstance(sink_config, NativeSinkConfig): + raise TypeError("sink_config must be a NativeSinkConfig") + if type(max_record_bytes) is not int or max_record_bytes <= 0: + raise ValueError("max_record_bytes must be a positive int") + from ..configuration.errors import ConfigurationError + + if sink_config.max_queue_bytes < max_record_bytes: + raise ConfigurationError( + f"a {max_record_bytes}-byte record exceeds the sink's " + f"max_queue_bytes ({sink_config.max_queue_bytes}); raise " + "max_queue_bytes to at least the largest record") + needed = max_record_bytes + PACK_FRAMING_RESERVE_BYTES + if sink_config.max_pack_bytes < needed: + raise ConfigurationError( + f"a {max_record_bytes}-byte record does not fit an empty pack of " + f"max_pack_bytes ({sink_config.max_pack_bytes}); raise " + f"max_pack_bytes to at least {needed} (the record plus " + f"{PACK_FRAMING_RESERVE_BYTES} bytes of pack framing)") + if (storage_config is not None and sink_config.max_pack_bytes + > storage_config.uploader_max_in_flight_bytes): + raise ConfigurationError( + f"the sink's max_pack_bytes ({sink_config.max_pack_bytes}) " + "exceeds the storage service's uploader_max_in_flight_bytes " + f"({storage_config.uploader_max_in_flight_bytes}), so a full " + "pack would never be uploaded; lower max_pack_bytes or raise " + "uploader_max_in_flight_bytes") + + class NativeCaptureStorage: """The in-process storage service: spool -> object store -> catalog.""" @@ -615,10 +767,14 @@ def read( __all__ = [ + "PACK_FRAMING_RESERVE_BYTES", + "SINK_OVERLOAD_POLICIES", + "NativeSinkConfig", "NativeCapture", "NativeCapturePage", "NativeCaptureReader", "NativeCaptureSelection", "NativeCaptureStorage", "NativeCaptureStorageConfig", + "validate_capture_bounds", ] diff --git a/src/dmi/transport/ring.py b/src/dmi/transport/ring.py index 27d61592f..7a57ccc78 100644 --- a/src/dmi/transport/ring.py +++ b/src/dmi/transport/ring.py @@ -420,6 +420,12 @@ def reserve_record(self, reservation_items: Any) -> int: return int(self._ring_engine.reserve_record(tuple(reservation_items))) + def begin_record_step(self) -> None: + """Start a record step: a fresh stall budget for the next + reservations (see RecordRuntime.begin_step).""" + + self._ring_engine.begin_record_step() + def push_record_descriptors(self, descriptors: Any) -> None: """Publish descriptors in the exact order of their producer tasks.""" diff --git a/tests/native/ring/test_record_consumer.cpp b/tests/native/ring/test_record_consumer.cpp index 5ab375efb..9fb2fffd9 100644 --- a/tests/native/ring/test_record_consumer.cpp +++ b/tests/native/ring/test_record_consumer.cpp @@ -1,4 +1,5 @@ #include "ring/record_consumer.h" +#include "ring/record_stall_budget.h" #include @@ -6,6 +7,7 @@ #include #include #include +#include #include #include #include @@ -192,6 +194,338 @@ static void test_submit_failure_precedes_durable_idle() { EXPECT(timeouts.load(std::memory_order_relaxed) == 0); } +template +static bool throws_runtime_error(Fn&& call) { + try { + call(); + } catch (const std::runtime_error&) { + return true; + } + return false; +} + +static void test_raise_policy_is_the_default_and_raises_at_the_producer() { + std::printf("[ TEST ] raise_at_producer is the default and raises on push\n"); + ring::RecordConsumer consumer(std::make_shared()); + EXPECT(consumer.snapshot().policy == + ring::RecordFailurePolicy::kRaiseAtProducer); + consumer.push_descriptor(descriptor("raise", "record")); + EXPECT(throws_runtime_error( + [&] { consumer.consume_payload(byte_payload({1})); })); + + EXPECT(throws_runtime_error( + [&] { consumer.push_descriptor(descriptor("raise", "later")); })); + const ring::RecordConsumerSnapshot snapshot = consumer.snapshot(); + EXPECT(snapshot.failed); + EXPECT(snapshot.failure.find("injected sink failure") != std::string::npos); + EXPECT(snapshot.discarded_descriptors == 0); + EXPECT(snapshot.discarded_payloads == 0); +} + +static void test_disable_capture_discards_after_a_latch_and_still_fails_flush() { + std::printf("[ TEST ] disable_capture discards after a latch; flush still fails\n"); + ring::RecordConsumer consumer( + std::make_shared(), + ring::RecordFailurePolicy::kDisableCapture); + // Three descriptors queued ahead of their payloads, as the forward + // publishes them before the drain delivers. + consumer.push_descriptors({ + descriptor("disable", "first"), + descriptor("disable", "second"), + descriptor("disable", "third"), + }); + + // The sink refuses the first: the consumer latches without throwing at + // the p2p worker, and the two descriptors still queued can never be + // stored, so they are dropped with it. + EXPECT(!throws_runtime_error( + [&] { consumer.consume_payload(byte_payload({1})); })); + // The forward keeps publishing and the drain keeps delivering: neither + // raises, and every one is counted. + EXPECT(!throws_runtime_error( + [&] { consumer.push_descriptor(descriptor("disable", "after")); })); + EXPECT(!throws_runtime_error([&] { + consumer.push_descriptors({descriptor("disable", "after-1"), + descriptor("disable", "after-2")}); + })); + for (int index = 0; index < 4; ++index) { + EXPECT(!throws_runtime_error( + [&] { consumer.consume_payload(byte_payload({2})); })); + } + + const ring::RecordConsumerSnapshot snapshot = consumer.snapshot(); + EXPECT(snapshot.policy == ring::RecordFailurePolicy::kDisableCapture); + EXPECT(snapshot.failed); + EXPECT(snapshot.failure.find("injected sink failure") != std::string::npos); + EXPECT(snapshot.discarded_descriptors == 5); + // The refused payload plus the four delivered after the latch. + EXPECT(snapshot.discarded_payloads == 5); + EXPECT(consumer.pending_descriptors() == 0); + + // The failure still surfaces at every checked completion. + EXPECT(throws_runtime_error([&] { consumer.rethrow_if_failed(); })); + EXPECT(throws_runtime_error( + [&] { consumer.wait_until_idle(std::chrono::milliseconds(10)); })); + EXPECT(throws_runtime_error([&] { consumer.finish(); })); +} + +static void test_disable_capture_latches_association_and_worker_failures() { + std::printf("[ TEST ] disable_capture latches association and worker failures\n"); + auto sink = std::make_shared(); + ring::RecordConsumer orphan( + sink, ring::RecordFailurePolicy::kDisableCapture); + EXPECT(!throws_runtime_error( + [&] { orphan.consume_payload(byte_payload({1})); })); + EXPECT(orphan.snapshot().failed); + EXPECT(orphan.snapshot().failure.find("without an encoded descriptor") != + std::string::npos); + EXPECT(orphan.snapshot().discarded_payloads == 1); + EXPECT(throws_runtime_error([&] { orphan.finish(); })); + + ring::RecordConsumer worker( + sink, ring::RecordFailurePolicy::kDisableCapture); + worker.push_descriptor(descriptor("worker", "queued")); + worker.record_failure(std::make_exception_ptr( + std::runtime_error("injected worker failure"))); + worker.push_descriptor(descriptor("worker", "after")); + EXPECT(!throws_runtime_error( + [&] { worker.consume_payload(byte_payload({1})); })); + const ring::RecordConsumerSnapshot snapshot = worker.snapshot(); + EXPECT(snapshot.failure == "injected worker failure"); + EXPECT(snapshot.discarded_descriptors == 2); + EXPECT(snapshot.discarded_payloads == 1); + EXPECT(sink->submitted.empty()); + EXPECT(throws_runtime_error([&] { worker.rethrow_if_failed(); })); +} + +// A step whose stall budget ran out: the records still queued for the sink +// and the rest of the step's records are skipped, and the next step is +// stored again. Payloads keep pairing with their own descriptors across +// the window, under either policy, and nothing latches. +static void test_a_discard_window_skips_one_steps_records( + ring::RecordFailurePolicy policy) { + std::printf("[ TEST ] a discard window skips one step's records (%s)\n", + policy == ring::RecordFailurePolicy::kDisableCapture + ? "disable_capture" : "raise"); + auto sink = std::make_shared(); + ring::RecordConsumer consumer(sink, policy); + + consumer.push_descriptors({descriptor("step", "queued-1"), + descriptor("step", "queued-2")}); + consumer.begin_discard_window(); + EXPECT(!throws_runtime_error([&] { + consumer.push_descriptors({descriptor("step", "rest-1"), + descriptor("step", "rest-2")}); + })); + consumer.end_discard_window(); + consumer.push_descriptor(descriptor("step", "next-step")); + + // Nothing is idle until every skipped payload has arrived. + EXPECT(!consumer.wait_until_idle(std::chrono::milliseconds(5))); + EXPECT(throws_runtime_error([&] { consumer.finish(); })); + for (float value = 1; value <= 5; ++value) { + EXPECT(!throws_runtime_error( + [&] { consumer.consume_payload(byte_payload({value})); })); + } + + EXPECT(sink->submitted.size() == 1); + if (sink->submitted.size() == 1) { + EXPECT(std::get( + sink->submitted[0].descriptor.rows[0].cells[0]) == + "next-step"); + EXPECT(at::equal(sink->submitted[0].payload.view(at::kFloat), + at::tensor({5.f}))); + } + const ring::RecordConsumerSnapshot snapshot = consumer.snapshot(); + EXPECT(!snapshot.failed); + EXPECT(snapshot.discarded_descriptors == 4); + EXPECT(snapshot.discarded_payloads == 4); + EXPECT(consumer.wait_until_idle(std::chrono::milliseconds(5))); + EXPECT(!throws_runtime_error([&] { consumer.finish(); })); + + // An empty window skips nothing. + consumer.begin_discard_window(); + consumer.end_discard_window(); + consumer.push_descriptor(descriptor("step", "after-empty")); + consumer.consume_payload(byte_payload({6})); + EXPECT(sink->submitted.size() == 2); +} + +// A window that dropped every queued descriptor leaves the queue empty while +// the drain still owes the dropped descriptors' payloads. A flush must wait +// for them: idle too early, and the next descriptor pushed after the flush +// would be paired with a skipped step's payload. +static void test_owed_payloads_hold_the_flush( + ring::RecordFailurePolicy policy) { + std::printf("[ TEST ] payloads owed to a discard window hold the flush (%s)\n", + policy == ring::RecordFailurePolicy::kDisableCapture + ? "disable_capture" : "raise"); + auto sink = std::make_shared(); + ring::RecordConsumer consumer(sink, policy); + + consumer.push_descriptors({descriptor("step", "queued-1"), + descriptor("step", "queued-2")}); + consumer.begin_discard_window(); + consumer.end_discard_window(); + EXPECT(consumer.pending_descriptors() == 0); + + // Only owed payloads remain: not idle, and finish() names them. + EXPECT(!consumer.wait_until_idle(std::chrono::milliseconds(5))); + std::string leftover; + try { + consumer.finish(); + } catch (const std::runtime_error& error) { + leftover = error.what(); + } + EXPECT(leftover.find("owed") != std::string::npos); + + consumer.consume_payload(byte_payload({1})); + EXPECT(!consumer.wait_until_idle(std::chrono::milliseconds(5))); + EXPECT(throws_runtime_error([&] { consumer.finish(); })); + + // A waiter already blocked wakes when the last owed payload arrives. + std::atomic idle{false}; + std::thread waiter([&] { + idle.store(consumer.wait_until_idle(std::chrono::seconds(5)), + std::memory_order_release); + }); + std::this_thread::sleep_for(std::chrono::milliseconds(20)); + consumer.consume_payload(byte_payload({2})); + waiter.join(); + EXPECT(idle.load(std::memory_order_acquire)); + EXPECT(!throws_runtime_error([&] { consumer.finish(); })); + EXPECT(sink->submitted.empty()); + EXPECT(!consumer.snapshot().failed); +} + +// A window drops every record still queued, including earlier steps' that +// the sink has not taken yet, so the steps that lost records can outnumber +// the windows. steps_with_discards counts each such step once; a window +// that dropped nothing counts no step. +static void test_steps_with_discards_counts_every_step_a_window_reaches( + ring::RecordFailurePolicy policy) { + std::printf("[ TEST ] steps_with_discards counts each step a window reaches (%s)\n", + policy == ring::RecordFailurePolicy::kDisableCapture + ? "disable_capture" : "raise"); + auto sink = std::make_shared(); + ring::RecordConsumer consumer(sink, policy); + + consumer.begin_step(); // 1 + consumer.push_descriptor(descriptor("step", "s1")); + consumer.begin_step(); // 2 + consumer.push_descriptors({descriptor("step", "s2a"), + descriptor("step", "s2b")}); + consumer.begin_step(); // 3: its budget runs out + consumer.push_descriptor(descriptor("step", "s3a")); + consumer.begin_discard_window(); + EXPECT(consumer.snapshot().steps_with_discards == 3); + consumer.push_descriptor(descriptor("step", "s3b")); + consumer.push_descriptors({}); + EXPECT(consumer.snapshot().steps_with_discards == 3); + consumer.end_discard_window(); + + consumer.begin_step(); // 4 + consumer.push_descriptor(descriptor("step", "s4")); + consumer.begin_step(); // 5: runs out too + consumer.begin_discard_window(); + EXPECT(consumer.snapshot().steps_with_discards == 4); + consumer.push_descriptors({descriptor("step", "s5")}); + EXPECT(consumer.snapshot().steps_with_discards == 5); + consumer.end_discard_window(); + + consumer.begin_step(); // 6: a window that finds nothing to drop + consumer.begin_discard_window(); + consumer.end_discard_window(); + consumer.begin_step(); // 7 + consumer.push_descriptor(descriptor("step", "s7")); + + for (float value = 1; value <= 8; ++value) { + consumer.consume_payload(byte_payload({value})); + } + const ring::RecordConsumerSnapshot snapshot = consumer.snapshot(); + EXPECT(!snapshot.failed); + EXPECT(snapshot.steps_with_discards == 5); + EXPECT(snapshot.discarded_descriptors == 7); + EXPECT(snapshot.discarded_payloads == 7); + EXPECT(sink->submitted.size() == 1); + if (sink->submitted.size() == 1) { + EXPECT(std::get( + sink->submitted[0].descriptor.rows[0].cells[0]) == "s7"); + EXPECT(at::equal(sink->submitted[0].payload.view(at::kFloat), + at::tensor({8.f}))); + } + EXPECT(consumer.wait_until_idle(std::chrono::milliseconds(5))); + EXPECT(!throws_runtime_error([&] { consumer.finish(); })); +} + +// The worker asks before copying a payload out of the pinned staging +// whether anyone will store it. A yes accounts the discard exactly as +// consume_payload would, so pairing and counters stay the same. +static void test_discard_next_payload_if_unwanted_matches_consume_payload() { + std::printf("[ TEST ] discard_next_payload_if_unwanted matches consume_payload\n"); + for (const auto policy : {ring::RecordFailurePolicy::kRaiseAtProducer, + ring::RecordFailurePolicy::kDisableCapture}) { + auto sink = std::make_shared(); + ring::RecordConsumer consumer(sink, policy); + // Nothing owed: the payload is wanted and nothing changes. + EXPECT(!consumer.discard_next_payload_if_unwanted()); + consumer.push_descriptors({descriptor("step", "skipped-1"), + descriptor("step", "skipped-2")}); + consumer.begin_discard_window(); + consumer.end_discard_window(); + consumer.push_descriptor(descriptor("step", "kept")); + EXPECT(consumer.discard_next_payload_if_unwanted()); + consumer.consume_payload(byte_payload({2})); + EXPECT(!consumer.discard_next_payload_if_unwanted()); + consumer.consume_payload(byte_payload({3})); + EXPECT(sink->submitted.size() == 1); + if (sink->submitted.size() == 1) { + EXPECT(at::equal(sink->submitted[0].payload.view(at::kFloat), + at::tensor({3.f}))); + } + EXPECT(consumer.snapshot().discarded_payloads == 2); + EXPECT(consumer.wait_until_idle(std::chrono::milliseconds(5))); + + consumer.record_failure(std::make_exception_ptr( + std::runtime_error("injected latch"))); + if (policy == ring::RecordFailurePolicy::kDisableCapture) { + // Capture is off: every later payload is unwanted. + EXPECT(consumer.discard_next_payload_if_unwanted()); + EXPECT(consumer.snapshot().discarded_payloads == 3); + } else { + // Left to consume_payload, which raises the latch. + EXPECT(!consumer.discard_next_payload_if_unwanted()); + EXPECT(consumer.snapshot().discarded_payloads == 2); + EXPECT(throws_runtime_error( + [&] { consumer.consume_payload(byte_payload({4})); })); + } + } +} + +// Past a spent budget the reservation may wait the sink's admission bound +// plus a grace that covers copying out the whole ring at 1 GB/s. +static void test_record_drain_grace_scales_with_the_ring() { + std::printf("[ TEST ] the post-budget drain grace scales with the ring\n"); + using std::chrono::nanoseconds; + using std::chrono::seconds; + EXPECT(ring::record_drain_grace(0, 0) == seconds(2)); + EXPECT(ring::record_drain_grace(4096, 4096) == + seconds(2) + nanoseconds(8192)); + // The engine's default 4 GiB ring and 4 GiB staging: about 10.6 s. + constexpr uint64_t kGiB = 1ull << 30; + EXPECT(ring::record_drain_grace(4 * kGiB, 4 * kGiB) == + seconds(2) + nanoseconds(8 * kGiB)); + EXPECT(ring::record_drain_grace(4 * kGiB, 0) > + ring::record_drain_grace(1 * kGiB, 0)); + // Saturates instead of wrapping, and a deadline built on it is sane. + const nanoseconds huge = ring::record_drain_grace( + std::numeric_limits::max(), + std::numeric_limits::max()); + EXPECT(huge > seconds(2)); + EXPECT(std::chrono::steady_clock::now() + huge > + std::chrono::steady_clock::now()); +} + int main() { setbuf(stdout, nullptr); std::printf("test_record_consumer\n"); @@ -199,6 +533,23 @@ int main() { test_zero_row_descriptor_consumes_without_sink_submission(); test_exact_association_failures(); test_submit_failure_precedes_durable_idle(); + test_raise_policy_is_the_default_and_raises_at_the_producer(); + test_disable_capture_discards_after_a_latch_and_still_fails_flush(); + test_disable_capture_latches_association_and_worker_failures(); + test_a_discard_window_skips_one_steps_records( + ring::RecordFailurePolicy::kRaiseAtProducer); + test_a_discard_window_skips_one_steps_records( + ring::RecordFailurePolicy::kDisableCapture); + test_owed_payloads_hold_the_flush( + ring::RecordFailurePolicy::kRaiseAtProducer); + test_owed_payloads_hold_the_flush( + ring::RecordFailurePolicy::kDisableCapture); + test_steps_with_discards_counts_every_step_a_window_reaches( + ring::RecordFailurePolicy::kRaiseAtProducer); + test_steps_with_discards_counts_every_step_a_window_reaches( + ring::RecordFailurePolicy::kDisableCapture); + test_discard_next_payload_if_unwanted_matches_consume_payload(); + test_record_drain_grace_scales_with_the_ring(); std::printf("Results: %d passed, %d failed\n", g_pass, g_fail); return g_fail == 0 ? 0 : 1; } diff --git a/tests/native/ring/test_ring_engine.cu b/tests/native/ring/test_ring_engine.cu index f3cb1ada2..23284a108 100644 --- a/tests/native/ring/test_ring_engine.cu +++ b/tests/native/ring/test_ring_engine.cu @@ -19,7 +19,9 @@ #include #include #include +#include #include +#include #include #include @@ -776,6 +778,391 @@ static void test_record_flush_reaches_sink_durability_boundary() { engine.stop(); } +// A sink that admits every record, each after `admission` -- a slow disk or +// a full queue with a bounded admission timeout. While it is inside +// submit(), the record worker holds every later payload's staging bytes. +// +// `bound` is the admission bound it declares to the ring (see +// RecordSink::admission_bound): by default the admission it really takes; a +// shorter one makes a sink that breaks its own bound, and nullopt a sink +// that cannot say. +class SlowRecordSink final : public ring::RecordSink { +public: + explicit SlowRecordSink(std::chrono::milliseconds admission) + : admission_(admission), bound_(admission) {} + SlowRecordSink(std::chrono::milliseconds admission, + std::optional bound) + : admission_(admission), bound_(bound) {} + + void submit(ring::RecordEnvelope) override { + std::this_thread::sleep_for(admission_); + submissions.fetch_add(1, std::memory_order_acq_rel); + } + + bool flush_and_wait(Duration) override { return true; } + void rethrow_if_failed() const override {} + std::optional admission_bound() const override { + return bound_; + } + + std::atomic submissions{0}; + +private: + std::chrono::milliseconds admission_; + std::optional bound_; +}; + +static ring::RecordDescriptor stall_descriptor() { + ring::RecordDescriptor descriptor; + descriptor.layout = "stall"; + descriptor.rows = {{std::vector{ + std::string("record"), ring::PayloadSlice{}}}}; + return descriptor; +} + +// Reserve, publish and launch one 1 KiB record, as RecordRuntime.emit_output +// and the hook it returns to do. +static int emit_stall_record(ring_py::RingEnginePy& engine, + const uint8_t* device) { + constexpr uint64_t kBytes = 1024; + const int reservation = engine.reserve_record({{kBytes, false}}); + engine.push_record_descriptors({stall_descriptor()}); + engine.record_no_notify( + reinterpret_cast(device), kBytes, 0, 0, + reinterpret_cast(at::cuda::getCurrentCUDAStream().stream())); + return reservation; +} + +struct StallOutcome { + bool threw{false}; + std::string error; + int reservation{-1}; + std::chrono::steady_clock::duration stall{}; + ring_py::RecordCaptureStatus status; + bool next_step_threw{false}; + std::string next_step_error; + bool later_reserve_threw{false}; + std::chrono::steady_clock::duration later_reserve{}; + bool flush_threw{false}; + std::string flush_error; + ring_py::RecordCaptureStatus flushed_status; + int submissions{0}; +}; + +// Fill a 4 KiB ring twice while the sink sits inside its first admission: +// records 1-4 reach the worker (which blocks in record 1), records 5-8 fill +// the ring again and cannot drain, because 2-4 still hold the staging. The +// ninth reservation has to wait for the sink, and the budget runs out. The +// step after it emits one more record. +static StallOutcome run_stalled_reservation(ring::RecordFailurePolicy policy) { + constexpr auto kBudget = std::chrono::milliseconds(50); + constexpr auto kAdmission = std::chrono::milliseconds(400); + ring_py::RingConfig cfg; + cfg.task_ring_entries = 16; + cfg.payload_ring_bytes = 4096; + cfg.pinned_staging_bytes = 4096; + cfg.drain_poll_timeout_us = 100; + ring_py::RecordRuntimeOptions options; + options.failure_policy = policy; + options.step_stall_budget_ms = kBudget.count(); + auto sink = std::make_shared(kAdmission); + ring_py::RingEnginePy engine(cfg, sink, options); + engine.init(); + engine.start(); + + const std::vector source = pattern(1024, 7); + uint8_t* device = upload(source, at::cuda::getCurrentCUDAStream().stream()); + + StallOutcome outcome; + for (int index = 0; index < 8; ++index) { + engine.begin_record_step(); + emit_stall_record(engine, device); + } + engine.begin_record_step(); + const auto started = std::chrono::steady_clock::now(); + try { + outcome.reservation = emit_stall_record(engine, device); + } catch (const std::runtime_error& error) { + outcome.threw = true; + outcome.error = error.what(); + } + outcome.stall = std::chrono::steady_clock::now() - started; + outcome.status = engine.record_capture_status(); + + try { + engine.begin_record_step(); + } catch (const std::runtime_error& error) { + outcome.next_step_threw = true; + outcome.next_step_error = error.what(); + } + const auto later = std::chrono::steady_clock::now(); + try { + emit_stall_record(engine, device); + } catch (const std::runtime_error&) { + outcome.later_reserve_threw = true; + } + outcome.later_reserve = std::chrono::steady_clock::now() - later; + + try { + engine.flush_records_and_wait(5000); + } catch (const std::runtime_error& error) { + outcome.flush_threw = true; + outcome.flush_error = error.what(); + } + outcome.flushed_status = engine.record_capture_status(); + engine.stop(); + outcome.submissions = sink->submissions.load(std::memory_order_acquire); + CUDA_CHECK(cudaFree(device)); + return outcome; +} + +static void test_raise_policy_skips_the_step_and_raises_at_the_next() { + banner("raise policy: a spent budget skips the step, raises at the next"); + const StallOutcome outcome = + run_stalled_reservation(ring::RecordFailurePolicy::kRaiseAtProducer); + // Nothing raises inside the forward: the step's capture is skipped. + EXPECT(!outcome.threw); + EXPECT(outcome.reservation == ring_py::RingEnginePy::STEP_RING_FLUSHED); + EXPECT(outcome.stall >= std::chrono::milliseconds(50)); + EXPECT(outcome.stall < std::chrono::milliseconds(50 + 400 + 250)); + EXPECT(outcome.status.failed); + EXPECT(outcome.status.failure.find("stall budget") != std::string::npos); + EXPECT(outcome.status.failure.find("from any step") != std::string::npos); + EXPECT(outcome.status.stall_budget_exhaustions == 1); + EXPECT(outcome.status.skipped_steps == 1); + // One skip, but records 2-8 were still queued: steps 2-9 lost records. + EXPECT(outcome.status.steps_with_discards == 8); + EXPECT(outcome.status.step_stall_budget_ms == 50); + // The next step boundary raises it, outside the forward. + EXPECT(outcome.next_step_threw); + EXPECT(outcome.next_step_error.find("stall budget") != std::string::npos); + // Latched from then on: a reservation raises before it waits. + EXPECT(outcome.later_reserve_threw); + EXPECT(outcome.later_reserve < std::chrono::milliseconds(50)); + EXPECT(outcome.flush_threw); + EXPECT(outcome.flush_error.find("stall budget") != std::string::npos); + // Only the record the sink was already admitting reached it. + EXPECT(outcome.submissions == 1); +} + +static void test_disable_capture_skips_the_step_and_resumes() { + banner("disable_capture: a spent budget skips the step; capture resumes"); + const StallOutcome outcome = + run_stalled_reservation(ring::RecordFailurePolicy::kDisableCapture); + EXPECT(!outcome.threw); + EXPECT(outcome.reservation == ring_py::RingEnginePy::STEP_RING_FLUSHED); + // At most the budget plus the one admission already in progress; the + // records queued behind it and the rest of the step are discarded. + EXPECT(outcome.stall >= std::chrono::milliseconds(50)); + EXPECT(outcome.stall < std::chrono::milliseconds(50 + 400 + 250)); + EXPECT(outcome.status.failure_policy == + ring::RecordFailurePolicy::kDisableCapture); + EXPECT(!outcome.status.failed); + EXPECT(outcome.status.stall_budget_exhaustions == 1); + EXPECT(outcome.status.skipped_steps == 1); + // One skip, but records 2-8 were still queued: steps 2-9 lost records. + EXPECT(outcome.status.steps_with_discards == 8); + EXPECT(outcome.status.max_step_wait_ns >= 50'000'000ull); + // The next step captures again, without raising or stalling long. + EXPECT(!outcome.next_step_threw); + EXPECT(!outcome.later_reserve_threw); + EXPECT(outcome.later_reserve < std::chrono::milliseconds(200)); + // A skipped step is not a failure: the flush succeeds. + EXPECT(!outcome.flush_threw); + EXPECT(!outcome.flushed_status.failed); + // Records 2-8, still queued at the exhaustion from steps 2-8, and the + // ninth: one skipped step, eight steps with discards. + EXPECT(outcome.flushed_status.discarded_descriptors == 8); + EXPECT(outcome.flushed_status.discarded_payloads == 8); + EXPECT(outcome.flushed_status.skipped_steps == 1); + EXPECT(outcome.flushed_status.steps_with_discards == 8); + // The record in the sink at the exhaustion, and the next step's. + EXPECT(outcome.submissions == 2); +} + +static ring_py::RingConfig small_record_ring() { + ring_py::RingConfig cfg; + cfg.task_ring_entries = 16; + cfg.payload_ring_bytes = 4096; + cfg.pinned_staging_bytes = 4096; + cfg.drain_poll_timeout_us = 100; + return cfg; +} + +static void test_a_stall_budget_needs_a_bounded_sink_admission() { + banner("a stall budget is refused for a sink with no admission bound"); + for (const auto policy : {ring::RecordFailurePolicy::kRaiseAtProducer, + ring::RecordFailurePolicy::kDisableCapture}) { + ring_py::RecordRuntimeOptions options; + options.failure_policy = policy; + options.step_stall_budget_ms = 50; + bool unbounded_refused = false; + try { + ring_py::RingEnginePy engine( + small_record_ring(), + std::make_shared( + std::chrono::milliseconds(1), std::nullopt), + options); + } catch (const std::invalid_argument& error) { + unbounded_refused = + std::string(error.what()).find("admission bound") != + std::string::npos; + } + EXPECT(unbounded_refused); + bool sinkless_refused = false; + try { + ring_py::RingEnginePy engine( + small_record_ring(), std::shared_ptr(), + options); + } catch (const std::invalid_argument&) { + sinkless_refused = true; + } + EXPECT(sinkless_refused); + // With no budget nothing is promised, so nothing is refused. + options.step_stall_budget_ms = 0; + bool unbudgeted_threw = false; + try { + ring_py::RingEnginePy engine( + small_record_ring(), + std::make_shared( + std::chrono::milliseconds(1), std::nullopt), + options); + } catch (const std::exception&) { + unbudgeted_threw = true; + } + EXPECT(!unbudgeted_threw); + } +} + +static void test_a_stall_budget_needs_begin_record_step() { + banner("a stall budget refuses reservations until a step has begun"); + ring_py::RecordRuntimeOptions options; + options.failure_policy = ring::RecordFailurePolicy::kDisableCapture; + options.step_stall_budget_ms = 50; + auto sink = std::make_shared(std::chrono::milliseconds(0)); + ring_py::RingEnginePy engine(small_record_ring(), sink, options); + engine.init(); + engine.start(); + bool refused = false; + try { + engine.reserve_record({{1024, false}}); + } catch (const std::logic_error& error) { + refused = std::string(error.what()).find("begin") != std::string::npos; + } + EXPECT(refused); + engine.begin_record_step(); + const std::vector source = pattern(1024, 3); + uint8_t* device = upload(source, at::cuda::getCurrentCUDAStream().stream()); + bool threw = false; + try { + emit_stall_record(engine, device); + EXPECT(engine.flush_records_and_wait(5000)); + } catch (const std::exception&) { + threw = true; + } + EXPECT(!threw); + EXPECT(sink->submissions.load(std::memory_order_acquire) == 1); + engine.stop(); + CUDA_CHECK(cudaFree(device)); +} + +static void test_a_sink_past_its_admission_bound_is_a_ring_failure() { + banner("a ring not drained by the admission bound plus the grace raises"); + // Declares 10 ms, takes 2.5 s: past the bound and the drain grace, which + // for this 4 KiB ring and 4 KiB staging is 2 s plus 8 us. + ring_py::RecordRuntimeOptions options; + options.failure_policy = ring::RecordFailurePolicy::kDisableCapture; + options.step_stall_budget_ms = 50; + auto sink = std::make_shared( + std::chrono::milliseconds(2500), std::chrono::milliseconds(10)); + ring_py::RingEnginePy engine(small_record_ring(), sink, options); + engine.init(); + engine.start(); + const std::vector source = pattern(1024, 5); + uint8_t* device = upload(source, at::cuda::getCurrentCUDAStream().stream()); + for (int index = 0; index < 8; ++index) { + engine.begin_record_step(); + emit_stall_record(engine, device); + } + engine.begin_record_step(); + const auto started = std::chrono::steady_clock::now(); + std::string error; + try { + emit_stall_record(engine, device); + } catch (const std::runtime_error& caught) { + error = caught.what(); + } + const auto stall = std::chrono::steady_clock::now() - started; + EXPECT(error.find("did not drain") != std::string::npos); + EXPECT(error.find("admission bound") != std::string::npos); + // Bounded by budget + declared bound + grace, not by the sink. + EXPECT(stall < std::chrono::milliseconds(2400)); + const ring_py::RecordCaptureStatus status = engine.record_capture_status(); + EXPECT(status.failed); + EXPECT(status.failure.find("admission bound") != std::string::npos); + CUDA_CHECK(cudaStreamSynchronize(at::cuda::getCurrentCUDAStream().stream())); + engine.stop(); + CUDA_CHECK(cudaFree(device)); +} + +// A sink whose checked flush reports lost records (NativePackSink's loss +// counters): the failure latches the record runtime, so capture stops and +// says why, instead of only failing that one flush. +class LossyFlushSink final : public ring::RecordSink { +public: + void submit(ring::RecordEnvelope) override { + submissions.fetch_add(1, std::memory_order_acq_rel); + } + bool flush_and_wait(Duration) override { + throw std::runtime_error("pipeline reported lost records"); + } + void rethrow_if_failed() const override {} + + std::atomic submissions{0}; +}; + +static void test_a_failed_sink_flush_latches_the_runtime() { + banner("a sink flush that fails latches the record runtime"); + ring_py::RecordRuntimeOptions options; + options.failure_policy = ring::RecordFailurePolicy::kDisableCapture; + auto sink = std::make_shared(); + ring_py::RingEnginePy engine(small_record_ring(), sink, options); + engine.init(); + engine.start(); + const std::vector source = pattern(1024, 9); + uint8_t* device = upload(source, at::cuda::getCurrentCUDAStream().stream()); + emit_stall_record(engine, device); + bool flush_threw = false; + try { + engine.flush_records_and_wait(5000); + } catch (const std::runtime_error&) { + flush_threw = true; + } + EXPECT(flush_threw); + const ring_py::RecordCaptureStatus status = engine.record_capture_status(); + EXPECT(status.failed); + EXPECT(status.failure.find("lost records") != std::string::npos); + // Capture is off: the next record is discarded, not submitted. + emit_stall_record(engine, device); + CUDA_CHECK(cudaStreamSynchronize(at::cuda::getCurrentCUDAStream().stream())); + try { + engine.flush_records_and_wait(5000); + } catch (const std::runtime_error&) { + } + EXPECT(sink->submissions.load(std::memory_order_acquire) == 1); + // The record worker discards the payload as the drain delivers it, so + // the count can trail the flush by a moment. + const auto discard_deadline = + std::chrono::steady_clock::now() + std::chrono::seconds(2); + while (engine.record_capture_status().discarded_payloads == 0 && + std::chrono::steady_clock::now() < discard_deadline) { + std::this_thread::sleep_for(std::chrono::milliseconds(5)); + } + EXPECT(engine.record_capture_status().discarded_payloads == 1); + engine.stop(); + CUDA_CHECK(cudaFree(device)); +} + int main() { setbuf(stdout, nullptr); ring::set_ring_null_mode(false); @@ -797,6 +1184,12 @@ int main() { test_drain_worker_binds_owner_device(); test_record_flush_bounds_current_stream_prefix_wait(); test_record_flush_reaches_sink_durability_boundary(); + test_raise_policy_skips_the_step_and_raises_at_the_next(); + test_disable_capture_skips_the_step_and_resumes(); + test_a_stall_budget_needs_a_bounded_sink_admission(); + test_a_stall_budget_needs_begin_record_step(); + test_a_sink_past_its_admission_bound_is_a_ring_failure(); + test_a_failed_sink_flush_latches_the_runtime(); std::printf("Results: %d passed, %d failed\n", g_pass, g_fail); return g_fail == 0 ? 0 : 1; diff --git a/tests/native/test_pack_sink_timeout.cpp b/tests/native/test_pack_sink_timeout.cpp index 795e7c2b8..7defd3fd8 100644 --- a/tests/native/test_pack_sink_timeout.cpp +++ b/tests/native/test_pack_sink_timeout.cpp @@ -21,6 +21,13 @@ // waits for the pipeline to visibly quiesce between submits), so the // finite timeout never touches them. // +// A second case pins the ENVELOPE deadline: the ring hands the sink one +// envelope of N rows per submit, and its admission as a whole must be bounded +// by one admission_timeout_s -- not one timeout per row, which let a steadily +// slow sink hold the record worker (and, through it, the forward) for N +// timeouts. There the stager is slowed rather than wedged, so each row after +// the pipeline fills waits a little under one timeout for room. +// // Built and run by tests/test_native_pack_sink_timeout.py. #include @@ -32,8 +39,10 @@ #include #include #include +#include #include "sink/pack_sink.h" +#include "sink/record_row.h" namespace fs = std::filesystem; @@ -183,10 +192,99 @@ void TestBlockedPipelineTimesOutAndCountsIt() { CHECK(final_snap.stage_packs == 0); } +std::string MetadataJson(int n) { + return std::string("{\"capture_id\": \"envelope-") + std::to_string(n) + + "\", \"tenant_id\": \"tenant-a\", \"experiment_id\": \"exp-1\", " + "\"run_id\": \"run-1\", \"session_id\": \"session-1\", " + "\"request_id\": \"req-1\", \"sequence_id\": \"seq-1\", " + "\"model_id\": \"model-1\", \"model_revision\": \"rev-1\", " + "\"adapter_revision\": null, \"capture_policy_version\": " + "\"policy-1\", \"hook_name\": \"hook.0\", \"layer_number\": 0, " + "\"producer_rank\": 0, \"step_number\": " + std::to_string(n) + + ", \"token_start\": 0, \"token_end\": 1, \"batch_position\": 0, " + "\"dtype\": \"float32\", \"shape\": [4], \"captured_at_ns\": " + + std::to_string(1700000000000000000ull + n) + "}"; +} + +void TestAnEnvelopeSharesOneAdmissionDeadline() { + const char* base = std::getenv("SPOOL_TEST_ROOT"); + const std::string root = + std::string(base != nullptr ? base : "/tmp") + "/sink-envelope"; + fs::remove_all(root); + + constexpr double kTimeoutS = 0.5; + // Each pack takes a little over half a timeout to stage, so every row + // that has to wait for room waits well under one timeout: row by row, + // none of them would ever time out. + constexpr auto kStage = std::chrono::milliseconds(300); + constexpr int kRows = 10; + + dmi_sink::SinkConfig config; + config.spool_root = root; + config.num_workers = 1; + config.max_queue_records = 1; + config.max_pack_records = 1; + config.stage_queue_packs = 1; + config.max_linger_ns = 3600ull * 1000 * 1000 * 1000; + config.overload = dmi_sink::Overload::kBlock; + config.admission_timeout_s = kTimeoutS; + + dmi_sink::PackSink sink(config); + const std::string start_error = sink.Start(); + CHECK(start_error.empty()); + if (!start_error.empty()) return; + sink.SpoolForTesting().SetStageHookForTesting( + [&] { std::this_thread::sleep_for(kStage); }); + + const uint8_t payload[16] = {}; + std::vector metadata; + for (int n = 0; n < kRows; ++n) metadata.push_back(MetadataJson(n)); + + const auto started = std::chrono::steady_clock::now(); + dmi_sink::EnvelopeAdmission envelope(sink); + dmi_sink::RowStatus status = dmi_sink::RowStatus::kOk; + std::string detail; + int admitted = 0; + for (int n = 0; n < kRows; ++n) { + dmi_sink::RowInput row; + row.metadata_json = metadata[n]; + row.payload = payload; + row.payload_bytes = sizeof(payload); + row.dtype_name = "float32"; + row.shape = {4}; + status = envelope.SubmitRow(row, &detail); + if (status != dmi_sink::RowStatus::kOk) break; + ++admitted; + } + const double elapsed_s = std::chrono::duration( + std::chrono::steady_clock::now() - started).count(); + + // One timeout for the whole envelope: a row times out once the + // envelope's deadline passes, well before every row got in. + CHECK(status == dmi_sink::RowStatus::kNotAccepted); + CHECK(detail == "timed_out"); + CHECK(admitted < kRows); + CHECK(elapsed_s >= kTimeoutS - 0.05); + CHECK(elapsed_s < kTimeoutS + 0.25); + if (status != dmi_sink::RowStatus::kNotAccepted || + elapsed_s >= kTimeoutS + 0.25) { + std::cerr << "envelope: " << admitted << " rows admitted in " + << elapsed_s << " s (status " << RowStatusName(status) + << ")\n"; + } + CHECK(sink.Snapshot().timed_out_records == 1); + + std::string close_error; + const dmi_sink::SinkSnapshot final_snap = sink.Close(-1.0, &close_error); + CHECK(close_error.empty()); + CHECK(final_snap.persisted_records == static_cast(admitted)); +} + } // namespace int main() { TestBlockedPipelineTimesOutAndCountsIt(); + TestAnEnvelopeSharesOneAdmissionDeadline(); if (g_failures != 0) { std::cerr << g_failures << " check(s) failed\n"; return 1; diff --git a/tests/test_capture_bounds.py b/tests/test_capture_bounds.py new file mode 100644 index 000000000..14acdd489 --- /dev/null +++ b/tests/test_capture_bounds.py @@ -0,0 +1,142 @@ +"""Capture bounds are checked at attach, never discovered in the forward. + +A record the sink can never admit -- larger than its queue, or than an empty +pack -- is refused on the record worker, which latches the record runtime: +under the raise policy the next forward raises. Likewise a pack larger than +the uploader's in-flight budget is staged and then never uploaded. All three +depend only on configuration and the largest record an adapter will emit, +so ``validate_capture_bounds`` checks them before any forward runs. +""" + +from __future__ import annotations + +import pytest + +from dmi.configuration.errors import ConfigurationError + +pytestmark = pytest.mark.cpu + +MiB = 1 << 20 + + +def _sink(tmp_path, **fields): + from dmi.storage.native_capture import NativeSinkConfig + + return NativeSinkConfig(spool_root=str(tmp_path), **fields) + + +def _storage(**fields): + from dmi.storage.native_capture import NativeCaptureStorageConfig + + base = dict(s3_endpoint="https://s3.example.test", s3_bucket="bucket", + s3_access_key="AKIA-test", s3_secret_key="secret-test") + base.update(fields) + return NativeCaptureStorageConfig(**base) + + +# --- uploader_max_in_flight_bytes --------------------------------------------- + + +def test_in_flight_budget_defaults_to_the_native_uploader_default(): + storage = _storage() + # native/csrc/store/uploader.h UploaderConfig::max_in_flight_bytes + assert storage.uploader_max_in_flight_bytes == 256 * MiB + assert storage._native_dict()["uploader_max_in_flight_bytes"] == 256 * MiB + + +def test_in_flight_budget_reaches_the_native_service_config(): + storage = _storage(uploader_max_in_flight_bytes=512 * MiB) + assert storage._native_dict()["uploader_max_in_flight_bytes"] == 512 * MiB + + +@pytest.mark.parametrize("value", [0, -1, 1.5, True]) +def test_in_flight_budget_must_be_a_positive_int(value): + with pytest.raises(ValueError, match="uploader_max_in_flight_bytes"): + _storage(uploader_max_in_flight_bytes=value) + + +# --- validate_capture_bounds ---------------------------------------------------- + + +def test_bounds_that_fit_pass(tmp_path): + from dmi.storage.native_capture import validate_capture_bounds + + validate_capture_bounds(_sink(tmp_path), 8 * MiB, storage_config=_storage()) + + +def test_a_record_larger_than_the_queue_is_refused(tmp_path): + from dmi.storage.native_capture import validate_capture_bounds + + # The default 16 MiB queue, and a 17 MiB prefill row. + with pytest.raises(ConfigurationError, match="max_queue_bytes"): + validate_capture_bounds(_sink(tmp_path), 17 * MiB) + + +def test_a_record_that_fills_an_empty_pack_is_refused(tmp_path): + from dmi.storage.native_capture import ( + PACK_FRAMING_RESERVE_BYTES, validate_capture_bounds, + ) + + # Admitted by the queue, but an empty pack still needs its header, row + # footer and trailer: the sink would admit the record and then count it + # oversized on the pack worker, storing nothing. + sink = _sink(tmp_path, max_queue_bytes=64 * MiB, max_pack_bytes=32 * MiB) + with pytest.raises(ConfigurationError, match="max_pack_bytes"): + validate_capture_bounds(sink, 32 * MiB) + validate_capture_bounds(sink, 32 * MiB - PACK_FRAMING_RESERVE_BYTES) + + +def test_a_pack_over_the_in_flight_budget_is_refused(tmp_path): + from dmi.storage.native_capture import validate_capture_bounds + + sink = _sink(tmp_path, max_pack_bytes=512 * MiB) + with pytest.raises(ConfigurationError, + match="uploader_max_in_flight_bytes"): + validate_capture_bounds(sink, MiB, storage_config=_storage()) + validate_capture_bounds( + sink, MiB, + storage_config=_storage(uploader_max_in_flight_bytes=512 * MiB)) + + +@pytest.mark.parametrize("value", [0, -1, 1.0, True]) +def test_the_record_size_must_be_a_positive_int(tmp_path, value): + from dmi.storage.native_capture import validate_capture_bounds + + with pytest.raises(ValueError, match="max_record_bytes"): + validate_capture_bounds(_sink(tmp_path), value) + + +# --- the engine's entry point ---------------------------------------------------- + + +def _engine(tmp_path, **sink_fields): + from dmi.config import MonitoringConfig + from dmi.engine import MonitoringEngine + + config = MonitoringConfig( + storage_backend="persistent", + capture_sink_config=_sink(tmp_path, **sink_fields), + capture_storage_config=_storage(), + ) + return MonitoringEngine(config=config, model_id="bounds", + enable_ring_transport=False) + + +def test_the_engine_checks_its_own_capture_configs(tmp_path): + engine = _engine(tmp_path) + engine.validate_capture_bounds(MiB) + with pytest.raises(ConfigurationError, match="max_queue_bytes"): + engine.validate_capture_bounds(17 * MiB) + engine = _engine(tmp_path, max_pack_bytes=512 * MiB, + max_queue_bytes=64 * MiB) + with pytest.raises(ConfigurationError, + match="uploader_max_in_flight_bytes"): + engine.validate_capture_bounds(MiB) + + +def test_the_engine_has_nothing_to_check_without_a_sink_config(): + from dmi.engine import MonitoringEngine + + engine = MonitoringEngine(config=None, model_id="bounds", + enable_ring_transport=False) + engine.validate_capture_bounds(1 << 40) diff --git a/tests/test_engine_runtime_api.py b/tests/test_engine_runtime_api.py index 963c82fe5..393c56129 100644 --- a/tests/test_engine_runtime_api.py +++ b/tests/test_engine_runtime_api.py @@ -285,7 +285,7 @@ def configure_record_schema(self, schema): class _Factory: @staticmethod - def create_record(config, host): + def create_record(config, host, **_options): created.append((config, host)) return new_ring @@ -444,7 +444,7 @@ def stop_and_release(): class _Factory: @staticmethod - def create_record(config, target): + def create_record(config, target, **_options): created.append((config, target)) return new_ring @@ -542,7 +542,7 @@ def failing_start(): class _Factory: @staticmethod - def create_record(config, target): + def create_record(config, target, **_options): return new_ring class _FakeTransport: @@ -822,7 +822,7 @@ def configure_record_schema(self, schema): "RingEngine", (), {"create_record": staticmethod( - lambda config, target: new_ring)}, + lambda config, target, **_options: new_ring)}, ) monkeypatch.setitem(sys.modules, "dmi.transport.ring", fake_transport_module) monkeypatch.setitem(sys.modules, "dmi.transport.native", fake_native_module) @@ -968,7 +968,7 @@ def test_second_record_runtime_is_refused_while_one_is_active(monkeypatch): activated = [] deactivated = [] - def create_record(config, target): + def create_record(config, target, **_options): created.append((config, target)) return new_ring @@ -1020,7 +1020,7 @@ def start(self): new_ring = _Boom() - def create_record(config, target): + def create_record(config, target, **_options): if failing_step == "create_record": raise RuntimeError("record ring create failed") return new_ring @@ -1068,7 +1068,7 @@ def boom_activate(_transport): _record_ring_fakes( monkeypatch, - create_record=lambda config, target: new_ring, + create_record=lambda config, target, **_options: new_ring, activate=boom_activate, deactivated=deactivated, ) diff --git a/tests/test_integration_api_v1.py b/tests/test_integration_api_v1.py index d1c1fd096..1272dcc82 100644 --- a/tests/test_integration_api_v1.py +++ b/tests/test_integration_api_v1.py @@ -237,8 +237,10 @@ def test_v1_public_surface_is_documented() -> None: assert missing == [] member_names = { "active_hook_specs", + "begin_step", "bind_hook", "capture_enabled", + "capture_status", "commit_step", "create_record_runtime", "emit_output", @@ -248,6 +250,7 @@ def test_v1_public_surface_is_documented() -> None: "prepare_replay", "ring_capacities", "set_capture_enabled", + "validate_capture_bounds", } missing_members = [ name diff --git a/tests/test_native_capture_chain_live.py b/tests/test_native_capture_chain_live.py index 134746806..5004db901 100644 --- a/tests/test_native_capture_chain_live.py +++ b/tests/test_native_capture_chain_live.py @@ -204,10 +204,14 @@ def _storage_config(endpoint, bucket, access, secret, prefix): table_prefix=prefix, poll_interval_s=0.05) -def _run_chain(config, spool_root: Path, envelopes, *, sink_overrides=None): +def _run_chain(config, spool_root: Path, envelopes, *, sink_overrides=None, + sink_config=None): """Service first, as the engine orders it (its start sweeps the spool, which is safe only while no sink writes there), then the sink; flush - both, and return the snapshots and what the reader reads back.""" + both, and return the snapshots and what the reader reads back. + + The sink is the raw binding with `sink_overrides`, or, given a + NativeSinkConfig, the one the engine builds from it.""" from dmi.storage.native_capture import ( NativeCaptureReader, NativeCaptureStorage, ) @@ -216,7 +220,15 @@ def _run_chain(config, spool_root: Path, envelopes, *, sink_overrides=None): spool_max_bytes=1 << 40, sweep_spool=True) service.start() try: - sink, _lease = _open_sink(spool_root, **(sink_overrides or {})) + if sink_config is None: + sink, _lease = _open_sink(spool_root, **(sink_overrides or {})) + else: + from dmi.storage.capture.native_sink import ( + create_native_pack_sink, + ) + + sink = create_native_pack_sink(sink_config).native_sink + _lease = sink.attach() for envelope in envelopes: sink.submit_envelope(LAYOUT, envelope.rows, envelope.payload()) assert sink.flush_and_wait(120.0) @@ -373,3 +385,78 @@ def test_a_multipart_pack_through_the_fake_s3(fake_s3, tmp_path): assert all(size >= S3_MIN_PART for size in parts[:-1]), parts assert STATE.objects[key]["etag"].endswith('-multipart"') _assert_read_back_exactly(captures, envelopes) + + +# --- B2: sizes and bursts are stored, or refused at configuration time ------- + + +def test_a_row_over_16_mib_is_refused_at_configuration_or_stored( + fake_s3, tmp_path): + """A 17 MiB prefill row against the default 16 MiB sink queue. + + The default bounds cannot admit it, and validate_capture_bounds says so + before any capture runs, naming the bound. With the queue raised the + same row goes through the real chain and reads back exactly.""" + from dmi.configuration.errors import ConfigurationError + from dmi.storage.native_capture import ( + NativeSinkConfig, validate_capture_bounds, + ) + + row_bytes = 17 * MiB + envelope = _Envelope() + envelope.add(0, torch_randn(row_bytes // 4, seed=17)) + spool_root = tmp_path / "spool" + with _catalog() as prefix: + config = _storage_config(fake_s3, BUCKET, ACCESS, SECRET, prefix) + with pytest.raises(ConfigurationError, match="max_queue_bytes"): + validate_capture_bounds( + NativeSinkConfig(spool_root=str(spool_root)), row_bytes, + storage_config=config) + + raised = NativeSinkConfig(spool_root=str(spool_root), + max_queue_bytes=64 * MiB, + max_linger_ns=600 * 10**9) + validate_capture_bounds(raised, row_bytes, storage_config=config) + sink_snapshot, snapshot, captures = _run_chain( + config, spool_root, [envelope], sink_config=raised) + + assert sink_snapshot["persisted_records"] == 1, sink_snapshot + assert sink_snapshot["oversized_records"] == 0, sink_snapshot + assert snapshot["indexed_rows"] == 1, snapshot + _assert_read_back_exactly(captures, [envelope]) + + +def test_a_burst_of_64_1_mib_rows_is_stored_under_the_default_sink_config( + fake_s3, tmp_path): + """One envelope of 64 x 1 MiB rows -- four times the default queue -- + through the sink built from NativeSinkConfig's defaults (block, 2 s). + The record worker waits for room rather than dropping: every row is + stored and reads back. Under the binding's old drop_newest default the + same envelope was refused after 18-25 rows.""" + from dmi.storage.native_capture import ( + NativeSinkConfig, validate_capture_bounds, + ) + + envelope = _Envelope() + for index in range(64): + envelope.add(index, torch_randn(MiB // 4, seed=index)) + spool_root = tmp_path / "spool" + sink_config = NativeSinkConfig(spool_root=str(spool_root)) + with _catalog() as prefix: + config = _storage_config(fake_s3, BUCKET, ACCESS, SECRET, prefix) + validate_capture_bounds(sink_config, MiB, storage_config=config) + sink_snapshot, snapshot, captures = _run_chain( + config, spool_root, [envelope], sink_config=sink_config) + + assert sink_snapshot["persisted_records"] == 64, sink_snapshot + assert sink_snapshot["dropped_records"] == 0, sink_snapshot + assert sink_snapshot["timed_out_records"] == 0, sink_snapshot + assert snapshot["indexed_rows"] == 64, snapshot + assert snapshot["pending_index"] == 0, snapshot + _assert_read_back_exactly(captures, [envelope]) + + +def torch_randn(elements: int, *, seed: int): + import torch + + return torch.randn(elements, generator=torch.Generator().manual_seed(seed)) diff --git a/tests/test_native_capture_storage_wiring.py b/tests/test_native_capture_storage_wiring.py index dbe9c8a20..3e38cd58c 100644 --- a/tests/test_native_capture_storage_wiring.py +++ b/tests/test_native_capture_storage_wiring.py @@ -490,7 +490,8 @@ class _RingEngine(_NewRing): def __init__(self, config, host): events.append(("ring", "create")) - create_record = staticmethod(lambda config, target: _NewRing()) + create_record = staticmethod( + lambda config, target, **_options: _NewRing()) native.RingEngine = _RingEngine ring = ModuleType("dmi.transport.ring") diff --git a/tests/test_native_pack_sink.py b/tests/test_native_pack_sink.py index 7034aa625..60f3815f0 100644 --- a/tests/test_native_pack_sink.py +++ b/tests/test_native_pack_sink.py @@ -283,6 +283,28 @@ def test_oversized_record_is_rejected_up_front(sink, tmp_path): assert snapshot["packs_persisted"] == 0 +def test_a_record_that_fits_no_empty_pack_is_oversized_not_a_failure( + sink, tmp_path): + """Admission screens the payload against max_pack_bytes; the pack adds a + header, the record's footer row and a trailer. A payload of exactly + max_pack_bytes clears admission and still fits no empty pack. The + oracle's PackAssembler raises OversizedRecordError for it and the + pipeline counts the record oversized and keeps going (pipeline.py + _run). The native worker instead sealed the empty pack it had just + opened for the record, and "cannot seal an empty pack" failed the whole + sink, losing everything queued behind it.""" + _open(sink, tmp_path / "spool", max_pack_bytes=1024) + edge = CaptureRecord( + metadata=_meta(0, dtype="uint8", shape=(1024,)), payload=bytes(1024)) + assert _submit(sink, edge) == "accepted" + assert _submit(sink, _record(1)) == "accepted" + assert sink.call(op="flush", timeout=30)["ok"] + snapshot = sink.call(op="close", timeout=30)["snapshot"] + assert snapshot["oversized_records"] == 1 + assert snapshot["failures"] == 0 + assert snapshot["persisted_records"] == 1 + + @pytest.mark.parametrize("overload, timeout", [ ("drop_newest", -1), # A timeout, so a regression cannot hang the suite: under BLOCK with diff --git a/tests/test_native_pack_sink_timeout.py b/tests/test_native_pack_sink_timeout.py index 1cae7f12f..ce78f76ba 100644 --- a/tests/test_native_pack_sink_timeout.py +++ b/tests/test_native_pack_sink_timeout.py @@ -9,6 +9,11 @@ with timed_out_records == 1. No admission along the way ever depends on where the packer thread happens to be. +A second case slows the stager instead of wedging it and submits one +envelope of rows through EnvelopeAdmission, the path NativePackSink::submit +takes: the whole envelope must time out within one admission_timeout_s, +not one per row. + This is the native-tier coverage of the deadline wait loop in PackSink::Submit; the conformance-driver suite (test_native_pack_sink.py) covers the admission bounds around it. @@ -53,6 +58,8 @@ def test_blocked_pipeline_times_out_and_counts_it(tmp_path): *extra, str(source), str(csrc / "sink" / "pack_sink.cpp"), + str(csrc / "sink" / "record_row.cpp"), + str(csrc / "common" / "json.cpp"), str(csrc / "sink" / "object_key.cpp"), str(csrc / "pack" / "pack_builder.cpp"), str(csrc / "store" / "spool.cpp"), diff --git a/tests/test_native_sink_admission.py b/tests/test_native_sink_admission.py new file mode 100644 index 000000000..e916614e4 --- /dev/null +++ b/tests/test_native_sink_admission.py @@ -0,0 +1,306 @@ +"""The native sink's admission policy, from NativeSinkConfig to the C++ sink. + +A ring-fed sink admits on the record worker, so its policy decides what a +burst does: under drop_newest a burst larger than the queue loses records +(and latches the record ring), under block the worker waits for room, up to +``admission_timeout_s``. The C++ ``SinkConfig`` keeps drop_newest with no +timeout, the oracle's default; ``NativeSinkConfig``, the ring-fed entry +point, defaults to block with a 2 s timeout. + +``NativeSinkConfig`` lives in :mod:`dmi.storage.native_capture`, live code, +so the public config no longer loads the backup capture package; the old +import path re-exports it. + +Build: make -C native cpu-goals PYTHON=/bin/python +""" + +from __future__ import annotations + +import json +import math +import subprocess +import sys +from pathlib import Path + +import pytest + +REPO = Path(__file__).resolve().parents[1] +BUILD = REPO / "native" / "build" + +pytestmark = pytest.mark.cpu + +LAYOUT = "capture_pack_reference_v1" +MiB = 1 << 20 + + +@pytest.fixture(scope="module") +def native_sink_module(): + # Not a skip: a CPU host builds this module (`make -C native cpu-goals`), + # and an unbuilt one is missing coverage, not missing hardware. + if not sorted(BUILD.glob("_dmi_native_sink*.so")): + pytest.fail("native/build/_dmi_native_sink*.so is not built; run " + "`make -C native cpu-goals PYTHON=/bin/python`") + from dmi.storage.capture.native_sink import _load_native_sink_extension + + return _load_native_sink_extension() + + +# --- NativeSinkConfig -------------------------------------------------------- + + +def test_ring_fed_default_is_block_with_a_two_second_timeout(tmp_path): + from dmi.storage.native_capture import NativeSinkConfig + + config = NativeSinkConfig(spool_root=str(tmp_path)) + assert config.overload == "block" + assert config.admission_timeout_s == 2.0 + + +@pytest.mark.parametrize("fields, error, match", [ + ({"overload": "drop_oldest"}, ValueError, "overload"), + ({"overload": None}, ValueError, "overload"), + ({"admission_timeout_s": 0}, ValueError, "admission_timeout_s"), + ({"admission_timeout_s": -1.0}, ValueError, "admission_timeout_s"), + ({"admission_timeout_s": math.nan}, ValueError, "admission_timeout_s"), + ({"admission_timeout_s": math.inf}, ValueError, "admission_timeout_s"), + ({"admission_timeout_s": "2"}, TypeError, "admission_timeout_s"), + ({"admission_timeout_s": True}, TypeError, "admission_timeout_s"), +]) +def test_admission_fields_are_validated(tmp_path, fields, error, match): + from dmi.storage.native_capture import NativeSinkConfig + + with pytest.raises(error, match=match): + NativeSinkConfig(spool_root=str(tmp_path), **fields) + + +def test_block_may_wait_without_bound_when_asked(tmp_path): + from dmi.storage.native_capture import NativeSinkConfig + + config = NativeSinkConfig(spool_root=str(tmp_path), + admission_timeout_s=None) + assert config.admission_timeout_s is None + + +def test_the_old_import_path_re_exports_the_same_type(): + from dmi.storage.capture import native_sink + from dmi.storage import native_capture + + assert native_sink.NativeSinkConfig is native_capture.NativeSinkConfig + + +def test_the_config_type_does_not_load_the_backup_capture_package(): + # A fresh interpreter: this process has long since imported everything. + probe = ( + "import sys\n" + "from dmi.storage.native_capture import NativeSinkConfig\n" + "from dmi.config import MonitoringConfig\n" + "MonitoringConfig(storage_backend='persistent',\n" + " capture_sink_config=NativeSinkConfig(spool_root='/x'))\n" + "print(sorted(m for m in sys.modules\n" + " if m.startswith('dmi.storage.capture')))\n" + ) + result = subprocess.run( + [sys.executable, "-c", probe], capture_output=True, text=True, + check=True, env={"PYTHONPATH": str(REPO / "src"), "PATH": ""}) + assert result.stdout.strip() == "[]", result.stdout + + +# A NativeSinkConfig pickled before the move and the admission fields: the +# old class at the old path (whose name now re-exports the new class), with +# spool_root="/spool/old" and max_queue_bytes=32 MiB, protocol 2. Its state +# is the eight fields the old class had. +_OLD_FORMAT_PICKLE = ( + b"\x80\x02cdmi.storage.capture.native_sink\nNativeSinkConfig\nq\x00)" + b"\x81q\x01]q\x02(X\n\x00\x00\x00/spool/oldq\x03\x8a\x06\x00\x00\x00" + b"\x00\x00\x01K\x01M\x00\x01J\x00\x00\x00\x02J\x00\x00\x00\x08M\x10'J" + b"\x00\xca\x9a;eb.") + + +def test_an_old_format_pickle_loads_with_the_new_defaults(): + import pickle + + from dmi.storage.native_capture import NativeSinkConfig + + config = pickle.loads(_OLD_FORMAT_PICKLE) + assert type(config) is NativeSinkConfig + assert config == NativeSinkConfig(spool_root="/spool/old", + max_queue_bytes=32 * MiB) + assert (config.overload, config.admission_timeout_s) == ("block", 2.0) + + +def test_a_new_pickle_round_trips(tmp_path): + import pickle + + from dmi.storage.native_capture import NativeSinkConfig + + config = NativeSinkConfig(spool_root=str(tmp_path), + overload="drop_newest") + assert pickle.loads(pickle.dumps(config)) == config + + +# --- the native binding ------------------------------------------------------ + + +def test_the_binding_keeps_the_cxx_default(native_sink_module, tmp_path): + sink = native_sink_module.NativePackSink( + spool_root=str(tmp_path), layout=LAYOUT) + assert sink.overload == "drop_newest" + assert sink.admission_timeout_s is None + + +def test_the_binding_accepts_block_with_a_timeout(native_sink_module, + tmp_path): + sink = native_sink_module.NativePackSink( + spool_root=str(tmp_path), layout=LAYOUT, overload="block", + admission_timeout_s=0.25) + assert sink.overload == "block" + assert sink.admission_timeout_s == 0.25 + + +@pytest.mark.parametrize("fields", [ + {"overload": "drop_oldest"}, + {"overload": "block", "admission_timeout_s": -0.5}, + {"overload": "block", "admission_timeout_s": math.nan}, +]) +def test_the_binding_refuses_a_bad_policy(native_sink_module, tmp_path, + fields): + with pytest.raises(ValueError): + native_sink_module.NativePackSink( + spool_root=str(tmp_path), layout=LAYOUT, **fields) + + +def test_the_handle_passes_the_policy_to_the_native_sink(native_sink_module, + tmp_path): + from dmi.storage.capture.native_sink import create_native_pack_sink + from dmi.storage.native_capture import NativeSinkConfig + + default = create_native_pack_sink( + NativeSinkConfig(spool_root=str(tmp_path / "a"))).native_sink + assert (default.overload, default.admission_timeout_s) == ("block", 2.0) + dropping = create_native_pack_sink(NativeSinkConfig( + spool_root=str(tmp_path / "b"), overload="drop_newest")).native_sink + assert dropping.overload == "drop_newest" + + +def _metadata(index: int) -> dict: + from dmi.storage.capture import CaptureMetadata + + return CaptureMetadata( + capture_id=f"burst-{index:04d}", tenant_id="t", experiment_id="e", + run_id="r", session_id="s", request_id=f"q{index}", + sequence_id=f"n{index}", model_id="m", model_revision="mr", + adapter_revision=None, capture_policy_version="v", + hook_name="resid_post", layer_number=0, producer_rank=0, + step_number=index, token_start=index, token_end=index + 1, + batch_position=0, dtype="float32", shape=(MiB // 4,), + captured_at_ns=1_700_000_000_000_000_000 + index, + ).to_mapping() + + +def test_a_burst_four_times_the_queue_is_admitted_under_the_default( + native_sink_module, tmp_path): + """64 x 1 MiB rows in one envelope, against the default 16 MiB queue: + the record worker waits for room instead of dropping.""" + import torch + + from dmi.storage.capture.native_sink import create_native_pack_sink + from dmi.storage.native_capture import NativeSinkConfig + + sink = create_native_pack_sink( + NativeSinkConfig(spool_root=str(tmp_path))).native_sink + lease = sink.attach() + rows = [{ + "metadata_json": json.dumps(_metadata(index)), + "offset": index * MiB, "length": MiB, "dtype": 6, # at::kFloat + "shape": [MiB // 4], + } for index in range(64)] + payload = torch.arange(64 * MiB // 4, dtype=torch.float32).view( + torch.uint8) + sink.submit_envelope(LAYOUT, rows, payload) + assert sink.flush_and_wait(60.0) + sink.rethrow_if_failed() + snapshot = sink.snapshot() + assert snapshot["persisted_records"] == 64, snapshot + assert snapshot["dropped_records"] == 0, snapshot + assert snapshot["timed_out_records"] == 0, snapshot + assert snapshot["rejected_closed_records"] == 0, snapshot + del lease + + +# --- the admission bound a record ring reads --------------------------------- + + +@pytest.mark.parametrize("fields, bound", [ + ({}, 0.0), # the C++ default, drop_newest: refused at once + ({"overload": "block", "admission_timeout_s": 0.25}, 0.25), + ({"overload": "block"}, None), # waits without bound +]) +def test_the_sink_reports_its_admission_bound(native_sink_module, tmp_path, + fields, bound): + """A record ring with a stall budget waits for one in-flight submit + after the budget is spent; it reads this bound to refuse a sink whose + admission has none.""" + sink = native_sink_module.NativePackSink( + spool_root=str(tmp_path), layout=LAYOUT, **fields) + assert sink.admission_bound_s == bound + + +# --- a record lost after admission ------------------------------------------- + + +def _one_row_envelope(index: int, nbytes: int): + import torch + + mapping = _metadata(index) + mapping["shape"] = [nbytes // 4] + rows = [{"metadata_json": json.dumps(mapping), "offset": 0, + "length": nbytes, "dtype": 6, "shape": [nbytes // 4]}] + return rows, torch.zeros(nbytes // 4, dtype=torch.float32).view( + torch.uint8) + + +def test_a_record_that_fits_no_pack_fails_the_flush(native_sink_module, + tmp_path): + """max_pack_bytes = one 1 MiB record: admission screens the payload + alone, so the record is admitted, and the pack worker then drops it as + oversized (it fits no empty pack with its framing). The reference + adapter counts oversized_records as a loss and raises at flush and + rethrow; the native sink used to report success with nothing + persisted.""" + from dmi.storage.capture.native_sink import create_native_pack_sink + from dmi.storage.native_capture import NativeSinkConfig + + sink = create_native_pack_sink(NativeSinkConfig( + spool_root=str(tmp_path), max_pack_bytes=MiB, + max_queue_bytes=4 * MiB)).native_sink + lease = sink.attach() + rows, payload = _one_row_envelope(0, MiB) + sink.submit_envelope(LAYOUT, rows, payload) + with pytest.raises(RuntimeError, match="oversized_records"): + sink.flush_and_wait(30.0) + with pytest.raises(RuntimeError, match="oversized_records"): + sink.rethrow_if_failed() + snapshot = sink.snapshot() + assert (snapshot["admitted_records"], snapshot["persisted_records"], + snapshot["oversized_records"]) == (1, 0, 1), snapshot + # The next envelope is refused too: capture stops and says why, on the + # record worker, instead of losing records until someone flushes. + rows, payload = _one_row_envelope(1, 1024) + with pytest.raises(RuntimeError, match="oversized_records"): + sink.submit_envelope(LAYOUT, rows, payload) + del lease + + +def test_a_stored_record_still_flushes_clean(native_sink_module, tmp_path): + from dmi.storage.capture.native_sink import create_native_pack_sink + from dmi.storage.native_capture import NativeSinkConfig + + sink = create_native_pack_sink( + NativeSinkConfig(spool_root=str(tmp_path))).native_sink + lease = sink.attach() + rows, payload = _one_row_envelope(0, 1024) + sink.submit_envelope(LAYOUT, rows, payload) + assert sink.flush_and_wait(30.0) + sink.rethrow_if_failed() + assert sink.snapshot()["persisted_records"] == 1 + del lease diff --git a/tests/test_record_failure_policy.py b/tests/test_record_failure_policy.py new file mode 100644 index 000000000..308fa02a1 --- /dev/null +++ b/tests/test_record_failure_policy.py @@ -0,0 +1,351 @@ +"""The record runtime's failure policy and capture status, from the engine. + +The native ring does the work (tests/native/ring: the consumer's discard +under disable_capture, the per-step stall budget). This suite pins the +Python surface around it with the native modules faked: what +create_record_runtime accepts and passes to the native ring, what it refuses +before touching the live ring, how a step is marked, and what +capture_status() and close() report. +""" + +from __future__ import annotations + +import logging + +import pytest + +from dmi.engine import MonitoringEngine +from tests.test_engine_runtime_api import ( + _FakeRingEngine, + _engine_with_fake_ring, + _explicit_sink_format, + _record_ring_fakes, +) + +pytestmark = pytest.mark.cpu + + +class _StatusRing(_FakeRingEngine): + """A record ring that reports the native capture status.""" + + def __init__(self, status=None): + super().__init__() + self.status = status or { + "failure_policy": "raise", "failed": False, "failure": "", + "discarded_descriptors": 0, "discarded_payloads": 0, + "step_stall_budget_ms": 0, "stall_budget_exhaustions": 0, + "skipped_steps": 0, "steps_with_discards": 0, + "reserve_wait_ns": 0, "max_step_wait_ns": 0, + } + self.steps = 0 + + def capture_status(self): + return dict(self.status) + + def begin_record_step(self): + self.steps += 1 + + +class _Sink: + """An explicit native sink, as far as create_record_runtime looks.""" + + def __init__(self, admission_bound_s=0.5): + self.admission_bound_s = admission_bound_s + + def _acquire_engine(self): + return object() + + +class _SinkWithoutABound: + """A native sink that cannot say how long its admission takes.""" + + def _acquire_engine(self): + return object() + + +def _record_engine(monkeypatch, ring, **runtime_options): + """An engine switched to `ring` through create_record_runtime.""" + engine, _old_transport, _old_ring = _engine_with_fake_ring() + engine._ring_config = object() + engine._host_engine = object() + created = [] + + def create_record(config, target, **options): + created.append(options) + return ring + + _record_ring_fakes(monkeypatch, create_record=create_record, + activate=lambda transport: None, deactivated=[]) + runtime = engine.create_record_runtime( + _explicit_sink_format(), **runtime_options) + return engine, runtime, created + + +def _refused_before_the_live_ring(monkeypatch, match, **runtime_options): + engine, _old_transport, old_ring = _engine_with_fake_ring() + engine._ring_config = object() + engine._host_engine = object() + _record_ring_fakes( + monkeypatch, + create_record=lambda *a, **k: pytest.fail("ring must not be built"), + activate=lambda transport: None, deactivated=[]) + with pytest.raises(ValueError, match=match): + engine.create_record_runtime( + _explicit_sink_format(), **runtime_options) + assert old_ring.stop_calls == 0 + assert engine._record_mode is False + + +# --- create_record_runtime ---------------------------------------------------- + + +def test_the_default_policy_is_raise_with_no_stall_budget(monkeypatch): + _engine, _runtime, created = _record_engine(monkeypatch, _StatusRing()) + assert created == [{"failure_policy": "raise", "step_stall_budget_ms": 0}] + + +def test_disable_capture_and_a_budget_reach_the_native_ring(monkeypatch): + _engine, _runtime, created = _record_engine( + monkeypatch, _StatusRing(), record_sink=_Sink(), + failure_policy="disable_capture", step_stall_budget_ms=2000) + assert created == [{"failure_policy": "disable_capture", + "step_stall_budget_ms": 2000}] + + +# A budget's promise -- a step stalls at most the budget plus one admission +# -- holds only when the sink bounds its admission: past the budget the +# forward still waits for the envelope the sink is admitting. Under both +# policies, since both skip the rest of the step the same way. + + +@pytest.mark.parametrize("policy", ["raise", "disable_capture"]) +def test_a_budget_is_refused_on_the_clickhouse_host_path(monkeypatch, policy): + """record_sink=None: ClickHouseRecordSink enqueues row by row, with + whatever wait the host engine's ingress policy has (by default, none) + and no bound per envelope.""" + _refused_before_the_live_ring( + monkeypatch, "ClickHouse host path", failure_policy=policy, + step_stall_budget_ms=50) + + +@pytest.mark.parametrize("sink", [_Sink(admission_bound_s=None), + _SinkWithoutABound()]) +@pytest.mark.parametrize("policy", ["raise", "disable_capture"]) +def test_a_budget_is_refused_for_a_sink_without_an_admission_bound( + monkeypatch, sink, policy): + """A NativePackSink that blocks with no admission timeout, the + reference bridge, or any sink that cannot say.""" + _refused_before_the_live_ring( + monkeypatch, "admission bound", record_sink=sink, + failure_policy=policy, step_stall_budget_ms=50) + + +def test_without_a_budget_nothing_is_promised_or_refused(monkeypatch): + _engine, _runtime, created = _record_engine( + monkeypatch, _StatusRing(), record_sink=_SinkWithoutABound(), + failure_policy="disable_capture") + assert created == [{"failure_policy": "disable_capture", + "step_stall_budget_ms": 0}] + + +@pytest.mark.parametrize("options, match", [ + ({"failure_policy": "ignore"}, "failure_policy"), + ({"failure_policy": None}, "failure_policy"), + ({"step_stall_budget_ms": 0}, "step_stall_budget_ms"), + ({"step_stall_budget_ms": -5}, "step_stall_budget_ms"), + ({"step_stall_budget_ms": 2.5}, "step_stall_budget_ms"), + ({"step_stall_budget_ms": True}, "step_stall_budget_ms"), +]) +def test_bad_options_are_refused_before_the_live_ring_is_touched( + monkeypatch, options, match): + engine, _old_transport, old_ring = _engine_with_fake_ring() + engine._ring_config = object() + engine._host_engine = object() + _record_ring_fakes( + monkeypatch, + create_record=lambda *a, **k: pytest.fail("ring must not be built"), + activate=lambda transport: None, deactivated=[]) + with pytest.raises(ValueError, match=match): + engine.create_record_runtime(_explicit_sink_format(), **options) + assert old_ring.stop_calls == 0 + assert engine._record_mode is False + + +@pytest.mark.parametrize("policy", ["raise", "disable_capture"]) +def test_a_budget_needs_a_bounded_persistent_sink(tmp_path, policy): + """The persistent path builds its NativePackSink from + capture_sink_config; block with no timeout has no admission bound.""" + from dmi.config import MonitoringConfig + from dmi.storage.native_capture import NativeSinkConfig + + config = MonitoringConfig( + storage_backend="persistent", + capture_sink_config=NativeSinkConfig( + spool_root=str(tmp_path), admission_timeout_s=None)) + engine = MonitoringEngine(config=config, model_id="policy", + enable_ring_transport=False) + with pytest.raises(ValueError, match="admission_timeout_s"): + engine.create_record_runtime( + _explicit_sink_format(), failure_policy=policy, + step_stall_budget_ms=50) + # Without a budget the same sink is accepted: past validation, this + # engine only lacks a ring. + with pytest.raises(RuntimeError, match="Ring transport"): + engine.create_record_runtime( + _explicit_sink_format(), failure_policy=policy) + + +@pytest.mark.parametrize("fields", [ + {}, # block with the 2 s default + {"overload": "drop_newest"}, +]) +def test_a_bounded_persistent_sink_passes_the_budget_check(tmp_path, fields): + from dmi.config import MonitoringConfig + from dmi.storage.native_capture import NativeSinkConfig + + config = MonitoringConfig( + storage_backend="persistent", + capture_sink_config=NativeSinkConfig( + spool_root=str(tmp_path), **fields)) + engine = MonitoringEngine(config=config, model_id="policy", + enable_ring_transport=False) + with pytest.raises(RuntimeError, match="Ring transport"): + engine.create_record_runtime( + _explicit_sink_format(), failure_policy="disable_capture", + step_stall_budget_ms=50) + + +# --- steps ----------------------------------------------------------------------- + + +def test_begin_step_starts_a_native_stall_budget_step(monkeypatch): + ring = _StatusRing() + _engine, runtime, _created = _record_engine(monkeypatch, ring) + runtime._transport.begin_record_step = ring.begin_record_step + runtime.begin_step() + runtime.begin_step() + assert ring.steps == 2 + + +# --- capture_status -------------------------------------------------------------- + + +def test_capture_status_without_a_record_runtime(): + engine = MonitoringEngine(enable_ring_transport=False) + assert engine.capture_status() == { + "record_mode": False, "capture_active": False, + "failure_policy": None, "failure": None, + "discarded_descriptors": 0, "discarded_payloads": 0, + "step_stall_budget_ms": None, "stall_budget_exhaustions": 0, + "skipped_steps": 0, "steps_with_discards": 0, + "reserve_wait_s": 0.0, "max_step_wait_s": 0.0, + "sink": None, "storage": None, + } + + +def test_capture_status_reports_a_disabled_capture(monkeypatch): + ring = _StatusRing({ + "failure_policy": "disable_capture", "failed": True, + "failure": "NativePackSink: sink refused durable admission: timed_out", + "discarded_descriptors": 7, "discarded_payloads": 9, + "step_stall_budget_ms": 2000, "stall_budget_exhaustions": 1, + "skipped_steps": 1, "steps_with_discards": 1, + "reserve_wait_ns": 2_500_000_000, "max_step_wait_ns": 2_100_000_000, + }) + engine, _runtime, _created = _record_engine( + monkeypatch, ring, record_sink=_Sink(), + failure_policy="disable_capture", step_stall_budget_ms=2000) + status = engine.capture_status() + assert status == { + "record_mode": True, "capture_active": False, + "failure_policy": "disable_capture", + "failure": "NativePackSink: sink refused durable admission: timed_out", + "discarded_descriptors": 7, "discarded_payloads": 9, + "step_stall_budget_ms": 2000, "stall_budget_exhaustions": 1, + "skipped_steps": 1, "steps_with_discards": 1, + "reserve_wait_s": 2.5, "max_step_wait_s": 2.1, + "sink": None, "storage": None, + } + + +def test_capture_status_reports_skipped_steps_with_capture_active( + monkeypatch): + """A spent budget under disable_capture skips that step and capture + goes on: active, nothing failed, the skip counted.""" + ring = _StatusRing({ + "failure_policy": "disable_capture", "failed": False, "failure": "", + "discarded_descriptors": 8, "discarded_payloads": 8, + "step_stall_budget_ms": 50, "stall_budget_exhaustions": 2, + "skipped_steps": 2, "steps_with_discards": 5, + "reserve_wait_ns": 900_000_000, "max_step_wait_ns": 450_000_000, + }) + engine, _runtime, _created = _record_engine( + monkeypatch, ring, record_sink=_Sink(), + failure_policy="disable_capture", step_stall_budget_ms=50) + status = engine.capture_status() + assert status["capture_active"] is True + assert status["failure"] is None + assert status["skipped_steps"] == 2 + assert status["stall_budget_exhaustions"] == 2 + # Each skip also dropped records queued from earlier steps. + assert status["steps_with_discards"] == 5 + + +def test_capture_status_includes_the_sink_and_storage_snapshots(monkeypatch): + engine, _runtime, _created = _record_engine(monkeypatch, _StatusRing()) + + class _Sink: + def snapshot(self): + return {"persisted_records": 3, "timed_out_records": 0} + + class _Storage: + def snapshot(self): + return {"indexed_packs": 1} + + engine._record_sink = _Sink() + engine._capture_storage = _Storage() + status = engine.capture_status() + assert status["capture_active"] is True + assert status["failure"] is None + assert status["step_stall_budget_ms"] is None + assert status["sink"] == {"persisted_records": 3, "timed_out_records": 0} + assert status["storage"] == {"indexed_packs": 1} + + +def test_close_reports_a_capture_that_stopped(monkeypatch, caplog): + ring = _StatusRing({ + "failure_policy": "disable_capture", "failed": True, + "failure": "NativePackSink: sink refused durable admission: dropped", + "discarded_descriptors": 4, "discarded_payloads": 5, + "step_stall_budget_ms": 50, "stall_budget_exhaustions": 0, + "skipped_steps": 0, "steps_with_discards": 0, + "reserve_wait_ns": 0, "max_step_wait_ns": 0, + }) + engine, _runtime, _created = _record_engine( + monkeypatch, ring, record_sink=_Sink(), + failure_policy="disable_capture", step_stall_budget_ms=50) + with caplog.at_level(logging.WARNING, logger="dmi.engine"): + engine.close() + assert ring.stop_calls == 1 + messages = [record.getMessage() for record in caplog.records] + assert any("durable admission" in message and "5 payloads" in message + for message in messages), messages + + +def test_close_reports_skipped_steps(monkeypatch, caplog): + ring = _StatusRing({ + "failure_policy": "disable_capture", "failed": False, "failure": "", + "discarded_descriptors": 8, "discarded_payloads": 8, + "step_stall_budget_ms": 50, "stall_budget_exhaustions": 3, + "skipped_steps": 3, "steps_with_discards": 4, + "reserve_wait_ns": 0, "max_step_wait_ns": 0, + }) + engine, _runtime, _created = _record_engine( + monkeypatch, ring, record_sink=_Sink(), + failure_policy="disable_capture", step_stall_budget_ms=50) + with caplog.at_level(logging.WARNING, logger="dmi.engine"): + engine.close() + messages = [record.getMessage() for record in caplog.records] + assert any("3 steps" in message and "stall budget" in message + and "4 steps lost records" in message + for message in messages), messages diff --git a/tests/test_record_failure_policy_gpu.py b/tests/test_record_failure_policy_gpu.py new file mode 100644 index 000000000..edbe8921d --- /dev/null +++ b/tests/test_record_failure_policy_gpu.py @@ -0,0 +1,377 @@ +"""Sink refusals and stalls on a real record ring: what the forward sees. + +The native suites pin the mechanism (tests/native/ring). This drives it the +way an integration does -- ``create_record_runtime``, a bound +``HookPointV1`` firing on a CUDA tensor, one ``begin_step`` per step -- and +measures what matters to serving: whether a hook call raises, and how long +the slowest step took. + +The stalling sink is a Python target behind the native reference bridge: +its admission sleeps, like a block-mode sink waiting for queue room, or +refuses, like one that dropped a record. The bridge is told the target's +admission time as its admission bound, standing in for a block-mode +NativePackSink's admission_timeout_s (a budget is refused for a sink with +no bound). The burst and lost-record cases use the real NativePackSink. + +Build: make -C native SM_ARCH=... PYTHON=... (and cpu-goals for the sink) +""" + +from __future__ import annotations + +import logging +import time +from pathlib import Path + +import pytest +import torch + +from tests._requirements import require_cuda, require_native_backend + +REPO_ROOT = Path(__file__).resolve().parents[1] +SINK_BUILT = bool( + sorted((REPO_ROOT / "native" / "build").glob("_dmi_native_sink*.so"))) + +pytestmark = [ + pytest.mark.gpu, + pytest.mark.native_backend, + require_cuda(), + require_native_backend(), +] + +BUDGET_MS = 50 +ADMISSION_S = 0.4 +RECORD_ELEMENTS = 256 # float32: 1 KiB per record +STEPS = 16 + + +class _Target: + """A capture target whose every admission takes ADMISSION_S, and which + refuses from record `refuse_from` on when set.""" + + def __init__(self, *, admission_s=0.0, refuse_from=None): + self.admission_s = admission_s + self.refuse_from = refuse_from + self.submitted = 0 + + def _attach(self): + pass + + def _detach(self): + pass + + def _submit_capture(self, metadata_json, payload): + if self.refuse_from is not None and self.submitted >= self.refuse_from: + raise RuntimeError("sink refused durable admission: dropped") + time.sleep(self.admission_s) + self.submitted += 1 + + def _flush_capture(self, timeout_s): + return True + + def _rethrow_capture(self): + pass + + +class _HookRuntime: + def __init__(self, runtime): + self._runtime = runtime + self.metadata = None + + def should_emit(self, hook): + return True + + def prepare_output(self, *, hook, output_index, output_id, output_spec, + output): + from dmi.api.v1 import ProducerPlanBuilder + + entry = ProducerPlanBuilder().record_output( + output_id=output_id, output_spec=output_spec, output=output) + return self._runtime.emit_output(entry, self.metadata, output) + + +def _metadata(step: int, tensor: torch.Tensor): + from dmi.storage.capture import CaptureMetadata + + return CaptureMetadata( + capture_id=f"policy-{step:04d}", tenant_id="t", experiment_id="e", + run_id="r", session_id="s", request_id=f"q{step}", + sequence_id="n", model_id="m", model_revision="mr", + adapter_revision=None, capture_policy_version="v", + hook_name="capture_tensor", layer_number=0, producer_rank=0, + step_number=step, token_start=step, token_end=step + 1, + batch_position=0, dtype=str(tensor.dtype).removeprefix("torch."), + shape=tuple(tensor.shape), + captured_at_ns=1_700_000_000_000_000_000 + step, + ) + + +def _ring_config(payload_bytes: int): + from dmi.api.v1 import RingConfig + + config = RingConfig() + config.task_ring_entries = 64 + config.payload_ring_bytes = payload_bytes + config.pinned_staging_bytes = payload_bytes + # Drain every record as soon as it is published, so a refusal reaches + # the worker while the steps are still running, as under serving load. + config.drain_flush_entry_threshold = 1 + return config + + +def _run_steps(record_sink, *, policy, budget_ms, payload_bytes=4096, + steps=STEPS, elements=RECORD_ELEMENTS, pace_s=0.0): + """Fire one hook per step; return per-step wall times, the first error + a hook call raised (if any), the status, and the flush outcome.""" + from dmi.api.v1 import ( + HookPointV1, HookSpecV1, MonitoringEngine, TransportSpec, + ) + from dmi.storage.capture import CaptureRecordFormat + + engine = MonitoringEngine(model_id="policy-gpu", + ring_config=_ring_config(payload_bytes)) + outcome = {"step_s": [], "errors": [], "begin_errors": [], + "flush_error": None} + try: + runtime = engine.create_record_runtime( + CaptureRecordFormat(), record_sink=record_sink, + failure_policy=policy, step_stall_budget_ms=budget_ms) + hook = HookPointV1( + HookSpecV1("capture_tensor", (TransportSpec("payload"),))) + hook_runtime = _HookRuntime(runtime) + runtime.bind_hook(hook, hook_runtime=hook_runtime) + tensor = torch.arange(elements, dtype=torch.float32, device="cuda") + for step in range(steps): + hook_runtime.metadata = _metadata(step, tensor) + try: + runtime.begin_step() + except RuntimeError as exc: + outcome["begin_errors"].append((step, str(exc))) + started = time.monotonic() + try: + hook(tensor + step) + torch.cuda.current_stream().synchronize() + except RuntimeError as exc: + outcome["errors"].append((step, str(exc))) + outcome["step_s"].append(time.monotonic() - started) + time.sleep(pace_s) + outcome["status"] = engine.capture_status() + try: + engine.flush_and_wait(30.0) + except RuntimeError as exc: + outcome["flush_error"] = str(exc) + outcome["flushed_status"] = engine.capture_status() + finally: + engine.close() + return outcome + + +def _slow_sink(target): + from dmi.storage.capture import CaptureRecordFormat + from dmi.transport import native + + return native.ReferencePythonCaptureSink( + target, CaptureRecordFormat.LAYOUT_NAME, + admission_bound_s=target.admission_s) + + +def test_disable_capture_skips_stalled_steps_and_keeps_capturing(caplog): + """A 4 KiB ring against a sink that takes 400 ms per record: every few + steps a reservation waits for the sink and the 50 ms budget runs out. + That step is skipped (its record, and those still queued for the sink, + are discarded), no step waits longer than the budget plus one + admission, and the next step captures again.""" + target = _Target(admission_s=ADMISSION_S) + with caplog.at_level(logging.WARNING, logger="dmi.engine"): + outcome = _run_steps(_slow_sink(target), policy="disable_capture", + budget_ms=BUDGET_MS) + + assert outcome["errors"] == [] + assert outcome["begin_errors"] == [] + worst = max(outcome["step_s"]) + assert worst < BUDGET_MS / 1000 + ADMISSION_S + 0.3, outcome["step_s"] + status = outcome["flushed_status"] + assert status["capture_active"] is True, status + assert status["failure"] is None + assert status["failure_policy"] == "disable_capture" + assert status["skipped_steps"] >= 1 + assert status["stall_budget_exhaustions"] == status["skipped_steps"] + # A skip also drops the records still queued from earlier steps. + assert status["steps_with_discards"] >= status["skipped_steps"] + assert status["step_stall_budget_ms"] == BUDGET_MS + assert status["max_step_wait_s"] >= BUDGET_MS / 1000 + # Capture resumed after a skip, and every record was either stored or + # discarded by one. + assert target.submitted >= 2 + assert target.submitted + status["discarded_payloads"] == STEPS + assert status["discarded_descriptors"] == status["discarded_payloads"] + # A skipped step is not a failure: the flush succeeds. + assert outcome["flush_error"] is None + assert any("stall budget" in r.getMessage() for r in caplog.records) + + +def test_raise_skips_the_spent_step_and_raises_at_the_next_begin_step(): + target = _Target(admission_s=ADMISSION_S) + outcome = _run_steps(_slow_sink(target), policy="raise", + budget_ms=BUDGET_MS) + + assert outcome["begin_errors"], "the next step must raise under 'raise'" + first_step, message = outcome["begin_errors"][0] + assert "stall budget" in message + # The step that ran out of budget did not raise inside its forward, and + # was bounded by the budget plus one admission. + assert all(step >= first_step for step, _ in outcome["errors"]) + spent = first_step - 1 + assert outcome["step_s"][spent] >= BUDGET_MS / 1000 + assert outcome["step_s"][spent] < BUDGET_MS / 1000 + ADMISSION_S + 0.3 + # Failed from then on: every later step boundary raises. + assert [step for step, _ in outcome["begin_errors"]] == list( + range(first_step, STEPS)) + assert outcome["status"]["capture_active"] is False + assert "stall budget" in outcome["flush_error"] + + +def test_a_budget_is_refused_until_a_step_has_begun(): + """Without begin_step the budget would span the runtime's life; the + first reservation says so instead of silently spending it.""" + from dmi.api.v1 import ( + HookPointV1, HookSpecV1, MonitoringEngine, TransportSpec, + ) + from dmi.storage.capture import CaptureRecordFormat + + engine = MonitoringEngine(model_id="policy-gpu", + ring_config=_ring_config(4096)) + try: + runtime = engine.create_record_runtime( + CaptureRecordFormat(), record_sink=_slow_sink(_Target()), + failure_policy="disable_capture", step_stall_budget_ms=BUDGET_MS) + hook = HookPointV1( + HookSpecV1("capture_tensor", (TransportSpec("payload"),))) + hook_runtime = _HookRuntime(runtime) + runtime.bind_hook(hook, hook_runtime=hook_runtime) + tensor = torch.arange(RECORD_ELEMENTS, dtype=torch.float32, + device="cuda") + hook_runtime.metadata = _metadata(0, tensor) + with pytest.raises(RuntimeError, match="begin_step"): + hook(tensor) + runtime.begin_step() + hook_runtime.metadata = _metadata(1, tensor) + hook(tensor) + engine.flush_and_wait(30.0) + finally: + engine.close() + + +def test_a_budget_is_refused_for_a_sink_with_no_admission_bound(): + from dmi.api.v1 import MonitoringEngine + from dmi.storage.capture import CaptureRecordFormat + from dmi.transport import native + + engine = MonitoringEngine(model_id="policy-gpu", + ring_config=_ring_config(4096)) + try: + unbounded = native.ReferencePythonCaptureSink( + _Target(), CaptureRecordFormat.LAYOUT_NAME) + assert unbounded.admission_bound_s is None + with pytest.raises(ValueError, match="admission bound"): + engine.create_record_runtime( + CaptureRecordFormat(), record_sink=unbounded, + failure_policy="disable_capture", + step_stall_budget_ms=BUDGET_MS) + finally: + engine.close() + + +def test_disable_capture_turns_a_sink_refusal_into_stopped_capture(): + target = _Target(refuse_from=3) + outcome = _run_steps(_slow_sink(target), policy="disable_capture", + budget_ms=None, payload_bytes=64 * 1024, + pace_s=0.02) + + assert outcome["errors"] == [] + status = outcome["status"] + assert status["capture_active"] is False + assert "sink refused durable admission" in status["failure"] + assert status["stall_budget_exhaustions"] == 0 + assert target.submitted == 3 + # The refused fourth record and the twelve after it. + assert outcome["flushed_status"]["discarded_payloads"] == STEPS - 3 + assert "sink refused durable admission" in outcome["flush_error"] + + +def test_raise_surfaces_a_sink_refusal_in_a_later_forward(): + target = _Target(refuse_from=3) + outcome = _run_steps(_slow_sink(target), policy="raise", budget_ms=None, + payload_bytes=64 * 1024, pace_s=0.02) + + assert outcome["errors"] + # The refusal happens on the worker after its record's forward returned, + # so it raises from a LATER step, and from every step after that. + first_step = outcome["errors"][0][0] + assert first_step > 3 + assert [step for step, _ in outcome["errors"]] == list( + range(first_step, STEPS)) + assert all("sink refused durable admission" in message + for _, message in outcome["errors"]) + # A step boundary after the latch raises it too, before any forward. + assert outcome["begin_errors"] + assert all("sink refused durable admission" in message + for _, message in outcome["begin_errors"]) + assert outcome["status"]["capture_active"] is False + + +@pytest.mark.skipif( + not SINK_BUILT, + reason="native/build/_dmi_native_sink*.so is not built; run " + "`make -C native cpu-goals PYTHON=/bin/python`") +def test_a_burst_four_times_the_sink_queue_is_stored_on_a_record_ring( + tmp_path): + """64 x 1 MiB records back to back through the real NativePackSink and + its default config (16 MiB queue, block with 2 s): the record worker + waits for room, nothing is refused, capture stays active. Under the + binding's old drop_newest default this burst dropped a record after + 18-25 and latched the ring.""" + from dmi.storage.capture.native_sink import create_native_pack_sink + from dmi.storage.native_capture import NativeSinkConfig + + handle = create_native_pack_sink( + NativeSinkConfig(spool_root=str(tmp_path / "spool"))) + sink = handle.native_sink + outcome = _run_steps(sink, policy="disable_capture", budget_ms=2000, + payload_bytes=128 << 20, steps=64, + elements=(1 << 20) // 4) + + assert outcome["errors"] == [] + assert outcome["flush_error"] is None + status = outcome["flushed_status"] + assert status["capture_active"] is True, status + assert status["sink"]["persisted_records"] == 64, status["sink"] + assert status["sink"]["dropped_records"] == 0 + assert status["sink"]["timed_out_records"] == 0 + + +@pytest.mark.skipif( + not SINK_BUILT, + reason="native/build/_dmi_native_sink*.so is not built; run " + "`make -C native cpu-goals PYTHON=/bin/python`") +def test_a_record_lost_after_admission_stops_capture_and_says_why(tmp_path): + """A 1 MiB record into 1 MiB packs: admitted, then dropped as oversized + on the pack worker. The sink reports it at its next submit and at the + flush, so under disable_capture capture stops with the reason, nothing + raises in a forward, and the flush raises.""" + from dmi.storage.capture.native_sink import create_native_pack_sink + from dmi.storage.native_capture import NativeSinkConfig + + handle = create_native_pack_sink(NativeSinkConfig( + spool_root=str(tmp_path / "spool"), max_pack_bytes=1 << 20, + max_queue_bytes=4 << 20)) + outcome = _run_steps(handle.native_sink, policy="disable_capture", + budget_ms=None, payload_bytes=8 << 20, steps=3, + elements=(1 << 20) // 4, pace_s=0.05) + + assert outcome["errors"] == [] + assert outcome["flush_error"] is not None + assert "oversized_records" in outcome["flush_error"] + status = outcome["flushed_status"] + assert status["capture_active"] is False + assert "oversized_records" in status["failure"] + assert status["sink"]["persisted_records"] == 0