Add `data_batch` telemetry by dhruv9vats · Pull Request #1068 · sirius-db/sirius

dhruv9vats · 2026-07-03T12:19:55Z

Add data_batch telemetry (engine emission → analyzer → quent UI)

Summary

Instruments the lifecycle of every GPU data_batch and surfaces it in the quent UI. Each batch now emits a telemetry FSM — Constructed → Stationary ⇄ InTransit → Destructed — attributed to the pipeline that produced it, so you can see per-batch memory residency and cross-tier transfers, and per-resource (memory/channel) occupancy over a query's execution.

Engine (C++ / Super Sirius)

quent_data_batch_probe (telemetry/data_batch_probe.hpp): an idata_batch_probe that forwards batch state transitions to quent, resolving memory/channel handles from a memory_context.
memory_context: registers a quent memory handle per memory space and a channel handle per space-pair, owned by the telemetry_context. The reservation manager is optional, so telemetry-free/no-GPU unit tests build a valid (empty-handle) context.
batch_telemetry_info POD (telemetry context + producing pipeline_uuid); quent_data_batch_probe::create(...) returns a real probe when a context is present, else a no-op base probe.
Every production batch-creation site now attaches a probe: the make_data_batch/make_data_batch_from_view factories, direct data_batch::make calls (top-n, limit, table-scan, ungrouped-aggregate, scan path, scan-manager cache provider), and the clone/convert sites (result-collector, partition, grouped-aggregate-merge). In-place convert_to reuses the existing probe.
Attribution is sourced from the operator's pipeline: both the telemetry context and pipeline UUID ride on pipeline_build_context (constructor-injected once at plan-convert time) and are read via the operator's _pipeline — no per-operator telemetry state, no per-call plumbing.

Analyzer (Rust)

Ingests the memory and channel resource streams (mirroring the existing task-queue/executor-thread handling).
Reconstructs DataBatch FSMs and exposes them alongside tasks: registered in fsm_types, fed into single/bulk/bulk-chunked resource timelines and long_fsms, filtered per-query by producing pipeline.
Per-state breakdowns are type-gated (a memory timeline broken down "by task" yields empty, and vice-versa), so task and data-batch state sets never mix.
Removed vestigial FsmCollection/Using impls on the model (unused by the timeline path).

Tests

Behavior unchanged. Call sites that build pipelines/batches without an engine pass a null telemetry context (→ no-op probe), so existing C++ unit tests are unaffected apart from the now-explicit argument.

Verification

C++ extension and Rust analyzer + server both build clean.
A GPU query run confirmed correct data_batch events with populated, per-pipeline producer_pipeline_uuid and resolved memory handles (no dropped events).
Pending: end-to-end UI runtime check against a live capture (confirm memory/channel timelines + FSM rendering) — not yet run.

🤖 PR description generated with @claude

mbrobbel · 2026-07-06T12:31:46Z

Should bump to main after NVIDIA/cuCascade#156 merge

mbrobbel · 2026-07-06T12:41:49Z

+    // auto maybe_memory_handle =
+    //   memory_context_->get_memory_handle(new_data.get_memory_space().get_id());
+    // if (not maybe_memory_handle) {
+    //   SIRIUS_LOG_WARN(fmt::format("No quent memory handle found for {}",
+    //                               new_data.get_memory_space().to_string()));
+    //   return;
+    // }
+    // handle_->stationary({
+    //   .memory_resource_id    = (*maybe_memory_handle).get().uuid(),
+    //   .memory_capacity_bytes = new_data.get_size_in_bytes(),
+    // });


Is this intended?

I was contemplating removing the set_data this corresponds to, but to reduce scope creep, lets just keep that API and this stationary -> stationary transition.

mbrobbel · 2026-07-06T12:43:33Z

