From 13ad1043272eea414c034f0d9fc6579f9d38e83f Mon Sep 17 00:00:00 2001 From: Alan Liu Date: Thu, 24 Sep 2026 11:56:44 -0400 Subject: [PATCH 01/14] Let a record ring stop capture instead of raising when its sink fails Decision 3 wants vLLM serving to keep running when capture storage refuses or fails: capture stops and says why. Until now the record consumer had one behaviour, raise_at_producer: the first sink refusal (a dropped, timed-out or oversized record, or a pipeline failure) latched, and the next descriptor push raised inside the model forward. RecordConsumer now takes a RecordFailurePolicy. raise_at_producer stays the default and is unchanged. Under disable_capture the latch drops the descriptors still queued, and every push and payload after it is discarded and counted instead of raised; the p2p worker keeps draining, so the ring keeps its protocol. The failure still surfaces at wait_until_idle, finish and rethrow_if_failed, and in a new snapshot() (policy, failure text, discarded descriptors and payloads). The policy is threaded through RingEngine and RecordP2PThread; nothing selects it yet. Test-first in tests/native/ring/test_record_consumer.cpp: with the snapshot API in place but no policy behaviour, the three new cases gave 12 FAILs and then an uncaught throw from a push after the latch; with the policy, 48 passed, 0 failed. --- native/csrc/ring/p2p_thread.cpp | 5 +- native/csrc/ring/p2p_thread.h | 4 +- native/csrc/ring/record_consumer.cpp | 83 +++++++++++++--- native/csrc/ring/record_consumer.h | 29 +++++- native/csrc/ring/record_failure_policy.h | 19 ++++ native/csrc/ring/ring_engine.cu | 5 +- native/csrc/ring/ring_engine.h | 4 +- tests/native/ring/test_record_consumer.cpp | 107 +++++++++++++++++++++ 8 files changed, 237 insertions(+), 19 deletions(-) create mode 100644 native/csrc/ring/record_failure_policy.h diff --git a/native/csrc/ring/p2p_thread.cpp b/native/csrc/ring/p2p_thread.cpp index 9c0001d9a..fcb66bcc3 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(); 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..fde2ec756 100644 --- a/native/csrc/ring/record_consumer.cpp +++ b/native/csrc/ring/record_consumer.cpp @@ -13,35 +13,77 @@ 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(); +} void RecordConsumer::push_descriptor(RecordDescriptor descriptor) { std::lock_guard lock(mu_); - if (failure_) std::rethrow_exception(failure_); + if (failure_) { + if (policy_ != RecordFailurePolicy::kDisableCapture) { + std::rethrow_exception(failure_); + } + ++discarded_descriptors_; + return; + } descriptors_.push_back(std::move(descriptor)); } 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; + } for (auto& descriptor : descriptors) { descriptors_.push_back(std::move(descriptor)); } } 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_) { + if (!disable) std::rethrow_exception(failure_); + // Capture is off: the drain still delivers what the forward + // produced, and the payload is dropped here, never submitted. + ++discarded_payloads_; + 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()); descriptors_.pop_front(); @@ -73,11 +115,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 +130,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. @@ -131,4 +176,20 @@ 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_; + return snapshot; +} + } // namespace ring diff --git a/native/csrc/ring/record_consumer.h b/native/csrc/ring/record_consumer.h index 721892dcb..f198fc748 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 @@ -20,9 +21,23 @@ namespace ring { +struct RecordConsumerSnapshot { + RecordFailurePolicy policy{RecordFailurePolicy::kRaiseAtProducer}; + bool failed{false}; + std::string failure; + // Descriptors dropped by the latch: those still queued when it happened + // plus every push after it. kDisableCapture only. + uint64_t discarded_descriptors{0}; + // Payloads not stored because of the latch: the one whose submission + // failed, plus every payload delivered after it. kDisableCapture only. + uint64_t discarded_payloads{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 +48,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 +60,24 @@ 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; private: + // Retain the first failure; under kDisableCapture also drop the queued + // descriptors. Caller holds mu_. + void latch_locked(std::exception_ptr failure); + std::shared_ptr sink_; + const RecordFailurePolicy policy_; mutable std::mutex mu_; mutable std::condition_variable idle_cv_; std::deque descriptors_; std::exception_ptr failure_; size_t active_payloads_{0}; + uint64_t discarded_descriptors_{0}; + uint64_t discarded_payloads_{0}; }; } // 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/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/tests/native/ring/test_record_consumer.cpp b/tests/native/ring/test_record_consumer.cpp index 5ab375efb..5ccef9bd3 100644 --- a/tests/native/ring/test_record_consumer.cpp +++ b/tests/native/ring/test_record_consumer.cpp @@ -192,6 +192,110 @@ 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(); })); +} + int main() { setbuf(stdout, nullptr); std::printf("test_record_consumer\n"); @@ -199,6 +303,9 @@ 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(); std::printf("Results: %d passed, %d failed\n", g_pass, g_fail); return g_fail == 0 ? 0 : 1; } From 0a280cd622014bdb439772af6d6f2e7000216c5c Mon Sep 17 00:00:00 2001 From: Alan Liu Date: Thu, 24 Sep 2026 11:56:53 -0400 Subject: [PATCH 02/14] Bound the forward's wait for record ring space by a per-step budget When a record reservation does not fit, reserve_record synchronised the stream and called the unbounded force_flush_and_wait. The drain is held back by the record worker and the worker by the sink, so a slow or stuck sink stalled the model forward for as long as the sink took. RingEnginePy takes RecordRuntimeOptions {failure_policy, step_stall_budget_ms}. Both reserve slow paths now wait through force_flush_and_wait_until, bounded by what is left of the step's budget (begin_record_step starts a step). Past it the policy applies: raise_at_producer raises from the reservation, with nothing reserved; disable_capture latches the failure in the consumer, which then discards instead of submitting, so the remaining wait is at most the one sink admission already in progress, and the reservation completes. Under raise_at_producer a latched failure now raises at the start of reserve_record, before a reservation no producer will use, rather than from the push after it. A budget of 0 keeps the old unbounded wait and is the default. record_capture_status() reports the policy, the failure, the discard counters, budget exhaustions and reserve wait time (total and worst step). Test-first in tests/native/ring/test_ring_engine.cu, on GPU 1: a 4 KiB ring filled twice while a sink sits in a 400 ms admission, 50 ms budget. Red was a compile failure (no options/status API); a mutation that ignores the budget fails 13 checks (the raise case waits the full admission queue and never raises; the disable case stalls past budget + one admission). Green: 138 passed, 0 failed. --- native/csrc/ring/ring_engine_py.cu | 140 +++++++++++++++++++++-- native/csrc/ring/ring_engine_py.h | 45 +++++++- tests/native/ring/test_ring_engine.cu | 159 ++++++++++++++++++++++++++ 3 files changed, 330 insertions(+), 14 deletions(-) diff --git a/native/csrc/ring/ring_engine_py.cu b/native/csrc/ring/ring_engine_py.cu index 78fe08187..08b9b471b 100644 --- a/native/csrc/ring/ring_engine_py.cu +++ b/native/csrc/ring/ring_engine_py.cu @@ -12,6 +12,7 @@ #include "ring/ring_debug.h" #include // at::cuda::getCurrentCUDAStream #include +#include #include #include #include @@ -138,6 +139,14 @@ struct RingEnginePy::Impl { at::Tensor payload_view; bool record_mode{false}; + // Record rings only. The step counter is written by the reserving + // thread; the atomics let record_capture_status read from any thread. + RecordRuntimeOptions record_options; + 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}; + Impl(ring::RingConfig cfg, SubmitFn sf) : engine(std::move(cfg), fifo, std::move(sf)) { @@ -150,8 +159,10 @@ 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) { const auto& state = engine.ring_state(); int dev_idx = 0; @@ -161,6 +172,75 @@ 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 failure policy + // applies. + 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(); + // Once capture is off the worker discards instead of submitting, so + // the sink is out of the path and only the device-to-host copy of + // what the ring already holds is left to wait for. + if (record_options.step_stall_budget_ms == 0 || + (disable && consumer.failed())) { + drain.force_flush_and_wait_until(FlushClock::time_point::max()); + account_record_wait(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); + const auto step_ms = + step_wait_ns.load(std::memory_order_relaxed) / 1'000'000ull; + const std::exception_ptr failure = + std::make_exception_ptr(std::runtime_error( + "record capture stall budget exhausted: reservations waited " + + std::to_string(step_ms) + " ms this step for the sink to " + "free ring space (step_stall_budget_ms=" + + std::to_string(record_options.step_stall_budget_ms) + ")")); + // Latched in the consumer either way, so the checked flush, the + // capture status and close report it. + consumer.record_failure(failure); + if (!disable) std::rethrow_exception(failure); + // Capture is off now, so the drain is held back only by the one sink + // admission already in progress (bounded by the sink's own admission + // timeout). The reservation must still complete: under CUDA-graph + // replay the producers launch whatever the host decides. + const auto resumed = FlushClock::now(); + drain.force_flush_and_wait_until(FlushClock::time_point::max()); + account_record_wait(resumed); + } }; // --------------------------------------------------------------------------- @@ -187,13 +267,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 +582,17 @@ 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 (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 +613,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 +625,35 @@ 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"); + } + impl_->step_wait_ns.store(0, std::memory_order_relaxed); +} + +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.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) { diff --git a/native/csrc/ring/ring_engine_py.h b/native/csrc/ring/ring_engine_py.h index 76fc307a5..ca5ae5235 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,36 @@ 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 (see + // begin_record_step) may wait for the drain to free ring space. Past + // it the policy applies: kRaiseAtProducer raises from the reservation, + // kDisableCapture latches the failure and stops capture. 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; + // 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 +204,17 @@ 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. + 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/tests/native/ring/test_ring_engine.cu b/tests/native/ring/test_ring_engine.cu index f3cb1ada2..3f6c2fcff 100644 --- a/tests/native/ring/test_ring_engine.cu +++ b/tests/native/ring/test_ring_engine.cu @@ -20,6 +20,7 @@ #include #include #include +#include #include #include @@ -776,6 +777,162 @@ 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. +class SlowRecordSink final : public ring::RecordSink { +public: + explicit SlowRecordSink(std::chrono::milliseconds admission) + : admission_(admission) {} + + 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::atomic submissions{0}; + +private: + std::chrono::milliseconds admission_; +}; + +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 later_reserve_threw{false}; + std::chrono::steady_clock::duration later_reserve{}; + bool flush_threw{false}; + 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. +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(); + + engine.begin_record_step(); + 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&) { + outcome.flush_threw = true; + } + engine.stop(); + outcome.submissions = sink->submissions.load(std::memory_order_acquire); + CUDA_CHECK(cudaFree(device)); + return outcome; +} + +static void test_raise_policy_bounds_a_stalled_reservation() { + banner("raise policy: a reservation stalled past the budget raises"); + const StallOutcome outcome = + run_stalled_reservation(ring::RecordFailurePolicy::kRaiseAtProducer); + EXPECT(outcome.threw); + EXPECT(outcome.error.find("stall budget") != std::string::npos); + // The budget, not the sink's 400 ms admission, bounds the wait. + EXPECT(outcome.stall < std::chrono::milliseconds(300)); + EXPECT(outcome.status.failed); + EXPECT(outcome.status.stall_budget_exhaustions == 1); + EXPECT(outcome.status.step_stall_budget_ms == 50); + // Latched: the next reservation raises before it waits for anything. + EXPECT(outcome.later_reserve_threw); + EXPECT(outcome.later_reserve < std::chrono::milliseconds(50)); + EXPECT(outcome.flush_threw); +} + +static void test_disable_capture_bounds_a_stalled_reservation() { + banner("disable_capture: a stalled reservation stops capture, not the forward"); + 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; once + // capture is off the worker discards instead of calling the sink. + 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.failure.find("stall budget") != std::string::npos); + EXPECT(outcome.status.stall_budget_exhaustions == 1); + EXPECT(outcome.status.max_step_wait_ns >= 50'000'000ull); + // Capture is off: the next step reserves without raising or stalling. + EXPECT(!outcome.later_reserve_threw); + EXPECT(outcome.later_reserve < std::chrono::milliseconds(200)); + // The failure still surfaces at the checked flush. + EXPECT(outcome.flush_threw); + // Only the record the sink was already admitting reached it. + EXPECT(outcome.submissions == 1); +} + int main() { setbuf(stdout, nullptr); ring::set_ring_null_mode(false); @@ -797,6 +954,8 @@ 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_bounds_a_stalled_reservation(); + test_disable_capture_bounds_a_stalled_reservation(); std::printf("Results: %d passed, %d failed\n", g_pass, g_fail); return g_fail == 0 ? 0 : 1; From b9731cc0c9dd466eba3373690f2f54512a2eab80 Mon Sep 17 00:00:00 2001 From: Alan Liu Date: Thu, 24 Sep 2026 11:59:53 -0400 Subject: [PATCH 03/14] Admit ring-fed bursts by blocking, and move NativeSinkConfig to live code The ring-fed native sink was built with the C++ SinkConfig admission default, drop_newest, which the Python binding did not let anyone change. Measured on this branch: one envelope of 64 x 1 MiB rows against the default 16 MiB queue dropped a record after 18-25 were admitted (three trials), and the refusal latches the record ring. bindings_sink.cpp now takes `overload` ("block" | "drop_newest") and `admission_timeout_s` (None waits without bound), exposes both as properties, and adds timed_out_records and rejected_closed_records to the snapshot. The binding's defaults and the C++ SinkConfig stay drop_newest with no timeout, so PackSink parity with the reference is untouched. NativeSinkConfig, the ring-fed entry point, moves to dmi.storage.native_capture (live code) with a re-export at the old path, gains the two fields with the plan's default of block with 2 s, and validates them. Its docstring no longer claims a field-by-field match with the reference pipeline config, which it never had. MonitoringConfig and the engine now import it from the live module, so the public config does not load the backup capture package (checked in a fresh interpreter). Test-first: tests/test_native_sink_admission.py was 19 failed (missing type, attributes and constructor arguments), then 19 passed; the burst test persists all 64 with no drop or timeout. The neighbouring cpu suites (rollback, storage wiring, backend choices, engine runtime API, adapter, capture storage, v1 API, pack sink, pack sink timeout): 333 passed, 1 skipped (no CUDA device). --- docs/integration-api-v1.md | 8 +- native/csrc/sink/bindings_sink.cpp | 50 ++++++- native/csrc/sink/pack_sink.h | 2 + src/dmi/config.py | 4 +- src/dmi/engine.py | 2 +- src/dmi/storage/capture/native_sink.py | 38 +---- src/dmi/storage/native_capture.py | 63 ++++++++ tests/test_native_sink_admission.py | 194 +++++++++++++++++++++++++ 8 files changed, 322 insertions(+), 39 deletions(-) create mode 100644 tests/test_native_sink_admission.py diff --git a/docs/integration-api-v1.md b/docs/integration-api-v1.md index e5e1ec0a9..a1002c4a9 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`: diff --git a/native/csrc/sink/bindings_sink.cpp b/native/csrc/sink/bindings_sink.cpp index 427edd63d..2c81ec862 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 { @@ -123,8 +146,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 +172,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 +207,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 +227,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/pack_sink.h b/native/csrc/sink/pack_sink.h index f9c48bb88..9c7a800c2 100644 --- a/native/csrc/sink/pack_sink.h +++ b/native/csrc/sink/pack_sink.h @@ -152,6 +152,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/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..a0a287efa 100644 --- a/src/dmi/engine.py +++ b/src/dmi/engine.py @@ -135,7 +135,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( 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 db2c90b6e..cef5c50e4 100644 --- a/src/dmi/storage/native_capture.py +++ b/src/dmi/storage/native_capture.py @@ -55,6 +55,67 @@ def _load_native_store_extension() -> Any: ) from exc +# 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; ``"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). + + 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 _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__}") @@ -415,6 +476,8 @@ def read( __all__ = [ + "SINK_OVERLOAD_POLICIES", + "NativeSinkConfig", "NativeCapture", "NativeCapturePage", "NativeCaptureReader", diff --git a/tests/test_native_sink_admission.py b/tests/test_native_sink_admission.py new file mode 100644 index 000000000..6fef5007e --- /dev/null +++ b/tests/test_native_sink_admission.py @@ -0,0 +1,194 @@ +"""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 + + +# --- 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 From 2b32283b7faacbc2a9d481e47e8595929ceeb73f Mon Sep 17 00:00:00 2001 From: Alan Liu Date: Thu, 24 Sep 2026 12:03:44 -0400 Subject: [PATCH 04/14] Count a record that fits no empty pack as oversized, not a sink failure Admission screens a record's payload against max_pack_bytes, but a pack also holds a header, the record's footer row and a trailer. A payload of exactly max_pack_bytes clears admission and fits no empty pack. The reference assembler raises OversizedRecordError for it and the pipeline counts it oversized and keeps going. The native worker opened a fresh builder, got kCapacity, and sealed that empty builder, so "seal failed: cannot seal an empty pack" failed the whole sink and every record queued behind it; on a ring that is a latched record runtime. Found while checking the framing reserve for validate_capture_bounds: through the torch binding, one 1 MiB float32 record into a 1 MiB pack failed the sink. Now a kCapacity on a builder with no records drops the record and counts it, as the reference does. Test-first in tests/test_native_pack_sink.py (conformance_sink driver): the new case failed with the next record refused as 'closed'; after the fix the sink suites (pack sink, pack sink timeout, sink admission, capture bounds, rollback) are 164 passed. --- native/csrc/sink/pack_sink.cpp | 14 ++++++++++++++ tests/test_native_pack_sink.py | 22 ++++++++++++++++++++++ 2 files changed, 36 insertions(+) diff --git a/native/csrc/sink/pack_sink.cpp b/native/csrc/sink/pack_sink.cpp index b015a2f6f..c0e98415c 100644 --- a/native/csrc/sink/pack_sink.cpp +++ b/native/csrc/sink/pack_sink.cpp @@ -642,6 +642,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/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 From b4c5da4fe4c29ef74b4376a64ccdce6ae9ddb19d Mon Sep 17 00:00:00 2001 From: Alan Liu Date: Thu, 24 Sep 2026 12:03:44 -0400 Subject: [PATCH 05/14] Check capture bounds at attach, and expose the uploader's in-flight budget Three configurations lose captures only once a forward runs: a record larger than the sink's max_queue_bytes is refused on the record worker (latching the record runtime), one that fits the queue but no empty pack is admitted and then dropped as oversized, and a pack larger than the uploader's in-flight budget is staged and never uploaded. All three are decided by configuration and the largest record an adapter emits. validate_capture_bounds(sink_config, max_record_bytes, storage_config=) in dmi.storage.native_capture refuses each with a ConfigurationError that names the bound to raise; max_pack_bytes must leave PACK_FRAMING_RESERVE_BYTES (64 KiB) of framing room above the record. MonitoringEngine.validate_capture_bounds(max_record_bytes) applies it to the engine's own configs, for adapters to call at attach (no adapter drives the persistent path yet; D2 wires it). With no sink config there is nothing to check. NativeCaptureStorageConfig gains uploader_max_in_flight_bytes (default 256 MiB, the native UploaderConfig default). bindings_store.cpp already read the key, but _native_dict never sent it. Test-first: tests/test_capture_bounds.py was 16 failed (missing field, function and method), then 16 passed; with the storage wiring and sink admission suites, 72 passed. --- src/dmi/engine.py | 24 +++++ src/dmi/storage/native_capture.py | 66 ++++++++++++++ tests/test_capture_bounds.py | 142 ++++++++++++++++++++++++++++++ 3 files changed, 232 insertions(+) create mode 100644 tests/test_capture_bounds.py diff --git a/src/dmi/engine.py b/src/dmi/engine.py index a0a287efa..6bd3ca6fb 100644 --- a/src/dmi/engine.py +++ b/src/dmi/engine.py @@ -266,6 +266,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')``.""" diff --git a/src/dmi/storage/native_capture.py b/src/dmi/storage/native_capture.py index cef5c50e4..c8f1164bc 100644 --- a/src/dmi/storage/native_capture.py +++ b/src/dmi/storage/native_capture.py @@ -169,6 +169,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 def __post_init__(self) -> None: for name in ("s3_endpoint", "s3_bucket", "s3_access_key", @@ -203,6 +207,9 @@ def __post_init__(self) -> None: _positive("clickhouse_request_timeout_s", self.clickhouse_request_timeout_s, float) _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): @@ -226,9 +233,66 @@ 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, } +# 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.""" @@ -476,6 +540,7 @@ def read( __all__ = [ + "PACK_FRAMING_RESERVE_BYTES", "SINK_OVERLOAD_POLICIES", "NativeSinkConfig", "NativeCapture", @@ -484,4 +549,5 @@ def read( "NativeCaptureSelection", "NativeCaptureStorage", "NativeCaptureStorageConfig", + "validate_capture_bounds", ] 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) From 3aca46bc29c47b5da7a095c73c4c9e49f08189ac Mon Sep 17 00:00:00 2001 From: Alan Liu Date: Thu, 24 Sep 2026 12:10:55 -0400 Subject: [PATCH 06/14] Choose a record runtime's failure policy, and report its capture status The native ring can now stop capture instead of raising and bound the forward's wait for a slow sink, but nothing above it could ask for either, and nothing reported a stopped capture short of flush raising. create_record_runtime(failure_policy="raise" | "disable_capture", step_stall_budget_ms=None | int) passes both to RingEngine.create_record. The default stays "raise" with no budget, so existing callers are unchanged; C2's YAML will choose (decision 3: disable_capture for vLLM serving). Bad values are refused before the live ring is touched, and disable_capture is refused with a NativeSinkConfig whose block admission has no timeout, since the post-latch wait is one sink admission. RecordRuntime.begin_step() starts a stall-budget step; integrations call it once per model step (it is explicit because piecewise CUDA graphs replay several plans per step). MonitoringEngine.capture_status() returns the policy, whether capture is still active and why not, discard counters, budget exhaustions, the reserve wait (total and worst step) and the native sink and storage snapshots. close(), and enable_ring_transport replacing a record ring, log a stopped capture at WARNING. The integration API document describes the policies, the stall bound (budget plus one admission timeout), the status fields and validate_capture_bounds. Test-first: tests/test_record_failure_policy.py was 14 failed, then 14 passed; the v1 doc-coverage test, extended with begin_step, capture_status and validate_capture_bounds, failed until they were documented. Seven create_record fakes in the engine and storage wiring suites now accept the options as keywords. Engine suites (runtime API, storage wiring, record runtime, v1 API, backend choices, HF refusal, policy): 162 passed. --- docs/integration-api-v1.md | 55 ++++- native/csrc/bindings.cpp | 47 ++++- src/dmi/engine.py | 132 +++++++++++- src/dmi/records.py | 11 + src/dmi/transport/ring.py | 5 + tests/test_engine_runtime_api.py | 14 +- tests/test_integration_api_v1.py | 3 + tests/test_native_capture_storage_wiring.py | 3 +- tests/test_record_failure_policy.py | 210 ++++++++++++++++++++ 9 files changed, 466 insertions(+), 14 deletions(-) create mode 100644 tests/test_record_failure_policy.py diff --git a/docs/integration-api-v1.md b/docs/integration-api-v1.md index a1002c4a9..6fcc58b1f 100644 --- a/docs/integration-api-v1.md +++ b/docs/integration-api-v1.md @@ -387,6 +387,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, @@ -394,14 +401,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 @@ -435,6 +445,49 @@ 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) 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 record reservation or descriptor push raises it, inside the forward that publishes the record. | +| `"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 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. `RecordRuntime.begin_step()` starts a step; call it once +per model step, before the step's first reservation, or the budget spans the +runtime's whole life. Past the budget the policy applies: `"raise"` raises +from the reservation with nothing reserved, and `"disable_capture"` latches +and completes the reservation after at most the one sink admission already +in progress. Under `"disable_capture"` the forward's stall per step is +therefore at most the budget plus one `admission_timeout_s`, and +`create_record_runtime()` refuses it with a `NativeSinkConfig` whose `block` +admission has no timeout. `None`, the default, waits without bound. Ring +failures that are not sink refusals, such as a CUDA error in the drain, raise +under both policies. + +`MonitoringEngine.capture_status()` returns a plain dict: `record_mode`, +`capture_active` (False once a failure latched), `failure_policy`, +`failure`, `discarded_descriptors`, `discarded_payloads`, +`step_stall_budget_ms`, `stall_budget_exhaustions`, `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..ebef1a9d2 100644 --- a/native/csrc/bindings.cpp +++ b/native/csrc/bindings.cpp @@ -157,6 +157,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) { @@ -678,7 +694,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 +715,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 +734,24 @@ 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["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/src/dmi/engine.py b/src/dmi/engine.py index 6bd3ca6fb..9901e4f66 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") @@ -339,12 +349,27 @@ 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, or a sink failure) does to the forward: + ``"raise"`` raises from the next 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 (see ``RecordRuntime.begin_step``) may wait for the sink to + free ring space; past it the policy applies. ``None`` waits without + bound. Under ``"disable_capture"`` the forward's stall per step is + then at most the budget plus one sink admission timeout. """ if getattr(self, "_storage_backend", "auto") == "none": @@ -353,6 +378,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) transport = self._ring_transport ring_config = self._ring_config if transport is None or ring_config is None: @@ -375,7 +402,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 @@ -383,6 +412,34 @@ def create_record_runtime( raise return runtime + def _validate_record_failure_options( + self, failure_policy: Any, step_stall_budget_ms: 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") + sink_config = self._capture_sink_config + if ( + failure_policy == "disable_capture" + and self._storage_backend == "persistent" + and sink_config is not None + and sink_config.overload == "block" + and sink_config.admission_timeout_s is None + ): + # Once capture latches, the forward still waits for the one sink + # admission in progress. Without a timeout that wait has no bound. + raise ValueError( + "failure_policy='disable_capture' needs a bounded sink " + "admission: set capture_sink_config.admission_timeout_s, or " + "overload='drop_newest'") + def _start_capture_storage(self, *, sweep_spool: bool) -> Optional[Any]: config = self._capture_storage_config if config is None or self._storage_backend != "persistent": @@ -405,6 +462,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 @@ -458,6 +518,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() @@ -466,6 +528,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: @@ -483,6 +546,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 @@ -508,6 +572,64 @@ 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 + ``"disable_capture"`` the forward keeps running and ``failure`` + says why capture stopped; ``discarded_*`` count what was dropped + after it. ``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, + "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"]), + 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, 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"] or 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( *, @@ -573,6 +695,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: @@ -589,6 +713,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) @@ -680,6 +805,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: @@ -699,6 +826,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) @@ -716,4 +844,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..314e3b6d4 100644 --- a/src/dmi/records.py +++ b/src/dmi/records.py @@ -309,6 +309,17 @@ 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 for the stall budget. + + The record reservations after this call share one fresh + ``step_stall_budget_ms`` (see ``create_record_runtime``). Call it + once per step, before the step's first ``emit_output`` or + ``prepare_replay``; without it the budget spans the runtime's life. + """ + + self._transport.begin_record_step() + def emit_output( self, entry: ProducerPlanEntry, diff --git a/src/dmi/transport/ring.py b/src/dmi/transport/ring.py index 27d61592f..503adf2d1 100644 --- a/src/dmi/transport/ring.py +++ b/src/dmi/transport/ring.py @@ -420,6 +420,11 @@ 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 fresh stall budget for the next reservations.""" + + 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/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_storage_wiring.py b/tests/test_native_capture_storage_wiring.py index ce59315b4..15d3efee7 100644 --- a/tests/test_native_capture_storage_wiring.py +++ b/tests/test_native_capture_storage_wiring.py @@ -269,7 +269,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_record_failure_policy.py b/tests/test_record_failure_policy.py new file mode 100644 index 000000000..73428a1ff --- /dev/null +++ b/tests/test_record_failure_policy.py @@ -0,0 +1,210 @@ +"""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, + "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 + + +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 + + +# --- 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(), failure_policy="disable_capture", + step_stall_budget_ms=2000) + assert created == [{"failure_policy": "disable_capture", + "step_stall_budget_ms": 2000}] + + +@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 + + +def test_disable_capture_needs_a_bounded_sink_admission(tmp_path): + """Once the budget latches, the forward still waits for the one sink + admission in progress; unbounded, that wait is the stall the policy + exists to prevent.""" + 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="disable_capture") + + +# --- 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, + "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, + "reserve_wait_ns": 2_500_000_000, "max_step_wait_ns": 2_100_000_000, + }) + engine, _runtime, _created = _record_engine( + monkeypatch, ring, 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, + "reserve_wait_s": 2.5, "max_step_wait_s": 2.1, + "sink": None, "storage": None, + } + + +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": "record capture stall budget exhausted", + "discarded_descriptors": 4, "discarded_payloads": 5, + "step_stall_budget_ms": 50, "stall_budget_exhaustions": 1, + "reserve_wait_ns": 0, "max_step_wait_ns": 0, + }) + engine, _runtime, _created = _record_engine( + monkeypatch, ring, 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("stall budget exhausted" in message and "5 payloads" in message + for message in messages), messages From fda8554ac60fc3d8b14f52875fa8892e240ebe3b Mon Sep 17 00:00:00 2001 From: Alan Liu Date: Thu, 24 Sep 2026 12:10:55 -0400 Subject: [PATCH 07/14] Pin a 17 MiB row and a 64 x 1 MiB burst through the live capture chain B2's live acceptance: each case is either stored or refused at configuration time, never lost in the forward. Through the real NativePackSink, the storage service, a live ClickHouse catalog and the fake S3: - a 17 MiB row is refused by validate_capture_bounds under the default 16 MiB queue, naming max_queue_bytes; with the queue raised it is stored and reads back exactly; - one envelope of 64 x 1 MiB rows, four times the default queue, is stored whole under NativeSinkConfig's defaults (block, 2 s): no drop, no timeout, every row read back. Written after the implementation; the burst's red was measured on the sink directly (the old drop_newest default refused it after 18-25 rows, three trials). tests/test_native_capture_chain_live.py against local ClickHouse 25.12: 4 passed, unique table prefixes, none left behind. --- tests/test_native_capture_chain_live.py | 93 ++++++++++++++++++++++++- 1 file changed, 90 insertions(+), 3 deletions(-) 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)) From 1f776009bdaec160e7d4a8d67a85ebcaa4109ee0 Mon Sep 17 00:00:00 2001 From: Alan Liu Date: Thu, 24 Sep 2026 12:12:02 -0400 Subject: [PATCH 08/14] Measure sink refusals and stalls from the forward on a real record ring The native suites pin the mechanism; this drives it as an integration does (create_record_runtime, a HookPointV1 on a CUDA tensor, begin_step per step) and checks what serving cares about: whether a hook call raises, and the slowest step's wall time. The stalling sink is a Python target behind the native reference bridge whose admission sleeps 400 ms or refuses; the burst uses the real NativePackSink. - disable_capture, 50 ms budget, 4 KiB ring: no step raises, the worst step is within budget + one admission (measured 0.401 s; 1.606 s with no budget), one record reached the sink, the other 15 were discarded, flush raises the stall-budget failure and close logs it. - raise, 50 ms budget: the stalled step raises at the budget (measured 0.050 s), every later step raises at once, flush raises. - a refusal on the 4th record: disable_capture never raises in a hook and reports "sink refused durable admission"; raise surfaces it from a later step and every step after. - 64 x 1 MiB records through NativePackSink's default config: all 64 persisted, capture active, no drop or timeout. On GPU 1, gated on it being idle: 5 passed, and 5 passed on each of five repeats. The first draft read counters before the drain had delivered anything (3 failures, test-side); the ring now drains each record promptly and the counters are read after the flush. --- tests/test_record_failure_policy_gpu.py | 278 ++++++++++++++++++++++++ 1 file changed, 278 insertions(+) create mode 100644 tests/test_record_failure_policy_gpu.py diff --git a/tests/test_record_failure_policy_gpu.py b/tests/test_record_failure_policy_gpu.py new file mode 100644 index 000000000..64a4c462e --- /dev/null +++ b/tests/test_record_failure_policy_gpu.py @@ -0,0 +1,278 @@ +"""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 burst case uses the real +NativePackSink with its default config. + +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": [], "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) + runtime.begin_step() + 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) + + +def test_disable_capture_keeps_the_forward_running_through_a_stalled_sink( + caplog): + """A 4 KiB ring against a sink that takes 400 ms per record: the ninth + record's reservation waits for the sink. With a 50 ms budget every step + returns, none waits longer than the budget plus one admission, and + capture reports why it stopped.""" + 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"] == [] + worst = max(outcome["step_s"]) + assert worst < BUDGET_MS / 1000 + ADMISSION_S + 0.3, outcome["step_s"] + status = outcome["status"] + assert status["capture_active"] is False + assert "stall budget" in status["failure"] + assert status["failure_policy"] == "disable_capture" + assert status["stall_budget_exhaustions"] == 1 + assert status["step_stall_budget_ms"] == BUDGET_MS + assert status["max_step_wait_s"] >= BUDGET_MS / 1000 + # Only the record the sink was already admitting reached it; everything + # after the latch was dropped on the worker, not submitted. + assert target.submitted == 1 + assert outcome["flushed_status"]["discarded_payloads"] == STEPS - 1 + assert "stall budget" in outcome["flush_error"] + assert any("stall budget" in r.getMessage() for r in caplog.records) + + +def test_raise_fails_fast_when_the_stall_budget_is_spent(): + target = _Target(admission_s=ADMISSION_S) + outcome = _run_steps(_slow_sink(target), policy="raise", + budget_ms=BUDGET_MS) + + assert outcome["errors"], "the stalled step must raise under 'raise'" + first_step, message = outcome["errors"][0] + assert "stall budget" in message + # Raised at the budget, not after the sink's 400 ms admission. + assert outcome["step_s"][first_step] < 0.3, outcome["step_s"] + # Latched: every later step raises at once. + assert [step for step, _ in outcome["errors"]] == list( + range(first_step, STEPS)) + assert max(outcome["step_s"][first_step + 1:]) < 0.1 + assert outcome["status"]["capture_active"] is False + assert "stall budget" in outcome["flush_error"] + + +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"]) + 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 From aeb6321f50c0c004e6d42e4dbd1bfd3c1510d3af Mon Sep 17 00:00:00 2001 From: Alan Liu Date: Thu, 24 Sep 2026 13:12:23 -0400 Subject: [PATCH 09/14] Give a sink envelope one admission deadline, and fail on lost records Three sink-side fixes the record ring's failure policy depends on. One deadline per envelope. The ring's record worker hands NativePackSink one envelope of N rows per submit, and each row's PackSink::Submit started its own admission_timeout_s, so a steadily slow sink held an N-row envelope (and through the drain, the forward) for up to N timeouts. EnvelopeAdmission (record_row.h) takes one deadline when the envelope arrives and admits every row against it, through the new PackSink::SubmitBy; Submit keeps its per-call deadline for its other callers and for parity with the reference. Red first: in tests/native/test_pack_sink_timeout.cpp a stager that takes 0.3 s per pack against a 0.5 s timeout let all 10 rows in over 1.99 s; now the envelope times out at ~0.5 s. Records lost after admission are failures. A record that fits the queue but no empty pack is admitted and then dropped as oversized on the pack worker; flush_and_wait returned True, rethrow_if_failed was silent, and the snapshot said admitted=1 persisted=0. The reference adapter counts oversized (and dropped, timed-out, duplicate, rejected-closed, failures) as losses and raises at flush and rethrow. NativePackSink now keeps its counters at construction and does the same: rethrow_if_failed and a completed flush raise naming the counters that moved, a flush also raises if fewer records were persisted than admitted, and submit checks first, so the ring's record runtime latches on the next envelope instead of storing around the hole. RecordSink::admission_bound(), read by the ring for the stall budget (next commit): zero under drop_newest, the timeout under block, none when block waits forever. The binding exposes it as admission_bound_s on every RecordSink. ClickHouseRecordSink has none, and the reference bridge has none unless its constructor is told what its Python target promises. tests/test_native_sink_admission.py: the bound per policy, the 1 MiB repro (flush, rethrow and the next envelope all raise naming oversized_records), and a stored record still flushing clean. Red: 4 failed; green with the sink and pack-sink suites, 149 passed. --- native/csrc/bindings.cpp | 30 +++++- native/csrc/reference_python_capture_sink.cpp | 6 +- native/csrc/reference_python_capture_sink.h | 12 ++- native/csrc/ring/record_sink.h | 9 ++ native/csrc/sink/bindings_sink.cpp | 11 ++- native/csrc/sink/native_pack_sink.cpp | 62 +++++++++++- native/csrc/sink/native_pack_sink.h | 21 ++++ native/csrc/sink/pack_sink.cpp | 16 ++- native/csrc/sink/pack_sink.h | 8 ++ native/csrc/sink/record_row.cpp | 28 +++++- native/csrc/sink/record_row.h | 16 +++ tests/native/test_pack_sink_timeout.cpp | 98 +++++++++++++++++++ tests/test_native_pack_sink_timeout.py | 7 ++ tests/test_native_sink_admission.py | 79 +++++++++++++++ 14 files changed, 386 insertions(+), 17 deletions(-) diff --git a/native/csrc/bindings.cpp b/native/csrc/bindings.cpp index ebef1a9d2..b8c0f583e 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 @@ -646,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") @@ -657,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); 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/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/sink/bindings_sink.cpp b/native/csrc/sink/bindings_sink.cpp index 2c81ec862..696256494 100644 --- a/native/csrc/sink/bindings_sink.cpp +++ b/native/csrc/sink/bindings_sink.cpp @@ -100,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 @@ -123,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; } 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 c0e98415c..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); diff --git a/native/csrc/sink/pack_sink.h b/native/csrc/sink/pack_sink.h index 9c7a800c2..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 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/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_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 index 6fef5007e..f4b26c692 100644 --- a/tests/test_native_sink_admission.py +++ b/tests/test_native_sink_admission.py @@ -192,3 +192,82 @@ def test_a_burst_four_times_the_queue_is_admitted_under_the_default( 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 From 7a55f5373a89a43742ba1d5893bd752ec547ff27 Mon Sep 17 00:00:00 2001 From: Alan Liu Date: Thu, 24 Sep 2026 13:25:17 -0400 Subject: [PATCH 10/14] Skip a step whose stall budget ran out, and bound the wait past it Four review findings on the record ring's stall budget. The budget was lifetime. Nothing in src/ calls begin_step, so the budget spanned the runtime's life: once spent, every later wait gave up at once and max_step_wait reported a lifetime total. The record ring has no reliable step boundary of its own: commit_step is the legacy adapter's and refuses record rings, eager hooks reserve once per output (and output ids repeat within a step), and a CUDA-graph step replays one plan, or several with piecewise graphs. So begin_record_step stays the boundary, and with a budget a reservation before the first one is refused with a logic_error naming RecordRuntime.begin_step(). Budget exhaustion latched capture off for good; the plan skips the rest of that step. RecordConsumer gains a discard window: its start drops the descriptors still queued, pushes until its end are dropped on arrival, and a count of owed payloads discards theirs as the drain delivers them, so descriptor/payload pairing and ring capacity are kept. Exhaustion opens it (counted in skipped_steps and discarded_*), begin_record_step closes it, and capture resumes. Under disable_capture a skipped step is not a failure: capture stays active and the flush succeeds. Under raise the forward no longer raises: the step is skipped the same way, the exhaustion is held and raised at the next begin_record_step (which then latches it) and at flush. Genuine sink failures still latch as before, and begin_record_step now also raises a failure latched during the previous step under raise, so it usually surfaces there rather than in the forward. The post-latch wait was unbounded, and the bound it relied on held only for a persistent-path NativeSinkConfig. The ring now reads the sink's admission_bound() and refuses a budget (either policy) for a sink with none: the ClickHouse host path, the reference bridge, a NativePackSink blocking without a timeout, or no sink. create_record_runtime refuses the same before tearing down the live ring, naming the case. Past the budget the wait is bounded by the sink's bound plus a 2 s drain grace; a sink that holds the ring longer is a ring failure raised from the reservation, rather than a forward that never returns. A sink failure found at the checked flush (NativePackSink's lost-record check) now latches the runtime like a refusal at submit. capture_status() reports skipped_steps, close logs skipped steps, and the integration doc and docstrings state the exact conditions. Test-first. tests/native/ring/test_record_consumer.cpp: a window under both policies (8 failed with a stub, 82 passed). test_ring_engine.cu (GPU, compile-red on the new status field): the skip under both policies, the refused unbounded and sinkless budgets, a reservation before begin_record_step, a sink past its bound, a failing sink flush. tests/test_record_failure_policy.py: 12 failed, then 26 passed. tests/test_record_failure_policy_gpu.py follows the new semantics and adds the begin_step, unbounded-sink and lost-record cases. CPU tier: 2367 passed. --- docs/integration-api-v1.md | 69 ++++-- native/csrc/bindings.cpp | 1 + native/csrc/ring/record_consumer.cpp | 38 ++- native/csrc/ring/record_consumer.h | 15 ++ native/csrc/ring/ring_engine_py.cu | 188 ++++++++++++--- native/csrc/ring/ring_engine_py.h | 28 ++- src/dmi/engine.py | 113 ++++++--- src/dmi/records.py | 17 +- src/dmi/transport/ring.py | 3 +- tests/native/ring/test_record_consumer.cpp | 57 +++++ tests/native/ring/test_ring_engine.cu | 264 +++++++++++++++++++-- tests/test_record_failure_policy.py | 166 +++++++++++-- tests/test_record_failure_policy_gpu.py | 153 +++++++++--- 13 files changed, 954 insertions(+), 158 deletions(-) diff --git a/docs/integration-api-v1.md b/docs/integration-api-v1.md index 6fcc58b1f..13d294811 100644 --- a/docs/integration-api-v1.md +++ b/docs/integration-api-v1.md @@ -445,36 +445,67 @@ 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) or fail, and it +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 record reservation or descriptor push raises it, inside the forward that publishes the record. | +| `"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 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. `RecordRuntime.begin_step()` starts a step; call it once -per model step, before the step's first reservation, or the budget spans the -runtime's whole life. Past the budget the policy applies: `"raise"` raises -from the reservation with nothing reserved, and `"disable_capture"` latches -and completes the reservation after at most the one sink admission already -in progress. Under `"disable_capture"` the forward's stall per step is -therefore at most the budget plus one `admission_timeout_s`, and -`create_record_runtime()` refuses it with a `NativeSinkConfig` whose `block` -admission has no timeout. `None`, the default, waits without bound. Ring -failures that are not sink refusals, such as a CUDA error in the drain, raise -under both policies. +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 records still queued for the +sink and every record the step reserves afterwards are discarded, counted in +`discarded_*` and `skipped_steps`, and capture resumes at the next step. +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 copy out and discard what the ring holds. A sink that holds the ring +past its bound plus a 2 s drain grace raises from the reservation under +either policy, rather than hanging the forward. `MonitoringEngine.capture_status()` returns a plain dict: `record_mode`, -`capture_active` (False once a failure latched), `failure_policy`, -`failure`, `discarded_descriptors`, `discarded_payloads`, -`step_stall_budget_ms`, `stall_budget_exhaustions`, `reserve_wait_s` and +`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`, `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 diff --git a/native/csrc/bindings.cpp b/native/csrc/bindings.cpp index b8c0f583e..2a75efcbe 100644 --- a/native/csrc/bindings.cpp +++ b/native/csrc/bindings.cpp @@ -770,6 +770,7 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { 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["reserve_wait_ns"] = status.reserve_wait_ns; out["max_step_wait_ns"] = status.max_step_wait_ns; return out; diff --git a/native/csrc/ring/record_consumer.cpp b/native/csrc/ring/record_consumer.cpp index fde2ec756..dfbb27413 100644 --- a/native/csrc/ring/record_consumer.cpp +++ b/native/csrc/ring/record_consumer.cpp @@ -36,6 +36,8 @@ void RecordConsumer::latch_locked(std::exception_ptr failure) { // at the latch can ever be paired again. discarded_descriptors_ += descriptors_.size(); descriptors_.clear(); + // Every later payload is discarded by the latch itself. + payloads_to_discard_ = 0; } void RecordConsumer::push_descriptor(RecordDescriptor descriptor) { @@ -47,6 +49,11 @@ void RecordConsumer::push_descriptor(RecordDescriptor descriptor) { ++discarded_descriptors_; return; } + if (discarding_) { + ++discarded_descriptors_; + ++payloads_to_discard_; + return; + } descriptors_.push_back(std::move(descriptor)); } @@ -60,6 +67,11 @@ void RecordConsumer::push_descriptors( discarded_descriptors_ += descriptors.size(); return; } + if (discarding_) { + discarded_descriptors_ += descriptors.size(); + payloads_to_discard_ += descriptors.size(); + return; + } for (auto& descriptor : descriptors) { descriptors_.push_back(std::move(descriptor)); } @@ -77,6 +89,14 @@ void RecordConsumer::consume_payload(at::Tensor payload) { ++discarded_payloads_; return; } + if (payloads_to_discard_ != 0) { + // 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; + } if (descriptors_.empty()) { latch_locked(std::make_exception_ptr(std::runtime_error( "record consumer: physical payload arrived without an encoded descriptor"))); @@ -145,7 +165,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; @@ -163,7 +184,7 @@ void RecordConsumer::rethrow_if_failed() const { void RecordConsumer::finish() const { std::lock_guard lock(mu_); if (failure_) std::rethrow_exception(failure_); - if (!descriptors_.empty()) { + if (!descriptors_.empty() || payloads_to_discard_ != 0) { invalid("durable completion found leftover encoded descriptors"); } if (active_payloads_ != 0) { @@ -192,4 +213,17 @@ RecordConsumerSnapshot RecordConsumer::snapshot() const { return snapshot; } +void RecordConsumer::begin_discard_window() { + std::lock_guard lock(mu_); + discarding_ = true; + 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 f198fc748..97446a58f 100644 --- a/native/csrc/ring/record_consumer.h +++ b/native/csrc/ring/record_consumer.h @@ -63,6 +63,15 @@ class RecordConsumer { bool failed() const; RecordConsumerSnapshot snapshot() const; + // Skip one step's records without latching (a spent stall budget). + // begin_discard_window() drops the descriptors still queued, 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 the + // discarded_* counters. A payload already in the sink is not recalled. + void begin_discard_window(); + void end_discard_window(); + private: // Retain the first failure; under kDisableCapture also drop the queued // descriptors. Caller holds mu_. @@ -76,6 +85,12 @@ class RecordConsumer { std::deque descriptors_; 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}; }; diff --git a/native/csrc/ring/ring_engine_py.cu b/native/csrc/ring/ring_engine_py.cu index 08b9b471b..2a80e25f2 100644 --- a/native/csrc/ring/ring_engine_py.cu +++ b/native/csrc/ring/ring_engine_py.cu @@ -15,6 +15,8 @@ #include #include #include +#include +#include #include #include #include @@ -30,6 +32,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( @@ -139,13 +151,24 @@ struct RingEnginePy::Impl { at::Tensor payload_view; bool record_mode{false}; - // Record rings only. The step counter is written by the reserving + // 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). + FlushClock::duration sink_admission_bound{}; 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)) @@ -164,6 +187,23 @@ struct RingEnginePy::Impl { : 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); + } const auto& state = engine.ring_state(); int dev_idx = 0; cudaGetDevice(&dev_idx); @@ -190,8 +230,10 @@ struct RingEnginePy::Impl { // 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 failure policy - // applies. + // is left of this step's stall budget. Past it the rest of the step is + // skipped: the records queued for the sink 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); @@ -200,15 +242,18 @@ struct RingEnginePy::Impl { const bool disable = record_options.failure_policy == ring::RecordFailurePolicy::kDisableCapture; const auto started = FlushClock::now(); - // Once capture is off the worker discards instead of submitting, so - // the sink is out of the path and only the device-to-host copy of - // what the ring already holds is left to wait for. - if (record_options.step_stall_budget_ms == 0 || - (disable && consumer.failed())) { + 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( @@ -221,25 +266,60 @@ struct RingEnginePy::Impl { if (drained) return; stall_budget_exhaustions.fetch_add(1, std::memory_order_relaxed); - const auto step_ms = - step_wait_ns.load(std::memory_order_relaxed) / 1'000'000ull; + 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 rest of that step was not captured")); + } + } + // 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 copying out what the ring + // holds. A sink that holds the worker past its own bound is a ring + // failure under either policy: the alternative is a forward that never + // returns. + void wait_past_the_budget(FlushClock::time_point started) { + auto& drain = engine.drain_thread(); + const auto deadline = started + sink_admission_bound + kDrainGrace; + const bool drained = drain.force_flush_and_wait_until(deadline); + account_record_wait(started); + if (drained) return; const std::exception_ptr failure = std::make_exception_ptr(std::runtime_error( - "record capture stall budget exhausted: reservations waited " + - std::to_string(step_ms) + " ms this step for the sink to " - "free ring space (step_stall_budget_ms=" + - std::to_string(record_options.step_stall_budget_ms) + ")")); - // Latched in the consumer either way, so the checked flush, the - // capture status and close report it. - consumer.record_failure(failure); - if (!disable) std::rethrow_exception(failure); - // Capture is off now, so the drain is held back only by the one sink - // admission already in progress (bounded by the sink's own admission - // timeout). The reservation must still complete: under CUDA-graph - // replay the producers launch whatever the host decides. - const auto resumed = FlushClock::now(); - drain.force_flush_and_wait_until(FlushClock::time_point::max()); - account_record_wait(resumed); + "record sink held the ring past its admission bound (" + + std::to_string(std::chrono::duration_cast< + std::chrono::milliseconds>(sink_admission_bound).count()) + + " ms, plus " + + std::to_string(std::chrono::duration_cast< + std::chrono::milliseconds>(kDrainGrace).count()) + + " ms to drain the ring)")); + engine.record_consumer().record_failure(failure); + std::rethrow_exception(failure); + } + + // What the drain may take, once the sink is out of the path, to copy out + // and discard everything the ring holds. + static constexpr FlushClock::duration kDrainGrace = + std::chrono::seconds(2); + + // 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; } }; @@ -590,6 +670,15 @@ int RingEnginePy::reserve_record( 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) { impl_->wait_for_record_space(); @@ -629,7 +718,21 @@ 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(); + } } RecordCaptureStatus RingEnginePy::record_capture_status() const { @@ -647,6 +750,14 @@ RecordCaptureStatus RingEnginePy::record_capture_status() const { 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); + 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 = @@ -681,6 +792,9 @@ bool RingEnginePy::flush_records_and_wait(uint64_t timeout_ms) { drain.rethrow_drain_failure(); drain.rethrow_record_reclaim_failure(); 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; @@ -702,14 +816,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 ca5ae5235..7c14dcb80 100644 --- a/native/csrc/ring/ring_engine_py.h +++ b/native/csrc/ring/ring_engine_py.h @@ -51,11 +51,19 @@ struct RingConfig { struct RecordRuntimeOptions { ring::RecordFailurePolicy failure_policy = ring::RecordFailurePolicy::kRaiseAtProducer; - // Cap on the time the record reservations of one step (see - // begin_record_step) may wait for the drain to free ring space. Past - // it the policy applies: kRaiseAtProducer raises from the reservation, - // kDisableCapture latches the failure and stops capture. 0 waits - // without bound, as before the budget existed. + // 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 rest of that step's records are skipped (discarded and + // counted in skipped_steps) 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. 0 waits without bound, as before the budget + // existed. uint64_t step_stall_budget_ms = 0; }; @@ -69,6 +77,8 @@ struct RecordCaptureStatus { uint64_t discarded_payloads = 0; uint64_t step_stall_budget_ms = 0; uint64_t stall_budget_exhaustions = 0; + // Steps whose remaining records were skipped because the budget ran out. + uint64_t skipped_steps = 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. @@ -210,7 +220,13 @@ class RingEnginePy { const std::vector>& reservation_items); // Start a new step for the stall budget: the reservations after this - // call share one fresh budget. + // 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; diff --git a/src/dmi/engine.py b/src/dmi/engine.py index 9901e4f66..92c077e6b 100644 --- a/src/dmi/engine.py +++ b/src/dmi/engine.py @@ -359,17 +359,37 @@ def create_record_runtime( 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, or a sink failure) does to the forward: - ``"raise"`` raises from the next record reservation or descriptor + 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 (see ``RecordRuntime.begin_step``) may wait for the sink to - free ring space; past it the policy applies. ``None`` waits without - bound. Under ``"disable_capture"`` the forward's stall per step is - then at most the budget plus one sink admission timeout. + 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 records still queued for the sink and the + ones the step reserves afterwards are discarded and counted in + ``skipped_steps`` -- 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 copy out what the ring holds; + a sink that holds the ring past its bound plus a 2 s drain grace + raises from the reservation under either policy. """ if getattr(self, "_storage_backend", "auto") == "none": @@ -379,7 +399,7 @@ def create_record_runtime( "or 'persistent' to capture" ) self._validate_record_failure_options( - failure_policy, step_stall_budget_ms) + 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: @@ -413,7 +433,8 @@ def create_record_runtime( return runtime def _validate_record_failure_options( - self, failure_policy: Any, step_stall_budget_ms: Any + self, failure_policy: Any, step_stall_budget_ms: Any, + record_sink: Optional[Any], ) -> None: if failure_policy not in RECORD_FAILURE_POLICIES: raise ValueError( @@ -425,20 +446,42 @@ def _validate_record_failure_options( 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 ( - failure_policy == "disable_capture" - and self._storage_backend == "persistent" - and sink_config is not None - and sink_config.overload == "block" - and sink_config.admission_timeout_s is None - ): - # Once capture latches, the forward still waits for the one sink - # admission in progress. Without a timeout that wait has no bound. - raise ValueError( - "failure_policy='disable_capture' needs a bounded sink " - "admission: set capture_sink_config.admission_timeout_s, or " - "overload='drop_newest'") + 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 @@ -575,10 +618,14 @@ def flush_and_wait(self, timeout_s: float = 600.0) -> None: 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 - ``"disable_capture"`` the forward keeps running and ``failure`` - says why capture stopped; ``discarded_*`` count what was dropped - after it. ``reserve_wait_s`` and ``max_step_wait_s`` are the time + ``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; under ``"disable_capture"`` capture stays active through them. + ``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. @@ -589,6 +636,7 @@ def capture_status(self) -> dict[str, Any]: "failure_policy": None, "failure": None, "discarded_descriptors": 0, "discarded_payloads": 0, "step_stall_budget_ms": None, "stall_budget_exhaustions": 0, + "skipped_steps": 0, "reserve_wait_s": 0.0, "max_step_wait_s": 0.0, "sink": None, "storage": None, } @@ -605,6 +653,7 @@ def capture_status(self) -> dict[str, Any]: 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"]), reserve_wait_s=int(native["reserve_wait_ns"]) / 1e9, max_step_wait_s=int(native["max_step_wait_ns"]) / 1e9, ) @@ -616,13 +665,23 @@ def capture_status(self) -> dict[str, Any]: return status def _report_capture_failure(self) -> None: - """Log, once, why capture stopped, as a record ring is retired.""" + """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"] or status["capture_active"]: + 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; %d descriptors and %d payloads " + "discarded", status["skipped_steps"], + status["step_stall_budget_ms"], + status["discarded_descriptors"], status["discarded_payloads"]) + if status["capture_active"]: return _LOG.warning( "record capture stopped before close (failure_policy=%s): %s; " diff --git a/src/dmi/records.py b/src/dmi/records.py index 314e3b6d4..748aaca13 100644 --- a/src/dmi/records.py +++ b/src/dmi/records.py @@ -310,12 +310,21 @@ def bind_hook( self._bound_hooks.add(id(hook)) def begin_step(self) -> None: - """Mark the start of one model step for the stall budget. + """Mark the start of one model step. The record reservations after this call share one fresh - ``step_stall_budget_ms`` (see ``create_record_runtime``). Call it - once per step, before the step's first ``emit_output`` or - ``prepare_replay``; without it the budget spans the runtime's life. + ``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() diff --git a/src/dmi/transport/ring.py b/src/dmi/transport/ring.py index 503adf2d1..7a57ccc78 100644 --- a/src/dmi/transport/ring.py +++ b/src/dmi/transport/ring.py @@ -421,7 +421,8 @@ 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 fresh stall budget for the next reservations.""" + """Start a record step: a fresh stall budget for the next + reservations (see RecordRuntime.begin_step).""" self._ring_engine.begin_record_step() diff --git a/tests/native/ring/test_record_consumer.cpp b/tests/native/ring/test_record_consumer.cpp index 5ccef9bd3..a11c92482 100644 --- a/tests/native/ring/test_record_consumer.cpp +++ b/tests/native/ring/test_record_consumer.cpp @@ -296,6 +296,59 @@ static void test_disable_capture_latches_association_and_worker_failures() { 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); +} + int main() { setbuf(stdout, nullptr); std::printf("test_record_consumer\n"); @@ -306,6 +359,10 @@ int main() { 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); 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 3f6c2fcff..4368f4f6b 100644 --- a/tests/native/ring/test_ring_engine.cu +++ b/tests/native/ring/test_ring_engine.cu @@ -19,6 +19,7 @@ #include #include #include +#include #include #include #include @@ -780,10 +781,18 @@ static void test_record_flush_reaches_sink_durability_boundary() { // 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) {} + : 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_); @@ -792,11 +801,15 @@ public: 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() { @@ -826,16 +839,21 @@ struct StallOutcome { 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. +// 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); @@ -871,7 +889,12 @@ static StallOutcome run_stalled_reservation(ring::RecordFailurePolicy policy) { outcome.stall = std::chrono::steady_clock::now() - started; outcome.status = engine.record_capture_status(); - engine.begin_record_step(); + 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); @@ -882,55 +905,244 @@ static StallOutcome run_stalled_reservation(ring::RecordFailurePolicy policy) { try { engine.flush_records_and_wait(5000); - } catch (const std::runtime_error&) { + } 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_bounds_a_stalled_reservation() { - banner("raise policy: a reservation stalled past the budget raises"); +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); - EXPECT(outcome.threw); - EXPECT(outcome.error.find("stall budget") != std::string::npos); - // The budget, not the sink's 400 ms admission, bounds the wait. - EXPECT(outcome.stall < std::chrono::milliseconds(300)); + // 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.stall_budget_exhaustions == 1); + EXPECT(outcome.status.skipped_steps == 1); EXPECT(outcome.status.step_stall_budget_ms == 50); - // Latched: the next reservation raises before it waits for anything. + // 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_bounds_a_stalled_reservation() { - banner("disable_capture: a stalled reservation stops capture, not the forward"); +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; once - // capture is off the worker discards instead of calling the sink. + // 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.failure.find("stall budget") != std::string::npos); + EXPECT(!outcome.status.failed); EXPECT(outcome.status.stall_budget_exhaustions == 1); + EXPECT(outcome.status.skipped_steps == 1); EXPECT(outcome.status.max_step_wait_ns >= 50'000'000ull); - // Capture is off: the next step reserves without raising or stalling. + // 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)); - // The failure still surfaces at the checked flush. - EXPECT(outcome.flush_threw); - // Only the record the sink was already admitting reached it. - EXPECT(outcome.submissions == 1); + // 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, and the ninth. + EXPECT(outcome.flushed_status.discarded_descriptors == 8); + EXPECT(outcome.flushed_status.discarded_payloads == 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 sink that holds the ring past its admission bound raises"); + // Declares 10 ms, takes 2.5 s: past the bound and the drain grace. + 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("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); + EXPECT(engine.record_capture_status().discarded_payloads == 1); + engine.stop(); + CUDA_CHECK(cudaFree(device)); } int main() { @@ -954,8 +1166,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_bounds_a_stalled_reservation(); - test_disable_capture_bounds_a_stalled_reservation(); + 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/test_record_failure_policy.py b/tests/test_record_failure_policy.py index 73428a1ff..c9a3e17ed 100644 --- a/tests/test_record_failure_policy.py +++ b/tests/test_record_failure_policy.py @@ -34,6 +34,7 @@ def __init__(self, status=None): "failure_policy": "raise", "failed": False, "failure": "", "discarded_descriptors": 0, "discarded_payloads": 0, "step_stall_budget_ms": 0, "stall_budget_exhaustions": 0, + "skipped_steps": 0, "reserve_wait_ns": 0, "max_step_wait_ns": 0, } self.steps = 0 @@ -45,6 +46,23 @@ 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() @@ -63,6 +81,21 @@ def create_record(config, target, **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 ---------------------------------------------------- @@ -73,12 +106,48 @@ def test_the_default_policy_is_raise_with_no_stall_budget(monkeypatch): def test_disable_capture_and_a_budget_reach_the_native_ring(monkeypatch): _engine, _runtime, created = _record_engine( - monkeypatch, _StatusRing(), failure_policy="disable_capture", - step_stall_budget_ms=2000) + 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"), @@ -102,10 +171,10 @@ def test_bad_options_are_refused_before_the_live_ring_is_touched( assert engine._record_mode is False -def test_disable_capture_needs_a_bounded_sink_admission(tmp_path): - """Once the budget latches, the forward still waits for the one sink - admission in progress; unbounded, that wait is the stall the policy - exists to prevent.""" +@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 @@ -117,7 +186,33 @@ def test_disable_capture_needs_a_bounded_sink_admission(tmp_path): enable_ring_transport=False) with pytest.raises(ValueError, match="admission_timeout_s"): engine.create_record_runtime( - _explicit_sink_format(), failure_policy="disable_capture") + _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 ----------------------------------------------------------------------- @@ -142,6 +237,7 @@ def test_capture_status_without_a_record_runtime(): "failure_policy": None, "failure": None, "discarded_descriptors": 0, "discarded_payloads": 0, "step_stall_budget_ms": None, "stall_budget_exhaustions": 0, + "skipped_steps": 0, "reserve_wait_s": 0.0, "max_step_wait_s": 0.0, "sink": None, "storage": None, } @@ -153,11 +249,12 @@ def test_capture_status_reports_a_disabled_capture(monkeypatch): "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, "reserve_wait_ns": 2_500_000_000, "max_step_wait_ns": 2_100_000_000, }) engine, _runtime, _created = _record_engine( - monkeypatch, ring, failure_policy="disable_capture", - step_stall_budget_ms=2000) + 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, @@ -165,11 +262,33 @@ def test_capture_status_reports_a_disabled_capture(monkeypatch): "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, "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, + "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 + + def test_capture_status_includes_the_sink_and_storage_snapshots(monkeypatch): engine, _runtime, _created = _record_engine(monkeypatch, _StatusRing()) @@ -194,17 +313,36 @@ def snapshot(self): def test_close_reports_a_capture_that_stopped(monkeypatch, caplog): ring = _StatusRing({ "failure_policy": "disable_capture", "failed": True, - "failure": "record capture stall budget exhausted", + "failure": "NativePackSink: sink refused durable admission: dropped", "discarded_descriptors": 4, "discarded_payloads": 5, - "step_stall_budget_ms": 50, "stall_budget_exhaustions": 1, + "step_stall_budget_ms": 50, "stall_budget_exhaustions": 0, + "skipped_steps": 0, "reserve_wait_ns": 0, "max_step_wait_ns": 0, }) engine, _runtime, _created = _record_engine( - monkeypatch, ring, failure_policy="disable_capture", - step_stall_budget_ms=50) + 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("stall budget exhausted" in message and "5 payloads" in message + 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, + "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 for message in messages), messages diff --git a/tests/test_record_failure_policy_gpu.py b/tests/test_record_failure_policy_gpu.py index 64a4c462e..f4433c533 100644 --- a/tests/test_record_failure_policy_gpu.py +++ b/tests/test_record_failure_policy_gpu.py @@ -8,8 +8,10 @@ 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 burst case uses the real -NativePackSink with its default config. +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) """ @@ -127,7 +129,8 @@ def _run_steps(record_sink, *, policy, budget_ms, payload_bytes=4096, engine = MonitoringEngine(model_id="policy-gpu", ring_config=_ring_config(payload_bytes)) - outcome = {"step_s": [], "errors": [], "flush_error": None} + outcome = {"step_s": [], "errors": [], "begin_errors": [], + "flush_error": None} try: runtime = engine.create_record_runtime( CaptureRecordFormat(), record_sink=record_sink, @@ -139,7 +142,10 @@ def _run_steps(record_sink, *, policy, budget_ms, payload_bytes=4096, tensor = torch.arange(elements, dtype=torch.float32, device="cuda") for step in range(steps): hook_runtime.metadata = _metadata(step, tensor) - runtime.begin_step() + try: + runtime.begin_step() + except RuntimeError as exc: + outcome["begin_errors"].append((step, str(exc))) started = time.monotonic() try: hook(tensor + step) @@ -164,56 +170,115 @@ def _slow_sink(target): from dmi.transport import native return native.ReferencePythonCaptureSink( - target, CaptureRecordFormat.LAYOUT_NAME) + target, CaptureRecordFormat.LAYOUT_NAME, + admission_bound_s=target.admission_s) -def test_disable_capture_keeps_the_forward_running_through_a_stalled_sink( - caplog): - """A 4 KiB ring against a sink that takes 400 ms per record: the ninth - record's reservation waits for the sink. With a 50 ms budget every step - returns, none waits longer than the budget plus one admission, and - capture reports why it stopped.""" +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["status"] - assert status["capture_active"] is False - assert "stall budget" in status["failure"] + status = outcome["flushed_status"] + assert status["capture_active"] is True, status + assert status["failure"] is None assert status["failure_policy"] == "disable_capture" - assert status["stall_budget_exhaustions"] == 1 + assert status["skipped_steps"] >= 1 + assert status["stall_budget_exhaustions"] == status["skipped_steps"] assert status["step_stall_budget_ms"] == BUDGET_MS assert status["max_step_wait_s"] >= BUDGET_MS / 1000 - # Only the record the sink was already admitting reached it; everything - # after the latch was dropped on the worker, not submitted. - assert target.submitted == 1 - assert outcome["flushed_status"]["discarded_payloads"] == STEPS - 1 - assert "stall budget" in outcome["flush_error"] + # 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_fails_fast_when_the_stall_budget_is_spent(): +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["errors"], "the stalled step must raise under 'raise'" - first_step, message = outcome["errors"][0] + assert outcome["begin_errors"], "the next step must raise under 'raise'" + first_step, message = outcome["begin_errors"][0] assert "stall budget" in message - # Raised at the budget, not after the sink's 400 ms admission. - assert outcome["step_s"][first_step] < 0.3, outcome["step_s"] - # Latched: every later step raises at once. - assert [step for step, _ in outcome["errors"]] == list( + # 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 max(outcome["step_s"][first_step + 1:]) < 0.1 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", @@ -245,6 +310,10 @@ def test_raise_surfaces_a_sink_refusal_in_a_later_forward(): 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 @@ -276,3 +345,31 @@ def test_a_burst_four_times_the_sink_queue_is_stored_on_a_record_ring( 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 From 13d39c33f870f8de135f191623b0e622da5cfd7a Mon Sep 17 00:00:00 2001 From: Alan Liu Date: Thu, 24 Sep 2026 13:26:02 -0400 Subject: [PATCH 11/14] Load a NativeSinkConfig pickled before the admission fields A config pickled from the old class (dmi.storage.capture.native_sink, eight fields) unpickles as the new class through the re-export, but the frozen slots dataclass's generated __setstate__ zips only the fields the pickle has, so .overload raised AttributeError on first use. The new __setstate__ fills the fields a pickle predates with their defaults (block, 2 s) and refuses state with too many values. Test-first in tests/test_native_sink_admission.py with the bytes of an old-class pickle (made from the pre-move module): AttributeError, then passing; a new pickle still round-trips. --- src/dmi/storage/native_capture.py | 25 +++++++++++++++++++++- tests/test_native_sink_admission.py | 33 +++++++++++++++++++++++++++++ 2 files changed, 57 insertions(+), 1 deletion(-) diff --git a/src/dmi/storage/native_capture.py b/src/dmi/storage/native_capture.py index c8f1164bc..dc3f18c69 100644 --- a/src/dmi/storage/native_capture.py +++ b/src/dmi/storage/native_capture.py @@ -29,7 +29,7 @@ import os import socket import uuid -from dataclasses import dataclass, field +from dataclasses import MISSING, dataclass, field, fields from typing import Any, Optional, Sequence @@ -116,6 +116,29 @@ def __post_init__(self) -> 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__}") diff --git a/tests/test_native_sink_admission.py b/tests/test_native_sink_admission.py index f4b26c692..e916614e4 100644 --- a/tests/test_native_sink_admission.py +++ b/tests/test_native_sink_admission.py @@ -105,6 +105,39 @@ def test_the_config_type_does_not_load_the_backup_capture_package(): 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 ------------------------------------------------------ From 67213f130acec1f76dd6bbc61f5a4bcec2f54d97 Mon Sep 17 00:00:00 2001 From: Alan Liu Date: Thu, 24 Sep 2026 13:27:41 -0400 Subject: [PATCH 12/14] Say in NativeSinkConfig what one admission timeout covers The rows of one ring record now share one admission deadline, a record lost after admission latches the runtime like a refusal, and a stall budget is refused for block with no timeout. The docstring said none of it. --- src/dmi/storage/native_capture.py | 7 +++++-- 1 file changed, 5 insertions(+), 2 deletions(-) diff --git a/src/dmi/storage/native_capture.py b/src/dmi/storage/native_capture.py index dc3f18c69..f393978c0 100644 --- a/src/dmi/storage/native_capture.py +++ b/src/dmi/storage/native_capture.py @@ -67,9 +67,12 @@ class NativeSinkConfig: 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; ``"drop_newest"`` refuses it at + 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). + 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 From b5c5a8193269bd8f3a0ed74997106d120abf55e2 Mon Sep 17 00:00:00 2001 From: Alan Liu Date: Thu, 24 Sep 2026 14:45:59 -0400 Subject: [PATCH 13/14] Count the steps a spent stall budget reaches, and scale its drain grace A spent budget's discard window drops every descriptor still queued, including records of earlier, finished steps that the sink had not taken yet; the stall bound needs that. The reporting said only "the rest of that step". The raise-mode error, the header comments, capture_status and the v1 contract now say the step's remaining records and every record still queued for the sink, from any step, are discarded. The consumer tags each queued descriptor with its step (begin_step, called from begin_record_step) and counts the distinct steps a window drops records from as steps_with_discards, reported beside skipped_steps (which always equals stall_budget_exhaustions) and logged at close. The drain grace past a spent budget was a fixed 2 s, while the drain still moves the whole ring and staging out before the reservation fits: with the default 4 GiB of each that could exceed it and fail the ring, blaming the sink. ring::record_drain_grace makes it 2 s plus the ring and staging bytes at 1 GB/s (about 10.6 s by default), and the error now says the ring did not drain and names both the sink and the drain as possible causes. The record worker also skips the pageable copy of a payload nobody will store (a skipped step's, or any after a disable_capture latch): discard_next_payload_if_unwanted accounts the discard through the same branch consume_payload uses, and the staging range is freed as before. A test now covers the payloads_to_discard_ term of wait_until_idle: with the queue empty and only owed payloads left, a flush must not go idle. finish() names the owed payloads in its error. The discarded_* comments no longer say kDisableCapture only. --- docs/integration-api-v1.md | 31 ++-- native/csrc/bindings.cpp | 1 + native/csrc/ring/p2p_thread.cpp | 11 ++ native/csrc/ring/record_consumer.cpp | 74 +++++--- native/csrc/ring/record_consumer.h | 62 +++++-- native/csrc/ring/record_stall_budget.h | 42 +++++ native/csrc/ring/ring_engine_py.cu | 54 +++--- native/csrc/ring/ring_engine_py.h | 23 ++- src/dmi/engine.py | 35 ++-- tests/native/ring/test_record_consumer.cpp | 187 +++++++++++++++++++++ tests/native/ring/test_ring_engine.cu | 16 +- tests/test_record_failure_policy.py | 17 +- tests/test_record_failure_policy_gpu.py | 2 + 13 files changed, 460 insertions(+), 95 deletions(-) create mode 100644 native/csrc/ring/record_stall_budget.h diff --git a/docs/integration-api-v1.md b/docs/integration-api-v1.md index 13d294811..a5cc45801 100644 --- a/docs/integration-api-v1.md +++ b/docs/integration-api-v1.md @@ -472,9 +472,14 @@ 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 records still queued for the -sink and every record the step reserves afterwards are discarded, counted in -`discarded_*` and `skipped_steps`, and capture resumes at the next step. +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 @@ -497,19 +502,23 @@ 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 copy out and discard what the ring holds. A sink that holds the ring -past its bound plus a 2 s drain grace raises from the reservation under -either policy, rather than hanging the forward. +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`, `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. +`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 diff --git a/native/csrc/bindings.cpp b/native/csrc/bindings.cpp index 2a75efcbe..feac37904 100644 --- a/native/csrc/bindings.cpp +++ b/native/csrc/bindings.cpp @@ -771,6 +771,7 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { 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; diff --git a/native/csrc/ring/p2p_thread.cpp b/native/csrc/ring/p2p_thread.cpp index fcb66bcc3..fb73ca8d1 100644 --- a/native/csrc/ring/p2p_thread.cpp +++ b/native/csrc/ring/p2p_thread.cpp @@ -363,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/record_consumer.cpp b/native/csrc/ring/record_consumer.cpp index dfbb27413..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 { @@ -36,10 +37,46 @@ void RecordConsumer::latch_locked(std::exception_ptr failure) { // at the latch can ever be paired again. discarded_descriptors_ += descriptors_.size(); descriptors_.clear(); - // Every later payload is discarded by the latch itself. + // 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_) { @@ -52,9 +89,10 @@ void RecordConsumer::push_descriptor(RecordDescriptor descriptor) { if (discarding_) { ++discarded_descriptors_; ++payloads_to_discard_; + note_discarded_step_locked(step_); return; } - descriptors_.push_back(std::move(descriptor)); + descriptors_.push_back({std::move(descriptor), step_}); } void RecordConsumer::push_descriptors( @@ -70,10 +108,11 @@ void RecordConsumer::push_descriptors( 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_}); } } @@ -82,21 +121,8 @@ void RecordConsumer::consume_payload(at::Tensor payload) { RecordDescriptor descriptor; { std::lock_guard lock(mu_); - if (failure_) { - if (!disable) std::rethrow_exception(failure_); - // Capture is off: the drain still delivers what the forward - // produced, and the payload is dropped here, never submitted. - ++discarded_payloads_; - return; - } - if (payloads_to_discard_ != 0) { - // 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; - } + if (failure_ && !disable) std::rethrow_exception(failure_); + if (take_unwanted_payload_locked()) return; if (descriptors_.empty()) { latch_locked(std::make_exception_ptr(std::runtime_error( "record consumer: physical payload arrived without an encoded descriptor"))); @@ -105,7 +131,7 @@ void RecordConsumer::consume_payload(at::Tensor payload) { idle_cv_.notify_all(); return; } - descriptor = std::move(descriptors_.front()); + descriptor = std::move(descriptors_.front().descriptor); descriptors_.pop_front(); ++active_payloads_; } @@ -185,7 +211,11 @@ void RecordConsumer::finish() const { std::lock_guard lock(mu_); if (failure_) std::rethrow_exception(failure_); if (!descriptors_.empty() || payloads_to_discard_ != 0) { - invalid("durable completion found leftover encoded descriptors"); + 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"); @@ -210,12 +240,16 @@ RecordConsumerSnapshot RecordConsumer::snapshot() const { 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(); diff --git a/native/csrc/ring/record_consumer.h b/native/csrc/ring/record_consumer.h index 97446a58f..4f65ddb3b 100644 --- a/native/csrc/ring/record_consumer.h +++ b/native/csrc/ring/record_consumer.h @@ -16,6 +16,7 @@ #include #include #include +#include #include #include @@ -25,12 +26,19 @@ struct RecordConsumerSnapshot { RecordFailurePolicy policy{RecordFailurePolicy::kRaiseAtProducer}; bool failed{false}; std::string failure; - // Descriptors dropped by the latch: those still queued when it happened - // plus every push after it. kDisableCapture only. + // 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 because of the latch: the one whose submission - // failed, plus every payload delivered after it. kDisableCapture only. + // 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 { @@ -63,26 +71,56 @@ class RecordConsumer { bool failed() const; RecordConsumerSnapshot snapshot() const; - // Skip one step's records without latching (a spent stall budget). - // begin_discard_window() drops the descriptors still queued, 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 the - // discarded_* counters. A payload already in the sink is not recalled. + // 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 @@ -93,6 +131,8 @@ class RecordConsumer { 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_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_py.cu b/native/csrc/ring/ring_engine_py.cu index 2a80e25f2..77a498d13 100644 --- a/native/csrc/ring/ring_engine_py.cu +++ b/native/csrc/ring/ring_engine_py.cu @@ -10,6 +10,7 @@ #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 @@ -154,8 +155,11 @@ struct RingEnginePy::Impl { // 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). + // 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}; @@ -203,6 +207,9 @@ struct RingEnginePy::Impl { } 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; @@ -231,9 +238,9 @@ struct RingEnginePy::Impl { // 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: the records queued for the sink and every record the step - // still reserves are discarded, which leaves the drain waiting only for - // the envelope the sink is admitting. + // 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); @@ -279,7 +286,8 @@ struct RingEnginePy::Impl { "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 rest of that step was not captured")); + "); 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 @@ -288,34 +296,34 @@ struct RingEnginePy::Impl { } // With the sink out of the path, what is left is the envelope it is - // admitting (at most its admission bound) and copying out what the ring - // holds. A sink that holds the worker past its own bound is a ring - // failure under either policy: the alternative is a forward that never - // returns. + // 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 + kDrainGrace; + 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 sink held the ring past its admission bound (" + - std::to_string(std::chrono::duration_cast< - std::chrono::milliseconds>(sink_admission_bound).count()) + - " ms, plus " + - std::to_string(std::chrono::duration_cast< - std::chrono::milliseconds>(kDrainGrace).count()) + - " ms to drain the ring)")); + "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); } - // What the drain may take, once the sink is out of the path, to copy out - // and discard everything the ring holds. - static constexpr FlushClock::duration kDrainGrace = - std::chrono::seconds(2); - // kRaiseAtProducer: the spent budget of an earlier step, if any. std::exception_ptr pending_step_failure() const { std::lock_guard lock(step_failure_mu); @@ -733,6 +741,7 @@ void RingEnginePy::begin_record_step() { if (impl_->skipping_step.exchange(false, std::memory_order_relaxed)) { consumer.end_discard_window(); } + consumer.begin_step(); } RecordCaptureStatus RingEnginePy::record_capture_status() const { @@ -752,6 +761,7 @@ RecordCaptureStatus RingEnginePy::record_capture_status() const { 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; diff --git a/native/csrc/ring/ring_engine_py.h b/native/csrc/ring/ring_engine_py.h index 7c14dcb80..a465f7c33 100644 --- a/native/csrc/ring/ring_engine_py.h +++ b/native/csrc/ring/ring_engine_py.h @@ -56,14 +56,16 @@ struct RecordRuntimeOptions { // 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 rest of that step's records are skipped (discarded and - // counted in skipped_steps) 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 + // 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. 0 waits without bound, as before the budget - // existed. + // 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; }; @@ -77,8 +79,13 @@ struct RecordCaptureStatus { uint64_t discarded_payloads = 0; uint64_t step_stall_budget_ms = 0; uint64_t stall_budget_exhaustions = 0; - // Steps whose remaining records were skipped because the budget ran out. + // 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. diff --git a/src/dmi/engine.py b/src/dmi/engine.py index 92c077e6b..211ff3ebf 100644 --- a/src/dmi/engine.py +++ b/src/dmi/engine.py @@ -371,9 +371,10 @@ def create_record_runtime( 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 records still queued for the sink and the - ones the step reserves afterwards are discarded and counted in - ``skipped_steps`` -- and capture resumes at the next step. Under + 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. @@ -387,9 +388,11 @@ def create_record_runtime( ``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 copy out what the ring holds; - a sink that holds the ring past its bound plus a 2 s drain grace - raises from the reservation under either policy. + 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": @@ -623,9 +626,12 @@ def capture_status(self) -> dict[str, Any]: 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; under ``"disable_capture"`` capture stays active through them. - ``discarded_*`` count the records dropped by a skip or after a - latch. ``reserve_wait_s`` and ``max_step_wait_s`` are the time + 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. @@ -636,7 +642,7 @@ def capture_status(self) -> dict[str, Any]: "failure_policy": None, "failure": None, "discarded_descriptors": 0, "discarded_payloads": 0, "step_stall_budget_ms": None, "stall_budget_exhaustions": 0, - "skipped_steps": 0, + "skipped_steps": 0, "steps_with_discards": 0, "reserve_wait_s": 0.0, "max_step_wait_s": 0.0, "sink": None, "storage": None, } @@ -654,6 +660,7 @@ def capture_status(self) -> dict[str, Any]: 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, ) @@ -677,9 +684,11 @@ def _report_capture_failure(self) -> None: if status["skipped_steps"]: _LOG.warning( "record capture skipped the rest of %d steps whose stall " - "budget (%s ms) ran out; %d descriptors and %d payloads " - "discarded", status["skipped_steps"], - status["step_stall_budget_ms"], + "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 diff --git a/tests/native/ring/test_record_consumer.cpp b/tests/native/ring/test_record_consumer.cpp index a11c92482..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 @@ -349,6 +351,181 @@ static void test_a_discard_window_skips_one_steps_records( 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"); @@ -363,6 +540,16 @@ int main() { 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 4368f4f6b..c9500aa83 100644 --- a/tests/native/ring/test_ring_engine.cu +++ b/tests/native/ring/test_ring_engine.cu @@ -927,8 +927,11 @@ static void test_raise_policy_skips_the_step_and_raises_at_the_next() { 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); @@ -957,6 +960,8 @@ static void test_disable_capture_skips_the_step_and_resumes() { 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); @@ -965,9 +970,12 @@ static void test_disable_capture_skips_the_step_and_resumes() { // 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, and the ninth. + // 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); } @@ -1059,8 +1067,9 @@ static void test_a_stall_budget_needs_begin_record_step() { } static void test_a_sink_past_its_admission_bound_is_a_ring_failure() { - banner("a sink that holds the ring past its admission bound raises"); - // Declares 10 ms, takes 2.5 s: past the bound and the drain grace. + 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; @@ -1084,6 +1093,7 @@ static void test_a_sink_past_its_admission_bound_is_a_ring_failure() { 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)); diff --git a/tests/test_record_failure_policy.py b/tests/test_record_failure_policy.py index c9a3e17ed..308fa02a1 100644 --- a/tests/test_record_failure_policy.py +++ b/tests/test_record_failure_policy.py @@ -34,7 +34,7 @@ def __init__(self, status=None): "failure_policy": "raise", "failed": False, "failure": "", "discarded_descriptors": 0, "discarded_payloads": 0, "step_stall_budget_ms": 0, "stall_budget_exhaustions": 0, - "skipped_steps": 0, + "skipped_steps": 0, "steps_with_discards": 0, "reserve_wait_ns": 0, "max_step_wait_ns": 0, } self.steps = 0 @@ -237,7 +237,7 @@ def test_capture_status_without_a_record_runtime(): "failure_policy": None, "failure": None, "discarded_descriptors": 0, "discarded_payloads": 0, "step_stall_budget_ms": None, "stall_budget_exhaustions": 0, - "skipped_steps": 0, + "skipped_steps": 0, "steps_with_discards": 0, "reserve_wait_s": 0.0, "max_step_wait_s": 0.0, "sink": None, "storage": None, } @@ -249,7 +249,7 @@ def test_capture_status_reports_a_disabled_capture(monkeypatch): "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, + "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( @@ -262,7 +262,7 @@ def test_capture_status_reports_a_disabled_capture(monkeypatch): "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, + "skipped_steps": 1, "steps_with_discards": 1, "reserve_wait_s": 2.5, "max_step_wait_s": 2.1, "sink": None, "storage": None, } @@ -276,7 +276,7 @@ def test_capture_status_reports_skipped_steps_with_capture_active( "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, + "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( @@ -287,6 +287,8 @@ def test_capture_status_reports_skipped_steps_with_capture_active( 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): @@ -316,7 +318,7 @@ def test_close_reports_a_capture_that_stopped(monkeypatch, caplog): "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, + "skipped_steps": 0, "steps_with_discards": 0, "reserve_wait_ns": 0, "max_step_wait_ns": 0, }) engine, _runtime, _created = _record_engine( @@ -335,7 +337,7 @@ def test_close_reports_skipped_steps(monkeypatch, caplog): "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, + "skipped_steps": 3, "steps_with_discards": 4, "reserve_wait_ns": 0, "max_step_wait_ns": 0, }) engine, _runtime, _created = _record_engine( @@ -345,4 +347,5 @@ def test_close_reports_skipped_steps(monkeypatch, caplog): 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 index f4433c533..edbe8921d 100644 --- a/tests/test_record_failure_policy_gpu.py +++ b/tests/test_record_failure_policy_gpu.py @@ -195,6 +195,8 @@ def test_disable_capture_skips_stalled_steps_and_keeps_capturing(caplog): 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 From 52b9627c7d612a3da1d6fd64687ff4edc108aac9 Mon Sep 17 00:00:00 2001 From: Alan Liu Date: Thu, 24 Sep 2026 19:28:52 -0400 Subject: [PATCH 14/14] Drain before reporting a latched failure at a disable_capture flush flush_records_and_wait rethrew the consumer's latched failure before draining, so under disable_capture a record emitted after the latch stayed in the ring: its payload was never delivered to be discarded, discarded_payloads stayed 0 and its ring space was not returned until something else forced a drain. The GPU case 'a sink flush that fails latches the record runtime' failed on exactly that (discarded_payloads == 1), and still failed with a 2 s poll. Under kDisableCapture a flush now waits for the stream prefix and forces a drain (bounded by the flush deadline) before rethrowing, so the discards are counted and the space is back when the caller hears why capture stopped. raise_at_producer is unchanged. test_ring_engine on GPU 1: 174 passed, 0 failed (was 173/1). CPU tier: 2369 passed. --- native/csrc/ring/ring_engine_py.cu | 12 ++++++++++++ tests/native/ring/test_ring_engine.cu | 8 ++++++++ 2 files changed, 20 insertions(+) diff --git a/native/csrc/ring/ring_engine_py.cu b/native/csrc/ring/ring_engine_py.cu index 77a498d13..541486193 100644 --- a/native/csrc/ring/ring_engine_py.cu +++ b/native/csrc/ring/ring_engine_py.cu @@ -801,6 +801,18 @@ 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); diff --git a/tests/native/ring/test_ring_engine.cu b/tests/native/ring/test_ring_engine.cu index c9500aa83..23284a108 100644 --- a/tests/native/ring/test_ring_engine.cu +++ b/tests/native/ring/test_ring_engine.cu @@ -1150,6 +1150,14 @@ static void test_a_failed_sink_flush_latches_the_runtime() { } 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));