From 453d2b260924c170470e67dc246ba5abe446423e Mon Sep 17 00:00:00 2001 From: zhaoyboo Date: Tue, 3 Nov 2026 17:03:36 -0500 Subject: [PATCH] [Core] Expose ring preparation, graph toggle, gather upload, and staging wait telemetry Make the legacy ring's existing waits and copies observable without adding any: host-only cumulative counters, read with no CUDA work. Native: - RingEngine.get_stats() returns RingFlushStats: reservation heads and drain tails; prepare_step outcomes (ok / flushed / oversized) and the time of its stream synchronizations and forced flushes; node-toggle pre-apply waits; drain waits for staging space; per-task ready-to-commit latency (first seen published to committed after D2H); D2H copy and synchronization failures; staging-to-pageable copies (calls, time, bytes). - HostWaitCounter (host_wait_counter.h) times a call that runs anyway and counts nonzero CUDA statuses and exceptions as errors; DrainTelemetry and PageableCopyTelemetry (pipeline_telemetry.h) hold the stage counters. An uncontended staging check reads no clock. Python: - dmi.sync_stats: synchronization_stats, pipeline_stats and gather_sync_stats (an integration's per-HookPoint row-table uploads, reuses and pre-write waits) return versioned dicts, report unavailable counters instead of zeros, and reject negative ones. A native RingEngine built before get_stats existed reports unavailable (rebuild required) rather than raising AttributeError; errors from an existing get_stats still propagate. - MonitoringEngine logs the ring's device payload, pinned staging and task entries when it builds a legacy ring: memory allocated before a serving framework profiles free device memory reduces its KV cache by the same amount. --- docs/ring-telemetry.md | 30 ++++ native/Makefile | 3 +- native/csrc/bindings.cpp | 32 ++++ native/csrc/ring/drain_thread.cpp | 20 ++- native/csrc/ring/drain_thread.h | 4 + native/csrc/ring/host_wait_counter.h | 57 +++++++ native/csrc/ring/p2p_thread.cpp | 8 +- native/csrc/ring/p2p_thread.h | 4 + native/csrc/ring/pipeline_telemetry.h | 58 +++++++ native/csrc/ring/ring_engine_py.cu | 58 ++++++- native/csrc/ring/ring_engine_py.h | 33 ++++ native/csrc/ring/toggle_registry.cu | 3 +- native/csrc/ring/toggle_registry.h | 5 + src/dmi/engine.py | 23 +++ src/dmi/sync_stats.py | 105 ++++++++++++ tests/test_ring_telemetry.py | 230 ++++++++++++++++++++++++++ 16 files changed, 656 insertions(+), 17 deletions(-) create mode 100644 docs/ring-telemetry.md create mode 100644 native/csrc/ring/host_wait_counter.h create mode 100644 native/csrc/ring/pipeline_telemetry.h create mode 100644 src/dmi/sync_stats.py create mode 100644 tests/test_ring_telemetry.py diff --git a/docs/ring-telemetry.md b/docs/ring-telemetry.md new file mode 100644 index 000000000..c6db1b2d2 --- /dev/null +++ b/docs/ring-telemetry.md @@ -0,0 +1,30 @@ +# Legacy ring telemetry + +`RingEngine.get_stats()` returns a `RingFlushStats` snapshot of host-only, +cumulative counters for waits and copies the legacy ring already performs. +Reading it does no CUDA work and adds no synchronization; each field is +monotonic on its own. + +| fields | what is counted | +|---|---| +| `cpu_payload_head`, `cpu_payload_tail_committed`, `cpu_task_head`, `cpu_task_tail_committed` | reservation heads and drain tails | +| `prepare_ok_calls`, `prepare_flushed_calls`, `prepare_oversized_calls` | `prepare_step` outcomes | +| `prepare_sync_{calls,ns,errors}` | stream synchronizations inside `prepare_step` (flushed and oversized steps) | +| `prepare_drain_{calls,ns,errors}` | forced flushes inside `prepare_step` | +| `toggle_wait_{calls,ns,errors}` | node-toggle pre-apply waits (`set_toggle_pre_apply_sync(True)` only) | +| `staging_wait_{calls,ns,errors}` | drain waits for pinned staging space (an uncontended check reads no clock) | +| `ready_commit_{tasks,us,invalid}` | per task, summed: first seen published to committed after its D2H copy | +| `d2h_copy_errors`, `d2h_sync_errors` | failed D2H copies and synchronizations | +| `pageable_copy_{calls,ns,bytes}` | staging-to-pageable copies in the P2P thread | + +`dmi.sync_stats` reads them as versioned dicts: +`synchronization_stats(engine)` (prepare outcomes and waits, toggle waits), +`pipeline_stats(engine)` (drain and P2P stages) and +`gather_sync_stats(source)` (an integration's per-HookPoint row-table +uploads, reuses and pre-write waits, from `source.row_table_stats()`). A reader returns `{"available": False, ...}` rather than +zeros when a counter is missing, and raises on a negative one. + +Building a legacy ring logs its device payload ring, pinned staging and task +entries at INFO (`dmi.engine`). The payload ring is allocated when the ring +is built, so a serving framework that sizes its KV cache from the free memory +it profiles afterwards loses KV capacity byte for byte. diff --git a/native/Makefile b/native/Makefile index efaa24d77..d456bb715 100644 --- a/native/Makefile +++ b/native/Makefile @@ -31,7 +31,8 @@ HEADERS := csrc/dmx_host_utils.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 + csrc/ring/buffer_pool.h csrc/ring/pooled_block.h \ + csrc/ring/host_wait_counter.h csrc/ring/pipeline_telemetry.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 6e26bfa50..6bddb5a6f 100644 --- a/native/csrc/bindings.cpp +++ b/native/csrc/bindings.cpp @@ -949,6 +949,36 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { }, py::arg("max_items") = 64, py::arg("timeout_ms") = 50); + // ---- RingFlushStats: host-only telemetry snapshot ---------------------- + py::class_(m, "RingFlushStats") + .def_readonly("cpu_payload_head", &ring_py::RingFlushStats::cpu_payload_head) + .def_readonly("cpu_payload_tail_committed", &ring_py::RingFlushStats::cpu_payload_tail_committed) + .def_readonly("cpu_task_head", &ring_py::RingFlushStats::cpu_task_head) + .def_readonly("cpu_task_tail_committed", &ring_py::RingFlushStats::cpu_task_tail_committed) + .def_readonly("prepare_ok_calls", &ring_py::RingFlushStats::prepare_ok_calls) + .def_readonly("prepare_flushed_calls", &ring_py::RingFlushStats::prepare_flushed_calls) + .def_readonly("prepare_oversized_calls", &ring_py::RingFlushStats::prepare_oversized_calls) + .def_readonly("prepare_sync_calls", &ring_py::RingFlushStats::prepare_sync_calls) + .def_readonly("prepare_sync_ns", &ring_py::RingFlushStats::prepare_sync_ns) + .def_readonly("prepare_sync_errors", &ring_py::RingFlushStats::prepare_sync_errors) + .def_readonly("prepare_drain_calls", &ring_py::RingFlushStats::prepare_drain_calls) + .def_readonly("prepare_drain_ns", &ring_py::RingFlushStats::prepare_drain_ns) + .def_readonly("prepare_drain_errors", &ring_py::RingFlushStats::prepare_drain_errors) + .def_readonly("toggle_wait_calls", &ring_py::RingFlushStats::toggle_wait_calls) + .def_readonly("toggle_wait_ns", &ring_py::RingFlushStats::toggle_wait_ns) + .def_readonly("toggle_wait_errors", &ring_py::RingFlushStats::toggle_wait_errors) + .def_readonly("staging_wait_calls", &ring_py::RingFlushStats::staging_wait_calls) + .def_readonly("staging_wait_ns", &ring_py::RingFlushStats::staging_wait_ns) + .def_readonly("staging_wait_errors", &ring_py::RingFlushStats::staging_wait_errors) + .def_readonly("ready_commit_tasks", &ring_py::RingFlushStats::ready_commit_tasks) + .def_readonly("ready_commit_us", &ring_py::RingFlushStats::ready_commit_us) + .def_readonly("ready_commit_invalid", &ring_py::RingFlushStats::ready_commit_invalid) + .def_readonly("d2h_copy_errors", &ring_py::RingFlushStats::d2h_copy_errors) + .def_readonly("d2h_sync_errors", &ring_py::RingFlushStats::d2h_sync_errors) + .def_readonly("pageable_copy_calls", &ring_py::RingFlushStats::pageable_copy_calls) + .def_readonly("pageable_copy_ns", &ring_py::RingFlushStats::pageable_copy_ns) + .def_readonly("pageable_copy_bytes", &ring_py::RingFlushStats::pageable_copy_bytes); + py::class_(m, "RingConfig") .def(py::init<>()) .def_readwrite("task_ring_entries", &ring_py::RingConfig::task_ring_entries) @@ -1256,6 +1286,8 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { .def("request_flush", &ring_py::RingEnginePy::request_flush, py::call_guard()) + .def("get_stats", &ring_py::RingEnginePy::get_stats, + py::call_guard()) .def_property_readonly("submit_exceptions", &ring_py::RingEnginePy::submit_exceptions) .def_property_readonly("last_sink_error", diff --git a/native/csrc/ring/drain_thread.cpp b/native/csrc/ring/drain_thread.cpp index 896635e49..6dff96676 100644 --- a/native/csrc/ring/drain_thread.cpp +++ b/native/csrc/ring/drain_thread.cpp @@ -416,7 +416,8 @@ void DrainThread::do_full_flush() { } { std::unique_lock lk(staging_mu_); - staging_cv_.wait(lk, [&] { return staging_.free_bytes() >= flush_bytes; }); + telemetry_.wait_for_space(staging_cv_, lk, + [&] { return staging_.free_bytes() >= flush_bytes; }); } enqueue_d2h(flush_bytes); sync_stream(); @@ -465,7 +466,8 @@ void DrainThread::flush_allowed() { } { std::unique_lock lk(staging_mu_); - staging_cv_.wait(lk, [&] { return staging_.free_bytes() >= flush_bytes; }); + telemetry_.wait_for_space(staging_cv_, lk, + [&] { return staging_.free_bytes() >= flush_bytes; }); } enqueue_d2h(flush_bytes); sync_stream(); @@ -597,7 +599,7 @@ void DrainThread::loop() { try { { std::unique_lock lk(staging_mu_); - staging_cv_.wait(lk, [&] { + telemetry_.wait_for_space(staging_cv_, lk, [&] { return staging_.free_bytes() >= flush_bytes; }); } @@ -749,8 +751,10 @@ void DrainThread::flush_state_update(uint64_t flush_count, uint64_t flush_bytes) } void DrainThread::sync_stream() { - throw_cuda_failure("cudaStreamSynchronize", - cudaStreamSynchronize(stream_)); + const cudaError_t synced = cudaStreamSynchronize(stream_); + if (synced != cudaSuccess) + telemetry_.d2h_sync_errors.fetch_add(1, std::memory_order_relaxed); + throw_cuda_failure("cudaStreamSynchronize", synced); } // --------------------------------------------------------------------------- @@ -776,6 +780,8 @@ void DrainThread::enqueue_d2h(uint64_t flush_bytes) { cudaError_t err = cudaMemcpyAsync(staging_.base() + stg_cursor, ring_.payload_buf + gpu_cursor, chunk, cudaMemcpyDeviceToHost, stream_); + if (err != cudaSuccess) + telemetry_.d2h_copy_errors.fetch_add(1, std::memory_order_relaxed); throw_cuda_failure("cudaMemcpyAsync", err); RING_DBG("[enqueue_d2h] chunk=%d enqueued OK\n", chunk_idx); @@ -793,6 +799,9 @@ void DrainThread::enqueue_d2h(uint64_t flush_bytes) { void DrainThread::submit_to_p2p(uint64_t flush_count, uint64_t flush_bytes) { uint64_t cumulative = 0; const uint64_t staging_batch_start = staging_.head(); + // Called after the batch's D2H copy completed: the commit instant of + // every task in it. + const int64_t committed_us = steady_us(std::chrono::steady_clock::now()); for (uint64_t i = 0; i < flush_count; ++i) { const uint64_t actual_bytes = scanned_[i]; @@ -803,6 +812,7 @@ void DrainThread::submit_to_p2p(uint64_t flush_count, uint64_t flush_bytes) { task.tensor_total_bytes = actual_bytes; task.alloc_bytes = alloc; task.capture_ts_us = scanned_ready_us_[i]; + telemetry_.record_commit(task.capture_ts_us, committed_us); if (data_len > 0) { uint64_t staging_logical = staging_batch_start + cumulative; diff --git a/native/csrc/ring/drain_thread.h b/native/csrc/ring/drain_thread.h index f71cc23a9..82f7ac882 100644 --- a/native/csrc/ring/drain_thread.h +++ b/native/csrc/ring/drain_thread.h @@ -16,6 +16,7 @@ #include "ring_config.h" #include "pinned_staging.h" #include "drain_task.h" +#include "pipeline_telemetry.h" #include #include @@ -85,6 +86,8 @@ class DrainThread { bool is_running() const { return running_.load(std::memory_order_relaxed); } + const DrainTelemetry& telemetry() const { return telemetry_; } + // 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 { @@ -110,6 +113,7 @@ class DrainThread { void rethrow_drain_failure(); private: + DrainTelemetry telemetry_; RingState& ring_; PinnedStaging& staging_; RingConfig cfg_; diff --git a/native/csrc/ring/host_wait_counter.h b/native/csrc/ring/host_wait_counter.h new file mode 100644 index 000000000..afe5b7d7c --- /dev/null +++ b/native/csrc/ring/host_wait_counter.h @@ -0,0 +1,57 @@ +// ring/host_wait_counter.h -- host elapsed-time counters for waits the ring +// already performs. No CUDA work, no device memory, no allocation, no +// polling: a counter only times a call that would run anyway. Snapshot +// fields are individually monotonic, not one atomic multi-field read while a +// measured call is completing. +#pragma once +#include +#include +#include +#include + +namespace ring_py { + +struct HostWaitSnapshot { + uint64_t calls{0}; + uint64_t ns{0}; + uint64_t errors{0}; +}; + +class HostWaitCounter { +public: + // Time f(). A non-void result is a CUDA status: nonzero counts as an + // error. An exception counts as an error and propagates. + template auto measure(F&& f) -> decltype(f()) { + const auto start = std::chrono::steady_clock::now(); + try { + if constexpr (std::is_void_v) { + f(); + record(start, false); + } else { + auto result = f(); + record(start, result != 0); + return result; + } + } catch (...) { + record(start, true); + throw; + } + } + HostWaitSnapshot snapshot() const { + return {calls_.load(std::memory_order_relaxed), + ns_.load(std::memory_order_relaxed), + errors_.load(std::memory_order_relaxed)}; + } + +private: + void record(std::chrono::steady_clock::time_point start, bool error) { + const auto elapsed = std::chrono::duration_cast( + std::chrono::steady_clock::now() - start).count(); + ns_.fetch_add(static_cast(elapsed), std::memory_order_relaxed); + if (error) errors_.fetch_add(1, std::memory_order_relaxed); + calls_.fetch_add(1, std::memory_order_relaxed); + } + std::atomic calls_{0}, ns_{0}, errors_{0}; +}; + +} // namespace ring_py diff --git a/native/csrc/ring/p2p_thread.cpp b/native/csrc/ring/p2p_thread.cpp index 5150da5f4..cbcf4bafe 100644 --- a/native/csrc/ring/p2p_thread.cpp +++ b/native/csrc/ring/p2p_thread.cpp @@ -244,12 +244,8 @@ void P2PThread::process(std::vector& tasks) { } uint8_t* dst = tensor.data_ptr(); - if (task.data_len1 > 0) { - std::memcpy(dst, task.data_ptr1, task.data_len1); - } - if (task.data_len2 > 0) { - std::memcpy(dst + task.data_len1, task.data_ptr2, task.data_len2); - } + copy_telemetry_.copy(dst, task.data_ptr1, task.data_len1, + task.data_ptr2, task.data_len2); // Release staging space if (task.alloc_bytes > 0) { diff --git a/native/csrc/ring/p2p_thread.h b/native/csrc/ring/p2p_thread.h index a252295c2..07c07ac19 100644 --- a/native/csrc/ring/p2p_thread.h +++ b/native/csrc/ring/p2p_thread.h @@ -9,6 +9,7 @@ #include "drain_thread.h" #include "drain_task.h" #include "ring_config.h" +#include "pipeline_telemetry.h" #include "record_consumer.h" #include "tensor_meta.h" @@ -119,7 +120,10 @@ class P2PThread { bool wait_until_processed_until(uint64_t target, std::chrono::steady_clock::time_point deadline); + const PageableCopyTelemetry& copy_telemetry() const { return copy_telemetry_; } + private: + PageableCopyTelemetry copy_telemetry_; DrainThread& drain_; ring_py::TensorMetaFifo& fifo_; RingConfig cfg_; diff --git a/native/csrc/ring/pipeline_telemetry.h b/native/csrc/ring/pipeline_telemetry.h new file mode 100644 index 000000000..5a4576d01 --- /dev/null +++ b/native/csrc/ring/pipeline_telemetry.h @@ -0,0 +1,58 @@ +// ring/pipeline_telemetry.h -- host-only accounting of the drain and P2P +// stages. No device reads and no added synchronization. +#pragma once +#include "host_wait_counter.h" + +#include +#include +#include +#include +#include + +namespace ring { + +class DrainTelemetry { +public: + // Wait for staging space under the caller's lock. An uncontended check + // reads no clock. + template + void wait_for_space(std::condition_variable& cv, + std::unique_lock& lock, Predicate ready) { + if (!ready()) staging_wait.measure([&] { cv.wait(lock, ready); }); + } + // One task's latency from first seen published to committed to the P2P + // queue (after its D2H copy completed). + void record_commit(int64_t ready_us, int64_t committed_us) { + if (ready_us <= 0 || committed_us < ready_us) { + ready_commit_invalid.fetch_add(1, std::memory_order_relaxed); + return; + } + ready_commit_us.fetch_add(static_cast(committed_us - ready_us), + std::memory_order_relaxed); + ready_commit_tasks.fetch_add(1, std::memory_order_relaxed); + } + ring_py::HostWaitCounter staging_wait; + std::atomic ready_commit_tasks{0}, ready_commit_us{0}, + ready_commit_invalid{0}; + std::atomic d2h_copy_errors{0}, d2h_sync_errors{0}; +}; + +class PageableCopyTelemetry { +public: + // Staging -> pageable copy of one task (one or two spans when the staging + // ring wraps; counted as one operation). Destination allocation and the + // staging release are outside the measurement. + void copy(uint8_t* dst, const uint8_t* a, uint64_t n, + const uint8_t* b, uint64_t m) { + if (n == 0 && m == 0) return; + elapsed.measure([&] { + if (n) std::memcpy(dst, a, n); + if (m) std::memcpy(dst + n, b, m); + }); + bytes.fetch_add(n + m, std::memory_order_relaxed); + } + ring_py::HostWaitCounter elapsed; + std::atomic bytes{0}; +}; + +} // namespace ring diff --git a/native/csrc/ring/ring_engine_py.cu b/native/csrc/ring/ring_engine_py.cu index fe40038b0..375c7a1b5 100644 --- a/native/csrc/ring/ring_engine_py.cu +++ b/native/csrc/ring/ring_engine_py.cu @@ -220,6 +220,10 @@ struct RingEnginePy::Impl { std::atomic gather_host_skips{0}; std::atomic head_launches{0}; + // prepare_step telemetry (get_stats). + std::atomic prepare_ok{0}, prepare_flushed{0}, prepare_oversized{0}; + HostWaitCounter prepare_sync, prepare_drain; + // True when the mirrored row at `table` holds the disabled sentinel. // A pointer outside the mirror is never treated as disabled. bool mirrored_row_disabled(uint64_t table) const { @@ -719,8 +723,9 @@ int RingEnginePy::prepare_step(uint64_t step_total_bytes, // net starts firing. if (step_total_bytes > effective_cap || num_hooks > tcap) { cudaStream_t ms = at::cuda::getCurrentCUDAStream().stream(); - cudaStreamSynchronize(ms); - drain.force_flush_and_wait(); + impl_->prepare_sync.measure([&] { return cudaStreamSynchronize(ms); }); + impl_->prepare_drain.measure([&] { drain.force_flush_and_wait(); }); + impl_->prepare_oversized.fetch_add(1, std::memory_order_relaxed); return STEP_OVERSIZED; } @@ -730,15 +735,17 @@ int RingEnginePy::prepare_step(uint64_t step_total_bytes, if (step_total_bytes <= payload_avail && num_hooks <= task_avail) { if (reserve) drain.reserve(step_total_bytes, num_hooks); + impl_->prepare_ok.fetch_add(1, std::memory_order_relaxed); return STEP_RING_OK; // fast path -- no CUDA or thread interaction } // Either payload or task ring full from prior steps. Sync main // stream so all producer kernels finish writing, then flush. cudaStream_t ms = at::cuda::getCurrentCUDAStream().stream(); - cudaStreamSynchronize(ms); - drain.force_flush_and_wait(); + impl_->prepare_sync.measure([&] { return cudaStreamSynchronize(ms); }); + impl_->prepare_drain.measure([&] { drain.force_flush_and_wait(); }); if (reserve) drain.reserve(step_total_bytes, num_hooks); + impl_->prepare_flushed.fetch_add(1, std::memory_order_relaxed); return STEP_RING_FLUSHED; } @@ -1356,6 +1363,49 @@ bool RingEnginePy::_test_applied_current(uint64_t graph) const { return impl_->toggle.applied_current(graph); } +RingFlushStats RingEnginePy::get_stats() const { + refuse_on_record_ring(impl_->record_mode, "legacy ring telemetry"); + RingFlushStats st{}; + auto& drain = impl_->engine.drain_thread(); + st.cpu_payload_head = drain.cpu_payload_head(); + st.cpu_payload_tail_committed = drain.cpu_payload_tail_committed(); + st.cpu_task_head = drain.cpu_task_head(); + st.cpu_task_tail_committed = drain.cpu_task_tail_committed(); + st.prepare_ok_calls = impl_->prepare_ok.load(std::memory_order_relaxed); + st.prepare_flushed_calls = impl_->prepare_flushed.load(std::memory_order_relaxed); + st.prepare_oversized_calls = impl_->prepare_oversized.load(std::memory_order_relaxed); + const auto sync = impl_->prepare_sync.snapshot(); + st.prepare_sync_calls = sync.calls; + st.prepare_sync_ns = sync.ns; + st.prepare_sync_errors = sync.errors; + const auto flush = impl_->prepare_drain.snapshot(); + st.prepare_drain_calls = flush.calls; + st.prepare_drain_ns = flush.ns; + st.prepare_drain_errors = flush.errors; + const auto toggle = impl_->toggle.replay_wait_stats(); + st.toggle_wait_calls = toggle.calls; + st.toggle_wait_ns = toggle.ns; + st.toggle_wait_errors = toggle.errors; + const auto& pipeline = drain.telemetry(); + const auto staging = pipeline.staging_wait.snapshot(); + st.staging_wait_calls = staging.calls; + st.staging_wait_ns = staging.ns; + st.staging_wait_errors = staging.errors; + st.ready_commit_tasks = pipeline.ready_commit_tasks.load(std::memory_order_relaxed); + st.ready_commit_us = pipeline.ready_commit_us.load(std::memory_order_relaxed); + st.ready_commit_invalid = pipeline.ready_commit_invalid.load(std::memory_order_relaxed); + st.d2h_copy_errors = pipeline.d2h_copy_errors.load(std::memory_order_relaxed); + st.d2h_sync_errors = pipeline.d2h_sync_errors.load(std::memory_order_relaxed); + if (const auto* p2p = impl_->engine.p2p_thread()) { + const auto& copy = p2p->copy_telemetry(); + const auto elapsed = copy.elapsed.snapshot(); + st.pageable_copy_calls = elapsed.calls; + st.pageable_copy_ns = elapsed.ns; + st.pageable_copy_bytes = copy.bytes.load(std::memory_order_relaxed); + } + return st; +} + bool RingEnginePy::request_flush() { return impl_->engine.drain_thread().request_flush(); } diff --git a/native/csrc/ring/ring_engine_py.h b/native/csrc/ring/ring_engine_py.h index bee8c8824..7e0893e5d 100644 --- a/native/csrc/ring/ring_engine_py.h +++ b/native/csrc/ring/ring_engine_py.h @@ -104,6 +104,36 @@ struct RecordCaptureStatus { uint64_t max_step_wait_ns = 0; }; +// Cumulative, host-only counters of the waits and copies the legacy ring +// already performs (get_stats()). Reading them does no CUDA work. Each +// field is monotonic on its own; a snapshot is not one atomic transaction. +// cpu_*_head / *_tail_committed reservation heads and drain tails +// prepare_{ok,flushed,oversized}_calls prepare_step outcomes +// prepare_sync_* stream synchronizations inside prepare_step +// prepare_drain_* forced flushes inside prepare_step +// toggle_wait_* node-toggle pre-apply waits (set_toggle_pre_apply_sync) +// staging_wait_* drain waits for pinned staging space +// ready_commit_* per task: first seen published -> committed after D2H +// d2h_*_errors failed D2H copies / synchronizations +// pageable_copy_* staging -> pageable copies in the P2P thread +// *_ns are nanoseconds and ready_commit_us microseconds of host time. +struct RingFlushStats { + uint64_t cpu_payload_head{0}; + uint64_t cpu_payload_tail_committed{0}; + uint64_t cpu_task_head{0}; + uint64_t cpu_task_tail_committed{0}; + uint64_t prepare_ok_calls{0}; + uint64_t prepare_flushed_calls{0}; + uint64_t prepare_oversized_calls{0}; + uint64_t prepare_sync_calls{0}, prepare_sync_ns{0}, prepare_sync_errors{0}; + uint64_t prepare_drain_calls{0}, prepare_drain_ns{0}, prepare_drain_errors{0}; + uint64_t toggle_wait_calls{0}, toggle_wait_ns{0}, toggle_wait_errors{0}; + uint64_t staging_wait_calls{0}, staging_wait_ns{0}, staging_wait_errors{0}; + uint64_t ready_commit_tasks{0}, ready_commit_us{0}, ready_commit_invalid{0}; + uint64_t d2h_copy_errors{0}, d2h_sync_errors{0}; + uint64_t pageable_copy_calls{0}, pageable_copy_ns{0}, pageable_copy_bytes{0}; +}; + // 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::functionsecond) { - cudaError_t se = cudaEventSynchronize(ev->second); + cudaError_t se = replay_wait_.measure( + [&] { return cudaEventSynchronize(ev->second); }); if (se != cudaSuccess) return static_cast(se); } diff --git a/native/csrc/ring/toggle_registry.h b/native/csrc/ring/toggle_registry.h index 4156477ec..7f2b59851 100644 --- a/native/csrc/ring/toggle_registry.h +++ b/native/csrc/ring/toggle_registry.h @@ -16,6 +16,7 @@ #pragma once #include +#include "host_wait_counter.h" #include "toggle_device.h" #include @@ -72,6 +73,9 @@ class ToggleRegistry { // pre_apply_sync_). Off by default. void set_pre_apply_sync(bool enabled); bool pre_apply_sync() const; + // Host time spent in the optional pre-apply wait (cumulative; clear() + // does not reset it). + HostWaitSnapshot replay_wait_stats() const { return replay_wait_.snapshot(); } void clear(); // -- device-side toggle (toggle_device.h) -------------------------------- @@ -178,6 +182,7 @@ class ToggleRegistry { std::set dev_nodes_; std::set dev_uploaded_; mutable std::mutex mu_; + HostWaitCounter replay_wait_; }; } // namespace ring_py diff --git a/src/dmi/engine.py b/src/dmi/engine.py index b508763ee..fe9fe1699 100644 --- a/src/dmi/engine.py +++ b/src/dmi/engine.py @@ -940,6 +940,7 @@ def enable_ring_transport( ring_engine.init() ring_engine.start() + _log_ring_footprint(ring_engine) transport = _rt.RingTransport(ring_engine) if model_shape is not None: @@ -1137,6 +1138,28 @@ def close(self) -> None: # Backend loader +def _log_ring_footprint(ring_engine: Any) -> None: + """Log the device and pinned host memory a legacy ring holds. + + The payload ring is allocated when the ring is built. A serving + framework that sizes its KV cache from the free device memory it + profiles afterwards loses KV capacity byte for byte, which shows up as + lower throughput with no transport involvement; the log line makes that + visible. + """ + try: + payload = int(ring_engine.payload_cap()) + staging = int(ring_engine.staging_cap()) + tasks = int(ring_engine.task_cap()) + except Exception: # pragma: no cover - partial native builds + return + _LOG.info( + "DMI legacy ring: %.1f MiB device payload ring, %.1f MiB pinned host " + "staging, %d task entries. Device memory allocated before a serving " + "framework profiles free memory reduces its KV cache by the same " + "amount.", payload / 2**20, staging / 2**20, tasks) + + class RingWithoutConsumerWarning(UserWarning): """A legacy ring was built with no consumer: captures are dropped.""" diff --git a/src/dmi/sync_stats.py b/src/dmi/sync_stats.py new file mode 100644 index 000000000..0a46f2e39 --- /dev/null +++ b/src/dmi/sync_stats.py @@ -0,0 +1,105 @@ +"""Host-only telemetry of the legacy ring's existing waits and copies. + +Every reader returns ``{"available": False, "reason": ...}`` instead of +inventing zeros when the native build or the adapter does not expose the +counters, and raises on a negative counter. Reading does no CUDA work: the +native snapshot (``RingEngine.get_stats()``) holds host counters only. +""" +from __future__ import annotations + +from typing import Any, Optional + +# prepare_step outcomes and its host waits, and node-toggle pre-apply waits. +FIELDS = ('prepare_ok_calls', 'prepare_flushed_calls', 'prepare_oversized_calls', + 'prepare_sync_calls', 'prepare_sync_ns', 'prepare_sync_errors', + 'prepare_drain_calls', 'prepare_drain_ns', 'prepare_drain_errors', + 'toggle_wait_calls', 'toggle_wait_ns', 'toggle_wait_errors') + +# Drain and P2P stages. +PIPELINE_FIELDS = ('staging_wait_calls', 'staging_wait_ns', 'staging_wait_errors', + 'ready_commit_tasks', 'ready_commit_us', 'ready_commit_invalid', + 'd2h_copy_errors', 'd2h_sync_errors', + 'pageable_copy_calls', 'pageable_copy_ns', 'pageable_copy_bytes') + +# Per-HookPoint row-table uploads an integration performs before a step +# (the integration's ``row_table_stats()`` mapping). +GATHER_FIELDS = ('hook_table_wait_calls', 'hook_table_wait_ns', + 'hook_table_wait_errors', 'hook_table_uploads', + 'hook_table_reuses') + + +def _native_counters(engine: Any, snapshot: Any, fields: tuple) -> dict: + if engine is None and snapshot is None: + return {'available': False, 'reason': 'no ring engine'} + if snapshot is None: + # A native build from before the snapshot API has no get_stats at + # all: report it like missing fields. Errors raised by an existing + # get_stats propagate. + getter = getattr(engine, 'get_stats', None) + if not callable(getter): + return {'available': False, + 'reason': 'native telemetry unavailable (RingEngine has ' + 'no get_stats); rebuild required', + 'missing_fields': list(fields)} + stats = getter() + else: + stats = snapshot + missing = [name for name in fields if not hasattr(stats, name)] + if missing: + return {'available': False, + 'reason': 'native telemetry unavailable; rebuild required', + 'missing_fields': missing} + counters = {name: int(getattr(stats, name)) for name in fields} + if any(value < 0 for value in counters.values()): + raise ValueError('negative native counter') + return {'available': True, 'schema_version': 1, 'cumulative': counters} + + +def synchronization_stats(engine: Any, *, snapshot: Optional[Any] = None) -> dict: + """prepare_step outcomes, its stream synchronizations and forced flushes, + and node-toggle pre-apply waits (host nanoseconds).""" + out = _native_counters(engine, snapshot, FIELDS) + if out['available']: + out.update(clock='host_steady_clock_ns', + scope='existing host waits only; not device execution time') + return out + + +def pipeline_stats(engine: Any, *, snapshot: Optional[Any] = None) -> dict: + """Drain staging waits, per-task ready-to-commit latency (summed + microseconds), D2H failures and staging-to-pageable copies.""" + out = _native_counters(engine, snapshot, PIPELINE_FIELDS) + if out['available']: + out.update(clock='host_steady_clock', + scope='host staging waits and memcpy ns; summed per-task ' + 'ready-to-commit us; not device time') + return out + + +def gather_sync_stats(source: Any) -> dict: + """An integration's per-HookPoint row-table upload counters: uploads, + reuses of an unchanged table, and the host event waits before + rewriting a pinned table. ``source.row_table_stats()`` returns them as + a mapping keyed by ``GATHER_FIELDS``. Never queries a CUDA event.""" + read = getattr(source, 'row_table_stats', None) + if not callable(read): + return {'available': False, + 'reason': 'integration row-table telemetry unavailable', + 'missing_fields': list(GATHER_FIELDS)} + stats = read() + missing = [name for name in GATHER_FIELDS if name not in stats] + if missing: + return {'available': False, + 'reason': 'integration row-table telemetry unavailable', + 'missing_fields': missing} + counters = {name: stats[name] for name in GATHER_FIELDS} + if any(type(value) is not int or value < 0 for value in counters.values()): + raise ValueError('invalid row-table synchronization counter') + return {'available': True, 'schema_version': 1, 'cumulative': counters, + 'clock': 'host_perf_counter_ns', + 'scope': 'row-table reuse/uploads and existing host event waits; ' + 'not GPU copy time'} + + +__all__ = ['FIELDS', 'GATHER_FIELDS', 'PIPELINE_FIELDS', 'gather_sync_stats', + 'pipeline_stats', 'synchronization_stats'] diff --git a/tests/test_ring_telemetry.py b/tests/test_ring_telemetry.py new file mode 100644 index 000000000..7d0b193f6 --- /dev/null +++ b/tests/test_ring_telemetry.py @@ -0,0 +1,230 @@ +"""Host-only telemetry of the legacy ring (RingEngine.get_stats and +dmi.sync_stats). + +The CPU tests check the readers' contracts on fake snapshots; the GPU tests +drive a real ring through ordinary, flushed and oversized steps, the +node-toggle pre-apply wait and the drain/P2P stages, and check the counters +move exactly where a wait or copy happened. +""" +from __future__ import annotations + +import logging +import time +from types import SimpleNamespace + +import pytest +import torch + +from dmi import sync_stats +from tests._requirements import require_cuda, require_native_backend + + +def _snapshot(fields, value=1): + return SimpleNamespace(**{name: value for name in fields}) + + +@pytest.mark.cpu +def test_readers_report_unavailable_instead_of_zeros(): + assert sync_stats.synchronization_stats(None) == { + 'available': False, 'reason': 'no ring engine'} + partial = SimpleNamespace(prepare_ok_calls=1) + out = sync_stats.synchronization_stats(None, snapshot=partial) + assert out['available'] is False and 'prepare_drain_ns' in out['missing_fields'] + assert sync_stats.gather_sync_stats(object())['available'] is False + + +class _LegacyRingEngine: + """The API shape of a native RingEngine built before get_stats existed.""" + + +@pytest.mark.cpu +@pytest.mark.parametrize("reader, fields", [ + (sync_stats.synchronization_stats, sync_stats.FIELDS), + (sync_stats.pipeline_stats, sync_stats.PIPELINE_FIELDS), +]) +def test_readers_report_a_native_build_without_get_stats(reader, fields): + out = reader(_LegacyRingEngine()) + assert out['available'] is False + assert 'rebuild required' in out['reason'] and 'get_stats' in out['reason'] + assert out['missing_fields'] == list(fields) + # A non-callable attribute of that name is no snapshot API either. + engine = _LegacyRingEngine() + engine.get_stats = None + assert reader(engine)['available'] is False + + +@pytest.mark.cpu +@pytest.mark.parametrize("reader", [sync_stats.synchronization_stats, + sync_stats.pipeline_stats]) +def test_readers_propagate_errors_from_an_existing_get_stats(reader): + class BrokenEngine: + def get_stats(self): + raise RuntimeError("snapshot failed") + + with pytest.raises(RuntimeError, match="snapshot failed"): + reader(BrokenEngine()) + + +@pytest.mark.cpu +def test_readers_copy_counters_and_reject_negative_ones(): + out = sync_stats.pipeline_stats(None, snapshot=_snapshot(sync_stats.PIPELINE_FIELDS, 3)) + assert out['available'] is True and out['schema_version'] == 1 + assert set(out['cumulative']) == set(sync_stats.PIPELINE_FIELDS) + assert all(v == 3 for v in out['cumulative'].values()) + with pytest.raises(ValueError, match="negative"): + sync_stats.synchronization_stats( + None, snapshot=_snapshot(sync_stats.FIELDS, -1)) + counters = {f: 2 for f in sync_stats.GATHER_FIELDS} + source = SimpleNamespace(row_table_stats=lambda: dict(counters)) + assert sync_stats.gather_sync_stats(source)['cumulative']['hook_table_uploads'] == 2 + del counters['hook_table_reuses'] + assert sync_stats.gather_sync_stats(source)['missing_fields'] == ['hook_table_reuses'] + counters['hook_table_reuses'] = -1 + with pytest.raises(ValueError, match="row-table"): + sync_stats.gather_sync_stats(source) + + +QLEN, HID = 4, 8 + + +def _ring(payload, consume, task_entries=64): + 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 + + cfg = ne.RingConfig() + cfg.payload_ring_bytes = payload + cfg.pinned_staging_bytes = payload + cfg.task_ring_entries = task_entries + cfg.drain_poll_timeout_us = 200 + engine = ne.RingEngine(cfg, consume) + engine.init(0) + engine.start() + transport = rt.RingTransport(engine) + rt.activate(transport) + 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=["r"], token_ranges=[(0, QLEN)], + dim0_offsets=[0], flattened=True) + return engine, transport, rt + + +@pytest.mark.gpu +@pytest.mark.native_backend +@require_cuda() +@require_native_backend() +def test_get_stats_counts_prepare_outcomes_and_pipeline_stages(): + from dmi.hooks.specs import HOOK_TYPE_RESID_PRE + + delivered = [] + step = QLEN * HID * 4 + engine, transport, rt = _ring(4 * step, lambda *a: delivered.append(a[2])) + try: + before = engine.get_stats() + assert before.prepare_ok_calls == before.pageable_copy_calls == 0 + x = torch.ones(QLEN, HID, device="cuda") + payload = engine.payload_tensor() + outcomes = [] + for _ in range(6): + outcomes.append(engine.prepare_step(step, 1)) + transport.pre_push_all_metas(batch=0, q_len=QLEN, kv_dim=0) + torch.ops.ring.producer(payload, x, HOOK_TYPE_RESID_PRE, 0) + assert engine.drain_to_sink_and_wait(0) + # Reservations with nothing published fill the 4-step ring: the + # fifth must synchronize and flush. + for _ in range(5): + outcomes.append(engine.prepare_step(step, 1)) + assert outcomes[-1] == 1 + assert engine.prepare_step(64 * step, 1) == 2 # oversized + stats = engine.get_stats() + assert stats.prepare_ok_calls == outcomes.count(0) + assert stats.prepare_flushed_calls == outcomes.count(1) >= 1 + assert stats.prepare_oversized_calls == 1 + # Every flushed or oversized step synchronized and flushed once. + waits = stats.prepare_flushed_calls + stats.prepare_oversized_calls + assert stats.prepare_sync_calls == stats.prepare_drain_calls == waits + assert stats.prepare_sync_ns > 0 and stats.prepare_sync_errors == 0 + assert stats.ready_commit_tasks == 6 and stats.ready_commit_invalid == 0 + assert stats.pageable_copy_calls == 6 + assert stats.pageable_copy_bytes == 6 * step + assert stats.d2h_copy_errors == stats.d2h_sync_errors == 0 + assert stats.cpu_task_head >= stats.cpu_task_tail_committed + out = sync_stats.synchronization_stats(engine) + assert out['available'] and out['cumulative']['prepare_oversized_calls'] == 1 + assert sync_stats.pipeline_stats(engine)['cumulative']['pageable_copy_calls'] == 6 + assert len(delivered) == 6 + finally: + torch.cuda.synchronize() + engine.flush_and_wait() + rt.deactivate() + engine.stop() + + +@pytest.mark.gpu +@pytest.mark.native_backend +@require_cuda() +@require_native_backend() +def test_toggle_pre_apply_waits_are_counted_only_when_enabled(): + from dmi.hooks.specs import HOOK_TYPE_RESID_PRE + + engine, transport, rt = _ring(1 << 20, lambda *a: None) + try: + x = torch.ones(QLEN, HID, device="cuda") + payload = engine.payload_tensor() + engine.enable_toggle_capture(True) + engine.set_null_mode(True) + graph = torch.cuda.CUDAGraph(keep_graph=True) + with torch.cuda.graph(graph): + torch.ops.ring.producer(payload, x, HOOK_TYPE_RESID_PRE, 0) + graph.instantiate() + engine.set_null_mode(False) + engine.enable_toggle_capture(False) + engine.bind_graph_exec(graph.raw_cuda_graph(), graph.raw_cuda_graph_exec()) + raw = graph.raw_cuda_graph() + + def reconfigure(enabled): + transport.set_active_hooks_lazy([(HOOK_TYPE_RESID_PRE, 0)] if enabled else []) + transport.pre_push_all_metas(batch=0, q_len=QLEN, kv_dim=0) + assert transport.ensure_graph_current(raw) == 0 + graph.replay() + assert transport.record_replay_event(raw) == 0 + + reconfigure(True) + reconfigure(False) + assert engine.get_stats().toggle_wait_calls == 0 # off by default + engine.set_toggle_pre_apply_sync(True) + reconfigure(True) # records the first replay event + reconfigure(False) # waits on it before mutating the exec + assert engine.get_stats().toggle_wait_calls == 1 + finally: + torch.cuda.synchronize() + engine.flush_and_wait() + rt.deactivate() + engine.stop() + + +@pytest.mark.gpu +@pytest.mark.native_backend +@require_cuda() +@require_native_backend() +def test_monitoring_engine_logs_the_ring_footprint(caplog): + from dmi.engine import MonitoringEngine + from dmi.transport import native as ne + + cfg = ne.RingConfig() + cfg.payload_ring_bytes = 2 << 20 + cfg.pinned_staging_bytes = 1 << 20 + engine = MonitoringEngine(model_id="m", ring_config=cfg, + ring_consumer=lambda *a: None) + try: + with caplog.at_level(logging.INFO, logger="dmi.engine"): + engine._ensure_ring_transport() + (record,) = [r for r in caplog.records if "legacy ring" in r.getMessage()] + assert "2.0 MiB device payload ring" in record.getMessage() + assert "1.0 MiB pinned host staging" in record.getMessage() + assert "KV cache" in record.getMessage() + finally: + engine.close()