@@ -161,7 +162,11 @@ void sirius_physical_materialized_collector::sink(const operator_data& input_dat

      // clone_to: creates new batch with data converted to host_data_representation
      auto result_batch = ro.clone_to<cucascade::host_data_representation>(


The probe is attached after registry.convert() finishes, so the GPU→host transfer never emits InTransit or channel usage. This omits result-collection transfers from the advertised telemetry.

There does not seem to be any straight forward way to achieve this, so lets take the "omits result-collection transfers from the advertised telemetry" compromise, for now. As this design is refined.

johanpel

Very nice, a few small comments about de-duplicating code.
Figuring those out for Sirius may provide some nice insights for analysis-related codegen in Quent later.

johanpel · 2026-07-07T08:58:52Z

+        for data_batch in view.data_batches() {
+            for usage in data_batch.usages() {
+                let resource_id = usage.resource_id();
+                if let Some(builder_indices) = plain_index.get(&resource_id) {
+                    for &builder_idx in builder_indices {
+                        let builder = &mut plain_builders[builder_idx];
+                        if data_batch.matches_filter(&builder.3) {
+                            builder.1.try_push(&usage)?;
+                        }
+                    }
+                }
+            }
+
+            for (state_name, usage) in data_batch.usages_with_state_names() {
+                let resource_id = usage.resource_id();
+                if let Some(builder_indices) = per_state_index.get(&resource_id) {
+                    for &builder_idx in builder_indices {
+                        let builder = &mut per_state_builders[builder_idx];
+                        if builder.4 == DATA_BATCH_TYPE_NAME
+                            && data_batch.matches_filter(&builder.3)
+                        {


Since this code will get quite repetitive the more FSM types we add, could we perhaps add some kind of umbrella FSM type that we can iterate over that just forwards trait implementations of e.g. usage_with_state_names() to its variants?

Might be possible in multiple places below as well.

Feel free to do in follow-up.

Yes, I started de-duplicating the code, but deferred that for a smaller follow-up.

johanpel

Approving Rust changes, thanks @dhruv9vats !

The rebase onto dev pulled in the data_batch entity type (sirius-db#1068) and the thread tree (sirius-db#1080), which interact with the attribute fields this branch adds to quent's FsmTransition: - populate the new attributes / derived_attributes fields when converting DataBatch transitions (no derived synthesis, so derived_attributes is empty) - dev removed the blanket FsmCollection impl on SiriusModel (two FSM types now), so list_entities dispatches on entity_type_name via TaskCollection / DataBatchCollection adapters, defaulting to tasks Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Signed-off-by: Felipe Aramburu <felipearamburu@gmail.com>

…UES source Brings in 73 commits including: - sirius-db#1080 GPU -> thread-group telemetry tree (telemetry_context now declares per-GPU device groups; created after the memory manager) - sirius-db#1068 DataBatch physical-residency telemetry (constructed/stationary/ in_transit/destructed with memory+channel usages) - coexists with our Batch placement FSM (queued/packaged/processing with memory_tier usages) - sirius-db#1144 GPU_VALUES source replacing CPU_SOURCE (drops our cpu_source publish hook; that path now flows through the standard operator sink) - sirius-db#1177 logging refactor (batch_telemetry switched fmt:: -> std::format) Merge notes: batch placement hooks re-seated on upstream's reworked lock_or_prepare_batch and weak-pointer _subscribed_batches; analyzer keeps upstream's string-keyed entity dispatch with 'batch' as a third entity kind; sirius-db#1112's per-transition tooltip attributes are disabled until quent's FsmTransition grows attribute fields (NOTE(merge) markers). Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>

…ew (sirius-db#1187) ## What Instruments **batch placements** — the lifecycle of each data batch published to a consuming pipeline's input port — and implements quent's new per-operator data-flow distribution endpoint, so quent's DAG view can scrub through time and show, per pipeline: how many batches (and bytes) are queued waiting for the scheduler, packaged into tasks, or actively processing, and **which memory tier / GPU device the data resides on**. Companion PR (merged): rapidsai/quent#393 — the generic categorical-timeline protocol + playhead UI. ## Model (`rust/crates/telemetry/model`) New `Batch` FSM: `batch_registered{batch_id, pipeline_uuid, port_uuid, origin} → batch_queued → batch_packaged{task_uuid} → batch_processing{task_uuid} → batch_consumed{reason}`, where every non-terminal state carries a tier usage on a new `MemoryTier` resource — **one per GPU device** (`GPU-0`, `GPU-1`, …) plus engine-wide `HOST`/`DISK` — weighted by the batch's bytes. Tier changes (downgrade/spill/prepare-time upgrade) are self-transitions with the new tier usage. Placements share the engine's `batch_id` across fan-out and OOM re-packaging, so batch identity survives task reschedules. Complementary to `DataBatch` (sirius-db#1068), which tracks physical residency; `Batch` tracks scheduling lifecycle per consumer. ## Engine instrumentation (`src/`) A process-global `batch_telemetry_registry` (sharded by batch id) maps port repositories to consumer pipelines (registered during plan telemetry), so publish sites only pass `(batch, repo)`: - Publishes (operator sink, partition consumer) emit `registered → queued` before the batch becomes poppable. - `gpu_pipeline_task`: ctor claims inputs (`queued → packaged`, with lazy registration for OOM-reschedule intermediates), execute emits `packaged → processing` post-prepare (tier re-read to capture upgrades; an id-based path covers merge/concat inputs consumed during materialization), dtor releases claims **by recorded batch id** (the weak batch refs are dead by then). OOM re-claims transfer ownership to the rescheduled task. - Tier moves reported from `convertible_data_batch::convert` and `lock_or_prepare_batch` while the exclusive lock is held (values passed, never re-locked). - Leftovers drained as `consumed{query_end}` in `QueryEnd`; gated by new `telemetry.enable_batch_events` (default on). Every placement in TPC-H SF1/SF300 (2-GPU) test runs completes `consumed{processed}` inside the query window. ## Analyzer (`rust/crates/telemetry/analyzer`) - Ingests Batch/MemoryTier events; implements `UiAnalyzer::data_flow_timeline` (states × tiers × count/bytes per pipeline; `Unsupported` for datasets without batch events, so old recordings keep working and the UI hides the view). - `batch` is a third entity kind in the per-state resource timelines and entity lists, beside `task`/`data_batch`; `batch_registered` (instantaneous bookkeeping entry state) is omitted from aggregated series. - Byte quantities declared with SI prefixes (GB rather than GiB). - 7 analyzer unit tests incl. hand-computed bin math with a mid-queue tier-change split. ## ⚠️ Do not merge yet (draft) - The quent crates are pinned to upstream `rapidsai/quent` main (the rapidsai/quent#393 merge commit); the temporary `[patch]` block used during development has been removed. - Includes a merge of latest `dev` (resolved against sirius-db#1068/sirius-db#1080/sirius-db#1112; the sirius-db#1112 per-transition tooltip attributes are wired for batches too). ## Testing - `pixi run make test` (1538 C++ cases) and `cargo test` green; full E2E: TPC-H SF1 + SF300 parquet on 2 GPUs (incl. a constrained-GPU-memory run), NDJSON lifecycle invariants verified, served through the embedded quent UI and scrubbed. 🤖 Generated with [Claude Code](https://claude.com/claude-code) --------- Co-authored-by: Claude Fable 5 <noreply@anthropic.com>

dhruv9vats added 8 commits July 1, 2026 17:36

Use standardized shared_ptr data_batch

d988ede

Merge dev

83820aa

Merge dev

4367b08

Bump cuCascade

0db572d

Add quent_data_batch_probe

64e04bb

Wire up data_batch_probes

d6d6d7d

cuCascade tracking

953698b

Merge dev

5669158

mbrobbel linked an issue Jul 3, 2026 that may be closed by this pull request

Add memory consumption and data movement instrumentation via Quent #1062

Closed

mbrobbel reviewed Jul 3, 2026

View reviewed changes

Comment thread experimental/starrocks/starrocks

dhruv9vats self-assigned this Jul 3, 2026

dhruv9vats added 2 commits July 3, 2026 20:27

Add data_batch analyzer

f1e5ac1

Revert unintended submodule bump

5eb18cb

dhruv9vats marked this pull request as ready for review July 3, 2026 15:16

Merge dev

ff97a85

mbrobbel reviewed Jul 6, 2026

View reviewed changes

Comment thread cucascade

mbrobbel Jul 6, 2026

Copy link
Copy Markdown

Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Should bump to main after NVIDIA/cuCascade#156 merge

dhruv9vats reacted with thumbs up emoji

mbrobbel added quent telemetry labels Jul 6, 2026

mbrobbel reviewed Jul 6, 2026

View reviewed changes

johanpel reviewed Jul 7, 2026

View reviewed changes

dhruv9vats added 3 commits July 8, 2026 14:23

use match for entity type matching, streamline

81611db

Merge dev

26050f7

Make batch_telemetry required in batch utils

963adb8

dhruv9vats requested review from johanpel and mbrobbel July 8, 2026 09:26

johanpel approved these changes Jul 8, 2026

View reviewed changes

mbrobbel reviewed Jul 8, 2026

View reviewed changes

Comment thread substrait

Revert accidental substrait submodule change

d2c8553

mbrobbel approved these changes Jul 8, 2026

View reviewed changes

dhruv9vats added this pull request to the merge queue Jul 8, 2026

github-merge-queue Bot removed this pull request from the merge queue due to failed status checks Jul 8, 2026

Update fmt paths

e02023d

dhruv9vats enabled auto-merge July 8, 2026 13:12

dhruv9vats added this pull request to the merge queue Jul 8, 2026

Merged via the queue into sirius-db:dev with commit c49301c Jul 8, 2026
10 checks passed

dhruv9vats deleted the data-batch-telemetry branch July 8, 2026 13:49

bwyogatama mentioned this pull request Jul 8, 2026

feat(scan): serve pinned entries through the split_provider machinery via a pinned-serving mode on the format ingestibles (#819 PR2) #1126

Closed

felipeblazing mentioned this pull request Jul 15, 2026

feat(telemetry): batch placement lifecycle + memory-tier data-flow view #1187

Merged

mike-wendt mentioned this pull request Aug 5, 2026

docs: add CODEOWNERS #1367

Merged

6 tasks

Navigation Menu

Add `data_batch` telemetry - #1068

Add `data_batch` telemetry#1068
dhruv9vats merged 16 commits into
sirius-db:devfrom
dhruv9vats:data-batch-telemetry

dhruv9vats commented Jul 3, 2026 •

edited

Loading

Uh oh!

Uh oh!

mbrobbel Jul 6, 2026

Uh oh!

mbrobbel Jul 6, 2026

Uh oh!

dhruv9vats Jul 8, 2026

Uh oh!

mbrobbel Jul 6, 2026

Uh oh!

dhruv9vats Jul 8, 2026

Uh oh!

Uh oh!

johanpel left a comment

Uh oh!

Uh oh!

Uh oh!

johanpel Jul 7, 2026

Uh oh!

dhruv9vats Jul 8, 2026

Uh oh!

johanpel left a comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

		@@ -161,7 +162,11 @@ void sirius_physical_materialized_collector::sink(const operator_data& input_dat

		// clone_to: creates new batch with data converted to host_data_representation
		auto result_batch = ro.clone_to<cucascade::host_data_representation>(

Conversation

dhruv9vats commented Jul 3, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

Uh oh!

mbrobbel Jul 6, 2026

Choose a reason for hiding this comment

Uh oh!

mbrobbel Jul 6, 2026

Choose a reason for hiding this comment

Uh oh!

dhruv9vats Jul 8, 2026

Choose a reason for hiding this comment

Uh oh!

mbrobbel Jul 6, 2026

Choose a reason for hiding this comment

Uh oh!

dhruv9vats Jul 8, 2026

Choose a reason for hiding this comment

Uh oh!

Uh oh!

johanpel left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

johanpel Jul 7, 2026

Choose a reason for hiding this comment

Uh oh!

dhruv9vats Jul 8, 2026

Choose a reason for hiding this comment

Uh oh!

johanpel left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

dhruv9vats commented Jul 3, 2026 •

edited

Loading