Skip to main content

query_runtime/observability/
mod.rs

1//! Native typed runtime observations and bounded query-local capture.
2
3#[cfg(any(test, feature = "test-hooks"))]
4use std::collections::BTreeSet;
5
6mod attempt_report;
7mod cache_fill_event;
8mod clock;
9pub mod graph_observation;
10mod jsonl;
11mod lifecycle;
12mod metrics_recorder;
13mod metrics_summary;
14mod operator_stage;
15mod recorder;
16mod runtime_observation;
17mod schedule_projection;
18mod task_identity;
19
20#[cfg(test)]
21pub(crate) use attempt_report::AttemptFailureOrigin;
22pub use attempt_report::{
23    AttemptAdmissionFact, AttemptAdmissionWaitFact, AttemptAllocatorEventSource, AttemptCapacityDisposition,
24    AttemptCleanupDisposition, AttemptDisposition, AttemptEvidenceProvenance, AttemptFailureFactKey,
25    AttemptFailureFactValue, AttemptFailureOperation, AttemptFailureStatus, AttemptFailureSummary,
26    AttemptFinalDomainUsage, AttemptHealthTransition, AttemptQueryId, AttemptRetryAdvice, AttemptRetryPrerequisite,
27    AttemptRetryRecovery, AttemptTimestamps, AttemptTraceContext, CleanupFailurePhase, CleanupFailureSummary,
28    MAX_ATTEMPT_QUERY_ID_BYTES, MAX_ATTEMPT_TRACE_CONTEXT_BYTES, QUERY_ATTEMPT_REPORT_SCHEMA, QueryAttemptReport,
29    RuntimeTerminalOutcome, SEALED_ATTEMPT_PROJECTION_SCHEMA, SealedAttemptLabels, SealedAttemptProjection,
30};
31pub(crate) use attempt_report::{
32    AttemptAllocatorEvent, AttemptAllocatorEventAction, AttemptAllocatorPoolUsage, AttemptAlternateDeviceEligibility,
33    AttemptDeviceHealth, AttemptDeviceHealthStatus, AttemptEvidenceField, AttemptFaultIsolation, AttemptNativeFault,
34    AttemptNativeFaultEvidence, AttemptNativeFaultKind, AttemptNativeStatus, AttemptResourceEvidence,
35    AttemptTeardownHealth, AttemptTerminalCommit, AttemptTimedDomainUsage, AttemptUnattributedAllocatorEvent,
36    has_proven_device_local_fault_evidence,
37};
38pub use cache_fill_event::LakehouseCacheFillEventSnapshot;
39#[doc(hidden)]
40pub use cache_fill_event::{LakehouseCacheFillErrorEvent, LakehouseCacheFillEvent, LakehouseCacheFillOutcome};
41pub use clock::DeterministicRuntimeObservationClock;
42pub(crate) use clock::{RuntimeObservationClock, SystemRuntimeObservationClock};
43pub use jsonl::{
44    RuntimeObservationExportError, RuntimeObservationExportErrorKind, RuntimeObservationExportErrorStatus,
45    write_runtime_observation_jsonl,
46};
47pub use lifecycle::{
48    RuntimeExecutionObservationScope, RuntimePlanExecutionLifecycle, RuntimeQueryLifecycle, RuntimeTaskLifecycle,
49    task_attempt_id,
50};
51pub(crate) use metrics_recorder::ExecutionProgress;
52pub use metrics_recorder::MetricsRecorder;
53#[cfg(test)]
54pub(crate) use metrics_summary::QueryMetricsSummary;
55pub(crate) use metrics_summary::{MetricsCaptureMode, MetricsCompleteness, QueryMetricsCheckpoint};
56pub use metrics_summary::{SourceReadFact, SourceReadTotals};
57pub(crate) use operator_stage::RuntimeExecutionStageContext;
58pub use operator_stage::{RuntimeOperatorStageGuard, RuntimeOperatorStageIdentity, stage_context_from_execution};
59#[cfg(test)]
60pub(crate) use recorder::RuntimeObservationDeliveryGate;
61pub(crate) use recorder::{
62    AttemptReportSealProvider, BudgetWindow, CleanupFailureTally, QueryCleanupSealDelivery,
63    in_runtime_observation_delivery,
64};
65pub use recorder::{
66    BudgetWindowPeaks, LiveRuntimeObservationSnapshot, RuntimeEventContext, RuntimeObservationConfig,
67    RuntimeObservationHandle, RuntimeObservationOwner, RuntimeObservationRecordDisposition, RuntimeObservationSession,
68    SealedRuntimeObservationSnapshot,
69};
70pub(crate) use runtime_observation::{
71    AdmissionLifecycleState, JobDefinitionRow, RuntimeAdmissionObservation, RuntimeClockDomain, RuntimeJobId,
72    RuntimeObservationClassLoss, RuntimeObservationFooter, RuntimeObservationManifest, RuntimeObservationReasonLoss,
73    RuntimeOperatorStageObservation, RuntimeSourceLocationId, RuntimeStreamLeaseObservation,
74    RuntimeTaskLifecycleObservation, StreamLeaseAcquisition,
75};
76pub use runtime_observation::{
77    AdmissionRequestId, BarrierDefinitionRow, CompiledEdgeRow, GpuJobLifecycleState, ObservationSequence, PlanEdgeRow,
78    PlanExecutionBoundaryKind, PlanExecutionId, PlanNodeRow, QueryBoundaryKind, RUNTIME_OBSERVATION_JSONL_SCHEMA,
79    RuntimeBarrierId, RuntimeCleanupFailureObservation, RuntimeGpuJobObservation, RuntimeObservationLimits,
80    RuntimeObservationLossReason, RuntimeObservationPayload, RuntimeObservationRecord, RuntimeObservationRecordClass,
81    RuntimeOperatorStageKind, RuntimeOperatorStageState, RuntimePipelineEdgeId, RuntimePlanNodeId,
82    RuntimeReservationAction, RuntimeReservationDisposition, RuntimeReservationId, RuntimeReservationObservation,
83    RuntimeResourceDomain, RuntimeResourceId, RuntimeResourceRow, RuntimeSinkId, RuntimeSourceReadObservation,
84    RuntimeSpanId, RuntimeStreamLeaseScope, RuntimeTimelineEvent, RuntimeTopologyRecord, RuntimeWorkerId,
85    SinkDefinitionRow, SourceReadLifecycleState, StreamLeaseLifecycleState, TaskDefinitionRow, TaskLifecycleState,
86    WorkerDefinitionRow,
87};
88pub use schedule_projection::{AttemptComposedScheduleProjection, AttemptId, COMPOSED_SCHEDULE_PROJECTION_SCHEMA};
89pub use task_identity::RuntimeTaskId;
90pub(crate) use task_identity::{RuntimeAttemptId, TaskAttemptId};
91
92#[cfg(any(test, feature = "test-hooks"))]
93#[doc(hidden)]
94pub fn assert_live_topology_references_resolve(snapshot: &LiveRuntimeObservationSnapshot) {
95    let mut plan_nodes = BTreeSet::new();
96    let mut tasks = BTreeSet::new();
97    let mut edges = BTreeSet::new();
98    let mut barriers = BTreeSet::new();
99    let mut sinks = BTreeSet::new();
100    let mut jobs = BTreeSet::new();
101    let mut workers = BTreeSet::new();
102    let mut resources = BTreeSet::new();
103    for record in snapshot.topology() {
104        let inserted = match record {
105            RuntimeTopologyRecord::PlanNode(row) => {
106                plan_nodes.insert((row.plan_execution_id.clone(), row.node_id.clone()))
107            },
108            RuntimeTopologyRecord::TaskDefinition(row) => {
109                tasks.insert((row.plan_execution_id.clone(), row.task_id.clone()))
110            },
111            RuntimeTopologyRecord::CompiledEdge(row) => {
112                edges.insert((row.plan_execution_id.clone(), row.edge_id.clone()))
113            },
114            RuntimeTopologyRecord::BarrierDefinition(row) => {
115                barriers.insert((row.plan_execution_id.clone(), row.barrier_id.clone()))
116            },
117            RuntimeTopologyRecord::SinkDefinition(row) => {
118                sinks.insert((row.plan_execution_id.clone(), row.sink_id.clone()))
119            },
120            RuntimeTopologyRecord::JobDefinition(row) => {
121                jobs.insert((row.plan_execution_id.clone(), row.job_id.clone()))
122            },
123            RuntimeTopologyRecord::WorkerDefinition(row) => {
124                workers.insert((row.plan_execution_id.clone(), row.worker_id.clone()))
125            },
126            RuntimeTopologyRecord::Resource(row) => {
127                resources.insert((row.plan_execution_id.clone(), row.resource_id.clone()))
128            },
129            RuntimeTopologyRecord::PlanEdge(_) => true,
130        };
131        assert!(inserted, "topology identity must be unique within its plan-execution scope: {record:?}");
132    }
133
134    for record in snapshot.topology() {
135        match record {
136            RuntimeTopologyRecord::PlanEdge(row) => {
137                assert!(plan_nodes.contains(&(row.plan_execution_id.clone(), row.from_node_id.clone())));
138                assert!(plan_nodes.contains(&(row.plan_execution_id.clone(), row.to_node_id.clone())));
139            },
140            RuntimeTopologyRecord::TaskDefinition(row) => {
141                if let Some(node_id) = &row.plan_node_id {
142                    assert!(plan_nodes.contains(&(row.plan_execution_id.clone(), node_id.clone())));
143                }
144            },
145            RuntimeTopologyRecord::CompiledEdge(row) => {
146                assert!(tasks.contains(&(row.plan_execution_id.clone(), row.from_task_id.clone())));
147                assert!(tasks.contains(&(row.plan_execution_id.clone(), row.to_task_id.clone())));
148                for resource_id in &row.resource_ids {
149                    assert!(resources.contains(&(Some(row.plan_execution_id.clone()), resource_id.clone())));
150                }
151            },
152            RuntimeTopologyRecord::BarrierDefinition(row) => {
153                for task_id in row.upstream_task_ids.iter().chain(&row.downstream_task_ids) {
154                    assert!(tasks.contains(&(row.plan_execution_id.clone(), task_id.clone())));
155                }
156            },
157            RuntimeTopologyRecord::SinkDefinition(row) => {
158                assert!(tasks.contains(&(row.plan_execution_id.clone(), row.input_task_id.clone())));
159            },
160            RuntimeTopologyRecord::Resource(row) => {
161                if let Some(plan_execution_id) = &row.plan_execution_id {
162                    if let Some(producer) = &row.producer_task_id {
163                        assert!(tasks.contains(&(plan_execution_id.clone(), producer.clone())));
164                    }
165                    for consumer in &row.consumer_task_ids {
166                        assert!(tasks.contains(&(plan_execution_id.clone(), consumer.clone())));
167                    }
168                }
169            },
170            RuntimeTopologyRecord::PlanNode(_)
171            | RuntimeTopologyRecord::JobDefinition(_)
172            | RuntimeTopologyRecord::WorkerDefinition(_) => {},
173        }
174    }
175
176    for event in snapshot.events() {
177        match &event.payload {
178            RuntimeObservationPayload::TaskLifecycle(observation) => {
179                let plan = event.plan_execution_id.clone().expect("task event has plan execution scope");
180                assert!(tasks.contains(&(plan, observation.attempt_id.task_id.clone())));
181            },
182            RuntimeObservationPayload::GpuJob(observation) => {
183                assert!(jobs.contains(&(event.plan_execution_id.clone(), observation.job_id.clone())));
184                if let Some(worker_id) = &observation.worker_id {
185                    assert!(workers.contains(&(event.plan_execution_id.clone(), worker_id.clone())));
186                }
187            },
188            RuntimeObservationPayload::StreamLease(observation) => {
189                assert!(resources.contains(&(event.plan_execution_id.clone(), observation.lease_id.clone())));
190                if let Some(task_id) = &observation.scope.task_id {
191                    let plan = event.plan_execution_id.clone().expect("task-correlated stream has plan scope");
192                    assert!(tasks.contains(&(plan, task_id.clone())));
193                }
194                if let Some(worker_id) = &observation.scope.source_worker_id {
195                    assert!(workers.iter().any(|(_, candidate)| candidate == worker_id));
196                }
197                if let Some(job_id) = &observation.scope.gpu_job_id {
198                    assert!(jobs.iter().any(|(_, candidate)| candidate == job_id));
199                }
200                if let Some(worker_id) = &observation.scope.gpu_worker_id {
201                    assert!(workers.iter().any(|(_, candidate)| candidate == worker_id));
202                }
203            },
204            RuntimeObservationPayload::ResourceMeasurement(observation) => {
205                assert!(resources.contains(&(event.plan_execution_id.clone(), observation.resource_id.clone())));
206            },
207            RuntimeObservationPayload::SourceRead(observation) => {
208                let plan = event.plan_execution_id.clone().expect("source event has plan execution scope");
209                assert!(tasks.contains(&(plan.clone(), observation.task_id.clone())));
210                assert!(workers.contains(&(Some(plan), observation.worker_id.clone())));
211            },
212            RuntimeObservationPayload::Reservation(observation) => {
213                if let Some(attempt) = &observation.task_attempt_id {
214                    let plan = event.plan_execution_id.clone().expect("task reservation has plan execution scope");
215                    assert!(tasks.contains(&(plan, attempt.task_id.clone())));
216                }
217            },
218            RuntimeObservationPayload::QueryBoundary { .. }
219            | RuntimeObservationPayload::PlanExecutionBoundary { .. }
220            | RuntimeObservationPayload::Admission(_)
221            | RuntimeObservationPayload::CleanupFailure(_)
222            | RuntimeObservationPayload::OperatorStage(_) => {},
223        }
224    }
225}