diff --git a/docs/ring-consumers.md b/docs/ring-consumers.md new file mode 100644 index 000000000..6edbb4ad2 --- /dev/null +++ b/docs/ring-consumers.md @@ -0,0 +1,128 @@ +# In-process consumers of the legacy ring + +The legacy ring delivers each captured tensor, sliced per request, to one +consumer. Besides the C++ ClickHouse host engine (`DMXHostEngine`), a ring can +deliver to an in-process consumer in the same Python process: + +```python +from dmi.engine import MonitoringEngine +from dmi.transport import native + +queue = native.BatchQueue(max_bytes=512 << 20) +engine = MonitoringEngine(model_id="m", ring_consumer=queue) +... +for (model_id, shard_rank, act_name, layer_no, node_id, step_seq, + last_in_step, tensor, req_ids, segments, *lease) in queue.drain_batches(64, 50): + ... +``` + +| consumer | handoff | GIL on the P2P thread | +|---|---|---| +| `DMXHostEngine` | one row per (request, hook, step) to ClickHouse | no | +| `SliceQueue` | one item per (request, hook, step); `drain_slices()` | no | +| `BatchQueue` | one item per (hook, step) with a segments table; `drain_batches()` | no | +| a Python callable | one call per (request, hook, step) | yes, per call | + +`MonitoringEngine(ring_consumer=...)` accepts a `SliceQueue`, a `BatchQueue` +(which switches `RingConfig.batch_submit` on) or a callable. A ring built with +no consumer warns (`RingWithoutConsumerWarning`): its captures are drained and +dropped. + +## The batch handoff + +With `RingConfig.batch_submit` and a `BatchQueue`, a packed step hands each +hook's whole tensor over once, with an int32 `[N, 4]` segments table of +`(dim0_offset, db_start, db_end, orig_position)` rows and the matching +request ids; requests with no rows are omitted. `final_logits` rows are +indexed by the request's original position. + +Batched (HF) steps are not packed, so they cannot be handed over whole. With a +`BatchQueue` they arrive as **one-request items**, one per request slice -- +the per-slice fallback, delivered through the same queue: + +| capture | `tensor` | `req_ids` | `segments` | +|---|---|---|---| +| HF token hook, `[batch, q_len, hidden...]` | the request's rows, `[rows, hidden...]`: its last `db_end - db_start` query positions (all `q_len` if fewer) | `[req_id]` | `[[0, db_start, db_end, orig_position]]` | +| HF attention matrix, `[batch, heads, q_len, kv_len]` | the request's slice, `[heads, rows, keys]`: its query rows and its real keys | `[req_id]` | `[[0, db_start, db_end, orig_position]]` | +| HF `final_logits` | `[logits_to_keep, vocab]` | `[req_id]` | `[[0, end_token - logits_to_keep, end_token, orig_position]]` | + +For token hooks the usual reconstruction `tensor[dim0_offset : dim0_offset + +(db_end - db_start)]` yields the request's rows (`dim0_offset` is 0). An +attention item is taken whole: its first dimension is heads, not rows. The +attention hooks are `blocks.attn.hook_pattern` and +`blocks.attn.hook_attn_scores` (C++ batch sinks also receive an `is_attn` +flag). The items of one hook share its `step_seq`; `last_in_step` is set on +the last item of the step's last hook. A consumer exception on a fallback item +is counted in `submit_exceptions` like any other. Packed attention matrices +are not delivered on either path. + +The tensor is a view of a pooled pageable buffer: consumers must not keep a +view past their handling of the item (clone what they keep). Pool blocks +return to the pool when the last view dies, on any thread, also after the +engine stopped. `RingConfig.p2p_pool_cap_bytes` bounds the idle bytes the pool +retains; `0` disables the pool, and every step's buffer is then allocated at +its exact size (pool blocks are size classes: powers of two, 64 KiB at +least). + +## Backpressure and drop attribution + +`BatchQueue` is bounded by retained bytes, `SliceQueue` by items and, +optionally, retained bytes (`SliceQueue(max_bytes=...)`, 0 = no byte cap). +Retained bytes are the allocations the queued items keep alive, not their +logical lengths: a 32-byte payload in a 64 KiB pool block is charged 64 KiB, +and items that are views of one buffer (the slices of one step) share a single +charge until the last of them is drained. Item metadata (segments table, +request ids) is not charged; `dropped_bytes` counts logical payload bytes. An +empty queue always admits one item. On overflow the item is shed, counted, +and attributed: a drop notice names the requests that lost data +(`drain_drop_notices()`). No notice means nothing was lost; a +saturated notice buffer raises a sticky flag for one poll window, meaning +every request in flight over that window may be incomplete. + +`BatchQueue.set_layer_class_map([(act_name, layer, rank), ...], default)` +installs drop ranks (0 = needed every step, 1 = needed when the request ends, +2 = no deadline). On overflow the queue evicts queued items strictly less +urgent than the arriving one (worst rank first, oldest first) until it fits; +when evicting all of them still would not make room, it evicts nothing and +refuses the newcomer. Evicting a view frees its shared buffer only once no +remaining item, of any rank, and not the arriving item itself still refers to +it. Without a map it is a FIFO that drops the newcomer. + +## Capture leases + +`RingTransport.set_step_context(capture_lease=native.CaptureLease(id))` +attaches an integration-defined identity to a step. The lease arrives with +every slice, batch and drop notice of the step, and `lease.idle` turns true +once every payload carrying it has been consumed or dropped. A lease built as +`CaptureLease(id, [(act_name, layer, rank), ...])` carries a drop-rank map +that orders eviction for the steps sharing it (`CaptureLease(id, policy)` +shares the map without copying it). + +## Delivery barrier and consumer failures + +`RingEngine.drain_to_sink_and_wait(timeout_ms)` flushes the ring and waits +until every flushed task has been processed by the P2P thread -- every +consumer call for every slice or batch captured so far has returned. It +returns False on timeout (0 waits forever). A finite `timeout_ms` bounds the +whole barrier: one deadline, taken at entry, covers the producer stream's +completion, the flush (which can wait behind staging saturated by a slow +consumer) and the delivery wait. `flush_and_wait()` only waits for the +ring-to-P2P handoff. + +Pending step metadata belongs to the ring: when an engine is torn down with +steps whose payloads never arrived (a forward that failed after publishing its +metadata), their contexts are released with it, so their capture leases turn +idle. + +An exception thrown by a consumer is logged once and counted: +`submit_exceptions`, `last_sink_error`. With `set_abort_on_sink_error(True)`, +`sink_failed` turns true at the first one and the next `prepare_step` raises. + +## Freshness and flushing + +Every delivered payload carries `capture_ts_us`, the steady-clock time the +drain first saw the task published, so `now - capture_ts_us` measures the +export latency including the drain's flush policy +(`drain_flush_entry_threshold`, `drain_flush_timeout_us`, ...). +`RingEngine.request_flush()` asks the drain to flush every task reserved so +far as soon as it is published, without waiting on the producer stream. diff --git a/native/Makefile b/native/Makefile index 1e8c3de03..efaa24d77 100644 --- a/native/Makefile +++ b/native/Makefile @@ -28,7 +28,10 @@ HEADERS := csrc/dmx_host_utils.h \ csrc/ring/p2p_thread.h csrc/ring/pinned_staging.h \ csrc/ring/ring_engine.h csrc/ring/ring_engine_py.h \ csrc/ring/ring_torch_op.h csrc/ring/tensor_meta.h \ - csrc/ring/toggle_device.h csrc/ring/toggle_registry.h + csrc/ring/toggle_device.h csrc/ring/toggle_registry.h \ + csrc/ring/capture_lease.h csrc/ring/drop_notice.h \ + csrc/ring/slice_queue.h csrc/ring/batch_queue.h \ + csrc/ring/buffer_pool.h csrc/ring/pooled_block.h BUILD_DIR := build HOST_BUILD_DIR := $(BUILD_DIR)/host OBJS := $(patsubst csrc/%.cpp,$(BUILD_DIR)/%.o,$(SRCS)) diff --git a/native/csrc/bindings.cpp b/native/csrc/bindings.cpp index 6e448a63d..6e26bfa50 100644 --- a/native/csrc/bindings.cpp +++ b/native/csrc/bindings.cpp @@ -23,8 +23,12 @@ namespace py = pybind11; #ifndef DMI_HOST_ONLY #include "clickhouse_record_sink.h" #include "reference_python_capture_sink.h" +#include "ring/batch_queue.h" +#include "ring/capture_lease.h" +#include "ring/drop_notice.h" #include "ring/ring_engine_py.h" #include "ring/ring_torch_op.h" +#include "ring/slice_queue.h" #include "ring/tensor_meta.h" #endif @@ -264,6 +268,62 @@ std::vector CopyRecordDescriptors( #endif #ifndef DMI_HOST_ONLY +// DropNoticeBuffer::Drained -> Python dict, shared by SliceQueue and +// BatchQueue so both shed paths report the same shape: +// {"notices": [(source, hook, layer_no, step_seq, start_token, end_token, +// nbytes, ts_us, [req_id, ...], rank[, capture_lease]), ...], +// "overflow": int, # notices lost since the last poll +// "sticky": bool, # any notice lost in this window +// "pending": int} # notices still buffered (max_items truncation) +// An empty "notices" with sticky False means nothing was lost. rank 255 is +// the unranked slice path; the lease is appended only when the step had one. +py::dict drop_notices_to_py(ring_py::DropNoticeBuffer::Drained&& d) { + py::list notices; + for (auto& n : d.notices) { + py::list rids; + for (auto& r : n.req_ids) rids.append(py::str(r)); + py::tuple notice = py::make_tuple(py::str(n.source), py::str(n.hook), n.layer_no, + n.step_seq, n.start_token, n.end_token, + n.nbytes, n.ts_us, rids, static_cast(n.rank)); + if (n.capture_lease.present()) + notice = py::cast(notice.attr("__add__")( + py::make_tuple(n.capture_lease.fork()))); + notices.append(notice); + } + py::dict out; + out["notices"] = notices; + out["overflow"] = d.overflow; + out["sticky"] = d.sticky; + out["pending"] = d.pending; + return out; +} + +ring_py::HookBatchItem make_batch_item( + const std::string& model_id, int32_t shard_rank, const std::string& act_name, + int32_t layer_no, uint64_t step_seq, bool last_in_step, at::Tensor tensor, + std::vector req_ids, at::Tensor segments, + ring_py::CaptureLease capture_lease, int64_t capture_ts_us, + bool is_attn = false) { + ring_py::HookBatchItem item; + item.model_id = model_id; + item.shard_rank = shard_rank; + item.act_name = act_name; + item.layer_no = layer_no; + item.node_id = 0; // in-process + item.step_seq = step_seq; + item.last_in_step = last_in_step; + item.is_attn = is_attn; + // Logical bytes; BatchQueue::push charges the allocation the tensor keeps + // alive (pooled_block.h), not this length. + item.nbytes = static_cast(tensor.nbytes()); + item.tensor = std::move(tensor); + item.req_ids = std::move(req_ids); + item.segments = std::move(segments); + item.capture_lease = std::move(capture_lease); + item.capture_ts_us = capture_ts_us; + return item; +} + using HookRowsPtr = std::shared_ptr>; // One hook's request layout for push_all_metas. Two encodings: @@ -710,6 +770,185 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { // Ring offload engine // ------------------------------------------------------------------------- #ifndef DMI_HOST_ONLY + // ---- Capture leases (capture_lease.h) ----------------------------------- + py::class_(m, "CaptureLease") + .def(py::init<>()) + .def(py::init(), py::arg("id")) + .def(py::init>&>(), + py::arg("id"), py::arg("ranks")) + .def(py::init(), + py::arg("id"), py::arg("policy")) + .def_property_readonly("id", &ring_py::CaptureLease::id) + .def_property_readonly("present", &ring_py::CaptureLease::present) + .def_property_readonly("idle", &ring_py::CaptureLease::idle) + .def("same_as", &ring_py::CaptureLease::same_as, py::arg("other")) + .def("rank_for", &ring_py::CaptureLease::rank_for, + py::arg("act_name"), py::arg("layer_no"), py::arg("fallback")) + .def("fork", &ring_py::CaptureLease::fork); + + // ---- SliceQueue: per-slice in-process consumer (slice_queue.h) --------- + py::class_>(m, "SliceQueue") + .def(py::init(), py::arg("max_items") = 200000, + py::arg("drop_notice_cap") = 4096, py::arg("max_bytes") = uint64_t{0}) + // Normally pushed by the P2P thread; exposed for tests. + .def("push", &ring_py::SliceQueue::push, + py::arg("model_id"), py::arg("shard_rank"), py::arg("req_id"), + py::arg("act_name"), py::arg("layer_no"), py::arg("start_token"), + py::arg("end_token"), py::arg("slice"), + py::arg("capture_lease") = ring_py::CaptureLease{}, + py::arg("capture_ts_us") = 0, + py::call_guard()) + .def("stop", &ring_py::SliceQueue::stop, + py::call_guard()) + .def("wake", &ring_py::SliceQueue::wake, + py::call_guard()) + .def_property_readonly("dropped", &ring_py::SliceQueue::dropped) + .def_property_readonly("qsize", &ring_py::SliceQueue::size) + .def_property_readonly("retained_bytes", &ring_py::SliceQueue::retained_bytes) + .def_property_readonly("max_bytes", &ring_py::SliceQueue::max_bytes) + .def("drain_drop_notices", + [](ring_py::SliceQueue& self, size_t max_items) { + ring_py::DropNoticeBuffer::Drained d; + { + py::gil_scoped_release rel; + d = self.drain_drop_notices(max_items); + } + return drop_notices_to_py(std::move(d)); + }, + py::arg("max_items") = 4096) + .def_property_readonly("drop_notices_pending", + [](const ring_py::SliceQueue& self) { return self.drop_notices().pending(); }) + .def_property_readonly("drop_notices_recorded", + [](const ring_py::SliceQueue& self) { return self.drop_notices().recorded(); }) + .def_property_readonly("drop_notices_lost", + [](const ring_py::SliceQueue& self) { return self.drop_notices().overflow_total(); }) + .def_property_readonly("drop_notice_cap", + [](const ring_py::SliceQueue& self) { return self.drop_notices().cap(); }) + // Wait up to timeout_ms with the GIL released; each item is + // (model_id, shard_rank, req_id, act_name, layer_no, start_token, + // end_token, tensor[, None, capture_lease, capture_ts_us]). + .def("drain_slices", + [](ring_py::SliceQueue& self, size_t max_items, int timeout_ms) { + std::vector batch; + { + py::gil_scoped_release rel; + batch = self.drain(max_items, timeout_ms); + } + py::list out; + for (auto& it : batch) { + py::tuple item = py::make_tuple(it.model_id, it.shard_rank, it.req_id, + it.act_name, it.layer_no, it.start_token, + it.end_token, it.slice); + if (it.capture_lease.present()) + item = py::cast(item.attr("__add__")(py::make_tuple( + py::none(), it.capture_lease.fork(), it.capture_ts_us))); + out.append(item); + } + return out; + }, + py::arg("max_items") = 256, py::arg("timeout_ms") = 50); + + // ---- BatchQueue: whole-hook batch handoff (batch_queue.h) -------------- + py::class_>(m, "BatchQueue") + .def(py::init(), + py::arg("max_bytes") = uint64_t{512} * 1024 * 1024, + py::arg("drop_notice_cap") = 4096) + .def("stop", &ring_py::BatchQueue::stop, + py::call_guard()) + .def("wake", &ring_py::BatchQueue::wake, + py::call_guard()) + .def_property_readonly("dropped", &ring_py::BatchQueue::dropped) + .def_property_readonly("dropped_bytes", &ring_py::BatchQueue::dropped_bytes) + .def_property_readonly("qsize", &ring_py::BatchQueue::size) + .def_property_readonly("retained_bytes", &ring_py::BatchQueue::retained_bytes) + .def_property_readonly("pushed", &ring_py::BatchQueue::pushed) + .def_property_readonly("max_bytes", &ring_py::BatchQueue::max_bytes) + .def("set_layer_class_map", &ring_py::BatchQueue::set_layer_class_map, + py::arg("entries"), + py::arg("default_rank") = static_cast(ring_py::kRankBestEffort), + py::call_guard()) + .def("set_class_eviction", &ring_py::BatchQueue::set_class_eviction, + py::arg("on"), py::call_guard()) + .def("rank_for", &ring_py::BatchQueue::rank_for, + py::arg("act_name"), py::arg("layer_no")) + .def_property_readonly("dropped_by_rank", &ring_py::BatchQueue::dropped_by_rank) + .def_property_readonly("dropped_bytes_by_rank", + &ring_py::BatchQueue::dropped_bytes_by_rank) + .def_property_readonly("qsize_by_rank", &ring_py::BatchQueue::qsize_by_rank) + .def_property_readonly("evicted", &ring_py::BatchQueue::evicted) + .def_property_readonly("class_eviction", &ring_py::BatchQueue::class_eviction) + .def_property_readonly("class_map_set", &ring_py::BatchQueue::class_map_set) + .def_property_readonly("class_map_size", &ring_py::BatchQueue::class_map_size) + .def_property_readonly("class_default_rank", + &ring_py::BatchQueue::class_default_rank) + .def("drain_drop_notices", + [](ring_py::BatchQueue& self, size_t max_items) { + ring_py::DropNoticeBuffer::Drained d; + { + py::gil_scoped_release rel; + d = self.drain_drop_notices(max_items); + } + return drop_notices_to_py(std::move(d)); + }, + py::arg("max_items") = 4096) + .def_property_readonly("drop_notices_pending", + [](const ring_py::BatchQueue& self) { return self.drop_notices().pending(); }) + .def_property_readonly("drop_notices_recorded", + [](const ring_py::BatchQueue& self) { return self.drop_notices().recorded(); }) + .def_property_readonly("drop_notices_lost", + [](const ring_py::BatchQueue& self) { return self.drop_notices().overflow_total(); }) + .def_property_readonly("drop_notice_cap", + [](const ring_py::BatchQueue& self) { return self.drop_notices().cap(); }) + // Normally pushed by the P2P thread's batch sink; exposed for tests and + // benchmarks that feed prebuilt items. + .def("push", + [](ring_py::BatchQueue& self, + const std::string& model_id, int32_t shard_rank, + const std::string& act_name, int32_t layer_no, + uint64_t step_seq, bool last_in_step, + at::Tensor tensor, std::vector req_ids, + at::Tensor segments, ring_py::CaptureLease capture_lease, + int64_t capture_ts_us) { + self.push(make_batch_item(model_id, shard_rank, act_name, layer_no, + step_seq, last_in_step, std::move(tensor), + std::move(req_ids), std::move(segments), + std::move(capture_lease), capture_ts_us)); + }, + py::arg("model_id"), py::arg("shard_rank"), py::arg("act_name"), + py::arg("layer_no"), py::arg("step_seq"), py::arg("last_in_step"), + py::arg("tensor"), py::arg("req_ids"), py::arg("segments"), + py::arg("capture_lease") = ring_py::CaptureLease{}, + py::arg("capture_ts_us") = 0, + py::call_guard()) + // Wait up to timeout_ms with the GIL released; each item is + // (model_id, shard_rank, act_name, layer_no, node_id, step_seq, + // last_in_step, tensor, req_ids, segments[, capture_lease, + // capture_ts_us]). + .def("drain_batches", + [](ring_py::BatchQueue& self, size_t max_items, int timeout_ms) { + std::vector batch; + { + py::gil_scoped_release rel; + batch = self.drain(max_items, timeout_ms); + } + py::list out; + for (auto& it : batch) { + py::list rids; + for (auto& r : it.req_ids) rids.append(py::str(r)); + py::tuple item = py::make_tuple( + it.model_id, it.shard_rank, it.act_name, it.layer_no, + it.node_id, it.step_seq, it.last_in_step, + it.tensor, rids, it.segments); + if (it.capture_lease.present()) + item = py::cast(item.attr("__add__")( + py::make_tuple(it.capture_lease.fork(), it.capture_ts_us))); + out.append(item); + } + return out; + }, + py::arg("max_items") = 64, py::arg("timeout_ms") = 50); + py::class_(m, "RingConfig") .def(py::init<>()) .def_readwrite("task_ring_entries", &ring_py::RingConfig::task_ring_entries) @@ -722,6 +961,8 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { .def_readwrite("drain_flush_byte_threshold", &ring_py::RingConfig::drain_flush_byte_threshold) .def_readwrite("drain_flush_timeout_us", &ring_py::RingConfig::drain_flush_timeout_us) .def_readwrite("clone_slices", &ring_py::RingConfig::clone_slices) + .def_readwrite("p2p_pool_cap_bytes", &ring_py::RingConfig::p2p_pool_cap_bytes) + .def_readwrite("batch_submit", &ring_py::RingConfig::batch_submit) .def_readwrite("insert_queue_max_bytes", &ring_py::RingConfig::insert_queue_max_bytes) .def_readwrite("insert_queue_max_items", &ring_py::RingConfig::insert_queue_max_items) .def_readwrite("drain_policy", &ring_py::RingConfig::drain_policy); @@ -773,13 +1014,53 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { py::class_>(m, "RingEngine") .def(py::init([](ring_py::RingConfig cfg, py::object host_engine_obj) { ring_py::SubmitFn submit_fn; - if (!host_engine_obj.is_none() && - py::isinstance(host_engine_obj)) { + ring_py::BatchSink batch_sink; + if (host_engine_obj.is_none()) { + // No consumer: captures are drained and dropped (null and + // benchmark runs). + } else if (py::isinstance(host_engine_obj)) { + // Batch handoff: one item per (hook x step), pushed in pure + // C++ (no GIL). Needs cfg.batch_submit for packed steps; + // a consumer drains it with BatchQueue.drain_batches(). + // HF batched steps (token and attention hooks) arrive as + // one-request items, one per request slice (p2p_thread.h). + if (!cfg.batch_submit) { + throw py::value_error( + "a BatchQueue consumer needs RingConfig.batch_submit = True"); + } + auto bq = host_engine_obj.cast>(); + batch_sink = [bq](const std::string& model_id, int32_t shard_rank, + const std::string& act_name, int32_t layer_no, + uint64_t step_seq, bool last_in_step, + at::Tensor tensor, std::vector req_ids, + at::Tensor segments, ring_py::CaptureLease capture_lease, + int64_t capture_ts_us, bool is_attn) { + bq->push(make_batch_item(model_id, shard_rank, act_name, layer_no, + step_seq, last_in_step, std::move(tensor), + std::move(req_ids), std::move(segments), + std::move(capture_lease), capture_ts_us, + is_attn)); + }; + } else if (py::isinstance(host_engine_obj)) { + // Per-slice queue, pushed in pure C++ (no GIL); a consumer + // drains it with SliceQueue.drain_slices(). + auto sq = host_engine_obj.cast>(); + submit_fn = [sq](const std::string& model_id, int32_t shard_rank, + const std::string& req_id, const std::string& act_name, + int32_t layer_no, int32_t start_token, int32_t end_token, + at::Tensor slice, int64_t capture_ts_us, + ring_py::CaptureLease capture_lease) { + sq->push(model_id, shard_rank, req_id, act_name, layer_no, + start_token, end_token, std::move(slice), + std::move(capture_lease), capture_ts_us); + }; + } else if (py::isinstance(host_engine_obj)) { auto host = host_engine_obj.cast>(); submit_fn = [host](const std::string& model_id, int32_t shard_rank, const std::string& req_id, const std::string& act_name, int32_t layer_no, int32_t start_token, int32_t end_token, - at::Tensor slice) { + at::Tensor slice, int64_t /*capture_ts_us*/, + ring_py::CaptureLease /*capture_lease*/) { dmx_host::ClickHouseRow row; row.emplace_back(model_id); row.emplace_back(req_id); @@ -792,17 +1073,18 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { row.emplace_back(std::move(slice)); host->submit_direct(std::move(row), nbytes); }; - } else if (!host_engine_obj.is_none()) { + } else { // An in-process consumer: any Python callable taking // (model_id, shard_rank, req_id, act_name, layer_no, - // start_token, end_token, tensor). The P2P thread calls - // it without the GIL, so take it for the call. An - // exception propagates to the P2P thread's submit-failure - // handling (logged once, the slice is dropped). + // start_token, end_token, tensor), plus (capture_lease, + // capture_ts_us) when the step carried a lease. The P2P + // thread calls it without the GIL, so take it for the call. + // An exception reaches the P2P thread's consumer-failure + // accounting (logged once, counted, slice dropped). if (!PyCallable_Check(host_engine_obj.ptr())) { throw py::type_error( "RingEngine host_engine must be a DMXHostEngine, a " - "callable consumer, or None"); + "BatchQueue, a SliceQueue, a callable consumer, or None"); } auto fn = std::shared_ptr( new py::object(host_engine_obj), @@ -814,19 +1096,26 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { submit_fn = [fn](const std::string& model_id, int32_t shard_rank, const std::string& req_id, const std::string& act_name, int32_t layer_no, int32_t start_token, int32_t end_token, - at::Tensor slice) { + at::Tensor slice, int64_t capture_ts_us, + ring_py::CaptureLease capture_lease) { py::gil_scoped_acquire gil; try { - (*fn)(model_id, shard_rank, req_id, act_name, layer_no, - start_token, end_token, std::move(slice)); + if (capture_lease.present()) { + (*fn)(model_id, shard_rank, req_id, act_name, layer_no, + start_token, end_token, std::move(slice), + capture_lease.fork(), capture_ts_us); + } else { + (*fn)(model_id, shard_rank, req_id, act_name, layer_no, + start_token, end_token, std::move(slice)); + } } catch (py::error_already_set& error) { std::string what = error.what(); throw std::runtime_error(what); } }; } - return MakeRingEngine( - std::move(cfg), std::move(submit_fn)); + return MakeRingEngine(std::move(cfg), std::move(submit_fn), + std::move(batch_sink)); }), py::arg("config"), py::arg("host_engine") = py::none()) .def_static( @@ -960,6 +1249,22 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { .def("flush_and_wait", &ring_py::RingEnginePy::flush_and_wait, py::call_guard()) + .def("drain_to_sink_and_wait", + &ring_py::RingEnginePy::drain_to_sink_and_wait, + py::arg("timeout_ms") = uint32_t{0}, + py::call_guard()) + .def("request_flush", + &ring_py::RingEnginePy::request_flush, + py::call_guard()) + .def_property_readonly("submit_exceptions", + &ring_py::RingEnginePy::submit_exceptions) + .def_property_readonly("last_sink_error", + &ring_py::RingEnginePy::last_sink_error) + .def_property_readonly("sink_failed", + &ring_py::RingEnginePy::sink_failed) + .def("set_abort_on_sink_error", + &ring_py::RingEnginePy::set_abort_on_sink_error, + py::arg("enabled")) .def("push_all_metas", [](ring_py::RingEnginePy& self, py::list hook_types_py, @@ -978,7 +1283,8 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { py::list dim0_offsets_py, py::list kv_offsets_py, py::list orig_positions_py, - py::list hook_requests_py) { + py::list hook_requests_py, + ring_py::CaptureLease capture_lease) { const size_t n_hooks = static_cast(py::len(hook_types_py)); if (py::len(hook_requests_py) && static_cast(py::len(hook_requests_py)) != n_hooks) @@ -986,6 +1292,7 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { // Build step context (heap-allocated, ownership to FIFO/p2p) auto owned_ctx = std::make_unique(); auto* ctx = owned_ctx.get(); + ctx->capture_lease = std::move(capture_lease); ctx->model_id = model_id; ctx->tp_rank = tp_rank; ctx->dp_rank = dp_rank; @@ -1053,7 +1360,8 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { py::arg("dim0_offsets"), py::arg("kv_offsets") = py::list(), py::arg("orig_positions") = py::list(), - py::arg("hook_requests") = py::list()) + py::arg("hook_requests") = py::list(), + py::arg("capture_lease") = ring_py::CaptureLease{}) .def_static("supports_hook_request_arrays", []() { return true; }) .def("set_gather_host_mirror", &ring_py::RingEnginePy::set_gather_host_mirror, diff --git a/native/csrc/ring/batch_queue.h b/native/csrc/ring/batch_queue.h new file mode 100644 index 000000000..3c531965d --- /dev/null +++ b/native/csrc/ring/batch_queue.h @@ -0,0 +1,410 @@ +// ring/batch_queue.h -- whole-HookPoint in-process consumer queue (the batch +// handoff). Companion to slice_queue.h. +// +// The per-slice path hands one item per (request x hook x step) across the +// C++ -> Python boundary: with many concurrent requests and many hooks that is +// thousands of handoffs per step, each paying a fixed cost (a locked push with +// string copies, a pybind tuple build, the consumer's routing) that is +// independent of what the consumer computes. +// +// The batch handoff hands one item per (hook x step) instead: the WHOLE +// contiguous tensor the producer wrote (a view of the P2P thread's pageable +// buffer -- never sliced, never cloned here), an int32 segments table and the +// request-id list. The consumer reconstructs the per-request views on its own +// thread, or consumes the batch as a whole. +// +// Lifetime: the tensor aliases a pooled pageable buffer (buffer_pool.h). +// Consumers must not retain a view of it past their handling of the item; +// anything kept must be cloned. clone_slices does not apply on this path. +// +// Backpressure: capped by RETAINED BYTES (item counts are meaningless when one +// item can be tens of MB). On overflow whole batches are shed, counted and +// attributed: every shed item files a DropNotice (drop_notice.h) naming the +// requests that were on it. +// +// Retained bytes are the payload ALLOCATIONS the queued items keep alive, not +// their logical lengths: a pooled 32-byte tensor pins its whole 64 KiB size +// class (pooled_block.h). Each distinct allocation is charged once, at +// max(logical bytes, allocated bytes), while any queued item refers to it -- +// one-request items cut from one step's buffer share one charge. Item +// metadata (segments table, request ids, act_name) is not charged: it is +// 16 bytes plus one id per request, small next to the payload. nbytes and +// dropped_bytes stay logical payload bytes. +// +// Rank-ordered eviction: each item carries the drop rank of its +// (act_name, layer) stream, resolved at enqueue from the installed map +// (set_layer_class_map, or the step's CaptureLease). On overflow the queue +// first makes room by EVICTING queued items strictly less urgent than the +// arriving one -- worst rank first, oldest first within a rank -- and refuses +// the newcomer only when nothing less urgent is left. Consequences: +// * a rank-0 batch is evicted only when every remaining item is rank 0; +// * with no map installed every item ranks best-effort, no candidate is +// ever strictly less urgent, and the queue behaves exactly like a FIFO +// that drops the newcomer; +// * set_class_eviction(false) keeps the map (per-rank loss is still +// attributed) but restores drop-the-newcomer. +// Every eviction files a DropNotice with source "batch_queue_evict". +// +// Header-only; included where ATen is already linked. + +#pragma once +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include "ring/drop_notice.h" +#include "ring/pooled_block.h" + +namespace ring_py { + +// One step x one hook. +struct HookBatchItem { + CaptureLease capture_lease; + int64_t capture_ts_us = 0; + std::string model_id; + int32_t shard_rank = 0; + std::string act_name; + int32_t layer_no = -1; + int32_t node_id = 0; // 0 in-process; wire placeholder + uint64_t step_seq = 0; // monotone step counter + bool last_in_step = false; + // The whole reshaped [total_tokens, ...] tensor -- a VIEW of the p2p pool + // buffer. Never sliced, never cloned here. Keeps the pool block alive + // until this item (and every view the consumer derives) is dropped. + at::Tensor tensor; + // Request ids, in segments row order (the step's request order, minus + // requests with no rows). + std::vector req_ids; + // int32 CPU [N,4]: (dim0_offset, db_start, db_end, orig_position). Row j + // gives request j's view: tensor[dim0_offset : dim0_offset + (db_end - + // db_start)] for token hooks; final_logits uses one row at dim0_offset + // (the request's original position). + at::Tensor segments; + // Logical payload bytes (tensor.nbytes()); reported in drop notices. + uint64_t nbytes = 0; + // A one-request item of the per-slice fallback holding an attention + // matrix slice [heads, rows, kv]: take the tensor whole. + bool is_attn = false; + // The allocation the tensor keeps alive (set by push; charged once per + // distinct allocation for the byte cap). + ring::BackingRef backing; + // Drop rank of this (act_name, layer) stream, resolved at enqueue. + // Best-effort by default: with no map every item ranks the same, which is + // FIFO drop-the-newcomer. + uint8_t rank = kRankBestEffort; +}; + +class BatchQueue { +public: + // max_bytes: retained-byte cap. When a push would exceed it, the whole + // item is dropped, counted AND attributed. Default 512 MiB. + // drop_notice_cap: notices retained between consumer polls. + explicit BatchQueue(uint64_t max_bytes = 512ULL * 1024 * 1024, + size_t drop_notice_cap = 4096) + : max_bytes_(max_bytes), drops_(drop_notice_cap) {} + + // Install the (act_name, layer) -> drop-rank map, normally once before the + // producer runs (safe later: guarded by mu_, affects later items only). + // `entries` is (act_name, layer_no, rank); `default_rank` covers every + // stream not named. Installing a map switches the queue from FIFO to + // rank-ordered eviction. + void set_layer_class_map( + const std::vector>& entries, + int default_rank = static_cast(kRankBestEffort)) { + std::lock_guard lk(mu_); + act_ids_.clear(); + rank_of_.clear(); + default_rank_ = clamp_rank(default_rank); + for (const auto& e : entries) { + const std::string& act = std::get<0>(e); + auto it = act_ids_.find(act); + if (it == act_ids_.end()) + it = act_ids_.emplace(act, static_cast(act_ids_.size())).first; + rank_of_[key_of(it->second, std::get<1>(e))] = clamp_rank(std::get<2>(e)); + } + class_map_set_ = true; + } + + // Keep the map installed (per-rank loss is still attributed and counted) + // but stop using rank to choose the victim: off => drop-the-newcomer. + void set_class_eviction(bool on) { + std::lock_guard lk(mu_); + class_eviction_ = on; + } + + // Pushed by the P2P thread in pure C++ (never the GIL). Bounded by + // retained bytes. On overflow: first try to make room by evicting queued + // items that are STRICTLY LESS URGENT than this one (worst rank first, FIFO + // within a rank); if that cannot free enough, refuse the newcomer. Either + // way every shed item is counted AND attributed (a DropNotice naming its + // req_ids). + void push(HookBatchItem item) { + bool shed = false; + std::vector evicted; // notices filed outside mu_ + item.backing = ring::backing_of(item.tensor); + { + std::lock_guard lk(mu_); + item.rank = item.capture_lease.rank_for(item.act_name, item.layer_no, + rank_locked(item.act_name, item.layer_no)); + if (retained_bytes_ + ledger_.cost(item.backing) > max_bytes_ && !q_.empty()) { + if (class_eviction_) + make_room_locked(item, evicted); + if (retained_bytes_ + ledger_.cost(item.backing) > max_bytes_ && !q_.empty()) { + // Nothing less urgent left to give up: drop the whole + // batch. An empty queue always admits one item, even one + // larger than the cap, so a tiny cap cannot wedge the + // pipeline (retained_bytes reports the overshoot). + dropped_++; + dropped_bytes_ += item.nbytes; + dropped_by_rank_[item.rank]++; + dropped_bytes_by_rank_[item.rank] += item.nbytes; + shed = true; + } + } + if (!shed) { + charge_locked(item); + pushed_++; + depth_by_rank_[item.rank]++; + q_.push_back(std::move(item)); + } + } + // Attribution runs outside mu_ (the notice buffer has its own lock: no + // nesting, no lock-order hazard). + for (auto& e : evicted) { + DropNotice n; + n.source = kDropSourceBatchEvict; + n.hook = std::move(e.act_name); + n.layer_no = e.layer_no; + n.step_seq = e.step_seq; + n.nbytes = e.nbytes; + n.ts_us = drop_notice_now_us(); + n.rank = e.rank; + n.capture_lease = std::move(e.capture_lease); + n.req_ids = std::move(e.req_ids); + drops_.record(std::move(n)); + } + if (shed) { + // `item` was not moved from on this branch, so its req_ids can be + // moved straight into the notice -- no copy on the shed path. + DropNotice n; + n.source = kDropSourceBatch; + n.hook = std::move(item.act_name); + n.layer_no = item.layer_no; + n.step_seq = item.step_seq; + n.nbytes = item.nbytes; + n.ts_us = drop_notice_now_us(); + n.rank = item.rank; + n.capture_lease = std::move(item.capture_lease); + n.req_ids = std::move(item.req_ids); + drops_.record(std::move(n)); + return; + } + cv_.notify_one(); + } + + // Wake any waiting consumer so it can observe shutdown / final items. + void stop() { + { std::lock_guard lk(mu_); stopped_ = true; } + cv_.notify_all(); + } + + // Control work lives in the host queue. Wake its consumer without a fake + // tensor or a permanent stop. Remember the wake if it arrives before wait. + void wake() { + { std::lock_guard lk(mu_); wake_pending_ = true; } + cv_.notify_one(); + } + + // Pop up to max_items, waiting up to timeout_ms for the first one. MUST be + // called with the GIL RELEASED (the binding wraps it in gil_scoped_release) + // so the blocking wait never holds the GIL. + std::vector drain(size_t max_items, int timeout_ms) { + std::vector out; + std::unique_lock lk(mu_); + if (q_.empty() && !stopped_) { + cv_.wait_for(lk, std::chrono::milliseconds(timeout_ms), + [this] { return !q_.empty() || stopped_ || wake_pending_; }); + } + wake_pending_ = false; + while (!q_.empty() && out.size() < max_items) { + const HookBatchItem& taken = out.emplace_back(std::move(q_.front())); + release_locked(taken); + if (depth_by_rank_[taken.rank]) depth_by_rank_[taken.rank]--; + q_.pop_front(); + } + return out; + } + + // Drain the drop notices filed since the last poll (see drop_notice.h). + // Non-blocking. Absence of notices == nothing was lost. + DropNoticeBuffer::Drained drain_drop_notices(size_t max_items = 4096) { + return drops_.drain(max_items); + } + const DropNoticeBuffer& drop_notices() const { return drops_; } + + uint64_t dropped() const { std::lock_guard lk(mu_); return dropped_; } + uint64_t dropped_bytes() const { std::lock_guard lk(mu_); return dropped_bytes_; } + size_t size() const { std::lock_guard lk(mu_); return q_.size(); } + uint64_t retained_bytes() const { std::lock_guard lk(mu_); return retained_bytes_; } + // Cumulative count of batches ever enqueued (monotone; survives drain): + // > 0 proves the batch path actually fired. + uint64_t pushed() const { std::lock_guard lk(mu_); return pushed_; } + uint64_t max_bytes() const { std::lock_guard lk(mu_); return max_bytes_; } + + // -- per-rank counters ------------------------------------------------ // + // Who lost what, by rank. dropped_by_rank sums to dropped() and + // dropped_bytes_by_rank to dropped_bytes(), evictions included. + std::vector dropped_by_rank() const { + std::lock_guard lk(mu_); + return {dropped_by_rank_, dropped_by_rank_ + kRankCount}; + } + std::vector dropped_bytes_by_rank() const { + std::lock_guard lk(mu_); + return {dropped_bytes_by_rank_, dropped_bytes_by_rank_ + kRankCount}; + } + // Current queue depth per rank. + std::vector qsize_by_rank() const { + std::lock_guard lk(mu_); + return {depth_by_rank_, depth_by_rank_ + kRankCount}; + } + // The subset of dropped() that was EVICTED from the queue to make room for a + // more urgent arrival (the rest were refused at the door). 0 under FIFO. + uint64_t evicted() const { std::lock_guard lk(mu_); return evicted_; } + bool class_eviction() const { std::lock_guard lk(mu_); return class_eviction_; } + bool class_map_set() const { std::lock_guard lk(mu_); return class_map_set_; } + size_t class_map_size() const { std::lock_guard lk(mu_); return rank_of_.size(); } + uint8_t class_default_rank() const { std::lock_guard lk(mu_); return default_rank_; } + // The rank an (act_name, layer) item would get under the installed map. + uint8_t rank_for(const std::string& act, int32_t layer) const { + std::lock_guard lk(mu_); return rank_locked(act, layer); + } + +private: + static uint64_t key_of(uint32_t act_id, int32_t layer) { + return (static_cast(act_id) << 32) | + static_cast(layer); // layer -1 keys fine, cast is total + } + + // Call with mu_ held. Charge / release the allocation an item keeps + // alive; an allocation shared by several queued items is charged once and + // released when the last of them leaves. The charge belongs to no rank: + // which items keep a shared allocation alive changes as they come and go, + // so eviction asks the ledger what a set of victims would actually free + // (make_room_locked) instead of keeping bytes per rank. + void charge_locked(const HookBatchItem& it) { + retained_bytes_ += ledger_.add(it.backing); + } + void release_locked(const HookBatchItem& it) { + retained_bytes_ -= ledger_.remove(it.backing); + } + + // Call with mu_ held. No map installed => every item is best-effort => + // all ranks equal => FIFO. + uint8_t rank_locked(const std::string& act, int32_t layer) const { + if (!class_map_set_) return kRankBestEffort; + auto a = act_ids_.find(act); + if (a == act_ids_.end()) return default_rank_; + auto r = rank_of_.find(key_of(a->second, layer)); + return r == rank_of_.end() ? default_rank_ : r->second; + } + + // Call with mu_ held. Evict queued items STRICTLY less urgent + // than `in` -- worst (highest) rank first, oldest-first within that rank -- + // the fewest that admit `in`. Moves the victims into `out` so the caller + // can file their notices after releasing mu_. + // + // "Rank 0 is evicted only when nothing else remains" falls out of the + // strict comparison: a rank-0 arrival can evict ranks 1 and 2, and a + // rank-0 item is never a candidate. Two rank-0 items contend by the FIFO + // rule: the newcomer is refused, never the resident. + // + // All or nothing: the victims are chosen first, without touching the + // queue, and evicted only if they admit `in`; otherwise nothing is + // evicted and the caller refuses `in`. What a victim frees is decided per + // allocation, from the references that would REMAIN: evicting a view of a + // shared allocation frees it only when no surviving item -- of any rank + // -- and not `in` itself still refers to it. So + // retained after = retained_bytes_ - freed + cost(in), + // where cost(in) is 0 when `in` aliases an allocation already queued (its + // charge is never counted as freed: `in` keeps it). Evicting everything + // queued empties the queue, which always admits. + void make_room_locked(const HookBatchItem& in, std::vector& out) { + const uint64_t in_cost = ledger_.cost(in.backing); + auto admits = [&](uint64_t freed) { + return retained_bytes_ - freed + in_cost <= max_bytes_; + }; + std::vector victims; // indices into q_ + std::unordered_map leaving; // refs the victims drop + uint64_t freed = 0; + bool fits = false; + for (int r = static_cast(kRankMax); + r > static_cast(in.rank) && !fits; --r) { + if (!depth_by_rank_[r]) continue; + for (size_t i = 0; i < q_.size() && !fits; ++i) { + const HookBatchItem& it = q_[i]; + if (it.rank != static_cast(r)) continue; + victims.push_back(i); + const ring::BackingRef& b = it.backing; + if (!b.key) { + freed += b.bytes; // charged per item + } else if (b.key != in.backing.key) { + const auto held = ledger_.held(b.key); + if (++leaving[b.key] == held.first) freed += held.second; + } + fits = admits(freed); + } + } + if (!fits && victims.size() < q_.size()) return; // evict nothing + if (victims.empty()) return; + // Commit, in the order chosen (notices follow eviction order). + for (size_t i : victims) { + HookBatchItem& it = q_[i]; + release_locked(it); + if (depth_by_rank_[it.rank]) depth_by_rank_[it.rank]--; + dropped_++; + evicted_++; + dropped_bytes_ += it.nbytes; + dropped_by_rank_[it.rank]++; + dropped_bytes_by_rank_[it.rank] += it.nbytes; + out.push_back(std::move(it)); + } + std::sort(victims.begin(), victims.end()); + for (auto i = victims.rbegin(); i != victims.rend(); ++i) + q_.erase(q_.begin() + static_cast(*i)); + } + + mutable std::mutex mu_; + std::condition_variable cv_; + std::deque q_; + uint64_t max_bytes_; + uint64_t retained_bytes_ = 0; // == ledger_.total() + ring::BackingLedger ledger_; + uint64_t dropped_ = 0; + uint64_t dropped_bytes_ = 0; + uint64_t pushed_ = 0; + uint64_t evicted_ = 0; + bool stopped_ = false; + bool wake_pending_ = false; + // -- rank map + per-rank ledger (all guarded by mu_) ------------------- // + std::unordered_map act_ids_; // act_name -> dense id + std::unordered_map rank_of_; // (act_id,layer) -> rank + uint8_t default_rank_ = kRankBestEffort; + bool class_map_set_ = false; + bool class_eviction_ = true; + uint64_t dropped_by_rank_[kRankCount] = {0, 0, 0}; + uint64_t dropped_bytes_by_rank_[kRankCount] = {0, 0, 0}; + uint64_t depth_by_rank_[kRankCount] = {0, 0, 0}; + // Own mutex; never taken while mu_ is held (see push()). + DropNoticeBuffer drops_; +}; + +} // namespace ring_py diff --git a/native/csrc/ring/buffer_pool.h b/native/csrc/ring/buffer_pool.h new file mode 100644 index 000000000..b8219055e --- /dev/null +++ b/native/csrc/ring/buffer_pool.h @@ -0,0 +1,173 @@ +// ring/buffer_pool.h -- refcount-recycled pageable buffers for the P2P thread. +// +// The P2P thread copies each step's tensor from the pinned staging ring into a +// pageable CPU buffer, then slices it per request or hands it over whole. +// PyTorch's default CPU allocator returns large blocks to the OS (glibc mmap +// threshold), so a fresh at::empty per step is freshly mmap'd and the copy +// takes cold page faults on every touched page. This pool keeps a small set +// of size-classed buffers resident so their pages stay warm across steps. +// +// Lifetime: a buffer is not consumed inside the P2P thread. Its views go to +// a consumer that may hold them for an unbounded time -- past the end of the +// step, and past the ring engine's own shutdown. A buffer is therefore +// wrapped with at::from_blob(..., deleter) and returned only when the LAST +// reference across all derived views is dropped, on whatever thread drops +// it. The deleter holds a reference to the pool's shared state, never to +// the BufferPool object, so a tensor outliving the pool frees its block +// instead of touching a destroyed pool. +// +// Accounting: a block is a size class (power of two, 64 KiB floor) larger +// than the tensor it backs. Its DataPtr context is a PooledBlockContext +// (pooled_block.h) carrying the real capacity, so a queue that retains the +// tensor -- or any view of it -- can charge the allocation it keeps alive +// rather than the logical length (backing_of()). +// +// Header-only; included only where ATen is already linked. + +#pragma once +#include "ring/pooled_block.h" + +#include +#include +#include +#include +#include +#include +#include + +namespace ring { + +class BufferPool { +public: + // idle_cap_bytes: maximum retained (free-list) bytes. A returned buffer + // that would push the free lists past it is freed instead. 0 disables + // retention (no pooling). + explicit BufferPool(uint64_t idle_cap_bytes = 256ULL * 1024 * 1024) + : state_(std::make_shared(idle_cap_bytes)) {} + + ~BufferPool() { + // Free the idle blocks now; blocks still referenced by live tensors + // are freed by their deleters when the last reference drops. + std::lock_guard lk(state_->mu); + state_->closed = true; + for (auto& kv : state_->free_lists) + for (void* p : kv.second) std::free(p); + state_->free_lists.clear(); + state_->idle_bytes = 0; + } + + BufferPool(const BufferPool&) = delete; + BufferPool& operator=(const BufferPool&) = delete; + + // A pageable buffer of at least nbytes, wrapped as a uint8 CPU tensor of + // exactly nbytes. When the tensor and every view derived from it die, + // the block returns to the pool. Returns an undefined tensor when + // allocation fails: the caller falls back to at::empty. + // pooled_block_capacity(tensor) reports the block's allocated size. + at::Tensor acquire(uint64_t nbytes) { + // Above 2^62 the power-of-two size class would overflow (and the + // byte count would not fit the tensor's int64 size): no buffer. + if (nbytes == 0 || nbytes > (1ULL << 62)) return at::Tensor{}; + const uint64_t cap = size_class(nbytes); + + void* ptr = nullptr; + { + std::lock_guard lk(state_->mu); + auto it = state_->free_lists.find(cap); + if (it != state_->free_lists.end() && !it->second.empty()) { + ptr = it->second.back(); + it->second.pop_back(); + state_->idle_bytes -= cap; + } + } + if (!ptr) { + ptr = aligned_alloc_bytes(cap); + if (!ptr) return at::Tensor{}; + } + + // Exactly one owner of the block at every point, so a throw anywhere + // returns it exactly once: + // 1. the raw block, until its context exists (a failed `new` + // returns it here); + // 2. `ctx`, until TensorMaker::context() takes it. context() is + // noexcept and only stores the pointer, so the handover cannot + // fail half way; + // 3. the TensorMaker, whose destructor deletes the context if + // make_tensor() throws, then the tensor's DataPtr. + // Nothing here deletes the context once TensorMaker owns it. + std::unique_ptr ctx; + try { + ctx.reset(new Block(ptr, cap, state_)); + } catch (...) { + state_->release(ptr, cap); + throw; + } + const int64_t sizes[1] = {static_cast(nbytes)}; + at::TensorMaker maker = at::for_blob(ptr, sizes); + maker.context(ctx.release(), &delete_pooled_block); + return maker + .options(at::TensorOptions().dtype(at::kByte).device(at::kCPU)) + .make_tensor(); + } + + // Size class acquire() allocates for nbytes (diagnostic and tests). + static uint64_t capacity_for(uint64_t nbytes) { return size_class(nbytes); } + + // Idle bytes currently retained (diagnostic). + uint64_t idle_bytes() const { + std::lock_guard lk(state_->mu); + return state_->idle_bytes; + } + +private: + struct State { + explicit State(uint64_t cap) : idle_cap_bytes(cap) {} + // Return a block to its size class, or free it when over the idle + // cap or after the pool closed. + void release(void* ptr, uint64_t cap) { + std::lock_guard lk(mu); + if (closed || idle_bytes + cap > idle_cap_bytes) { + std::free(ptr); + return; + } + free_lists[cap].push_back(ptr); + idle_bytes += cap; + } + mutable std::mutex mu; + std::unordered_map> free_lists; // class -> blocks + uint64_t idle_bytes = 0; + const uint64_t idle_cap_bytes; + bool closed = false; + }; + + // DataPtr context of one handed-out block: returns it on destruction. + // Holds the pool's shared state, never the BufferPool object. + struct Block final : PooledBlockContext { + Block(void* p, uint64_t cap, std::shared_ptr st) + : PooledBlockContext(p, cap), state(std::move(st)) {} + ~Block() override { state->release(ptr, capacity); } + std::shared_ptr state; + }; + + // Round up to the next power of two, with a floor so small steps do not + // fragment the free lists into many singleton classes. + static uint64_t size_class(uint64_t nbytes) { + uint64_t c = kMinClass; + while (c < nbytes) c <<= 1; + return c; + } + + static void* aligned_alloc_bytes(uint64_t cap) { + // cap is a power of two >= kMinClass: a valid alignment multiple. + void* p = nullptr; + if (posix_memalign(&p, kAlign, cap) != 0) return nullptr; + return p; + } + + static constexpr uint64_t kMinClass = 64ULL * 1024; // 64 KiB floor + static constexpr size_t kAlign = 64; + + std::shared_ptr state_; +}; + +} // namespace ring diff --git a/native/csrc/ring/capture_lease.h b/native/csrc/ring/capture_lease.h new file mode 100644 index 000000000..80ed7477d --- /dev/null +++ b/native/csrc/ring/capture_lease.h @@ -0,0 +1,75 @@ +// ring/capture_lease.h -- an integration-defined identity carried with a +// captured step through the legacy ring to its consumer. +// +// A consumer that changes what it captures over time (for example, which +// HookPoints or requests a step selects) attaches a lease to each step when it +// publishes the step's metadata. The lease rides with the StepContext through +// the drain and P2P threads and arrives with every slice or batch the step +// produces, so the consumer can tell which configuration produced a payload, +// and its reference count tells the integration when every payload of a +// configuration has been delivered or dropped (idle()). +// +// A lease may also carry an immutable (act_name, layer) -> drop-rank map, +// shared by every step of one configuration; the batch queue uses it to order +// load shedding (drop_notice.h). +// +// Host metadata only: no Python objects and no GIL on the data path. Copies +// share the configuration until the FIFO, the queues and the consumer release +// it. +#pragma once +#include +#include +#include +#include +#include +#include +#include + +namespace ring_py { + +class CaptureLease { + using Ranks = std::map, uint8_t>; + struct Configuration { + uint64_t id; + std::shared_ptr ranks; + }; +public: + CaptureLease() = default; + explicit CaptureLease(uint64_t id) + : config_(std::make_shared(Configuration{id, {}})) {} + // entries: (act_name, layer_no, rank), ranks clamped to [0, 2]. + CaptureLease(uint64_t id, + const std::vector>& entries) { + Ranks ranks; + for (const auto& [hook, layer, rank] : entries) + ranks[{hook, layer}] = + static_cast(rank < 0 ? 0 : rank > 2 ? 2 : rank); + config_ = std::make_shared(Configuration{ + id, std::make_shared(std::move(ranks))}); + } + // A new lease (own lifetime accounting) sharing `policy`'s immutable rank + // map: no map allocation or copy per captured step. + CaptureLease(uint64_t id, const CaptureLease& policy) + : config_(std::make_shared(Configuration{ + id, policy.config_ ? policy.config_->ranks : nullptr})) {} + + uint64_t id() const { return config_ ? config_->id : 0; } + bool present() const { return bool(config_); } + // No copy other than this one is alive: every payload of the step that + // carried this lease has been consumed or dropped. + bool idle() const { return config_ && config_.use_count() == 1; } + bool same_as(const CaptureLease& other) const { return config_ == other.config_; } + // Rank of (hook, layer) under this lease's map; `fallback` without a map, + // the lowest priority (2) for a stream the map does not name. + uint8_t rank_for(const std::string& hook, int32_t layer, uint8_t fallback) const { + if (!config_ || !config_->ranks) return fallback; + auto it = config_->ranks->find({hook, layer}); + return it == config_->ranks->end() ? 2 : it->second; + } + CaptureLease fork() const { return *this; } + +private: + std::shared_ptr config_; +}; + +} // namespace ring_py diff --git a/native/csrc/ring/drain_task.h b/native/csrc/ring/drain_task.h index 075287898..3a559baa0 100644 --- a/native/csrc/ring/drain_task.h +++ b/native/csrc/ring/drain_task.h @@ -20,6 +20,13 @@ struct DrainTask { // CPU-direct path: tensor already in pageable CPU memory. // When defined(), skip staging copy in p2p thread. at::Tensor cpu_paged_tensor; + + // steady_clock microseconds when the drain thread first saw this task's + // publication on the host: the capture exists and is host-visible. + // Stamped before the flush wait and the D2H copy, so a consumer's + // (now - capture_ts_us) includes the drain-flush latency the drain_flush + // settings control. 0 = unstamped. + int64_t capture_ts_us = 0; }; } // namespace ring diff --git a/native/csrc/ring/drain_thread.cpp b/native/csrc/ring/drain_thread.cpp index f40e4418f..896635e49 100644 --- a/native/csrc/ring/drain_thread.cpp +++ b/native/csrc/ring/drain_thread.cpp @@ -21,6 +21,13 @@ namespace ring { +namespace { +int64_t steady_us(std::chrono::steady_clock::time_point t) { + return std::chrono::duration_cast( + t.time_since_epoch()).count(); +} +} // namespace + namespace { static inline bool publication_try_acquire( @@ -115,6 +122,18 @@ void DrainThread::notify() { cv_.notify_one(); } +bool DrainThread::request_flush() { + if (!is_running()) return false; + { + std::lock_guard lk(mgmt_mu_); + if (cpu_task_head_ <= cpu_task_tail_ || + cpu_task_head_ <= async_flush_until_task_) return false; + async_flush_until_task_ = cpu_task_head_; + } + notify(); + return true; +} + // --------------------------------------------------------------------------- // force_flush_and_wait -- signal drain thread to flush, block until done. // @@ -361,6 +380,8 @@ void DrainThread::submit_cpu_direct(at::Tensor cpu_tensor, uint64_t tensor_bytes DrainTask task{}; task.tensor_total_bytes = tensor_bytes; task.cpu_paged_tensor = std::move(cpu_tensor); + // Already in host memory when handed over: host-visible as of now. + task.capture_ts_us = steady_us(std::chrono::steady_clock::now()); { std::lock_guard lk(queue_mu_); @@ -370,6 +391,7 @@ void DrainThread::submit_cpu_direct(at::Tensor cpu_tensor, uint64_t tensor_bytes std::lock_guard lk(pop_mu_); can_pop_count_ += 1; } + tasks_enqueued_.fetch_add(1, std::memory_order_release); pop_cv_.notify_one(); } @@ -641,13 +663,15 @@ void DrainThread::scan_ready() { account_record_task(visible_head_, actual_bytes); + const auto now = std::chrono::steady_clock::now(); scanned_.push_back(actual_bytes); + scanned_ready_us_.push_back(steady_us(now)); pending_entries_++; pending_bytes_ += align_up(actual_bytes, PAYLOAD_ALIGN); visible_head_++; if (!has_complete_time_) { - first_complete_time_ = std::chrono::steady_clock::now(); + first_complete_time_ = now; has_complete_time_ = true; } } @@ -696,6 +720,7 @@ bool DrainThread::should_flush() const { if (fe >= task_cap) return true; if (fb >= payload_cap) return true; + if (async_flush_until_task_ > cpu_task_tail_) return true; if (fc.task_ratio > 0.0f && fe >= static_cast(fc.task_ratio * task_cap)) return true; if (fc.payload_ratio > 0.0f && @@ -777,6 +802,7 @@ void DrainThread::submit_to_p2p(uint64_t flush_count, uint64_t flush_bytes) { DrainTask task{}; task.tensor_total_bytes = actual_bytes; task.alloc_bytes = alloc; + task.capture_ts_us = scanned_ready_us_[i]; if (data_len > 0) { uint64_t staging_logical = staging_batch_start + cumulative; @@ -803,6 +829,7 @@ void DrainThread::submit_to_p2p(uint64_t flush_count, uint64_t flush_bytes) { std::lock_guard lk(pop_mu_); can_pop_count_ += 1; } + tasks_enqueued_.fetch_add(1, std::memory_order_release); pop_cv_.notify_one(); } @@ -816,6 +843,7 @@ void DrainThread::submit_to_p2p(uint64_t flush_count, uint64_t flush_bytes) { void DrainThread::trim_scanned(uint64_t flush_count, uint64_t flush_bytes) { for (uint64_t i = 0; i < flush_count; ++i) { scanned_.pop_front(); + scanned_ready_us_.pop_front(); } pending_entries_ -= flush_count; pending_bytes_ -= flush_bytes; diff --git a/native/csrc/ring/drain_thread.h b/native/csrc/ring/drain_thread.h index 9bcb7d7d3..f71cc23a9 100644 --- a/native/csrc/ring/drain_thread.h +++ b/native/csrc/ring/drain_thread.h @@ -48,6 +48,13 @@ class DrainThread { void notify(); + // Expedite the tasks already reserved, without waiting on the producer + // stream: the drain flushes everything up to the current reservation + // head as soon as it is published, whatever the flush thresholds say. + // The target survives a wake-up that comes before the GPU publishes those + // tasks. Returns false when nothing is reserved beyond the drain tail. + bool request_flush(); + // Request the drain thread to flush all pending entries, then block // until it finishes. Caller must have done cudaStreamSynchronize on // the main stream first so all GPU writes are visible. @@ -78,6 +85,12 @@ class DrainThread { bool is_running() const { return running_.load(std::memory_order_relaxed); } + // Cumulative count of DrainTasks pushed to the P2P queue. Paired with + // P2PThread::tasks_processed() for a barrier through the consumer. + uint64_t tasks_enqueued() const { + return tasks_enqueued_.load(std::memory_order_acquire); + } + // Capacity query accessors (called from RingEnginePy::prepare_step). uint64_t cpu_payload_head() const; uint64_t cpu_payload_tail_committed() const; @@ -116,6 +129,10 @@ class DrainThread { uint64_t pending_entries_{0}; uint64_t pending_bytes_{0}; std::deque scanned_; + // Parallel to scanned_: when each entry was first seen published. + std::deque scanned_ready_us_; + // request_flush() target: flush every task below this sequence. + uint64_t async_flush_until_task_{0}; uint64_t cpu_task_head_{0}; uint64_t cpu_payload_head_{0}; @@ -157,6 +174,7 @@ class DrainThread { std::deque task_queue_; std::mutex queue_mu_; uint64_t can_pop_count_{0}; + std::atomic tasks_enqueued_{0}; // cumulative tasks pushed std::mutex pop_mu_; std::condition_variable pop_cv_; bool p2p_stop_requested_{false}; diff --git a/native/csrc/ring/drop_notice.h b/native/csrc/ring/drop_notice.h new file mode 100644 index 000000000..aa04d71c2 --- /dev/null +++ b/native/csrc/ring/drop_notice.h @@ -0,0 +1,156 @@ +// ring/drop_notice.h -- attribution for the consumer queues' load shedding. +// +// Both in-process consumer queues shed load on overflow: BatchQueue drops +// whole batches when its retained-byte cap would be exceeded, SliceQueue drops +// single slices when its item cap is reached. Counting a drop is not enough +// for a consumer that aggregates per request: it has to know WHICH requests +// lost data, or it reports results over data it silently lacks. Every drop +// therefore files a DropNotice naming the request ids that were on the dropped +// item; the consumer drains the notices and marks those requests incomplete. +// +// The notice buffer is itself bounded. When it saturates, overflowing +// notices are counted and a sticky flag is raised; the consumer contract for +// that flag is pessimistic -- treat every request in flight since the previous +// poll as incomplete -- because the identity of the losers is then unknown. +// +// Invariant: no notice means nothing was lost. A quiet buffer must never mean +// "we forgot to look". +// +// Header-only, ATen-free, and lock-order safe: the buffer has its own mutex +// and never calls back into the owning queue, so a queue may record a notice +// while holding or after releasing its own lock. + +#pragma once +#include +#include +#include +#include +#include +#include +#include + +#include "ring/capture_lease.h" + +namespace ring_py { + +// Which queue shed the load. Stable strings: consumers key on them. +inline constexpr const char* kDropSourceBatch = "batch_queue"; +inline constexpr const char* kDropSourceSlice = "slice_queue"; +// The shed item was not the arriving one: a queued, lower-priority batch was +// evicted to make room for it (BatchQueue rank-ordered eviction). +inline constexpr const char* kDropSourceBatchEvict = "batch_queue_evict"; + +// --------------------------------------------------------------------------- +// Drop rank -- the shedding order. +// +// The transport knows nothing about the consumer's analyses; it only knows +// that an (act_name, layer) stream has a priority, and priorities are totally +// ordered. Lower rank = more urgent = shed later. The consumer installs the +// map (BatchQueue::set_layer_class_map, or per configuration through a +// CaptureLease); a stream the map does not name gets the default rank. +// Three ranks cover a deadline per token, a deadline per finished request, +// and no deadline. +// --------------------------------------------------------------------------- +inline constexpr uint8_t kRankStreaming = 0; // needed every step +inline constexpr uint8_t kRankRequestBoundary = 1; // needed when the request ends +inline constexpr uint8_t kRankBestEffort = 2; // no deadline +inline constexpr uint8_t kRankMax = 2; // ranks are [0, kRankMax] +inline constexpr size_t kRankCount = 3; +// A path that does not rank (SliceQueue). Never stored on a queued item, +// only reported on a notice. +inline constexpr uint8_t kRankUnset = 255; + +// Clamp a rank crossing from Python: an invented rank becomes the documented +// default for the unknown, never "more sheddable than everything". +inline uint8_t clamp_rank(int r) { + if (r < 0 || r > static_cast(kRankMax)) return kRankBestEffort; + return static_cast(r); +} + +// steady_clock microseconds (the clock of DrainTask::capture_ts_us). +inline int64_t drop_notice_now_us() { + return std::chrono::duration_cast( + std::chrono::steady_clock::now().time_since_epoch()).count(); +} + +// One shed item, with the identity of everything that was on it. +struct DropNotice { + CaptureLease capture_lease; + const char* source = kDropSourceBatch; // static string, never freed + std::string hook; // act_name of the lost capture + int32_t layer_no = -1; + uint64_t step_seq = 0; // batch path: the step that lost data + int32_t start_token = -1; // slice path only (token ranges are + int32_t end_token = -1; // per request in a batch's segments) + uint64_t nbytes = 0; // bytes of capture lost + int64_t ts_us = 0; // steady-clock time of the drop + // Rank of the lost stream; kRankUnset on the slice path. + uint8_t rank = kRankUnset; + // Every request that had rows on the dropped item: a batch's req_ids, or + // exactly one on the slice path. + std::vector req_ids; +}; + +class DropNoticeBuffer { +public: + explicit DropNoticeBuffer(size_t cap = 4096) : cap_(cap) {} + + // Called on the drop path (P2P thread): no GIL, no allocation beyond the + // notice itself. + void record(DropNotice n) { + std::lock_guard lk(mu_); + recorded_++; + if (q_.size() >= cap_) { + overflow_++; + overflow_total_++; + sticky_ = true; + return; + } + q_.push_back(std::move(n)); + } + + struct Drained { + std::vector notices; + // Notices lost since the previous drain. + uint64_t overflow = 0; + // Any notice lost since the previous drain: the consumer must treat + // every request in flight over that window as incomplete. + bool sticky = false; + // Notices still buffered after this drain (max_items truncation). + size_t pending = 0; + }; + + // Drain up to max_items notices and the overflow/sticky state, clearing + // both. Non-blocking. + Drained drain(size_t max_items) { + Drained out; + std::lock_guard lk(mu_); + while (!q_.empty() && out.notices.size() < max_items) { + out.notices.push_back(std::move(q_.front())); + q_.pop_front(); + } + out.overflow = overflow_; + out.sticky = sticky_; + out.pending = q_.size(); + overflow_ = 0; + sticky_ = false; // sticky covers one inter-poll window + return out; + } + + size_t pending() const { std::lock_guard lk(mu_); return q_.size(); } + // Cumulative, never reset by drain. + uint64_t recorded() const { std::lock_guard lk(mu_); return recorded_; } + uint64_t overflow_total() const { std::lock_guard lk(mu_); return overflow_total_; } + size_t cap() const { return cap_; } + +private: + mutable std::mutex mu_; + std::deque q_; + size_t cap_; + uint64_t recorded_ = 0; // notices offered (kept + lost) + uint64_t overflow_ = 0; // lost since the last drain + uint64_t overflow_total_ = 0; // lost, cumulative + bool sticky_ = false; +}; + +} // namespace ring_py diff --git a/native/csrc/ring/p2p_thread.cpp b/native/csrc/ring/p2p_thread.cpp index 4278d6645..5150da5f4 100644 --- a/native/csrc/ring/p2p_thread.cpp +++ b/native/csrc/ring/p2p_thread.cpp @@ -2,13 +2,16 @@ // Compiled with g++ (not nvcc); uses ATen for tensor operations. #include "p2p_thread.h" +#include "buffer_pool.h" #include "pinned_staging.h" #include +#include #include #include #include #include +#include namespace ring { @@ -121,8 +124,15 @@ static at::Tensor slice_flattened( // --------------------------------------------------------------------------- P2PThread::P2PThread(DrainThread& drain, ring_py::TensorMetaFifo& fifo, - const RingConfig& cfg, SubmitFn submit_fn) - : drain_(drain), fifo_(fifo), cfg_(cfg), submit_fn_(std::move(submit_fn)) + const RingConfig& cfg, SubmitFn submit_fn, + BatchSink batch_sink) + : drain_(drain), fifo_(fifo), cfg_(cfg), submit_fn_(std::move(submit_fn)), + batch_sink_(std::move(batch_sink)), + // p2p_pool_cap_bytes == 0: no pool at all -- exact-size at::empty + // buffers, as before the pool existed (no size-class rounding). + pool_(cfg.p2p_pool_cap_bytes > 0 + ? std::make_unique(cfg.p2p_pool_cap_bytes) + : nullptr) {} P2PThread::~P2PThread() noexcept { @@ -131,6 +141,10 @@ P2PThread::~P2PThread() noexcept { } void P2PThread::start() { + { + std::lock_guard lk(barrier_mu_); + stopped_ = false; + } thread_ = std::thread([this] { loop(); }); } @@ -140,6 +154,18 @@ void P2PThread::stop() { // --------------------------------------------------------------------------- void P2PThread::loop() { + // However the loop ends (stop, or an escaping exception), release every + // delivery barrier: no further task will be processed. + struct ExitNotice { + P2PThread* self; + ~ExitNotice() { + { + std::lock_guard lk(self->barrier_mu_); + self->stopped_ = true; + } + self->barrier_cv_.notify_all(); + } + } exit_notice{this}; while (true) { uint64_t n = drain_.wait_for_tasks(); if (n == 0) break; @@ -147,9 +173,49 @@ void P2PThread::loop() { std::vector local; drain_.pop_tasks(n, local); process(local); + + // Publish progress under barrier_mu_ so a waiter cannot miss the + // wake-up (it checks the predicate under the same mutex). + { + std::lock_guard lk(barrier_mu_); + tasks_processed_.fetch_add(local.size(), std::memory_order_release); + } + barrier_cv_.notify_all(); } } +// --------------------------------------------------------------------------- +bool P2PThread::wait_until_processed(uint64_t target, uint32_t timeout_ms) { + std::unique_lock lk(barrier_mu_); + auto reached = [&] { + return tasks_processed_.load(std::memory_order_acquire) >= target; + }; + // A stopped thread processes nothing more: return its final verdict. + auto done = [&] { return reached() || stopped_; }; + if (timeout_ms == 0) + barrier_cv_.wait(lk, done); + else + barrier_cv_.wait_for(lk, std::chrono::milliseconds(timeout_ms), done); + return reached(); +} + +bool P2PThread::wait_until_processed_until( + uint64_t target, std::chrono::steady_clock::time_point deadline) { + std::unique_lock lk(barrier_mu_); + auto reached = [&] { + return tasks_processed_.load(std::memory_order_acquire) >= target; + }; + barrier_cv_.wait_until(lk, deadline, [&] { return reached() || stopped_; }); + return reached(); +} + +// --------------------------------------------------------------------------- +void P2PThread::note_submit_error(const char* what) { + submit_exceptions_.fetch_add(1, std::memory_order_relaxed); + std::lock_guard lk(err_mu_); + last_error_ = what ? what : "unknown consumer error"; +} + // --------------------------------------------------------------------------- void P2PThread::process(std::vector& tasks) { for (size_t i = 0; i < tasks.size(); ++i) { @@ -162,11 +228,20 @@ void P2PThread::process(std::vector& tasks) { continue; } - // Normal tensor: copy from pinned staging to pageable + // Normal tensor: copy from pinned staging to pageable. The + // destination comes from the warm-page pool; the tensor (and every + // view derived from it) keeps its pool block alive until the last + // reference drops (buffer_pool.h). at::empty, exact size, without + // a pool (p2p_pool_cap_bytes == 0) or when it cannot serve the + // request. uint64_t total_bytes = task.tensor_total_bytes; - auto tensor = at::empty({static_cast(total_bytes)}, - at::TensorOptions().dtype(at::kByte).device(at::kCPU)); + at::Tensor tensor; + if (pool_) tensor = pool_->acquire(total_bytes); + if (!tensor.defined()) { + tensor = at::empty({static_cast(total_bytes)}, + at::TensorOptions().dtype(at::kByte).device(at::kCPU)); + } uint8_t* dst = tensor.data_ptr(); if (task.data_len1 > 0) { @@ -194,16 +269,19 @@ void P2PThread::do_post_processing(at::Tensor& tensor, const DrainTask& first_ta // Get step context -- pop from context queue if this is the first // hook in a new step (current_ctx_ is null). + // Each new context starts a step: every batch of the step carries the + // same step_seq. if (!current_ctx_) { current_ctx_ = fifo_.pop_context(); if (!current_ctx_) return; // no context available + step_seq_++; } if (meta.shape.empty() || first_task.tensor_total_bytes == 0) { if (meta.last_in_step) { delete current_ctx_; current_ctx_ = nullptr; } return; } - if (!submit_fn_) { + if (!submit_fn_ && !batch_sink_) { if (meta.last_in_step) { delete current_ctx_; current_ctx_ = nullptr; } return; } @@ -254,8 +332,146 @@ void P2PThread::do_post_processing(at::Tensor& tensor, const DrainTask& first_ta // An independently gathered hook carries its own request layout; every // other hook uses the step's. const auto& requests = meta.requests ? *meta.requests : current_ctx_->requests; + + // Batch handoff: hand the WHOLE reshaped tensor and a segments table to + // the consumer once, instead of narrowing, copying and submitting one + // slice per request. Packed steps only; batched (HF) steps and + // attention matrices use the per-slice path (with only a batch sink, + // delivered as one-request items, below). The tensor is a view of + // the pooled buffer, never sliced or cloned here: consumers clone what + // they keep. + if (cfg_.batch_submit && batch_sink_ && current_ctx_->flattened && !is_attn) { + // One int32 row per request that yields a non-empty view: + // (dim0_offset, db_start, db_end, orig_position). Rows the + // per-slice path would skip are omitted, so a step with no rows + // yields an empty table. + const int64_t n_req = static_cast(requests.size()); + auto segments = at::empty({n_req, 4}, + at::TensorOptions().dtype(at::kInt).device(at::kCPU)); + int32_t* seg = segments.data_ptr(); + std::vector req_ids; + req_ids.reserve(requests.size()); + int64_t kept = 0; + for (size_t j = 0; j < requests.size(); ++j) { + const ring_py::RequestMeta& req = requests[j]; + int32_t db_start = req.start_token; + int32_t db_end = req.end_token; + int64_t dim0; + if (meta.hook_type == ring_py::HOOK_TYPE_FINAL_LOGITS) { + // One logit row per request, indexed by the request's + // original position; it covers the request's last token. + dim0 = (req.orig_position >= 0) + ? static_cast(req.orig_position) + : static_cast(j); + if (dim0 < 0 || dim0 + 1 > tensor.size(0) || req.end_token <= 0) + continue; + db_start = req.end_token - 1; + db_end = req.end_token; + } else { + // Widen before subtracting: int32 token positions. + const int64_t num_tokens = static_cast(req.end_token) - + static_cast(req.start_token); + if (num_tokens <= 0 || req.dim0_offset < 0 || + req.dim0_offset + num_tokens > tensor.size(0)) + continue; + dim0 = req.dim0_offset; + } + seg[kept * 4 + 0] = static_cast(dim0); + seg[kept * 4 + 1] = db_start; + seg[kept * 4 + 2] = db_end; + seg[kept * 4 + 3] = req.orig_position; + req_ids.push_back(req.req_id); + kept++; + } + if (kept < n_req) segments = segments.narrow(0, 0, kept); + + // Nothing to deliver: skip the push, but keep the context pairing. + if (kept > 0) { + try { + batch_sink_(current_ctx_->model_id, shard_rank, act_name, + meta.layer_no, step_seq_, meta.last_in_step, + tensor, std::move(req_ids), std::move(segments), + current_ctx_->capture_lease, first_task.capture_ts_us, + /*is_attn=*/false); + } catch (const std::exception& e) { + log_submit_failure_once(current_ctx_->model_id, "", + act_name, meta.layer_no, shard_rank, + 0, 0, e.what()); + note_submit_error(e.what()); + } catch (...) { + log_submit_failure_once(current_ctx_->model_id, "", + act_name, meta.layer_no, shard_rank, + 0, 0, "unknown non-std exception"); + note_submit_error("unknown non-std exception"); + } + } + if (meta.last_in_step) { delete current_ctx_; current_ctx_ = nullptr; } + return; + } + // Per-slice path: the legacy consumer (submit_fn_), or -- with only a + // batch consumer installed -- every capture the whole-batch branch above + // does not take (HF batched steps: tokens [batch, q_len, ...] and + // attention matrices [batch, heads, q_len, kv_len]). Those go to the + // batch sink as one-request items, one per request slice, so a + // BatchQueue consumer receives them instead of losing them silently. + // Packed attention matrices are not sliceable per request on either + // path and stay skipped. + const bool to_batch_sink = !submit_fn_; bool should_clone = cfg_.clone_slices && requests.size() > 1; + struct Slice { + const ring_py::RequestMeta* req; + int32_t db_start; + int32_t db_end; + at::Tensor tensor; + }; + auto deliver = [&](Slice& sl, bool last_item) { + const ring_py::RequestMeta& req = *sl.req; + const int32_t db_start = sl.db_start; + const int32_t db_end = sl.db_end; + try { + if (!to_batch_sink) { + submit_fn_(current_ctx_->model_id, shard_rank, + req.req_id, act_name, meta.layer_no, + db_start, db_end, + std::move(sl.tensor), first_task.capture_ts_us, + current_ctx_->capture_lease); + } else { + // One-request batch item. The tensor IS the request's + // slice, so its segments row has dim0_offset 0: + // token hooks: [rows, ...], rows = the request's last + // (db_end - db_start) query positions + // (fewer if q_len is shorter); + // attention hooks: [heads, rows, kv] (is_attn), taken whole. + auto segments = at::empty( + {1, 4}, at::TensorOptions().dtype(at::kInt).device(at::kCPU)); + int32_t* seg = segments.data_ptr(); + seg[0] = 0; + seg[1] = db_start; + seg[2] = db_end; + seg[3] = req.orig_position; + batch_sink_(current_ctx_->model_id, shard_rank, act_name, + meta.layer_no, step_seq_, last_item, + std::move(sl.tensor), std::vector{req.req_id}, + std::move(segments), current_ctx_->capture_lease, + first_task.capture_ts_us, is_attn); + } + } catch (const std::exception& e) { + log_submit_failure_once(current_ctx_->model_id, req.req_id, + act_name, meta.layer_no, shard_rank, + db_start, db_end, e.what()); + note_submit_error(e.what()); + } catch (...) { + log_submit_failure_once(current_ctx_->model_id, req.req_id, + act_name, meta.layer_no, shard_rank, + db_start, db_end, + "unknown non-std exception"); + note_submit_error("unknown non-std exception"); + } + }; + // Batch sink: hold back one slice so the hook's last delivered item can + // carry last_in_step. The legacy consumer gets each slice as it is cut. + std::optional pending; for (size_t j = 0; j < requests.size(); ++j) { const ring_py::RequestMeta& req = requests[j]; at::Tensor slice; @@ -311,23 +527,15 @@ void P2PThread::do_post_processing(at::Tensor& tensor, const DrainTask& first_ta if (should_clone) { slice = slice.clone(); } - - try { - submit_fn_(current_ctx_->model_id, shard_rank, - req.req_id, act_name, meta.layer_no, - db_start, db_end, - std::move(slice)); - } catch (const std::exception& e) { - log_submit_failure_once(current_ctx_->model_id, req.req_id, - act_name, meta.layer_no, shard_rank, - db_start, db_end, e.what()); - } catch (...) { - log_submit_failure_once(current_ctx_->model_id, req.req_id, - act_name, meta.layer_no, shard_rank, - db_start, db_end, - "unknown non-std exception"); + Slice sl{&req, db_start, db_end, std::move(slice)}; + if (!to_batch_sink) { + deliver(sl, false); + } else { + if (pending) deliver(*pending, false); + pending = std::move(sl); } } + if (pending) deliver(*pending, meta.last_in_step); // Last hook in step -- free context if (meta.last_in_step) { diff --git a/native/csrc/ring/p2p_thread.h b/native/csrc/ring/p2p_thread.h index a7e87dc10..a252295c2 100644 --- a/native/csrc/ring/p2p_thread.h +++ b/native/csrc/ring/p2p_thread.h @@ -12,31 +12,71 @@ #include "record_consumer.h" #include "tensor_meta.h" +#include +#include +#include #include #include +#include #include #include +#include namespace at { class Tensor; } namespace ring { +class BufferPool; // buffer_pool.h, included only by the .cpp + // Called by the p2p thread for each per-request tensor slice. -// Runs without the GIL -- all arguments are C++ types. +// Runs without the GIL -- all arguments are C++ types. capture_ts_us is +// DrainTask::capture_ts_us; capture_lease is the step's lease. using SubmitFn = std::function; + const std::string& model_id, + int32_t shard_rank, + const std::string& req_id, + const std::string& act_name, + int32_t layer_no, + int32_t start_token, + int32_t end_token, + at::Tensor slice, + int64_t capture_ts_us, + ring_py::CaptureLease capture_lease)>; + +// Called by the p2p thread ONCE per (hook x step) on the batch handoff +// (RingConfig::batch_submit, packed steps only): the whole reshaped tensor +// plus a segments table; the consumer does the per-request slicing. +// segments: int32 CPU [N, 4] rows (dim0_offset, db_start, db_end, +// orig_position) +// req_ids: N request ids, in segments row order +// last_in_step marks the last hook of a step. Runs without the GIL. +// +// Per-slice fallback: when the batch sink is the only consumer, captures the +// whole-batch handoff does not cover (HF batched steps, token or attention) +// arrive here as ONE-REQUEST items, one call per request slice: N == 1, +// req_ids = {req_id}, segments = {(0, db_start, db_end, orig_position)} and +// `tensor` is that request's slice -- [rows, ...] for token hooks, +// [heads, rows, kv] for attention matrices (is_attn true; take it whole). +// last_in_step is set on the last item of the step's last hook. +using BatchSink = std::function req_ids, + at::Tensor segments, + ring_py::CaptureLease capture_lease, + int64_t capture_ts_us, + bool is_attn)>; class P2PThread { public: P2PThread(DrainThread& drain, ring_py::TensorMetaFifo& fifo, - const RingConfig& cfg, SubmitFn submit_fn); + const RingConfig& cfg, SubmitFn submit_fn, + BatchSink batch_sink = nullptr); ~P2PThread() noexcept; P2PThread(const P2PThread&) = delete; @@ -45,18 +85,71 @@ class P2PThread { void start(); void stop(); + // -- consumer failures --------------------------------------------------- + // Exceptions thrown by the consumer are logged once and counted here + // instead of disappearing; the most recent message is kept. + uint64_t submit_exceptions() const { + return submit_exceptions_.load(std::memory_order_relaxed); + } + std::string last_error() const { + std::lock_guard lk(err_mu_); + return last_error_; + } + void set_abort_on_sink_error(bool v) { + abort_on_sink_error_.store(v, std::memory_order_relaxed); + } + // True once a consumer error occurred AND abort-on-sink-error is armed. + bool sink_failed() const { + return abort_on_sink_error_.load(std::memory_order_relaxed) + && submit_exceptions_.load(std::memory_order_relaxed) > 0; + } + + // -- delivery barrier ----------------------------------------------------- + // DrainTasks fully processed by this thread: every consumer call for the + // task has returned. Paired with DrainThread::tasks_enqueued(). + uint64_t tasks_processed() const { + return tasks_processed_.load(std::memory_order_acquire); + } + // Block until tasks_processed() >= target, timeout_ms elapses (0 = no + // timeout), or the thread has exited. Returns true when the target was + // reached. + bool wait_until_processed(uint64_t target, uint32_t timeout_ms); + // Same, bounded by an absolute deadline (a deadline already passed + // checks once and returns). + bool wait_until_processed_until(uint64_t target, + std::chrono::steady_clock::time_point deadline); + private: DrainThread& drain_; ring_py::TensorMetaFifo& fifo_; RingConfig cfg_; SubmitFn submit_fn_; + BatchSink batch_sink_; + // Incremented for every new StepContext; carried in each batch. + uint64_t step_seq_{0}; std::thread thread_; ring_py::StepContext* current_ctx_{nullptr}; // owned, freed on last_in_step + // Warm pageable buffers for the staging -> pageable copy, recycled by + // reference count: it must outlive every slice it hands out, so it is + // destroyed last, in ~P2PThread, after stop() joins this thread. + std::unique_ptr pool_; + + std::atomic submit_exceptions_{0}; + mutable std::mutex err_mu_; + std::string last_error_; + std::atomic abort_on_sink_error_{false}; + + std::atomic tasks_processed_{0}; + std::mutex barrier_mu_; + std::condition_variable barrier_cv_; + bool stopped_{false}; // loop exited; under barrier_mu_ + void loop(); void process(std::vector& tasks); void do_post_processing(at::Tensor& tensor, const DrainTask& first_task); + void note_submit_error(const char* what); }; // Schema-driven record path. It intentionally has a distinct worker class diff --git a/native/csrc/ring/pooled_block.h b/native/csrc/ring/pooled_block.h new file mode 100644 index 000000000..00c9c4af0 --- /dev/null +++ b/native/csrc/ring/pooled_block.h @@ -0,0 +1,121 @@ +// ring/pooled_block.h -- identify the backing allocation a CPU tensor keeps +// alive, and charge it once per queue. +// +// The P2P thread's buffer pool (buffer_pool.h) hands out size-classed blocks +// (power of two, 64 KiB floor) wrapped as tensors of the REQUESTED length, so +// tensor.nbytes() and even storage().nbytes() understate what a queued item +// keeps alive. A pooled block carries a PooledBlockContext as its DataPtr +// context; backing_of() reads the real capacity from it. Any other tensor is +// charged its storage's byte size. +// +// Several queued items may alias one allocation (per-request views of one +// step's buffer). BackingLedger charges each distinct allocation once, at +// max(storage bytes, block capacity), while at least one queued item refers +// to it, and releases the charge when the last such item leaves the queue. +// +// Header-only; included where ATen is already linked. + +#pragma once +#include +#include +#include +#include +#include +#include + +namespace ring { + +// DataPtr context of a pooled block. The deleter (delete_pooled_block) +// destroys it through the virtual destructor, which returns the block to its +// pool (or frees it). +struct PooledBlockContext { + PooledBlockContext(void* p, uint64_t cap) : ptr(p), capacity(cap) {} + virtual ~PooledBlockContext() = default; + void* ptr; + uint64_t capacity; // bytes actually allocated for the block +}; + +inline void delete_pooled_block(void* ctx) { + delete static_cast(ctx); +} + +// Allocated bytes behind a pooled tensor (or any view of it); 0 when the +// tensor's storage is not a pool block. +inline uint64_t pooled_block_capacity(const at::Tensor& t) { + if (!t.defined() || !t.has_storage()) return 0; + const c10::DataPtr& dp = t.storage().data_ptr(); + if (dp.get_deleter() != &delete_pooled_block) return 0; + const auto* ctx = static_cast(dp.get_context()); + return ctx ? ctx->capacity : 0; +} + +// The allocation a tensor keeps alive: an identity shared by every view of +// it, and the bytes it pins. +struct BackingRef { + const void* key = nullptr; // storage base address; nullptr = nothing + uint64_t bytes = 0; +}; + +inline BackingRef backing_of(const at::Tensor& t) { + BackingRef r; + if (!t.defined()) return r; + const uint64_t logical = static_cast(t.nbytes()); + if (!t.has_storage()) { + r.bytes = logical; + return r; + } + const c10::Storage& s = t.storage(); + r.key = s.data(); + r.bytes = std::max({logical, static_cast(s.nbytes()), + pooled_block_capacity(t)}); + if (!r.key) r.bytes = logical; // no allocation to share + return r; +} + +// Retained-byte ledger: each distinct allocation is charged once while any +// queued item refers to it. Not thread-safe; the owning queue guards it. +class BackingLedger { +public: + // Bytes the ledger would grow by if `r` were added. + uint64_t cost(const BackingRef& r) const { + if (!r.key) return r.bytes; + return blocks_.count(r.key) ? 0 : r.bytes; + } + // Add one reference. Returns the bytes newly charged. + uint64_t add(const BackingRef& r) { + if (!r.key) { total_ += r.bytes; return r.bytes; } + auto it = blocks_.find(r.key); + if (it != blocks_.end()) { it->second.refs++; return 0; } + blocks_.emplace(r.key, Entry{r.bytes, 1}); + total_ += r.bytes; + return r.bytes; + } + // Drop one reference. Returns the bytes released: the allocation's + // charge when its last queued reference leaves, else 0. + uint64_t remove(const BackingRef& r) { + if (!r.key) { total_ -= r.bytes; return r.bytes; } + auto it = blocks_.find(r.key); + if (it == blocks_.end()) return 0; + if (--it->second.refs > 0) return 0; + const uint64_t bytes = it->second.bytes; + blocks_.erase(it); + total_ -= bytes; + return bytes; + } + // Queued references to `key` and the bytes charged for it ({0, 0} when + // the allocation is not held). Lets a queue work out what removing a + // set of items would release before it removes any. + std::pair held(const void* key) const { + auto it = blocks_.find(key); + if (it == blocks_.end()) return {0, 0}; + return {it->second.refs, it->second.bytes}; + } + uint64_t total() const { return total_; } + +private: + struct Entry { uint64_t bytes; uint64_t refs; }; + std::unordered_map blocks_; + uint64_t total_ = 0; +}; + +} // namespace ring diff --git a/native/csrc/ring/ring_config.h b/native/csrc/ring/ring_config.h index e39975b72..391cb912c 100644 --- a/native/csrc/ring/ring_config.h +++ b/native/csrc/ring/ring_config.h @@ -75,6 +75,16 @@ struct RingConfig { // that keep the full tensor alive until consumed. bool clone_slices = false; + // Maximum idle bytes the P2P thread's pageable buffer pool retains + // between steps (buffer_pool.h). 0 disables the pool entirely: every + // step gets an exact-size at::empty buffer (no size-class rounding). + uint64_t p2p_pool_cap_bytes = 256ULL * 1024 * 1024; + + // Batch handoff: with a batch consumer (BatchQueue), packed (flattened) + // steps hand each hook's whole tensor and a segments table to the + // consumer once, instead of one slice per request. + bool batch_submit = false; + // ClickHouse insert queue limits (host engine). // P2p thread blocks on submit_direct() when queue is full. uint64_t insert_queue_max_bytes = 4096ULL * 1024 * 1024; // 4 GiB diff --git a/native/csrc/ring/ring_engine.cu b/native/csrc/ring/ring_engine.cu index ea43b7500..d9e9911e2 100644 --- a/native/csrc/ring/ring_engine.cu +++ b/native/csrc/ring/ring_engine.cu @@ -8,7 +8,7 @@ namespace ring { RingEngine::RingEngine(const RingConfig& cfg, ring_py::TensorMetaFifo& fifo, - SubmitFn submit_fn) + SubmitFn submit_fn, BatchSink batch_sink) : cfg_(cfg), ring_(cfg) { if (cfg_.payload_ring_bytes % PAYLOAD_ALIGN != 0) { @@ -32,7 +32,35 @@ RingEngine::RingEngine(const RingConfig& cfg, ring_py::TensorMetaFifo& fifo, staging_.init(staging_bytes); drain_ = std::make_unique(ring_.state(), staging_, cfg_); - p2p_ = std::make_unique(*drain_, fifo, cfg_, std::move(submit_fn)); + p2p_ = std::make_unique(*drain_, fifo, cfg_, std::move(submit_fn), + std::move(batch_sink)); +} + +bool RingEngine::drain_to_sink_and_wait(uint32_t timeout_ms) { + if (timeout_ms == 0) return drain_to_sink_and_wait_until(std::nullopt); + return drain_to_sink_and_wait_until( + std::chrono::steady_clock::now() + std::chrono::milliseconds(timeout_ms)); +} + +bool RingEngine::drain_to_sink_and_wait_until( + std::optional deadline) { + if (!p2p_) { + throw std::logic_error( + "drain_to_sink_and_wait is a legacy-ring barrier; record rings " + "use flush_records_and_wait"); + } + // Everything published is handed to the P2P queue by the flush, so the + // enqueued count read afterwards covers every task produced so far. + // One deadline bounds both phases: a flush stuck behind saturated + // staging (slow consumer) times out like the final delivery wait. + if (!deadline) { + drain_->force_flush_and_wait(); + drain_->rethrow_drain_failure(); + return p2p_->wait_until_processed(drain_->tasks_enqueued(), 0); + } + if (!drain_->force_flush_and_wait_until(*deadline)) return false; + drain_->rethrow_drain_failure(); + return p2p_->wait_until_processed_until(drain_->tasks_enqueued(), *deadline); } RingEngine::RingEngine(const RingConfig& cfg, diff --git a/native/csrc/ring/ring_engine.h b/native/csrc/ring/ring_engine.h index 6675d0bb8..582fdf6e4 100644 --- a/native/csrc/ring/ring_engine.h +++ b/native/csrc/ring/ring_engine.h @@ -8,7 +8,9 @@ #include "record_sink.h" #include "tensor_meta.h" +#include #include +#include #include namespace ring { @@ -16,7 +18,7 @@ namespace ring { class RingEngine { public: RingEngine(const RingConfig& cfg, ring_py::TensorMetaFifo& fifo, - SubmitFn submit_fn); + SubmitFn submit_fn, BatchSink batch_sink = nullptr); RingEngine(const RingConfig& cfg, std::shared_ptr lease, RecordFailurePolicy failure_policy = RecordFailurePolicy::kRaiseAtProducer); @@ -35,6 +37,22 @@ class RingEngine { const RecordConsumer& record_consumer() const; std::shared_ptr record_sink() const { return record_sink_; } bool record_mode() const { return static_cast(record_p2p_); } + // Legacy rings only (nullptr on a record ring). + P2PThread* p2p_thread() { return p2p_.get(); } + const P2PThread* p2p_thread() const { return p2p_.get(); } + + // Flush the ring, then wait until the P2P thread has processed every task + // the drain has handed it: every consumer call for every captured slice + // or batch has returned. Returns true when delivered, false on timeout + // (0 = wait forever). The caller synchronizes the producer stream first. + // timeout_ms == 0 waits forever; otherwise ONE deadline, taken at entry, + // bounds the flush and the delivery wait together. + bool drain_to_sink_and_wait(uint32_t timeout_ms); + // Same with an absolute deadline (nullopt = no deadline). Lets a caller + // that already spent part of its budget (producer-stream completion) + // pass the remainder on. + bool drain_to_sink_and_wait_until( + std::optional deadline); uint64_t payload_cap() const { return cfg_.payload_ring_bytes; } bool drain_step_gated() const { return cfg_.drain_step_gated; } diff --git a/native/csrc/ring/ring_engine_py.cu b/native/csrc/ring/ring_engine_py.cu index b9c29f139..fe40038b0 100644 --- a/native/csrc/ring/ring_engine_py.cu +++ b/native/csrc/ring/ring_engine_py.cu @@ -152,6 +152,39 @@ uint64_t legacy_ring_room(uint64_t cap, uint64_t head, uint64_t tail) { return outstanding >= cap ? 0 : cap - outstanding; } +// Wait for the work queued on `stream` so far, or until `deadline`. Records +// a timing-free event and polls it with a short, growing sleep instead of an +// unbounded cudaStreamSynchronize. Returns false on timeout (the stream's +// work is still pending), and when the event API fails: completion cannot be +// proven then, which the caller reports as not delivered. +bool wait_stream_until(cudaStream_t stream, + std::chrono::steady_clock::time_point deadline) { + cudaEvent_t ev = nullptr; + if (cudaEventCreateWithFlags(&ev, cudaEventDisableTiming) != cudaSuccess) { + cudaGetLastError(); + return false; + } + bool done = false; + if (cudaEventRecord(ev, stream) == cudaSuccess) { + auto backoff = std::chrono::microseconds(20); + while (true) { + const cudaError_t q = cudaEventQuery(ev); + if (q == cudaSuccess) { done = true; break; } + if (q != cudaErrorNotReady) { cudaGetLastError(); break; } + const auto now = std::chrono::steady_clock::now(); + if (now >= deadline) break; + const auto left = std::chrono::duration_cast( + deadline - now); + std::this_thread::sleep_for(std::min(backoff, left)); + backoff = std::min(backoff * 2, std::chrono::microseconds(1000)); + } + } else { + cudaGetLastError(); + } + cudaEventDestroy(ev); + return done; +} + } // namespace // --------------------------------------------------------------------------- @@ -225,8 +258,8 @@ struct RingEnginePy::Impl { 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)) + Impl(ring::RingConfig cfg, SubmitFn sf, BatchSink bs = nullptr) + : engine(std::move(cfg), fifo, std::move(sf), std::move(bs)) { const auto& state = engine.ring_state(); int dev_idx = 0; @@ -395,6 +428,8 @@ static ring::RingConfig convert(const RingConfig& c) { r.drain_flush.byte_threshold = c.drain_flush_byte_threshold; r.drain_flush.timeout_us = c.drain_flush_timeout_us; r.clone_slices = c.clone_slices; + r.p2p_pool_cap_bytes = c.p2p_pool_cap_bytes; + r.batch_submit = c.batch_submit; r.insert_queue_max_bytes = c.insert_queue_max_bytes; r.insert_queue_max_items = c.insert_queue_max_items; if (c.drain_policy == "step_gated") { @@ -408,12 +443,14 @@ static ring::RingConfig convert(const RingConfig& c) { } // --------------------------------------------------------------------------- -RingEnginePy::RingEnginePy(RingConfig cfg, SubmitFn submit_fn) { +RingEnginePy::RingEnginePy(RingConfig cfg, SubmitFn submit_fn, + BatchSink batch_sink) { if (cfg.drain_policy == "step_gated") { throw std::invalid_argument( "drain_policy 'step_gated' requires a record ring"); } - impl_ = std::make_unique(convert(cfg), std::move(submit_fn)); + impl_ = std::make_unique(convert(cfg), std::move(submit_fn), + std::move(batch_sink)); } RingEnginePy::RingEnginePy( @@ -632,6 +669,11 @@ int RingEnginePy::prepare_step(uint64_t step_total_bytes, { // Record rings reserve through reserve_record. refuse_on_record_ring(impl_->record_mode, "legacy step reservation"); + if (sink_failed()) { + throw std::runtime_error( + "ring consumer failed (abort-on-sink-error is armed): " + + last_sink_error()); + } impl_->current_hook_idx = 0; if (step_total_bytes % ring::PAYLOAD_ALIGN != 0) { throw std::invalid_argument( @@ -1314,6 +1356,46 @@ bool RingEnginePy::_test_applied_current(uint64_t graph) const { return impl_->toggle.applied_current(graph); } +bool RingEnginePy::request_flush() { + return impl_->engine.drain_thread().request_flush(); +} + +bool RingEnginePy::drain_to_sink_and_wait(uint32_t timeout_ms) { + refuse_on_record_ring(impl_->record_mode, "legacy delivery barrier"); + cudaStream_t ms = at::cuda::getCurrentCUDAStream().stream(); + if (timeout_ms == 0) { + // No timeout: wait for the producer stream, the flush and delivery. + cudaStreamSynchronize(ms); + return impl_->engine.drain_to_sink_and_wait_until(std::nullopt); + } + // A finite timeout bounds the WHOLE barrier: one deadline, taken here, + // covers producer-stream completion, the flush and the delivery wait. + const auto deadline = + std::chrono::steady_clock::now() + std::chrono::milliseconds(timeout_ms); + if (!wait_stream_until(ms, deadline)) return false; + return impl_->engine.drain_to_sink_and_wait_until(deadline); +} + +uint64_t RingEnginePy::submit_exceptions() const { + const auto* p2p = impl_->engine.p2p_thread(); + return p2p ? p2p->submit_exceptions() : 0; +} + +std::string RingEnginePy::last_sink_error() const { + const auto* p2p = impl_->engine.p2p_thread(); + return p2p ? p2p->last_error() : std::string{}; +} + +void RingEnginePy::set_abort_on_sink_error(bool enabled) { + auto* p2p = impl_->engine.p2p_thread(); + if (p2p) p2p->set_abort_on_sink_error(enabled); +} + +bool RingEnginePy::sink_failed() const { + const auto* p2p = impl_->engine.p2p_thread(); + return p2p && p2p->sink_failed(); +} + uint64_t RingEnginePy::available_capacity() const { auto& drain = impl_->engine.drain_thread(); const uint64_t head = drain.cpu_payload_head(); diff --git a/native/csrc/ring/ring_engine_py.h b/native/csrc/ring/ring_engine_py.h index 2baccf519..bee8c8824 100644 --- a/native/csrc/ring/ring_engine_py.h +++ b/native/csrc/ring/ring_engine_py.h @@ -13,6 +13,7 @@ #include #include +#include "ring/capture_lease.h" #include "ring/tensor_meta.h" // TensorMeta, TensorMetaFifo #include "ring/record_descriptor.h" #include "ring/record_failure_policy.h" @@ -41,6 +42,12 @@ struct RingConfig { uint64_t drain_flush_timeout_us = 0; // Clone per-request slices bool clone_slices = false; + // Idle bytes the P2P thread's pageable buffer pool retains (0 = no pool: + // exact-size buffers, no size-class rounding). + uint64_t p2p_pool_cap_bytes = 256ULL * 1024 * 1024; + // Batch handoff for packed steps: with a batch consumer (BatchQueue), + // each hook's whole tensor and a segments table are handed over once. + bool batch_submit = false; // ClickHouse insert queue limits uint64_t insert_queue_max_bytes = 4096ULL * 1024 * 1024; uint64_t insert_queue_max_items = 65536; @@ -97,7 +104,8 @@ struct RecordCaptureStatus { uint64_t max_step_wait_ns = 0; }; -// Called by the p2p thread for each per-request tensor slice. +// Called by the p2p thread for each per-request tensor slice; see +// ring/p2p_thread.h for capture_ts_us and capture_lease. using SubmitFn = std::function; + at::Tensor slice, + int64_t capture_ts_us, + CaptureLease capture_lease)>; + +// Called by the p2p thread once per (hook x step) on the batch handoff, and +// once per request slice for HF batched steps when it is the only consumer; +// see ring/p2p_thread.h BatchSink. +using BatchSink = std::function req_ids, + at::Tensor segments, + CaptureLease capture_lease, + int64_t capture_ts_us, + bool is_attn)>; // Opaque RAII engine. class RingEnginePy { public: - explicit RingEnginePy(RingConfig cfg, SubmitFn submit_fn); + explicit RingEnginePy(RingConfig cfg, SubmitFn submit_fn, + BatchSink batch_sink = nullptr); explicit RingEnginePy(RingConfig cfg, std::shared_ptr sink, RecordRuntimeOptions options = {}); @@ -247,6 +275,11 @@ class RingEnginePy { // Lightweight wake-up for the drain thread. void notify_drain(); + // Nonblocking hint: flush the tasks already reserved as soon as they are + // published, whatever the flush thresholds say. Returns false when + // nothing is reserved beyond the drain tail. + bool request_flush(); + // Pre-forward capacity check + conditional flush. // // Called once per generate() step from _prepare_wrapper (GIL released). @@ -367,6 +400,24 @@ class RingEnginePy { // ensure FIFO ordering before consuming the next meta out-of-band. void flush_and_wait(); + // Like flush_and_wait(), and additionally wait until the P2P thread has + // delivered everything flushed: every consumer call for every slice or + // batch captured so far has returned. Returns true when delivered, + // false on timeout (timeout_ms = 0 waits forever). A finite timeout is + // ONE deadline taken at entry; it bounds producer-stream completion + // (polled through an event, no unbounded stream synchronize), the flush + // and the delivery wait together. Legacy rings only. + bool drain_to_sink_and_wait(uint32_t timeout_ms); + + // Consumer failures: exceptions thrown by the consumer are counted and + // the latest message kept, instead of vanishing. With + // abort-on-sink-error armed, sink_failed() turns true at the first one + // and prepare_step() raises it. + uint64_t submit_exceptions() const; + std::string last_sink_error() const; + void set_abort_on_sink_error(bool enabled); + bool sink_failed() const; + // ---- Runtime CUDA-graph node toggling (legacy rings) ----------------- // cudaGraph_t / cudaGraphNode_t / cudaGraphExec_t travel as opaque // uint64_t so this header stays CUDA-free (bindings.cpp is compiled with diff --git a/native/csrc/ring/slice_queue.h b/native/csrc/ring/slice_queue.h new file mode 100644 index 000000000..d08caaf8a --- /dev/null +++ b/native/csrc/ring/slice_queue.h @@ -0,0 +1,157 @@ +// ring/slice_queue.h -- per-slice in-process consumer queue. +// +// A Python-callable consumer makes the P2P thread take the GIL for every +// slice, which serializes extraction against a GIL-bound serving thread under +// load (ring backpressure, then producer stalls). With this queue the P2P +// thread pushes finished slices in pure C++ (no GIL), and a Python consumer +// pulls them in batches with drain(), releasing the GIL while it blocks. +// +// Bounds: max_items, and optionally max_bytes (0 = no byte cap). Retained +// bytes are the payload allocations the queued slices keep alive, each +// distinct allocation charged once at max(logical bytes, allocated bytes) +// (pooled_block.h): per-request views of one step's pooled buffer share one +// charge, and a 32-byte pooled slice is charged its 64 KiB block. +// +// Header-only; included where ATen is already linked. + +#pragma once +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include "ring/drop_notice.h" +#include "ring/pooled_block.h" + +namespace ring_py { + +struct SliceItem { + std::string model_id; + int32_t shard_rank; + std::string req_id; + std::string act_name; + int32_t layer_no; + int32_t start_token; + int32_t end_token; + at::Tensor slice; + CaptureLease capture_lease; + int64_t capture_ts_us = 0; + ring::BackingRef backing; // the allocation `slice` keeps alive +}; + +class SliceQueue { +public: + explicit SliceQueue(size_t max_items = 200000, size_t drop_notice_cap = 4096, + uint64_t max_bytes = 0) + : max_items_(max_items), max_bytes_(max_bytes), drops_(drop_notice_cap) {} + + // Pushed by the P2P thread in pure C++ (never the GIL). Bounded: when + // full, the slice is dropped, counted AND attributed -- a DropNotice + // naming this req_id is filed (drop_notice.h). + void push(const std::string& model_id, int32_t shard_rank, + const std::string& req_id, const std::string& act_name, + int32_t layer_no, int32_t start_token, int32_t end_token, + at::Tensor slice, CaptureLease capture_lease = {}, + int64_t capture_ts_us = 0) { + bool shed = false; + const ring::BackingRef backing = ring::backing_of(slice); + { + std::lock_guard lk(mu_); + // The byte cap, like BatchQueue's, always admits into an empty + // queue so a tiny cap cannot wedge the pipeline. + const bool over_bytes = max_bytes_ > 0 && !q_.empty() && + ledger_.total() + ledger_.cost(backing) > max_bytes_; + if (q_.size() >= max_items_ || over_bytes) { + dropped_++; + shed = true; + } else { + ledger_.add(backing); + q_.push_back(SliceItem{model_id, shard_rank, req_id, act_name, + layer_no, start_token, end_token, + std::move(slice), std::move(capture_lease), capture_ts_us, + backing}); + } + } + if (shed) { + // Outside mu_: the notice buffer has its own lock (no nesting). + DropNotice n; + n.source = kDropSourceSlice; + n.hook = act_name; + n.layer_no = layer_no; + n.start_token = start_token; + n.end_token = end_token; + n.nbytes = static_cast(slice.defined() ? slice.nbytes() : 0); + n.ts_us = drop_notice_now_us(); + n.req_ids.push_back(req_id); + n.rank = capture_lease.rank_for(act_name, layer_no, n.rank); + n.capture_lease = std::move(capture_lease); + drops_.record(std::move(n)); + return; + } + cv_.notify_one(); + } + + // Wake any waiting consumer so it can observe shutdown / final items. + void stop() { + { std::lock_guard lk(mu_); stopped_ = true; } + cv_.notify_all(); + } + + // Coalesced notification for host-side control work, not a queue shutdown. + void wake() { + { std::lock_guard lk(mu_); wake_pending_ = true; } + cv_.notify_one(); + } + + // Pop up to max_items, waiting up to timeout_ms for the first one. MUST be + // called with the GIL RELEASED (the binding wraps it in gil_scoped_release) + // so the blocking wait never holds the GIL. + std::vector drain(size_t max_items, int timeout_ms) { + std::vector out; + std::unique_lock lk(mu_); + if (q_.empty() && !stopped_) { + cv_.wait_for(lk, std::chrono::milliseconds(timeout_ms), + [this] { return !q_.empty() || stopped_ || wake_pending_; }); + } + wake_pending_ = false; + while (!q_.empty() && out.size() < max_items) { + ledger_.remove(q_.front().backing); + out.push_back(std::move(q_.front())); + q_.pop_front(); + } + return out; + } + + // Drain the drop notices filed since the last poll (see drop_notice.h). + // Non-blocking. Absence of notices == nothing was lost. + DropNoticeBuffer::Drained drain_drop_notices(size_t max_items = 4096) { + return drops_.drain(max_items); + } + const DropNoticeBuffer& drop_notices() const { return drops_; } + + uint64_t dropped() const { std::lock_guard lk(mu_); return dropped_; } + size_t size() const { std::lock_guard lk(mu_); return q_.size(); } + // Allocated bytes the queued slices keep alive (see the header comment). + uint64_t retained_bytes() const { std::lock_guard lk(mu_); return ledger_.total(); } + uint64_t max_bytes() const { std::lock_guard lk(mu_); return max_bytes_; } + +private: + mutable std::mutex mu_; + std::condition_variable cv_; + std::deque q_; + size_t max_items_; + uint64_t max_bytes_; // 0 = no byte cap + ring::BackingLedger ledger_; + uint64_t dropped_ = 0; + bool stopped_ = false; + bool wake_pending_ = false; + // Own mutex; never taken while mu_ is held (see push()). + DropNoticeBuffer drops_; +}; + +} // namespace ring_py diff --git a/native/csrc/ring/tensor_meta.h b/native/csrc/ring/tensor_meta.h index ba56afca9..5fd39ef00 100644 --- a/native/csrc/ring/tensor_meta.h +++ b/native/csrc/ring/tensor_meta.h @@ -15,6 +15,8 @@ #include #include +#include "ring/capture_lease.h" + namespace ring_py { // --------------------------------------------------------------------------- @@ -199,6 +201,9 @@ struct RequestMeta { // Per-step shared context. Heap-allocated by push_step caller. // Ownership transfers to p2p thread via pop_context; p2p deletes it. struct StepContext { + // Integration-defined identity of the configuration that captured this + // step; delivered with every slice or batch of the step (capture_lease.h). + CaptureLease capture_lease; std::string model_id; int32_t tp_rank = 0; int32_t dp_rank = 0; @@ -234,17 +239,52 @@ struct TensorMeta { }; // Thread-safe dual FIFO for TensorMeta + StepContext. +// +// Ownership: the FIFO owns every StepContext it holds. push_step stores ONE +// context per step (the step's metas do not point at it), so each pending +// context is owned exactly once; pop_context hands that ownership to the +// caller (the P2P thread, which deletes it after the step's last hook). +// Contexts still queued when the FIFO is cleared or destroyed -- metadata +// published for a forward that never produced its payloads -- are deleted +// then, releasing their CaptureLease references so lease.idle() can turn +// true after teardown. class TensorMetaFifo { public: + TensorMetaFifo() = default; + TensorMetaFifo(const TensorMetaFifo&) = delete; + TensorMetaFifo& operator=(const TensorMetaFifo&) = delete; + ~TensorMetaFifo() { clear(); } + // Push context + all hook metas in one lock. // ctx is heap-allocated by caller; ownership transfers here. void push_step(StepContext* ctx, std::vector& metas) { + std::unique_ptr owned(ctx); std::lock_guard lk(mu_); - ctx_q_.push_back(ctx); + if (owned) ctx_q_.push_back(std::move(owned)); for (auto& m : metas) q_.push_back(std::move(m)); } + // Drop every pending meta and delete every pending context. Only for a + // FIFO whose consumer has stopped (teardown): a context already popped + // by the consumer is the consumer's to delete. + void clear() { + std::deque> ctxs; + std::deque metas; + { + std::lock_guard lk(mu_); + ctxs.swap(ctx_q_); + metas.swap(q_); + } + // Destroyed outside mu_: a lease's last release runs no callbacks + // today, but nothing here needs the lock. + } + + size_t pending_contexts() const { + std::lock_guard lk(mu_); + return ctx_q_.size(); + } + bool pop(TensorMeta& out) { std::lock_guard lk(mu_); if (q_.empty()) return false; @@ -257,15 +297,15 @@ class TensorMetaFifo { StepContext* pop_context() { std::lock_guard lk(mu_); if (ctx_q_.empty()) return nullptr; - StepContext* p = ctx_q_.front(); + StepContext* p = ctx_q_.front().release(); ctx_q_.pop_front(); return p; } private: - std::mutex mu_; - std::deque q_; - std::deque ctx_q_; + mutable std::mutex mu_; + std::deque q_; + std::deque> ctx_q_; }; } // namespace ring_py diff --git a/src/dmi/adapters/types.py b/src/dmi/adapters/types.py index ae85155b1..84b991513 100644 --- a/src/dmi/adapters/types.py +++ b/src/dmi/adapters/types.py @@ -64,6 +64,8 @@ class StepContext: # (hook_type, layer_no). A planned hook reserves and publishes its # plan's row_count rows; None = every hook uses the step layout. hook_rows: Optional[dict] = None + # Integration-defined CaptureLease delivered with this step's payloads. + capture_lease: Optional[object] = None def transport_kwargs(self) -> dict: """Return kwargs accepted by ``RingTransport.set_step_context``.""" @@ -83,6 +85,8 @@ def transport_kwargs(self) -> dict: kwargs["orig_positions"] = self.orig_positions if self.hook_rows is not None: kwargs["hook_rows"] = self.hook_rows + if self.capture_lease is not None: + kwargs["capture_lease"] = self.capture_lease return kwargs diff --git a/src/dmi/engine.py b/src/dmi/engine.py index 5191e19d3..b508763ee 100644 --- a/src/dmi/engine.py +++ b/src/dmi/engine.py @@ -6,6 +6,7 @@ import importlib import logging import time +import warnings from typing import Any, Optional, Sequence, TYPE_CHECKING, TypeVar from .config import MonitoringConfig @@ -105,6 +106,7 @@ def __init__( model_id: Optional[str] = None, host_engine: Optional[Any] = None, db_config: Optional[HostEngineConfig] = None, + ring_consumer: Optional[Any] = None, enable_ring_transport: bool = True, ring_config: Optional[Any] = None, ring_payload_mb: int = 4096, @@ -129,6 +131,14 @@ def __init__( if host_engine is not None and db_config is not None: raise ValueError("Provide either host_engine or db_config, not both") + if ring_consumer is not None and (host_engine is not None + or db_config is not None): + raise ValueError( + "Provide either ring_consumer or host_engine/db_config, not " + "both: the legacy ring delivers to exactly one consumer") + # An in-process consumer of the legacy ring: a Python callable, a + # SliceQueue or a BatchQueue (see enable_ring_transport). + self._ring_consumer: Optional[Any] = ring_consumer # The canonical name: a deprecated one ("native", "capture") is # replaced here, for a MonitoringConfig and any config-like object @@ -167,7 +177,8 @@ def __init__( # before the service starts until the sink and the service are done. self._spool_claim: Optional[Any] = None host_configured = host_engine is not None or db_config is not None - if self._storage_backend == "in-memory" and not host_configured: + if (self._storage_backend == "in-memory" and not host_configured + and ring_consumer is None): raise ValueError( f"config.storage_backend={self._backend_label()} delivers " "records through " @@ -181,7 +192,8 @@ def __init__( "engine allocates no ring, but a ring_config was given. Pick " "'in-memory' or 'persistent' to capture, or drop ring_config" ) - if self._storage_backend in ("persistent", "none") and host_configured: + if self._storage_backend in ("persistent", "none") and ( + host_configured or ring_consumer is not None): raise ValueError( f"config.storage_backend={self._backend_label()} does not " "use the C++ ClickHouse host, but host_engine/db_config was " @@ -901,16 +913,30 @@ def enable_ring_transport( self._retire_capture_storage(storage, drain_deadline, sink_sealed=sealed) - # Pass the DMXHostEngine C++ object directly; RingEngine builds a - # SubmitFn that calls submit_direct without touching Python/GIL. - # Pass None for null/benchmark mode (no DB writes). - host_cpp = None - if self._host_engine is not None and isinstance( + # The consumer the ring delivers to. A DMXHostEngine is passed as + # the C++ object, so RingEngine calls submit_direct without touching + # Python or the GIL; SliceQueue/BatchQueue consumers are pushed in + # C++ as well; a callable takes the GIL per slice. None is the + # null/benchmark mode: captures are drained and dropped. + consumer = getattr(self, "_ring_consumer", None) + if consumer is None and self._host_engine is not None and isinstance( self._host_engine, _native_engine.DMXHostEngine ): - host_cpp = self._host_engine - - ring_engine = _native_engine.RingEngine(ring_config, host_cpp) + consumer = self._host_engine + if consumer is None: + # Never run without a consumer by accident: an attribute that + # exists but holds None looks exactly like a configured consumer + # everywhere except here, where every capture is dropped. + warnings.warn( + "MonitoringEngine builds a legacy ring with no consumer: " + "captures are drained and dropped. Pass host_engine=, " + "db_config= or ring_consumer= to deliver them.", + RingWithoutConsumerWarning, stacklevel=2) + if consumer is not None and isinstance( + consumer, _native_engine.BatchQueue): + ring_config.batch_submit = True + + ring_engine = _native_engine.RingEngine(ring_config, consumer) ring_engine.init() ring_engine.start() @@ -1102,10 +1128,18 @@ def close(self) -> None: except Exception: pass self._host_engine = None + # An in-process consumer belongs to its caller; drop the reference + # only, after the ring that delivered to it has stopped. + self._ring_consumer = None # --------------------------------------------------------------------------- # Backend loader -__all__ = ["MonitoringEngine", "RECORD_FAILURE_POLICIES", "RingCapacities"] +class RingWithoutConsumerWarning(UserWarning): + """A legacy ring was built with no consumer: captures are dropped.""" + + +__all__ = ["MonitoringEngine", "RECORD_FAILURE_POLICIES", "RingCapacities", + "RingWithoutConsumerWarning"] diff --git a/src/dmi/transport/native.py b/src/dmi/transport/native.py index 0af261988..315556101 100644 --- a/src/dmi/transport/native.py +++ b/src/dmi/transport/native.py @@ -27,11 +27,14 @@ ) _RING_EXPORTS = frozenset( { + "BatchQueue", + "CaptureLease", "ClickHouseRecordSink", "RecordSink", "ReferencePythonCaptureSink", "RingConfig", "RingEngine", + "SliceQueue", "ring_set_active_engine", "ring_clear_active_engine", "parse_hook_requests", diff --git a/src/dmi/transport/ring.py b/src/dmi/transport/ring.py index dd68ff34b..498589eb1 100644 --- a/src/dmi/transport/ring.py +++ b/src/dmi/transport/ring.py @@ -253,6 +253,7 @@ def __init__(self, ring_engine: Any) -> None: # keyed by (hook_type, layer_no); None = every hook uses the step # layout. self._current_hook_rows: Optional[dict] = None + self._current_capture_lease: Any = None # When True: meta pushes are skipped so the FIFO stays empty. # Producer kernel still fires (for CUDA graph capture) but as no-ops. @@ -338,6 +339,7 @@ def set_step_context( flattened: bool = False, orig_positions: Optional[List[int]] = None, hook_rows: Optional[dict] = None, + capture_lease: Any = None, ) -> None: """Called before each forward pass to provide per-step batch metadata. @@ -361,6 +363,9 @@ def set_step_context( keyed by ``(hook_type, layer_no)``. A hook with a plan publishes ``row_count`` rows and its own request layout; others use the step layout. + capture_lease: an integration-defined ``CaptureLease`` delivered + with every slice or batch of this step; it is handed to the + next metadata publication only. """ self._current_model_id = model_id self._current_tp_rank = tp_rank @@ -380,6 +385,7 @@ def set_step_context( ) self._current_orig_positions = orig_positions self._current_hook_rows = hook_rows + self._current_capture_lease = capture_lease def set_model_cfg(self, cfg: ModelShapeConfig) -> None: """Set the model shape config for analytical shape computation.""" @@ -491,6 +497,9 @@ def pre_push_all_metas(self, batch: int, q_len: int, kv_dim: int, padding-strip mode. Other specs and the no-strip case use ``q_len`` (today's behavior). """ + # The lease belongs to this submission, never to a later step. + lease = getattr(self, "_current_capture_lease", None) + self._current_capture_lease = None if self.null_offload: return # kernel launches happen; metas are intentionally skipped if self._model_cfg is None or not self._active_specs: @@ -537,6 +546,8 @@ def pre_push_all_metas(self, batch: int, q_len: int, kv_dim: int, if hook_types: extra = {} + if lease is not None: + extra["capture_lease"] = lease orig_positions = getattr(self, "_current_orig_positions", None) if orig_positions is not None: extra["orig_positions"] = list(orig_positions) diff --git a/tests/native/ring/cpu_cuda_stub/cuda_runtime.h b/tests/native/ring/cpu_cuda_stub/cuda_runtime.h index ba02b4660..e2550b918 100644 --- a/tests/native/ring/cpu_cuda_stub/cuda_runtime.h +++ b/tests/native/ring/cpu_cuda_stub/cuda_runtime.h @@ -1,5 +1,6 @@ // CPU-only stand-in for the CUDA runtime API used by the host half of -// ring/toggle_registry.cu and by ring/ring_torch_op.cpp's producer ops. No +// ring/toggle_registry.cu, ring/ring_torch_op.cpp's producer ops and the +// host-side P2P handoff (ring/p2p_thread.cpp, ring/pinned_staging.h). No // CUDA library is linked or initialized: each call records what a real exec // would hold so tests can compare the registry's view with the exec's actual // node state. @@ -79,6 +80,12 @@ inline int cudaMallocManaged(void** p, size_t n) { return *p ? 0 : 1; } inline int cudaFree(void* p) { std::free(p); return 0; } +// Pinned host memory (ring/pinned_staging.h). Declared only, with the C +// linkage of the real runtime API: a host-only test that reaches them must +// define them itself (or never call them, and drop them at link time). +constexpr unsigned cudaHostAllocDefault = 0; +extern "C" cudaError_t cudaHostAlloc(void** p, size_t n, unsigned flags); +extern "C" cudaError_t cudaFreeHost(void* p); inline int cudaStreamCreateWithFlags(cudaStream_t* p, unsigned) { *p = reinterpret_cast(1); return 0; diff --git a/tests/native/ring/test_ring_engine.cu b/tests/native/ring/test_ring_engine.cu index fc905b62f..e039b0eee 100644 --- a/tests/native/ring/test_ring_engine.cu +++ b/tests/native/ring/test_ring_engine.cu @@ -1,6 +1,7 @@ // CUDA integration tests for producer -> drain -> pinned-staging delivery. #include "ring/drain_thread.h" +#include "ring/p2p_thread.h" #include "ring/pinned_staging.h" #include "ring/producer.cuh" #include "ring/ring_alloc.h" @@ -326,7 +327,7 @@ static void check_eager_delivery_past_the_task_ring(uint32_t task_slots) { cfg, [&](const std::string&, int32_t, const std::string&, const std::string& act_name, int32_t layer_no, int32_t, int32_t, - at::Tensor slice) { + at::Tensor slice, int64_t, ring_py::CaptureLease) { const uint8_t* data = slice.data_ptr(); Delivery delivery{act_name, layer_no, slice.sizes().vec(), std::vector(data, data + slice.nbytes())}; @@ -473,6 +474,31 @@ static void test_record_ring_refuses_every_legacy_producer_entry() { CUDA_CHECK(cudaFree(source)); } +static void test_delivery_barrier_returns_once_the_consumer_thread_exits() { + banner("delivery barrier returns when the consumer thread has exited"); + DrainHarness harness(make_config()); + ring_py::TensorMetaFifo fifo; + ring::P2PThread p2p(*harness.drain, fifo, harness.cfg, ring::SubmitFn{}); + p2p.start(); + harness.drain->stop(); + harness.drain->signal_p2p_stop(); // the consumer loop exits + // An unbounded wait for a task that will never be processed must not + // hang: a watchdog fails the test instead of the whole suite stalling. + auto verdict = std::make_shared>(); + auto result = verdict->get_future(); + std::thread([&p2p, verdict] { + verdict->set_value(p2p.wait_until_processed(1, 0)); + }).detach(); + if (result.wait_for(std::chrono::seconds(10)) != std::future_status::ready) { + std::fprintf(stderr, "FAIL: wait_until_processed hung after thread exit\n"); + std::_Exit(1); + } + EXPECT(result.get() == false); + EXPECT(p2p.wait_until_processed(0, 0) == true); + EXPECT(p2p.wait_until_processed(1, 50) == false); + p2p.stop(); +} + static void test_static_force_flush() { banner("static producer drains exact data"); DrainHarness harness(make_config()); @@ -1369,6 +1395,7 @@ int main() { test_ring_room_saturates_past_capacity(); test_eager_safety_net_delivers_past_the_task_ring(); test_record_ring_refuses_every_legacy_producer_entry(); + test_delivery_barrier_returns_once_the_consumer_thread_exits(); test_static_force_flush(); test_prefix_force_flush(); test_chunked_short_then_static_preserves_payload_offset(); diff --git a/tests/native/ring/test_ring_torch_op_host.cpp b/tests/native/ring/test_ring_torch_op_host.cpp index 1e3bf1adf..91fe88f48 100644 --- a/tests/native/ring/test_ring_torch_op_host.cpp +++ b/tests/native/ring/test_ring_torch_op_host.cpp @@ -59,7 +59,7 @@ namespace ring_py { struct RingEnginePy::Impl {}; -RingEnginePy::RingEnginePy(RingConfig, SubmitFn) : impl_(new Impl) {} +RingEnginePy::RingEnginePy(RingConfig, SubmitFn, BatchSink) : impl_(new Impl) {} RingEnginePy::~RingEnginePy() = default; void RingEnginePy::hook_no_notify(uint64_t, uint64_t, uint32_t, uint64_t) { diff --git a/tests/native/test_batch_queue_shared_eviction.cpp b/tests/native/test_batch_queue_shared_eviction.cpp new file mode 100644 index 000000000..1e8c3d0df --- /dev/null +++ b/tests/native/test_batch_queue_shared_eviction.cpp @@ -0,0 +1,178 @@ +// BatchQueue rank-ordered eviction with views of a shared allocation. +// +// The byte cap charges each distinct allocation once while any queued item +// refers to it. Eviction must judge what a set of victims would actually +// free from the references that REMAIN -- surviving aliases of any rank and +// the arriving item's own -- and evict nothing unless the arrival is then +// admitted. +// +// Compiles batch_queue.h against ATen (CPU only). Built and run by +// tests/test_native_ring_consumer_fixes.py. + +#include + +#include +#include +#include +#include + +#include "ring/batch_queue.h" + +namespace { + +int g_failures = 0; + +void check(bool ok, const std::string& what) { + if (!ok) { + std::cerr << "FAIL: " << what << '\n'; + ++g_failures; + } +} + +// Ranks: "urgent*" 0, "mid" 1, anything else best-effort (2). +void install_map(ring_py::BatchQueue& q) { + q.set_layer_class_map( + {{"urgent", 0, 0}, {"newurgent", 0, 0}, {"mid", 0, 1}}, 2); +} + +void push(ring_py::BatchQueue& q, const char* act, const at::Tensor& t) { + ring_py::HookBatchItem item; + item.act_name = act; + item.layer_no = 0; + item.tensor = t; + item.nbytes = t.nbytes(); + item.req_ids = {act}; + q.push(std::move(item)); +} + +struct Outcome { + size_t size; + uint64_t retained, dropped, evicted; + std::string items; + uint64_t retained_after_drain; +}; + +Outcome finish(ring_py::BatchQueue& q) { + Outcome o{q.size(), q.retained_bytes(), q.dropped(), q.evicted(), "", 0}; + for (const auto& item : q.drain(100, 0)) o.items += item.act_name + ","; + o.retained_after_drain = q.retained_bytes(); + return o; +} + +void expect(const char* label, const Outcome& o, size_t size, uint64_t retained, + uint64_t dropped, uint64_t evicted, const std::string& items) { + std::cout << label << " size=" << o.size << " retained=" << o.retained + << " dropped=" << o.dropped << " evicted=" << o.evicted + << " items=" << o.items + << " retained_after_drain=" << o.retained_after_drain << '\n'; + const std::string l = label; + check(o.size == size, l + ": size"); + check(o.retained == retained, l + ": retained"); + check(o.dropped == dropped, l + ": dropped"); + check(o.evicted == evicted, l + ": evicted"); + check(o.items == items, l + ": items"); + check(o.retained_after_drain == 0, l + ": nothing retained after drain"); +} + +} // namespace + +int main() { + // A drained urgent view leaves its 80-byte allocation pinned only by a + // low view. Evicting that low view frees all 80 bytes: the new urgent + // item fits (10 + 80 <= 100). + { + ring_py::BatchQueue q(100); + install_map(q); + auto shared = at::zeros({80}, at::kByte); + push(q, "urgent", shared.narrow(0, 0, 40)); + push(q, "low", shared.narrow(0, 40, 40)); + q.drain(1, 0); + push(q, "urgent", at::zeros({10}, at::kByte)); + push(q, "newurgent", at::zeros({80}, at::kByte)); + expect("underestimated_evictable_capacity", finish(q), + 2, 90, 1, 1, "urgent,newurgent,"); + } + // A low and an urgent view share 80 bytes. Evicting low frees nothing + // (urgent keeps the allocation), so the newcomer cannot fit: nothing is + // evicted and only the newcomer is refused. + { + ring_py::BatchQueue q(100); + install_map(q); + auto shared = at::zeros({80}, at::kByte); + push(q, "low", shared.narrow(0, 0, 40)); + push(q, "urgent", shared.narrow(0, 40, 40)); + push(q, "newurgent", at::zeros({80}, at::kByte)); + expect("overestimated_evictable_capacity", finish(q), + 2, 80, 1, 0, "low,urgent,"); + check(q.qsize_by_rank() == std::vector({0, 0, 0}), + "depth by rank after drain"); + } + // A shared allocation is freed only when its LAST view goes: the low view + // alone frees nothing, the mid view then frees all 80 bytes. Victims are + // taken worst rank first, oldest first within a rank. + { + ring_py::BatchQueue q(120); + install_map(q); + auto shared = at::zeros({80}, at::kByte); + push(q, "low", shared.narrow(0, 0, 40)); + push(q, "mid", shared.narrow(0, 40, 40)); + push(q, "low2", at::zeros({10}, at::kByte)); + push(q, "newurgent", at::zeros({100}, at::kByte)); + // 90 + 100 > 120. low frees 0, low2 frees 10 (180 > 120), mid frees + // the shared 80: 100 retained. + expect("last_view_frees_the_allocation", finish(q), + 1, 100, 3, 3, "newurgent,"); + } + // The arriving item aliases a queued allocation: it adds no bytes, so it + // is admitted at the cap with nothing evicted. + { + ring_py::BatchQueue q(100); + install_map(q); + auto shared = at::zeros({80}, at::kByte); + push(q, "low", shared.narrow(0, 0, 40)); + push(q, "low2", at::zeros({20}, at::kByte)); + push(q, "newurgent", shared.narrow(0, 40, 40)); + expect("arrival_aliases_a_queued_allocation", finish(q), + 3, 100, 0, 0, "low,low2,newurgent,"); + } + // A less urgent arrival never evicts a more urgent resident. + { + ring_py::BatchQueue q(100); + install_map(q); + auto big = at::zeros({100}, at::kByte); + push(q, "urgent", at::zeros({10}, at::kByte)); // 10 + push(q, "low", big.narrow(0, 0, 50)); // 110: refused + expect("oversized_low_is_refused", finish(q), 1, 10, 1, 0, "urgent,"); + } + { + // Queue over its cap through the empty-queue rule (one item larger + // than the cap); an urgent alias of that same allocation arrives. + // Evicting the low view frees nothing, but it empties the queue, + // which always admits: the urgent alias replaces it and the + // allocation stays charged once. + ring_py::BatchQueue q(50); + install_map(q); + auto big = at::zeros({80}, at::kByte); + push(q, "low", big.narrow(0, 0, 40)); // 80 > 50, admitted + push(q, "newurgent", big.narrow(0, 40, 40)); + expect("alias_over_cap_empties_the_queue", finish(q), + 1, 80, 1, 1, "newurgent,"); + } + // Reclassifying streams while items are queued: an item keeps the rank it + // was enqueued with, and the per-rank depths stay consistent. + { + ring_py::BatchQueue q(100); + install_map(q); + push(q, "a", at::zeros({60}, at::kByte)); // best-effort + q.set_layer_class_map({{"a", 0, 0}, {"b", 0, 0}}, 2); + check(q.qsize_by_rank() == std::vector({0, 0, 1}), + "queued item keeps its enqueue rank"); + push(q, "b", at::zeros({60}, at::kByte)); // urgent now + check(q.qsize_by_rank() == std::vector({1, 0, 0}), + "depth by rank after eviction"); + expect("reclassified_while_queued", finish(q), 1, 60, 1, 1, "b,"); + } + if (g_failures) return 1; + std::cout << "batch queue shared eviction ok\n"; + return 0; +} diff --git a/tests/native/test_buffer_pool_failure.cpp b/tests/native/test_buffer_pool_failure.cpp new file mode 100644 index 000000000..9344982cd --- /dev/null +++ b/tests/native/test_buffer_pool_failure.cpp @@ -0,0 +1,106 @@ +// BufferPool::acquire under allocation failure: every allocation acquire() +// makes -- the block's context, then ATen's StorageImpl, TensorImpl, ... -- +// is failed in turn with std::bad_alloc. Each failure must propagate the +// exception and return the block to the pool exactly once; the first count +// that does not fail must hand out a working tensor whose death returns the +// block exactly once. +// +// The pool is given an idle cap, so a returned block sits on its free list +// and idle_bytes() counts returns: one return is exactly one size class, a +// double return (the context deleted twice) is two, a leak is zero. +// +// Compiles buffer_pool.h against ATen (CPU only). Built and run by +// tests/test_native_ring_consumer_fixes.py. + +#include + +#include +#include +#include +#include +#include +#include + +#include "ring/buffer_pool.h" +#include "ring/pooled_block.h" + +namespace { + +// > 0: fail the n-th operator new from now on, once. +std::atomic g_fail_after{0}; +std::atomic g_injected{0}; + +int g_failures = 0; + +void check(bool ok, const std::string& what) { + if (!ok) { + std::cerr << "FAIL: " << what << '\n'; + ++g_failures; + } +} + +} // namespace + +void* operator new(std::size_t n) { + if (g_fail_after.load() > 0 && g_fail_after.fetch_sub(1) == 1) { + g_injected.fetch_add(1); + throw std::bad_alloc(); + } + if (void* p = std::malloc(n ? n : 1)) return p; + throw std::bad_alloc(); +} +void operator delete(void* p) noexcept { std::free(p); } +void operator delete(void* p, std::size_t) noexcept { std::free(p); } + +int main() { + at::zeros({1}, at::kByte); // initialize ATen before injecting + constexpr uint64_t kBytes = 32; + const uint64_t cls = ring::BufferPool::capacity_for(kBytes); + bool succeeded = false; + int failures_seen = 0; + for (int k = 1; k <= 64 && !succeeded; ++k) { + ring::BufferPool pool(1ULL << 30); + const int injected0 = g_injected.load(); + bool threw = false; + at::Tensor t; + g_fail_after.store(k); + try { + t = pool.acquire(kBytes); + } catch (const std::bad_alloc&) { + threw = true; + } + g_fail_after.store(0); + const std::string at_k = " (allocation " + std::to_string(k) + ")"; + if (threw) { + ++failures_seen; + check(g_injected.load() == injected0 + 1, "one injected failure" + at_k); + check(pool.idle_bytes() == cls, + "a failed acquire returns its block exactly once" + at_k + + ": idle_bytes=" + std::to_string(pool.idle_bytes())); + // The returned block is reused, then returned once more. + at::Tensor again = pool.acquire(kBytes); + check(again.defined() && again.nbytes() == kBytes, + "the pool still serves after a failure" + at_k); + check(pool.idle_bytes() == 0, "the returned block is reused" + at_k); + again.reset(); + check(pool.idle_bytes() == cls, "reuse returns it once" + at_k); + } else { + succeeded = true; + check(t.defined() && t.nbytes() == kBytes, "acquire succeeds" + at_k); + check(ring::pooled_block_capacity(t) == cls, "block capacity"); + check(pool.idle_bytes() == 0, "a live tensor holds its block"); + at::Tensor view = t.narrow(0, 8, 8); + t.reset(); + check(pool.idle_bytes() == 0, "a live view holds the block"); + view.reset(); + check(pool.idle_bytes() == cls, "the last view returns it once"); + } + } + check(succeeded, "acquire eventually succeeds"); + // The context, StorageImpl and TensorImpl allocations at least. + check(failures_seen >= 3, "every allocation of acquire was failed: " + + std::to_string(failures_seen)); + if (g_failures) return 1; + std::cout << "buffer pool failures=" << failures_seen << " ok\n"; + return 0; +} diff --git a/tests/native/test_p2p_batch_handoff.cpp b/tests/native/test_p2p_batch_handoff.cpp new file mode 100644 index 000000000..b6d9c891d --- /dev/null +++ b/tests/native/test_p2p_batch_handoff.cpp @@ -0,0 +1,373 @@ +// P2P post-processing with a BatchQueue-style consumer, and pooled-buffer +// retention accounting. +// +// 1. With only a batch sink installed (BatchQueue), HF batched captures -- +// token [batch, q_len, hidden] and attention [batch, heads, q_len, kv] -- +// used to fall through the whole-batch branch (packed, non-attention only) +// and return at the empty per-slice consumer: nothing delivered, nothing +// counted. They now arrive as one-request items, one per request slice. +// 2. A pooled 32-byte payload pins its whole 64 KiB size class. BatchQueue +// and SliceQueue charge the allocation a queued item keeps alive, once per +// distinct allocation, not the logical length. +// 3. p2p_pool_cap_bytes == 0 allocates exact-size buffers (no pool, no +// size-class rounding). +// 4. The P2P delivery wait honours an absolute deadline. +// +// Compiles the real p2p_thread.cpp against ATen (CPU only). The DrainThread +// is never started: its constructor is replaced below so no CUDA API runs, +// and cudaFreeHost aborts if anything reaches it. Built and run by +// tests/test_native_ring_consumer_fixes.py. + +#include + +#include +#include +#include +#include + +#include "ring/batch_queue.h" +#include "ring/buffer_pool.h" +#include "ring/drain_thread.h" +#include "ring/pooled_block.h" +#include "ring/slice_queue.h" +#include "ring/tensor_meta.h" +// The test drives the private post-processing entry points directly. +#define private public +#include "ring/p2p_thread.h" +#undef private + +extern "C" cudaError_t cudaFreeHost(void*) { std::abort(); } + +namespace ring { +DrainThread::DrainThread(RingState& rs, PinnedStaging& st, const RingConfig& cfg) + : ring_(rs), staging_(st), cfg_(cfg) {} +DrainThread::~DrainThread() noexcept = default; +// Staging is never used here (every task has alloc_bytes == 0). +void DrainThread::notify_staging_freed_bytes(uint64_t) { std::abort(); } +} // namespace ring + +namespace { + +int g_failures = 0; + +void check(bool ok, const std::string& what) { + if (!ok) { + std::cerr << "FAIL: " << what << '\n'; + ++g_failures; + } +} + +struct Delivered { + std::string act_name; + int32_t layer_no; + uint64_t step_seq; + bool last_in_step; + at::Tensor tensor; + std::vector req_ids; + at::Tensor segments; + ring_py::CaptureLease lease; + int64_t ts; + bool is_attn; +}; + +struct Fixture { + ring::RingState state{}; + ring::PinnedStaging staging; + ring::RingConfig cfg; + ring_py::TensorMetaFifo fifo; + std::unique_ptr drain; + std::unique_ptr p2p; + std::vector got; + bool throw_in_sink = false; + + explicit Fixture(bool batch_submit = true, uint64_t pool_cap = 256ULL << 20) { + cfg.batch_submit = batch_submit; + cfg.p2p_pool_cap_bytes = pool_cap; + drain = std::make_unique(state, staging, cfg); + ring::BatchSink sink = [this](const std::string&, int32_t, + const std::string& act, int32_t layer, + uint64_t seq, bool last, at::Tensor t, + std::vector ids, at::Tensor seg, + ring_py::CaptureLease lease, int64_t ts, + bool is_attn) { + if (throw_in_sink) throw std::runtime_error("consumer refused"); + got.push_back(Delivered{act, layer, seq, last, std::move(t), std::move(ids), + std::move(seg), std::move(lease), ts, is_attn}); + }; + p2p = std::make_unique(*drain, fifo, cfg, ring::SubmitFn{}, + std::move(sink)); + } + + // Publish one single-hook step and post-process `tensor` for it. + void run(bool flattened, int hook, std::vector shape, + std::vector reqs, const at::Tensor& tensor, + const ring_py::CaptureLease& lease) { + auto* ctx = new ring_py::StepContext; + ctx->model_id = "m"; + ctx->flattened = flattened; + ctx->requests = std::move(reqs); + ctx->capture_lease = lease; + ring_py::TensorMeta meta; + meta.hook_type = hook; + meta.layer_no = 3; + meta.shape = shape; + meta.dtype = int(at::kFloat); + meta.last_in_step = true; + std::vector metas{meta}; + fifo.push_step(ctx, metas); + // The P2P thread hands post-processing a flat byte buffer. + at::Tensor bytes = tensor.contiguous().view(at::kByte).reshape({-1}).clone(); + ring::DrainTask task; + task.tensor_total_bytes = bytes.nbytes(); + task.capture_ts_us = 1234; + p2p->do_post_processing(bytes, task); + } +}; + +ring_py::RequestMeta req(const char* id, int32_t start, int32_t end, int64_t dim0, + int32_t kv_offset = 0, int32_t orig = -1) { + ring_py::RequestMeta r; + r.req_id = id; + r.start_token = start; + r.end_token = end; + r.dim0_offset = dim0; + r.kv_offset = kv_offset; + r.orig_position = orig; + return r; +} + +std::vector seg_row(const at::Tensor& seg) { + auto c = seg.contiguous(); + const int32_t* p = c.data_ptr(); + return {p[0], p[1], p[2], p[3]}; +} + +void test_hf_token_fallback() { + Fixture f; + ring_py::CaptureLease lease(9); + // [batch=2, q_len=3, hidden=4], left-padded: request a has 2 real rows, + // request b 3. + auto src = at::arange(24, at::kFloat).reshape({2, 3, 4}); + f.run(false, ring_py::HOOK_TYPE_RESID_PRE, {2, 3, 4}, + {req("a", 0, 2, 0, 0, 0), req("b", 5, 8, 1, 0, 1)}, src, lease); + check(f.p2p->submit_exceptions() == 0, "hf token: no consumer error"); + check(f.got.size() == 2, "hf token: one item per request slice"); + if (f.got.size() != 2) return; + const auto& a = f.got[0]; + const auto& b = f.got[1]; + check(a.req_ids == std::vector{"a"}, "hf token: item 0 is request a"); + check(b.req_ids == std::vector{"b"}, "hf token: item 1 is request b"); + check(a.tensor.sizes().vec() == std::vector{2, 4}, "hf token: a is [2, 4]"); + check(at::equal(a.tensor, src[0].narrow(0, 1, 2)), "hf token: a holds its last 2 rows"); + check(at::equal(b.tensor, src[1]), "hf token: b holds all 3 rows"); + check(seg_row(a.segments) == std::vector{0, 0, 2, 0}, "hf token: a segments row"); + check(seg_row(b.segments) == std::vector{0, 5, 8, 1}, "hf token: b segments row"); + check(a.segments.sizes().vec() == std::vector{1, 4}, "hf token: [1, 4] segments"); + check(!a.is_attn && !b.is_attn, "hf token: not attention"); + check(a.act_name == "blocks.hook_resid_pre" && a.layer_no == 3, "hf token: act/layer"); + check(a.step_seq == b.step_seq && a.step_seq > 0, "hf token: one step_seq per step"); + check(!a.last_in_step && b.last_in_step, "hf token: last_in_step on the last item only"); + check(a.lease.same_as(lease) && b.lease.same_as(lease), "hf token: lease carried"); + check(a.ts == 1234, "hf token: capture_ts_us carried"); +} + +void test_hf_attention_fallback() { + Fixture f; + ring_py::CaptureLease lease(10); + // [batch=2, heads=2, q=3, kv=4]. Request a: 2 query rows, keys [1, 3) + // (left-padded dynamic cache, kv_offset 1). Request b: 1 decode row. + auto src = at::arange(48, at::kFloat).reshape({2, 2, 3, 4}); + f.run(false, ring_py::HOOK_TYPE_PATTERN, {2, 2, 3, 4}, + {req("a", 0, 2, 0, 1), req("b", 3, 4, 1, 0)}, src, lease); + check(f.got.size() == 2, "hf attention: one item per request slice"); + if (f.got.size() != 2) return; + const auto& a = f.got[0]; + const auto& b = f.got[1]; + check(a.is_attn && b.is_attn, "hf attention: is_attn set"); + check(a.act_name == "blocks.attn.hook_pattern", "hf attention: act_name"); + check(a.tensor.sizes().vec() == std::vector{2, 2, 2}, + "hf attention: a is [heads, rows, kv]"); + check(at::equal(a.tensor, src[0].narrow(1, 1, 2).narrow(2, 1, 2)), + "hf attention: a holds its query rows and real keys"); + check(b.tensor.sizes().vec() == std::vector{2, 1, 4}, + "hf attention: b is [heads, 1, kv]"); + check(at::equal(b.tensor, src[1].narrow(1, 2, 1)), "hf attention: b holds the last row"); + check(seg_row(a.segments) == std::vector{0, 0, 2, -1}, "hf attention: a segments"); + check(seg_row(b.segments) == std::vector{0, 3, 4, -1}, "hf attention: b segments"); + check(b.last_in_step && !a.last_in_step, "hf attention: last_in_step"); +} + +void test_fallback_consumer_errors_are_counted() { + Fixture f; + f.throw_in_sink = true; + auto src = at::zeros({2, 3, 4}, at::kFloat); + f.run(false, ring_py::HOOK_TYPE_RESID_PRE, {2, 3, 4}, + {req("a", 0, 3, 0), req("b", 0, 3, 1)}, src, {}); + check(f.p2p->submit_exceptions() == 2, "fallback: every refused slice is counted"); + check(f.p2p->last_error() == "consumer refused", "fallback: last error kept"); +} + +void test_packed_paths_unchanged() { + Fixture f; + auto src = at::arange(20, at::kFloat).reshape({5, 4}); + f.run(true, ring_py::HOOK_TYPE_RESID_PRE, {5, 4}, + {req("a", 0, 2, 0), req("b", 0, 3, 2)}, src, {}); + check(f.got.size() == 1, "packed token: one whole-batch item"); + if (!f.got.empty()) { + check(f.got[0].req_ids == std::vector{"a", "b"}, "packed: req ids"); + check(f.got[0].tensor.sizes().vec() == std::vector{5, 4}, "packed: whole tensor"); + check(!f.got[0].is_attn && f.got[0].last_in_step, "packed: flags"); + } + // Packed attention is not sliceable per request on any path. + Fixture g; + auto attn = at::zeros({1, 2, 2}, at::kFloat); + g.run(true, ring_py::HOOK_TYPE_PATTERN, {1, 2, 2}, {req("a", 0, 2, 0)}, attn, {}); + check(g.got.empty() && g.p2p->submit_exceptions() == 0, "packed attention: skipped"); +} + +void test_pooled_retention_is_charged_by_allocation() { + const uint64_t block = ring::BufferPool::capacity_for(32); + check(block == 64 * 1024, "the minimum size class is 64 KiB"); + ring::BufferPool pool(256ULL << 20); + + // The reviewer's case: 32 independent 32-byte pooled items, logical sum + // exactly max_bytes. Each pins a whole block: only the first fits (an + // empty queue always admits one item), the rest are shed and counted. + { + ring_py::BatchQueue queue(1024); + for (int i = 0; i < 32; ++i) { + ring_py::HookBatchItem item; + item.tensor = pool.acquire(32); + item.nbytes = item.tensor.nbytes(); + item.act_name = "blocks.hook_resid_pre"; + item.req_ids = {"r"}; + check(ring::pooled_block_capacity(item.tensor) == block, + "a pooled tensor reports its block capacity"); + queue.push(std::move(item)); + } + check(queue.retained_bytes() == block, "batch queue: charged one block"); + check(queue.size() == 1 && queue.dropped() == 31, "batch queue: 31 shed"); + check(queue.dropped_bytes() == 31 * 32, "batch queue: dropped bytes are logical"); + queue.drain(64, 0); + check(queue.retained_bytes() == 0, "batch queue: drained to zero"); + } + + // Views of ONE block (one-request items of one step) share one charge, + // released when the last of them leaves the queue. + { + ring_py::BatchQueue queue(1ULL << 20); + at::Tensor whole = pool.acquire(4096); + for (int i = 0; i < 4; ++i) { + ring_py::HookBatchItem item; + item.tensor = whole.narrow(0, i * 1024, 1024); + item.nbytes = item.tensor.nbytes(); + item.act_name = "blocks.hook_resid_pre"; + queue.push(std::move(item)); + } + check(queue.retained_bytes() == block, "shared block: charged once"); + auto first = queue.drain(3, 0); + check(queue.retained_bytes() == block, "shared block: held while one item remains"); + queue.drain(1, 0); + check(queue.retained_bytes() == 0, "shared block: released with the last item"); + } + + // An unpooled tensor is charged its storage. + { + ring_py::BatchQueue queue(1ULL << 20); + ring_py::HookBatchItem item; + item.tensor = at::zeros({100}, at::kByte); + item.nbytes = 100; + queue.push(std::move(item)); + check(queue.retained_bytes() == 100, "unpooled: charged its storage bytes"); + } + + // SliceQueue: retained bytes by allocation, optional byte cap. + { + ring_py::SliceQueue queue(1000, 64, /*max_bytes=*/block + 1); + at::Tensor a = pool.acquire(64); + at::Tensor b = pool.acquire(64); + queue.push("m", 0, "a", "x", 0, 0, 1, a.narrow(0, 0, 32)); + queue.push("m", 0, "b", "x", 0, 0, 1, a.narrow(0, 32, 32)); // same block + check(queue.retained_bytes() == block, "slice queue: shared block charged once"); + queue.push("m", 0, "c", "x", 0, 0, 1, b); // a second block + check(queue.size() == 2 && queue.dropped() == 1, "slice queue: byte cap sheds"); + queue.drain(10, 0); + check(queue.retained_bytes() == 0, "slice queue: drained to zero"); + ring_py::SliceQueue unbounded; // max_bytes 0: no byte cap + unbounded.push("m", 0, "a", "x", 0, 0, 1, pool.acquire(32)); + unbounded.push("m", 0, "b", "x", 0, 0, 1, pool.acquire(32)); + check(unbounded.size() == 2 && unbounded.retained_bytes() == 2 * block, + "slice queue: no byte cap by default, retention still reported"); + } +} + +void test_pool_cap_zero_allocates_exact_size() { + for (uint64_t cap : {uint64_t{0}, uint64_t{256} << 20}) { + Fixture f(true, cap); + check((f.p2p->pool_ != nullptr) == (cap != 0), "pool exists iff cap > 0"); + std::vector staged(48, 7); + auto* ctx = new ring_py::StepContext; + ctx->model_id = "m"; + ctx->flattened = true; + ctx->requests = {req("a", 0, 3, 0)}; + ring_py::TensorMeta meta; + meta.hook_type = ring_py::HOOK_TYPE_RESID_PRE; + meta.layer_no = 0; + meta.shape = {3, 4}; + meta.dtype = int(at::kFloat); + meta.last_in_step = true; + std::vector metas{meta}; + f.fifo.push_step(ctx, metas); + std::vector tasks(1); + tasks[0].tensor_total_bytes = staged.size(); + tasks[0].data_ptr1 = staged.data(); + tasks[0].data_len1 = staged.size(); + tasks[0].alloc_bytes = 0; + f.p2p->process(tasks); + check(f.got.size() == 1, "staged task delivered"); + if (f.got.empty()) continue; + const at::Tensor& t = f.got[0].tensor; + if (cap == 0) { + check(ring::pooled_block_capacity(t) == 0, "cap 0: not a pool block"); + check(t.storage().nbytes() == staged.size(), "cap 0: exact-size buffer"); + check(ring::backing_of(t).bytes == staged.size(), "cap 0: charged exact size"); + } else { + check(ring::pooled_block_capacity(t) == 64 * 1024, "pooled: 64 KiB block"); + check(ring::backing_of(t).bytes == 64 * 1024, "pooled: charged the block"); + } + } +} + +void test_delivery_wait_honours_deadline() { + // Nothing is ever processed (the thread is not started): a finite wait + // must return false at its deadline, not block. + Fixture f; + const auto t0 = std::chrono::steady_clock::now(); + const bool ok = f.p2p->wait_until_processed_until( + 1, t0 + std::chrono::milliseconds(50)); + const auto ms = std::chrono::duration_cast( + std::chrono::steady_clock::now() - t0).count(); + check(!ok, "deadline wait: reports not delivered"); + check(ms >= 45 && ms < 5000, "deadline wait: returns at the deadline"); + // A deadline already in the past checks once. + check(!f.p2p->wait_until_processed_until(1, t0), "past deadline: immediate false"); + check(f.p2p->wait_until_processed_until(0, t0), "reached target: true"); +} + +} // namespace + +int main() { + test_hf_token_fallback(); + test_hf_attention_fallback(); + test_fallback_consumer_errors_are_counted(); + test_packed_paths_unchanged(); + test_pooled_retention_is_charged_by_allocation(); + test_pool_cap_zero_allocates_exact_size(); + test_delivery_wait_honours_deadline(); + if (g_failures) { + std::cerr << g_failures << " failure(s)\n"; + return 1; + } + std::cout << "ok\n"; + return 0; +} diff --git a/tests/native/test_ring_meta_fifo_ownership.cpp b/tests/native/test_ring_meta_fifo_ownership.cpp new file mode 100644 index 000000000..d4e153571 --- /dev/null +++ b/tests/native/test_ring_meta_fifo_ownership.cpp @@ -0,0 +1,134 @@ +// TensorMetaFifo owns the StepContexts it holds. +// +// Metadata can be published for a forward that then fails before any +// producer publishes a payload; the step's context stays queued until +// teardown. The context carries the step's CaptureLease, so a FIFO that +// dropped its raw pointers without deleting them kept the lease non-idle +// forever and an integration could never retire that configuration. +// +// Covered: destruction and clear() release every pending context exactly +// once -- including multi-hook steps, whose metas share ONE context -- and a +// context popped by the consumer is the consumer's to delete (the FIFO no +// longer frees it). +// +// Standard library only (no ATen, no CUDA). Built and run by +// tests/test_native_ring_consumer_fixes.py. + +#include +#include + +#include "ring/tensor_meta.h" + +namespace { + +int g_failures = 0; + +void check(bool ok, const char* what) { + if (!ok) { + std::cerr << "FAIL: " << what << '\n'; + ++g_failures; + } +} + +ring_py::StepContext* make_ctx(const ring_py::CaptureLease& lease) { + auto* ctx = new ring_py::StepContext; + ctx->capture_lease = lease; + return ctx; +} + +std::vector metas(size_t n) { + std::vector out(n); + if (n) out.back().last_in_step = true; + return out; +} + +} // namespace + +int main() { + // Destroying the FIFO with a pending single-hook step frees its context. + { + ring_py::CaptureLease lease(17); + check(lease.idle(), "a fresh lease is idle"); + { + ring_py::TensorMetaFifo fifo; + auto m = metas(1); + fifo.push_step(make_ctx(lease), m); + check(!lease.idle(), "a queued context holds the lease"); + } + check(lease.idle(), "destroying the FIFO releases the pending lease"); + } + + // A multi-hook step pushes ONE context for all of its metas; several + // pending steps, each released once (no double free under ASan/glibc). + { + ring_py::CaptureLease a(1), b(2); + ring_py::TensorMetaFifo fifo; + auto ma = metas(5); + auto mb = metas(3); + fifo.push_step(make_ctx(a), ma); + fifo.push_step(make_ctx(b), mb); + check(fifo.pending_contexts() == 2, "one context per step, not per meta"); + check(!a.idle() && !b.idle(), "pending steps hold their leases"); + fifo.clear(); + check(fifo.pending_contexts() == 0, "clear() empties the context queue"); + ring_py::TensorMeta m; + check(!fifo.pop(m), "clear() empties the meta queue"); + check(a.idle() && b.idle(), "clear() releases every pending lease"); + } + + // The pop path: a popped context belongs to the consumer. The FIFO must + // not free it again on destruction, and the consumer's delete releases + // the lease (the P2P thread deletes it after the step's last hook). + { + ring_py::CaptureLease lease(3); + ring_py::StepContext* popped = nullptr; + { + ring_py::TensorMetaFifo fifo; + auto m = metas(2); + fifo.push_step(make_ctx(lease), m); + ring_py::TensorMeta out; + check(fifo.pop(out) && fifo.pop(out) && out.last_in_step, + "both metas pop in order"); + popped = fifo.pop_context(); + check(popped != nullptr, "the context pops"); + check(fifo.pop_context() == nullptr, "only one context per step"); + } + check(!lease.idle(), "the consumer still owns the popped context"); + delete popped; + check(lease.idle(), "the consumer's delete releases the lease"); + } + + // A context with a payload-less step still pending behind a popped one: + // only the pending one is freed by the FIFO. + { + ring_py::CaptureLease first(4), second(5); + ring_py::StepContext* popped = nullptr; + { + ring_py::TensorMetaFifo fifo; + auto m1 = metas(1); + auto m2 = metas(1); + fifo.push_step(make_ctx(first), m1); + fifo.push_step(make_ctx(second), m2); + popped = fifo.pop_context(); + } + check(second.idle(), "the still-pending context is freed with the FIFO"); + check(!first.idle(), "the popped context is not freed by the FIFO"); + delete popped; + check(first.idle(), "and is released by its owner"); + } + + // A null context (defensive) is accepted and not queued. + { + ring_py::TensorMetaFifo fifo; + auto m = metas(1); + fifo.push_step(nullptr, m); + check(fifo.pending_contexts() == 0, "a null context is not queued"); + } + + if (g_failures) { + std::cerr << g_failures << " failure(s)\n"; + return 1; + } + std::cout << "ok\n"; + return 0; +} diff --git a/tests/test_native_ring_consumer_fixes.py b/tests/test_native_ring_consumer_fixes.py new file mode 100644 index 000000000..74d15efb7 --- /dev/null +++ b/tests/test_native_ring_consumer_fixes.py @@ -0,0 +1,174 @@ +"""Legacy-ring consumer handoff: CPU-only native coverage. + +* tests/native/test_ring_meta_fifo_ownership.cpp -- TensorMetaFifo owns its + pending StepContexts: clearing or destroying it releases their capture + leases exactly once (standard library only). +* tests/native/test_p2p_batch_handoff.cpp -- the real p2p_thread.cpp with a + batch-only consumer delivers HF batched token/attention captures as + one-request items; queues charge the pooled allocation a payload keeps + alive; p2p_pool_cap_bytes == 0 allocates exact-size buffers; the delivery + wait honours a deadline. Needs ATen only: the CUDA types come from the + CPU stand-in tests/native/ring/cpu_cuda_stub/ (no CUDA API is called and + the drain is never started), used even where a CUDA toolkit is installed + so every machine builds what the CUDA-free CI runner builds. +* tests/native/test_buffer_pool_failure.cpp -- BufferPool::acquire with each + of its allocations failed in turn: std::bad_alloc propagates and the block + returns to the pool exactly once. +* tests/native/test_batch_queue_shared_eviction.cpp -- BatchQueue eviction + with views of a shared allocation frees only what the remaining references + release, and evicts nothing unless the arrival is then admitted. +* The delivery barrier's finite timeout covers producer-stream completion, + the flush and the delivery wait with one deadline (source check: the + CUDA path itself needs a GPU). +""" + +from __future__ import annotations + +import os +import re +import shutil +import subprocess +from pathlib import Path + +import pytest + +ROOT = Path(__file__).resolve().parents[1] +CSRC = ROOT / "native" / "csrc" +CPU_CUDA_STUB = ROOT / "tests" / "native" / "ring" / "cpu_cuda_stub" + + +def _compiler(): + compiler = shutil.which("g++") or shutil.which("c++") + if compiler is None: + pytest.skip("a C++17 compiler is required") + return compiler + + +def _run(executable: Path) -> None: + env = {**os.environ, "CUDA_VISIBLE_DEVICES": ""} + result = subprocess.run([str(executable)], capture_output=True, text=True, + check=False, env=env, timeout=120) + assert result.returncode == 0, result.stdout + result.stderr + assert result.stdout.strip().endswith("ok") + + +@pytest.mark.cpu +def test_meta_fifo_owns_pending_step_contexts(tmp_path): + executable = tmp_path / "test_ring_meta_fifo_ownership" + result = subprocess.run( + [_compiler(), "-std=c++17", "-O0", "-g", "-pthread", f"-I{CSRC}", + str(ROOT / "tests" / "native" / "test_ring_meta_fifo_ownership.cpp"), + "-o", str(executable)], + capture_output=True, text=True, check=False) + assert result.returncode == 0, result.stdout + result.stderr + _run(executable) + + +def _build_with_aten(tmp_path: Path, name: str, *sources: Path) -> Path: + """Compile ``sources`` against CPU-only ATen and the CPU CUDA stub (never + a CUDA toolkit's headers), C++20 as the Torch headers require.""" + torch = pytest.importorskip("torch") + torch_root = Path(torch.__file__).resolve().parent + torch_lib = torch_root / "lib" + abi = int(torch.compiled_with_cxx11_abi()) + executable = tmp_path / name + result = subprocess.run( + [_compiler(), "-std=c++20", "-O0", "-w", "-pthread", + f"-D_GLIBCXX_USE_CXX11_ABI={abi}", + f"-I{CSRC}", f"-I{CPU_CUDA_STUB}", + f"-I{torch_root / 'include'}", + f"-I{torch_root / 'include' / 'torch' / 'csrc' / 'api' / 'include'}", + *[str(src) for src in sources], + "-ffunction-sections", "-fdata-sections", "-Wl,--gc-sections", + # RPATH, not RUNPATH: the interpreter's torch libraries win over any + # other torch on LD_LIBRARY_PATH. + f"-L{torch_lib}", f"-Wl,-rpath,{torch_lib}", "-Wl,--disable-new-dtags", + "-ltorch_cpu", "-lc10", "-o", str(executable)], + capture_output=True, text=True, check=False) + assert result.returncode == 0, result.stdout[-4000:] + result.stderr[-4000:] + return executable + + +@pytest.mark.cpu +def test_p2p_batch_handoff_and_pooled_retention(tmp_path): + _run(_build_with_aten( + tmp_path, "test_p2p_batch_handoff", + ROOT / "tests" / "native" / "test_p2p_batch_handoff.cpp", + CSRC / "ring" / "p2p_thread.cpp")) + + +@pytest.mark.cpu +def test_buffer_pool_acquire_failure_releases_the_block_once(tmp_path): + _run(_build_with_aten( + tmp_path, "test_buffer_pool_failure", + ROOT / "tests" / "native" / "test_buffer_pool_failure.cpp")) + + +@pytest.mark.cpu +def test_batch_queue_eviction_with_shared_allocations(tmp_path): + _run(_build_with_aten( + tmp_path, "test_batch_queue_shared_eviction", + ROOT / "tests" / "native" / "test_batch_queue_shared_eviction.cpp")) + + +@pytest.mark.cpu +def test_cpu_cuda_stub_stands_in_for_the_toolkit(): + """The host-only builds above must not depend on a CUDA toolkit: none of + their sources may reach a CUDA header the stub does not provide.""" + provided = {p.relative_to(CPU_CUDA_STUB).as_posix() + for p in CPU_CUDA_STUB.rglob("*.h")} + assert "cuda_runtime.h" in provided + seen, pending = set(), [ + ROOT / "tests" / "native" / "test_p2p_batch_handoff.cpp", + ROOT / "tests" / "native" / "test_buffer_pool_failure.cpp", + ROOT / "tests" / "native" / "test_batch_queue_shared_eviction.cpp", + CSRC / "ring" / "p2p_thread.cpp", + ] + include = re.compile(r'^\s*#\s*include\s*[<"]([^>"]+)[>"]', re.M) + while pending: + path = pending.pop() + if path in seen: + continue + seen.add(path) + for name in include.findall(path.read_text()): + if name.startswith(("cuda", "cub/", "thrust/", "ATen/cuda/", + "c10/cuda/")) or name.endswith(".cuh"): + assert name in provided, f"{path.name} includes {name}" + for base in (path.parent, CSRC): + if (base / name).is_file(): + pending.append(base / name) + break + + +def _function_body(source: str, signature: str) -> str: + start = source.index(signature) + brace = source.index("{", start) + depth = 0 + for i in range(brace, len(source)): + depth += {"{": 1, "}": -1}.get(source[i], 0) + if depth == 0: + return source[brace:i + 1] + raise AssertionError(f"unterminated body for {signature}") + + +@pytest.mark.cpu +def test_delivery_barrier_timeout_bounds_every_phase(): + py = (CSRC / "ring" / "ring_engine_py.cu").read_text() + body = _function_body(py, "bool RingEnginePy::drain_to_sink_and_wait(") + finite = body[body.index("const auto deadline"):] + # A finite timeout never synchronizes the stream unboundedly: it polls + # an event against the deadline, then hands the same deadline on. + assert "cudaStreamSynchronize" not in finite + assert "wait_stream_until(ms, deadline)" in finite + assert "drain_to_sink_and_wait_until(deadline)" in finite + poll = _function_body(py, "bool wait_stream_until(") + assert "cudaEventDisableTiming" in poll and "cudaEventQuery" in poll + assert "cudaEventDestroy" in poll and "deadline" in poll + + engine = (CSRC / "ring" / "ring_engine.cu").read_text() + until = _function_body(engine, "bool RingEngine::drain_to_sink_and_wait_until(") + bounded = until[until.index("force_flush_and_wait_until"):] + assert re.search(r"if \(!drain_->force_flush_and_wait_until\(\*deadline\)\)\s*return false;", + until) + assert "wait_until_processed_until(drain_->tasks_enqueued(), *deadline)" in bounded + assert "force_flush_and_wait()" not in bounded diff --git a/tests/test_ring_consumer_queues.py b/tests/test_ring_consumer_queues.py new file mode 100644 index 000000000..0c3de9652 --- /dev/null +++ b/tests/test_ring_consumer_queues.py @@ -0,0 +1,240 @@ +"""In-process consumer queues of the legacy ring: SliceQueue and BatchQueue. + +Drives the queues directly through their push bindings (no ring, no GPU): +retained-byte backpressure, rank-ordered eviction, drop attribution through +drop notices, capture-lease identity and lifetime, and the drain/wake +protocol. tests/test_ring_consumers_gpu.py feeds them from a real ring. +""" +from __future__ import annotations + +import threading +import time + +import pytest +import torch + +from tests._requirements import require_native_backend + +pytestmark = [pytest.mark.native_backend, require_native_backend()] + + +@pytest.fixture +def ne(): + from dmi.transport import native + return native + + +def _push_batch(queue, act, layer, nbytes, req_ids, step=0, lease=None): + tensor = torch.zeros(nbytes, dtype=torch.uint8) + segments = torch.zeros((len(req_ids), 4), dtype=torch.int32) + kwargs = {} if lease is None else {"capture_lease": lease, "capture_ts_us": 7} + queue.push("m", 0, act, layer, step, True, tensor, list(req_ids), segments, **kwargs) + + +def _notices(queue): + drained = queue.drain_drop_notices() + return [(n[0], n[1], n[2], tuple(n[8]), n[9]) for n in drained["notices"]], drained + + +def test_batch_queue_is_fifo_and_sheds_whole_batches_with_attribution(ne): + queue = ne.BatchQueue(max_bytes=300) + _push_batch(queue, "blocks.hook_resid_pre", 0, 100, ["a", "b"], step=1) + _push_batch(queue, "blocks.hook_resid_pre", 1, 100, ["a"], step=1) + _push_batch(queue, "blocks.hook_resid_pre", 2, 150, ["c"], step=1) # over the cap + assert queue.qsize == 2 and queue.retained_bytes == 200 + assert queue.dropped == 1 and queue.dropped_bytes == 150 + notices, drained = _notices(queue) + assert notices == [("batch_queue", "blocks.hook_resid_pre", 2, ("c",), 2)] + assert drained["sticky"] is False and drained["overflow"] == 0 + batches = queue.drain_batches(10, 0) + assert [(b[3], b[5], tuple(b[8])) for b in batches] == [(0, 1, ("a", "b")), (1, 1, ("a",))] + assert queue.retained_bytes == 0 and queue.pushed == 2 + # No notice means nothing was lost. + assert _notices(queue)[0] == [] + + +def test_an_empty_queue_admits_one_item_larger_than_the_cap(ne): + queue = ne.BatchQueue(max_bytes=10) + _push_batch(queue, "x", 0, 100, ["a"]) + assert queue.qsize == 1 and queue.dropped == 0 + + +def test_rank_ordered_eviction_sheds_less_urgent_streams_first(ne): + queue = ne.BatchQueue(max_bytes=300) + queue.set_layer_class_map([("urgent", 0, 0), ("boundary", 0, 1)], 2) + assert queue.class_map_set and queue.class_map_size == 2 + assert queue.rank_for("urgent", 0) == 0 + assert queue.rank_for("boundary", 0) == 1 + assert queue.rank_for("unnamed", 3) == 2 + _push_batch(queue, "unnamed", 3, 100, ["a"]) # rank 2 + _push_batch(queue, "boundary", 0, 100, ["b"]) # rank 1 + _push_batch(queue, "unnamed", 3, 100, ["c"]) # rank 2 + _push_batch(queue, "urgent", 0, 150, ["d"]) # rank 0: evicts two rank-2 + assert queue.evicted == 2 and queue.dropped == 2 + assert queue.dropped_by_rank == [0, 0, 2] + assert queue.qsize_by_rank == [1, 1, 0] + notices, _ = _notices(queue) + assert [(n[0], n[3]) for n in notices] == [("batch_queue_evict", ("a",)), + ("batch_queue_evict", ("c",))] + # A rank-0 arrival never evicts another rank-0 item, and evicting the + # rank-1 batch would not make room (150 + 200 > 300): the newcomer alone + # is refused and the rank-1 batch survives. + _push_batch(queue, "urgent", 0, 200, ["e"]) + assert queue.evicted == 2 and queue.qsize_by_rank == [1, 1, 0] + notices, _ = _notices(queue) + assert notices == [("batch_queue", "urgent", 0, ("e",), 0)] + # Evicting it does make room for a smaller one. + _push_batch(queue, "urgent", 0, 140, ["f"]) + assert queue.evicted == 3 and queue.qsize_by_rank == [2, 0, 0] + notices, _ = _notices(queue) + assert [(n[0], n[3]) for n in notices] == [("batch_queue_evict", ("b",))] + assert sum(queue.dropped_by_rank) == queue.dropped + + +def test_eviction_that_cannot_admit_the_newcomer_evicts_nothing(ne): + queue = ne.BatchQueue(max_bytes=100) + queue.set_layer_class_map([("urgent", 0, 0), ("boundary", 0, 1)], 2) + _push_batch(queue, "urgent", 0, 80, ["a"]) # rank 0 + _push_batch(queue, "other", 0, 15, ["b"]) # rank 2 + _push_batch(queue, "boundary", 0, 50, ["c"]) # rank 1: 80 + 50 > 100 anyway + assert queue.evicted == 0 and queue.dropped == 1 + assert queue.dropped_by_rank == [0, 1, 0] + assert queue.qsize_by_rank == [1, 0, 1] and queue.retained_bytes == 95 + notices, _ = _notices(queue) + assert notices == [("batch_queue", "boundary", 0, ("c",), 1)] + # Draining the queue also clears the per-rank byte accounting. + queue.drain_batches(10, 0) + _push_batch(queue, "other", 0, 60, ["d"]) # rank 2 + _push_batch(queue, "boundary", 0, 50, ["e"]) # rank 1 evicts d: fits + assert queue.evicted == 1 and queue.qsize_by_rank == [0, 1, 0] + + +def test_class_eviction_off_restores_drop_the_newcomer(ne): + queue = ne.BatchQueue(max_bytes=200) + queue.set_layer_class_map([("urgent", 0, 0)], 2) + queue.set_class_eviction(False) + assert queue.class_eviction is False + _push_batch(queue, "other", 0, 150, ["a"]) + _push_batch(queue, "urgent", 0, 100, ["b"]) + assert queue.evicted == 0 and queue.dropped == 1 + assert queue.dropped_by_rank == [1, 0, 0] + + +def test_drop_notice_buffer_overflow_is_sticky_for_one_window(ne): + queue = ne.BatchQueue(max_bytes=1, drop_notice_cap=2) + _push_batch(queue, "x", 0, 10, ["keep"]) + for i in range(5): + _push_batch(queue, "x", 0, 10, [f"r{i}"]) + assert queue.drop_notices_recorded == 5 + drained = queue.drain_drop_notices() + assert len(drained["notices"]) == 2 + assert drained["overflow"] == 3 and drained["sticky"] is True + assert queue.drop_notices_lost == 3 + again = queue.drain_drop_notices() + assert again["sticky"] is False and again["notices"] == [] + + +def test_slice_queue_caps_items_and_attributes_each_dropped_slice(ne): + queue = ne.SliceQueue(max_items=2) + lease = ne.CaptureLease(42) + for i in range(3): + queue.push("m", 0, f"r{i}", "blocks.hook_resid_pre", 5, i, i + 1, + torch.zeros(4), lease, 1000 + i) + assert queue.qsize == 2 and queue.dropped == 1 + notices = queue.drain_drop_notices()["notices"] + assert len(notices) == 1 + source, hook, layer, _, start, end, nbytes, _, rids, rank, kept = notices[0] + assert (source, hook, layer, start, end, rids, rank) == ( + "slice_queue", "blocks.hook_resid_pre", 5, 2, 3, ["r2"], 255) + assert kept.same_as(lease) and nbytes == 16 + items = queue.drain_slices(10, 0) + assert [i[2] for i in items] == ["r0", "r1"] + assert items[0][9].id == 42 and items[0][10] == 1000 + + +def test_capture_lease_identity_ranks_and_lifetime(ne): + policy = ne.CaptureLease(1, [("a", 0, 0), ("b", 1, 7)]) + step = ne.CaptureLease(2, policy) + assert step.id == 2 and step.present and not ne.CaptureLease().present + assert step.rank_for("a", 0, 9) == 0 + assert step.rank_for("b", 1, 9) == 2 # clamped + assert step.rank_for("c", 3, 9) == 2 # unnamed stream + assert ne.CaptureLease(3).rank_for("a", 0, 1) == 1 # no map: fallback + assert step.idle + queue = ne.BatchQueue() + _push_batch(queue, "a", 0, 8, ["r"], lease=step) + assert not step.idle # the queued item holds a copy + (item,) = queue.drain_batches(1, 0) + assert item[10].same_as(step) and item[11] == 7 + del item + assert step.idle + + +def test_lease_ranks_order_eviction_per_configuration(ne): + queue = ne.BatchQueue(max_bytes=200) + lease = ne.CaptureLease(5, [("x", 0, 0)]) + _push_batch(queue, "y", 0, 150, ["low"], lease=lease) # rank 2 under the lease + _push_batch(queue, "x", 0, 100, ["high"], lease=lease) # rank 0: evicts it + assert queue.evicted == 1 + (notice,) = queue.drain_drop_notices()["notices"] + assert notice[8] == ["low"] and notice[10].same_as(lease) + + +def test_drain_waits_without_the_gil_and_wake_ends_a_wait(ne): + queue = ne.BatchQueue() + started = time.monotonic() + threading.Timer(0.05, queue.wake).start() + assert queue.drain_batches(4, 5000) == [] + assert time.monotonic() - started < 2.0 + threading.Timer(0.05, lambda: _push_batch(queue, "x", 0, 8, ["r"])).start() + assert len(queue.drain_batches(4, 5000)) == 1 + queue.stop() + assert queue.drain_batches(4, 5000) == [] + + +def test_ring_config_exposes_the_handoff_settings(ne): + cfg = ne.RingConfig() + assert cfg.batch_submit is False + assert cfg.p2p_pool_cap_bytes == 256 * 1024 * 1024 + with pytest.raises(ValueError, match="batch_submit"): + ne.RingEngine(cfg, ne.BatchQueue()) + + +def test_queues_charge_each_backing_allocation_once(ne): + # Views of one buffer (the per-request items of one step) keep ONE + # allocation alive: it is charged once, at its full size, until the last + # of them leaves the queue. + whole = torch.zeros(400, dtype=torch.uint8) + queue = ne.BatchQueue(max_bytes=1000) + for i in range(4): + segments = torch.zeros((1, 4), dtype=torch.int32) + queue.push("m", 0, "x", 0, 1, i == 3, whole[i * 100:(i + 1) * 100], ["r"], segments) + assert queue.qsize == 4 and queue.retained_bytes == 400 and queue.dropped == 0 + queue.drain_batches(3, 0) + assert queue.retained_bytes == 400 + queue.drain_batches(1, 0) + assert queue.retained_bytes == 0 + # A small view of a large allocation is charged the allocation. + big = torch.zeros(800, dtype=torch.uint8) + queue.push("m", 0, "x", 0, 2, True, big[:8], ["r"], torch.zeros((1, 4), dtype=torch.int32)) + assert queue.retained_bytes == 800 + _push_batch(queue, "x", 1, 300, ["s"], step=2) # 800 + 300 > 1000 + assert queue.dropped == 1 and queue.dropped_bytes == 300 + + +def test_slice_queue_reports_retention_and_honours_an_optional_byte_cap(ne): + unbounded = ne.SliceQueue() + assert unbounded.max_bytes == 0 + whole = torch.zeros(256, dtype=torch.uint8) + for i in range(4): + unbounded.push("m", 0, f"r{i}", "x", 0, 0, 1, whole[i * 64:(i + 1) * 64]) + assert unbounded.qsize == 4 and unbounded.retained_bytes == 256 + + capped = ne.SliceQueue(max_items=100, drop_notice_cap=16, max_bytes=300) + capped.push("m", 0, "a", "x", 0, 0, 1, torch.zeros(200, dtype=torch.uint8)) + capped.push("m", 0, "b", "x", 0, 0, 1, torch.zeros(200, dtype=torch.uint8)) # over + assert capped.qsize == 1 and capped.dropped == 1 and capped.retained_bytes == 200 + (notice,) = capped.drain_drop_notices()["notices"] + assert notice[8] == ["b"] + capped.drain_slices(10, 0) + assert capped.retained_bytes == 0 diff --git a/tests/test_ring_consumers_gpu.py b/tests/test_ring_consumers_gpu.py new file mode 100644 index 000000000..6d77f9e1a --- /dev/null +++ b/tests/test_ring_consumers_gpu.py @@ -0,0 +1,310 @@ +"""In-process consumers of a real legacy ring. + +Producers write through the ring; a SliceQueue, a BatchQueue or a Python +callable receives the result. Checks the batch handoff's segments table, +capture leases and timestamps, the delivery barrier through the consumer, +consumer-failure accounting, the nonblocking flush hint, and the +MonitoringEngine consumer wiring. + +Needs CUDA and the full native backend. Memory use is a few MiB. +""" +from __future__ import annotations + +import time +import warnings + +import pytest +import torch + +from tests._requirements import require_cuda, require_native_backend + +pytestmark = [ + pytest.mark.gpu, + pytest.mark.native_backend, + require_cuda(), + require_native_backend(), +] + +HID = 8 + + +def _config(ne, *, batch=False, timeout_us=200, poll_us=200): + cfg = ne.RingConfig() + cfg.payload_ring_bytes = 4 << 20 + cfg.pinned_staging_bytes = 4 << 20 + cfg.drain_poll_timeout_us = poll_us + cfg.drain_flush_timeout_us = timeout_us + cfg.batch_submit = batch + return cfg + + +class Rig: + """Ring engine + transport + one hook type over three requests.""" + + def __init__(self, consumer, *, batch=False, timeout_us=200, specs=None): + from dmi.hooks.specs import (HOOK_TYPE_RESID_PRE, HookSpec, + ModelShapeConfig) + from dmi.transport import native as ne + from dmi.transport import ring as rt + + self.ne, self.rt = ne, rt + self.ht = HOOK_TYPE_RESID_PRE + self.engine = ne.RingEngine(_config(ne, batch=batch, timeout_us=timeout_us), + consumer) + self.engine.init(0) + self.engine.start() + self.transport = rt.RingTransport(self.engine) + rt.activate(self.transport) + self.transport.set_model_cfg(ModelShapeConfig( + hidden_dim=HID, num_heads=2, num_kv_heads=2, head_dim=HID // 2, + dtype=torch.float32, vocab_size=4)) + self.transport._active_specs = specs or [ + HookSpec(hook_type=self.ht, module=None, layer_no=layer) + for layer in range(2)] + self.payload = self.engine.payload_tensor() + self.x = torch.arange(6 * HID, dtype=torch.float32, + device="cuda").reshape(6, HID) + + def step(self, lease=None): + # Requests a (2 rows), b (3 rows), c (1 row) packed into 6 rows. + self.transport.set_step_context( + model_id="m", req_ids=["a", "b", "c"], + token_ranges=[(5, 7), (0, 3), (9, 10)], dim0_offsets=[0, 2, 5], + flattened=True, capture_lease=lease) + self.transport.pre_push_all_metas(batch=0, q_len=6, kv_dim=0) + for spec in self.transport._active_specs: + torch.ops.ring.producer(self.payload, self.x, spec.hook_type, spec.layer_no) + + def close(self): + torch.cuda.synchronize() + try: + self.engine.flush_and_wait() + finally: + self.rt.deactivate() + self.engine.stop() + + +def _wait_for(predicate, timeout=5.0): + deadline = time.monotonic() + timeout + while not predicate() and time.monotonic() < deadline: + time.sleep(0.005) + return predicate() + + +def test_slice_queue_delivers_slices_with_lease_and_capture_time(): + from dmi.transport import native as ne + + queue = ne.SliceQueue() + rig = Rig(queue) + try: + lease = ne.CaptureLease(11) + before = time.monotonic_ns() // 1000 + rig.step(lease) + torch.cuda.synchronize() + rig.engine.flush_and_wait() + items = [] + assert _wait_for(lambda: items.extend(queue.drain_slices(64, 10)) or len(items) >= 6) + got = {(i[4], i[2]): i for i in items} + assert sorted(got) == [(layer, rid) for layer in (0, 1) for rid in "abc"] + torch.testing.assert_close(got[(0, "b")][7], rig.x[2:5].cpu()) + assert got[(1, "c")][5:7] == (9, 10) + for item in items: + assert item[9].same_as(lease) + assert before <= item[10] <= time.monotonic_ns() // 1000 + 1 + assert not lease.idle + del got, item + items.clear() + assert _wait_for(lambda: lease.idle) + finally: + rig.close() + + +def test_batch_queue_hands_each_hook_over_once_with_a_segments_table(): + from dmi.transport import native as ne + + queue = ne.BatchQueue() + rig = Rig(queue, batch=True) + try: + rig.step(ne.CaptureLease(3)) + rig.step() + torch.cuda.synchronize() + rig.engine.flush_and_wait() + batches = [] + assert _wait_for(lambda: batches.extend(queue.drain_batches(64, 10)) or len(batches) >= 4) + assert queue.pushed == 4 + model, shard, act, layer, node, step_seq, last, tensor, rids, segments = batches[0][:10] + assert (act, layer, node, last) == ("blocks.hook_resid_pre", 0, 0, False) + assert batches[1][6] is True # last hook of the step + assert rids == ["a", "b", "c"] + assert segments.tolist() == [[0, 5, 7, -1], [2, 0, 3, -1], [5, 9, 10, -1]] + torch.testing.assert_close(tensor, rig.x.cpu()) + assert batches[0][10].id == 3 and batches[0][11] > 0 + assert len(batches[2]) == 10 # second step carried no lease + assert batches[2][5] == batches[0][5] + 1 # one step_seq per step + finally: + rig.close() + + +def test_batch_handoff_indexes_final_logits_by_original_position(): + from dmi.hooks.specs import HOOK_TYPE_FINAL_LOGITS, HookSpec + from dmi.transport import native as ne + + queue = ne.BatchQueue() + rig = Rig(queue, batch=True, specs=[HookSpec(hook_type=HOOK_TYPE_FINAL_LOGITS, + module=None, layer_no=-1)]) + try: + logits = torch.arange(4 * 4, dtype=torch.float32, device="cuda").reshape(4, 4) + rig.transport.set_step_context(model_id="m", req_ids=["q1", "q3"], + token_ranges=[(0, 5), (0, 2)], + dim0_offsets=[3, 9], flattened=True, + orig_positions=[1, 3]) + rig.transport.pre_push_all_metas(batch=0, q_len=11, kv_dim=0, + logits_to_keep=4) + torch.ops.ring.producer(rig.payload, logits, HOOK_TYPE_FINAL_LOGITS, -1) + torch.cuda.synchronize() + rig.engine.flush_and_wait() + batches = [] + assert _wait_for(lambda: batches.extend(queue.drain_batches(8, 10)) or batches) + segments = batches[0][9] + assert segments.tolist() == [[1, 4, 5, 1], [3, 1, 2, 3]] + finally: + rig.close() + + +def test_drain_to_sink_and_wait_returns_after_the_consumer_has_everything(): + delivered = [] + + def slow(model_id, shard_rank, req_id, act_name, layer_no, start, end, tensor): + time.sleep(0.02) + delivered.append((layer_no, req_id)) + + rig = Rig(slow) + try: + rig.step() + rig.step() + assert rig.engine.drain_to_sink_and_wait(0) is True + assert len(delivered) == 12 # 2 steps x 2 hooks x 3 requests + rig.step() + # A barrier that cannot finish in time reports it. + assert rig.engine.drain_to_sink_and_wait(1) is False + assert rig.engine.drain_to_sink_and_wait(0) is True + assert len(delivered) == 18 + finally: + rig.close() + + +def test_consumer_failures_are_counted_and_can_fail_the_next_step(): + calls = [] + + def failing(*args): + calls.append(args[2]) + raise ValueError("consumer rejected " + args[2]) + + rig = Rig(failing) + try: + rig.step() + assert rig.engine.drain_to_sink_and_wait(0) is True + assert rig.engine.submit_exceptions == 6 + assert "consumer rejected" in rig.engine.last_sink_error + assert rig.engine.sink_failed is False # not armed + assert rig.engine.prepare_step(256, 1) in (0, 1) + rig.engine.flush_and_wait() + rig.engine.set_abort_on_sink_error(True) + assert rig.engine.sink_failed is True + with pytest.raises(RuntimeError, match="consumer failed"): + rig.engine.prepare_step(256, 1) + finally: + rig.engine.set_abort_on_sink_error(False) + rig.close() + + +def test_request_flush_expedites_reserved_tasks(): + delivered = [] + rig = Rig(lambda *args: delivered.append(args[2]), timeout_us=60_000_000) + try: + assert rig.engine.request_flush() is False # nothing reserved + assert rig.engine.prepare_step(2 * 6 * HID * 4, 2) == 0 + rig.step() + torch.cuda.synchronize() + time.sleep(0.05) + assert delivered == [] # thresholds not met + assert rig.engine.request_flush() is True + assert _wait_for(lambda: len(delivered) == 6) + finally: + rig.close() + + +def test_monitoring_engine_delivers_to_a_ring_consumer_and_warns_without_one(): + from dmi.engine import MonitoringEngine, RingWithoutConsumerWarning + from dmi.transport import native as ne + + delivered = [] + cfg = _config(ne) + engine = MonitoringEngine(model_id="m", ring_config=cfg, + ring_consumer=lambda *args: delivered.append(args[2])) + with warnings.catch_warnings(): + warnings.simplefilter("error", RingWithoutConsumerWarning) + transport = engine._ensure_ring_transport() + try: + from dmi.hooks.specs import HOOK_TYPE_RESID_PRE, HookSpec, ModelShapeConfig + transport.set_model_cfg(ModelShapeConfig(hidden_dim=HID, num_heads=2, + num_kv_heads=2, head_dim=HID // 2, + dtype=torch.float32)) + transport._active_specs = [HookSpec(hook_type=HOOK_TYPE_RESID_PRE, + module=None, layer_no=0)] + transport.set_step_context(model_id="m", req_ids=["a"], token_ranges=[(0, 2)], + dim0_offsets=[0], flattened=True) + transport.pre_push_all_metas(batch=0, q_len=2, kv_dim=0) + torch.ops.ring.producer(transport._ring_payload, + torch.ones(2, HID, device="cuda"), HOOK_TYPE_RESID_PRE, 0) + assert engine._ring_engine.drain_to_sink_and_wait(0) + assert delivered == ["a"] + finally: + engine.close() + + with pytest.raises(ValueError, match="ring_consumer or host_engine"): + MonitoringEngine(model_id="m", ring_config=_config(ne), + ring_consumer=print, host_engine=object()) + quiet = MonitoringEngine(model_id="m", ring_config=_config(ne)) + try: + with pytest.warns(RingWithoutConsumerWarning, match="no consumer"): + quiet._ensure_ring_transport() + finally: + quiet.close() + + +def test_monitoring_engine_switches_a_batch_queue_consumer_to_the_batch_handoff(): + from dmi.engine import MonitoringEngine + from dmi.transport import native as ne + + queue = ne.BatchQueue() + engine = MonitoringEngine(model_id="m", ring_config=_config(ne), + ring_consumer=queue) + try: + engine._ensure_ring_transport() + assert engine._ring_config.batch_submit is True + finally: + engine.close() + + +def test_delivered_tensors_outlive_the_ring_engine(): + """Pooled pageable buffers stay valid after the engine that filled them + stops, and are freed when the consumer drops them.""" + from dmi.transport import native as ne + + queue = ne.BatchQueue() + rig = Rig(queue, batch=True) + batches = [] + try: + rig.step() + torch.cuda.synchronize() + rig.engine.flush_and_wait() + assert _wait_for(lambda: batches.extend(queue.drain_batches(8, 10)) or len(batches) >= 2) + finally: + rig.close() + del rig + import gc + gc.collect() + torch.testing.assert_close(batches[0][7], torch.arange(6 * HID, dtype=torch.float32).reshape(6, HID)) + batches.clear() + gc.collect()