Sink refusals and stalls never kill the forward: failure policy, per-step stall budget, bounded admission

docs/integration-api-v1.md

@@ -136,8 +136,12 @@ The predicates apply warmup, then offset, then stride. Step selection also
 honors `capture_prefill`/`capture_decode`; an unknown phase raises `ValueError`.
 `MonitoringConfig` carries this schedule plus three storage fields:
 `storage_backend`, `capture_sink_config` (a `NativeSinkConfig`, or `None`) and
-`capture_storage_config` (a `NativeCaptureStorageConfig` from
-`dmi.storage.native_capture`, or `None`).
+`capture_storage_config` (a `NativeCaptureStorageConfig`, or `None`), both
+from `dmi.storage.native_capture`. `NativeSinkConfig` bounds the native pack
+sink's queue and packs, and sets its admission policy: `overload="block"`
+(the default) waits up to `admission_timeout_s` (2 s; `None` waits without
+bound) for queue room before refusing a record, and `"drop_newest"` refuses
+it at once.
 
 `storage_backend` is the user's storage choice, one of
 `dmi.config.USER_STORAGE_CHOICES`:
@@ -408,21 +412,31 @@ 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,
     gate_tensor=None,
     gate_value=0,
 )
 
+runtime.begin_step()
 runtime.emit_output(entry, metadata, output)
 runtime.prepare_replay(plan, metadata)
 
 engine.flush_and_wait(timeout_s=600.0)
+engine.capture_status()
 ```
 
 The public operations are `create_record_runtime()`, `bind_hook()`,
-`emit_output()`, `prepare_replay()`, and `flush_and_wait()`.
+`begin_step()`, `emit_output()`, `prepare_replay()`, `flush_and_wait()`,
+`capture_status()`, and `validate_capture_bounds()`.
 
 `RecordRuntime.bind_hook()` assigns stable output IDs and binds the hook to the
 record ring. `RecordRuntime.emit_output()` reserves and publishes one eager
@@ -456,6 +470,89 @@ sink reached that sink's durability boundary within one timeout. The current
 ClickHouse adapter waits for acknowledged inserts. A timeout raises
 `TimeoutError`; asynchronous failures propagate unchanged.
 
+A sink can refuse a record (dropped, timed out, too large), lose one it
+admitted (a record that fits no empty pack, a duplicate id), or fail, and it
+does so on the record worker, after the forward that produced the record has
+moved on. The refusal latches the runtime, and `failure_policy` decides what
+the forward sees next:
+
+| `failure_policy` | After a latch |
+| --- | --- |
+| `"raise"` (default) | The next `RecordRuntime.begin_step()`, record reservation or descriptor push raises it. |
+| `"disable_capture"` | Capture stops; the forward keeps running. Descriptors and payloads that follow are discarded and counted. |
+
+Under both, `flush_and_wait()` raises the failure, `capture_status()` reports
+it, and `close()` logs it at WARNING. A `NativePackSink` reports a record lost
+after admission at its next submit and at the checked flush, as the reference
+adapter does, so such a loss also latches the runtime. Ring failures that are
+not sink refusals, such as a CUDA error in the drain, raise under both
+policies.
+
+A slow sink reaches the forward as a stall instead: a reservation that does
+not fit waits for the drain, which waits for the sink.
+`step_stall_budget_ms` caps the time the reservations of one step may wait.
+A step is what lies between two `RecordRuntime.begin_step()` calls. The
+runtime cannot find step boundaries itself (eager hooks reserve once per
+output, output ids repeat within a step, and a CUDA-graph step replays one
+plan, or several with piecewise graphs), so with a budget the integration
+calls `begin_step()` once per model step, before the step's first
+reservation, and a reservation before the first call is refused. Past the
+budget the rest of that step is skipped: the step's remaining records and
+every record still queued for the sink, from any step, are discarded, and
+capture resumes at the next step. The stall bound needs that whole queue out
+of the drain's way, so records of earlier, already finished steps that the
+sink had not taken yet are lost too. The discards are counted in
+`discarded_*`; `skipped_steps` counts the skips (always equal to
+`stall_budget_exhaustions`), and `steps_with_discards` the distinct steps
+that lost at least one record to them, which can be more.
+Under `"disable_capture"` a skipped step is not a failure: capture stays
+active and `flush_and_wait()` succeeds. Under `"raise"` the exhaustion is
+raised at the next `begin_step()` and at `flush_and_wait()`, never inside
+the forward that ran out of budget, and the runtime is failed from then on.
+`None`, the default, waits without bound.
+
+The skip still waits for the one envelope the sink is admitting, so a budget
+needs a sink whose admission is bounded, and `create_record_runtime()`
+refuses a budget otherwise, under either policy:
+
+- the ClickHouse host path (`record_sink=None` without the persistent
+  backend), which enqueues row by row under the host engine's ingress
+  policy, with no bound per envelope;
+- a `record_sink` whose `admission_bound_s` is None: a `NativePackSink` with
+  `overload="block"` and no `admission_timeout_s`, or the reference bridge;
+- a persistent `capture_sink_config` with `overload="block"` and
+  `admission_timeout_s=None`.
+
+A `NativePackSink` admits all rows of one envelope against one
+`admission_timeout_s`, so its `admission_bound_s` is that timeout under
+`"block"` and zero under `"drop_newest"`. With such a sink, a step stalls the
+forward for at most `step_stall_budget_ms` plus `admission_bound_s`, plus the
+time to move what the ring holds out to the record worker, which discards it
+without copying it to pageable memory. A ring that has not drained by the
+bound plus a drain grace fails from the reservation under either policy,
+rather than hanging the forward: the error cannot tell a sink that held its
+admission past its bound from a slow drain, and names both. The grace is
+2 s plus the bytes of the payload ring and its pinned staging at 1 GB/s,
+about 10.6 s with the engine's default 4 GiB of each.
+
+`MonitoringEngine.capture_status()` returns a plain dict: `record_mode`,
+`capture_active` (False once a failure latched, and under `"raise"` once a
+spent budget is waiting to be raised), `failure_policy`, `failure`,
+`discarded_descriptors`, `discarded_payloads`, `step_stall_budget_ms`,
+`stall_budget_exhaustions`, `skipped_steps`, `steps_with_discards`,
+`reserve_wait_s` and `max_step_wait_s` (time reservations waited for the
+sink, in total and in the worst step), and the `sink` and `storage`
+snapshots when the engine holds a native pack sink or storage service.
+Without a record runtime every field is empty.
+
+`MonitoringEngine.validate_capture_bounds(max_record_bytes)` refuses, with
+`ConfigurationError`, capture bounds under which the largest record an
+integration emits could not be stored: a `max_queue_bytes` below it, a
+`max_pack_bytes` below it plus 64 KiB of pack framing, or a `max_pack_bytes`
+above the storage service's `uploader_max_in_flight_bytes`. Integrations call
+it at attach, so these refusals never reach the forward. Without a
+`capture_sink_config` there is nothing to check.
+
 ### `deactivate_ring_transport`
 
 ```python

