1#[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}