native/csrc/bindings.cpp

@@ -141,6 +141,14 @@ std::shared_ptr<dmx_host::ClickHouseRecordSink> 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<double> AdmissionBoundSeconds(const ring::RecordSink& sink) {
+  const auto bound = sink.admission_bound();
+  if (!bound) return std::nullopt;
+  return std::chrono::duration<double>(*bound).count();
+}
+
 template <typename... Args>
 std::shared_ptr<ring_py::RingEnginePy> MakeRingEngine(Args&&... args) {
   // Ring destruction may join a worker that is completing a callback.  Never
@@ -157,6 +165,22 @@ std::shared_ptr<ring_py::RingEnginePy> MakeRingEngine(Args&&... args) {
       });
 }
 
+ring::RecordFailurePolicy ParseRecordFailurePolicy(const std::string& name) {
+  if (name == "raise") return ring::RecordFailurePolicy::kRaiseAtProducer;
+  if (name == "disable_capture") {
+    return ring::RecordFailurePolicy::kDisableCapture;
+  }
+  throw py::value_error(
+      "failure_policy must be 'raise' or 'disable_capture', got '" + name +
+      "'");
+}
+
+const char* RecordFailurePolicyName(ring::RecordFailurePolicy policy) {
+  return policy == ring::RecordFailurePolicy::kDisableCapture
+             ? "disable_capture"
+             : "raise";
+}
+
 ring::RecordDescriptor CopyRecordDescriptor(
     const py::handle& descriptor_py,
     const dmx_host::RecordSchema& schema) {
@@ -630,7 +654,8 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
       .def("_acquire_engine",
            [](std::shared_ptr<ring::RecordSink> sink) {
              return ring::RecordSinkLease::acquire(std::move(sink));
-           });
+           })
+      .def_property_readonly("admission_bound_s", &AdmissionBoundSeconds);
   py::class_<dmx_host::ClickHouseRecordSink, ring::RecordSink,
              std::shared_ptr<dmx_host::ClickHouseRecordSink>>(
       m, "ClickHouseRecordSink")
@@ -641,12 +666,25 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
   py::class_<dmi_capture::ReferencePythonCaptureSink, ring::RecordSink,
              std::shared_ptr<dmi_capture::ReferencePythonCaptureSink>>(
       m, "ReferencePythonCaptureSink")
-      .def(py::init([](py::object target, std::string layout) {
+      .def(py::init([](py::object target, std::string layout,
+                       std::optional<double> admission_bound_s) {
+             std::optional<ring::RecordSink::Duration> 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<ring::RecordSink::Duration>(
+                   std::chrono::duration<double>(*admission_bound_s));
+             }
              return std::make_shared<
                  dmi_capture::ReferencePythonCaptureSink>(
-                     target.ptr(), std::move(layout));
+                     target.ptr(), std::move(layout), bound);
            }),
-           py::arg("target"), py::arg("layout"))
+           py::arg("target"), py::arg("layout"),
+           py::arg("admission_bound_s") = py::none())
       .def_property_readonly(
           "attached", &dmi_capture::ReferencePythonCaptureSink::engine_owned);
 
@@ -678,7 +716,12 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
            py::arg("config"), py::arg("host_engine") = py::none())
       .def_static(
           "create_record",
-          [](ring_py::RingConfig cfg, py::object sink_or_host) {
+          [](ring_py::RingConfig cfg, py::object sink_or_host,
+             const std::string& failure_policy,
+             uint64_t step_stall_budget_ms) {
+            ring_py::RecordRuntimeOptions options;
+            options.failure_policy = ParseRecordFailurePolicy(failure_policy);
+            options.step_stall_budget_ms = step_stall_budget_ms;
             std::shared_ptr<ring::RecordSinkLease> lease;
             if (!sink_or_host.is_none()) {
               if (py::isinstance<ring::RecordSinkLease>(sink_or_host)) {
@@ -694,9 +737,11 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
                     MakeClickHouseRecordSink(std::move(host)));
               }
             }
-            return MakeRingEngine(std::move(cfg), std::move(lease));
+            return MakeRingEngine(std::move(cfg), std::move(lease), options);
           },
-          py::arg("config"), py::arg("sink_or_host") = py::none())
+          py::arg("config"), py::arg("sink_or_host") = py::none(),
+          py::arg("failure_policy") = "raise",
+          py::arg("step_stall_budget_ms") = uint64_t{0})
       .def("init",  &ring_py::RingEnginePy::init,
            py::arg("stream_handle") = uint64_t{0})
       .def("start", &ring_py::RingEnginePy::start)
@@ -711,6 +756,26 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
            &ring_py::RingEnginePy::reserve_record,
            py::arg("reservation_items"),
            py::call_guard<py::gil_scoped_release>())
+      .def("begin_record_step", &ring_py::RingEnginePy::begin_record_step)
+      .def("capture_status",
+           [](const ring_py::RingEnginePy& self) {
+             const ring_py::RecordCaptureStatus status =
+                 self.record_capture_status();
+             py::dict out;
+             out["failure_policy"] =
+                 RecordFailurePolicyName(status.failure_policy);
+             out["failed"] = status.failed;
+             out["failure"] = status.failure;
+             out["discarded_descriptors"] = status.discarded_descriptors;
+             out["discarded_payloads"] = status.discarded_payloads;
+             out["step_stall_budget_ms"] = status.step_stall_budget_ms;
+             out["stall_budget_exhaustions"] = status.stall_budget_exhaustions;
+             out["skipped_steps"] = status.skipped_steps;
+             out["steps_with_discards"] = status.steps_with_discards;
+             out["reserve_wait_ns"] = status.reserve_wait_ns;
+             out["max_step_wait_ns"] = status.max_step_wait_ns;
+             return out;
+           })
       .def("push_record_descriptors",
            [](ring_py::RingEnginePy& self, py::sequence descriptors,
               py::object schema) {

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<Duration> 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");

native/csrc/reference_python_capture_sink.h

@@ -6,13 +6,19 @@
 
 #include <Python.h>
 
+#include <optional>
 #include <string>
 
 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<Duration> 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<Duration> 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<Duration> admission_bound_;
 };
 
 }  // namespace dmi_capture

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<RecordSink> sink)
-    : drain_(drain), consumer_(std::move(sink)) {}
+    DrainThread& drain, std::shared_ptr<RecordSink> sink,
+    RecordFailurePolicy failure_policy)
+    : drain_(drain), consumer_(std::move(sink), failure_policy) {}
 
 RecordP2PThread::~RecordP2PThread() noexcept {
     stop();
@@ -362,6 +363,17 @@ void RecordP2PThread::loop() {
 void RecordP2PThread::process(std::vector<DrainTask>& 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);

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<RecordSink> sink);
+    RecordP2PThread(DrainThread& drain, std::shared_ptr<RecordSink> sink,
+                    RecordFailurePolicy failure_policy =
+                        RecordFailurePolicy::kRaiseAtProducer);
     ~RecordP2PThread() noexcept;
 
     RecordP2PThread(const RecordP2PThread&) = delete;