1use serde::{Deserialize, Serialize};
2use std::{
3 fmt,
4 time::{Duration, Instant},
5};
6
7use super::super::{
8 BatchInvocationId, CompletionHandle, DefinitelyNotSubmittedRetryAuthority,
9 DefinitelyNotSubmittedWaveRetryAuthority, DeviceCommandPhase, DeviceComputePathRequirement,
10 DeviceRuntime, DeviceSubmissionAttribution, DeviceSubmissionExecutionTiming,
11 DeviceSubmissionStage, DeviceSubmissionTimingSink, DeviceTimingMeasurement, HostTransferLayout,
12 IdentifiedFailure, IndeterminateSubmissionHandle, NodeId, VNextError,
13};
14use super::foundation::invalid_operation;
15use super::{BatchOperationIdentity, OperationFailure};
16
17pub trait DispatchRetryAuthority: fmt::Debug {
18 fn prior_attempt(&self) -> BatchInvocationId;
19}
20
21impl<R: DeviceRuntime> DispatchRetryAuthority for DefinitelyNotSubmittedRetryAuthority<R> {
22 fn prior_attempt(&self) -> BatchInvocationId {
23 self.prior_attempt()
24 }
25}
26
27impl<R: DeviceRuntime> DispatchRetryAuthority for DefinitelyNotSubmittedWaveRetryAuthority<R> {
28 fn prior_attempt(&self) -> BatchInvocationId {
29 self.prior_attempt()
30 }
31}
32
33pub enum OperationDispatchError<R, Retry = DefinitelyNotSubmittedRetryAuthority<R>>
34where
35 R: DeviceRuntime,
36 Retry: DispatchRetryAuthority,
37{
38 Contract(VNextError),
39 Provider(OperationFailure),
40 Initialization(IdentifiedFailure),
41 InputUpload(IdentifiedFailure),
42 DefinitelyNotSubmitted {
43 failures: Vec<IdentifiedFailure>,
44 retry: Retry,
45 },
46 SubmissionIndeterminate {
47 recovery: IndeterminateSubmissionHandle<R>,
48 },
49 PostSubmitContract {
50 error: VNextError,
51 completion: CompletionHandle<R>,
52 },
53}
54
55pub type SubmissionWaveDispatchError<R> =
56 OperationDispatchError<R, DefinitelyNotSubmittedWaveRetryAuthority<R>>;
57
58#[derive(Debug, Clone, PartialEq, Eq, Serialize)]
59pub struct BoundDeviceSubmissionAttribution {
60 batch_identity: BatchOperationIdentity,
61 submission_fingerprint: String,
62 device: DeviceSubmissionAttribution,
63 terminal_timing: DeviceTimingMeasurement<DeviceSubmissionExecutionTiming>,
64}
65
66impl BoundDeviceSubmissionAttribution {
67 pub(super) fn new(
68 batch_identity: BatchOperationIdentity,
69 submission_fingerprint: String,
70 device: DeviceSubmissionAttribution,
71 ) -> Result<Self, VNextError> {
72 for command in device.commands() {
73 let Some(node_index) = command.node_index() else {
74 continue;
75 };
76 let node_index_usize = usize::try_from(node_index).map_err(|_| {
77 invalid_operation(format!(
78 "device command {} node index exceeds host address space",
79 command.command_index()
80 ))
81 })?;
82 let node_participant_count = batch_identity
83 .node_participant_count(node_index_usize)
84 .and_then(|count| u32::try_from(count).ok())
85 .ok_or_else(|| {
86 invalid_operation(format!(
87 "device command {} references absent node {}",
88 command.command_index(),
89 node_index
90 ))
91 })?;
92 let participant_range_is_valid = command.participant_start() < node_participant_count
93 && command.participant_end() <= node_participant_count;
94 let command_requires_full_node =
95 command.command_phase() != DeviceCommandPhase::Initialization;
96 if !participant_range_is_valid
97 || command_requires_full_node
98 && (command.participant_start() != 0
99 || command.participant_count() != node_participant_count)
100 {
101 return Err(invalid_operation(format!(
102 "device command {} phase {:?} participant range {}..{} differs from node {} participant count {}",
103 command.command_index(),
104 command.command_phase(),
105 command.participant_start(),
106 command.participant_end(),
107 node_index,
108 node_participant_count
109 )));
110 }
111 }
112 for replayed_segment in device.replayed_segments() {
113 let program_id = replayed_segment.program_id();
114 if program_id.plan_hash() != batch_identity.plan_hash()
115 || program_id.runtime_implementation_fingerprint()
116 != batch_identity.runtime_implementation_fingerprint()
117 || program_id.lane_id() != batch_identity.lane_id()
118 || replayed_segment.logical_commands().iter().any(|command| {
119 let Ok(node_index) = usize::try_from(command.node_index()) else {
120 return true;
121 };
122 batch_identity.node_id_at(node_index).is_none()
123 || u32::try_from(
124 batch_identity
125 .node_participant_count(node_index)
126 .unwrap_or_default(),
127 )
128 .map_or(true, |count| count != command.participant_count())
129 })
130 || replayed_segment.logical_commands().first().is_none_or(|_| {
131 usize::try_from(replayed_segment.physical_command_index())
132 .ok()
133 .and_then(|index| device.commands().get(index))
134 .is_none_or(|physical| {
135 physical.participant_start() != 0
136 || program_id.immediate_sequences() != physical.participant_count()
137 || program_id.immediate_tokens() != physical.token_count()
138 })
139 })
140 {
141 return Err(invalid_operation(
142 "replayed segment attribution differs from its batch or sealed program identity",
143 ));
144 }
145 }
146 Ok(Self {
147 batch_identity,
148 submission_fingerprint,
149 device,
150 terminal_timing: DeviceTimingMeasurement::NotRequested,
151 })
152 }
153
154 pub fn bind_terminal_timing(
155 mut self,
156 terminal_timing: DeviceTimingMeasurement<DeviceSubmissionExecutionTiming>,
157 ) -> Result<Self, VNextError> {
158 if let DeviceTimingMeasurement::Measured(timing) = &terminal_timing {
159 if !terminal_timing_matches_submission_attribution(timing, &self.device) {
160 return Err(invalid_operation(
161 "terminal device timing coverage differs from submission command attribution",
162 ));
163 }
164 }
165 self.terminal_timing = terminal_timing;
166 Ok(self)
167 }
168
169 pub fn batch_identity(&self) -> &BatchOperationIdentity {
170 &self.batch_identity
171 }
172
173 pub fn submission_fingerprint(&self) -> &str {
174 &self.submission_fingerprint
175 }
176
177 pub fn device(&self) -> &DeviceSubmissionAttribution {
178 &self.device
179 }
180
181 pub const fn terminal_timing(
182 &self,
183 ) -> &DeviceTimingMeasurement<DeviceSubmissionExecutionTiming> {
184 &self.terminal_timing
185 }
186}
187
188fn terminal_timing_matches_submission_attribution(
189 timing: &DeviceSubmissionExecutionTiming,
190 attribution: &DeviceSubmissionAttribution,
191) -> bool {
192 let commands = attribution.commands();
193 let attributed = |command_index| {
194 commands
195 .binary_search_by_key(&command_index, |command| command.command_index())
196 .is_ok()
197 };
198 commands
199 .iter()
200 .all(|command| timing.span_for_command(command.command_index()).is_some())
201 && timing.spans().iter().all(|span| {
202 span.measurement().elapsed_ns().is_none()
203 || (span.start_command_index()..span.end_command_index())
204 .all(|command_index| attributed(command_index))
205 })
206}
207
208#[must_use = "profiled submission evidence and completion must be consumed together"]
209pub struct ProfiledSubmissionHandle<R: DeviceRuntime> {
210 completion: CompletionHandle<R>,
211 attribution: Option<BoundDeviceSubmissionAttribution>,
212}
213
214impl<R: DeviceRuntime> ProfiledSubmissionHandle<R> {
215 pub(super) fn new(
216 completion: CompletionHandle<R>,
217 attribution: Option<BoundDeviceSubmissionAttribution>,
218 ) -> Self {
219 Self {
220 completion,
221 attribution,
222 }
223 }
224
225 pub fn into_parts(
226 self,
227 ) -> (
228 CompletionHandle<R>,
229 Option<BoundDeviceSubmissionAttribution>,
230 ) {
231 (self.completion, self.attribution)
232 }
233}
234
235#[derive(Debug, Clone, Copy, PartialEq, Eq)]
238pub enum SubmissionWaveDispatchStage {
239 ContractValidateAndReserve,
240 BackingAndInputEncode,
241 ProviderNodeEncode,
242 LaneReserve,
243 DeviceRuntimeSubmit,
244 CompletionArm,
245 LaneReserveSubmitAndArm,
246}
247
248pub trait SubmissionWaveDispatchTimingSink: DeviceSubmissionTimingSink {
256 fn record(&self, stage: SubmissionWaveDispatchStage, elapsed: Duration);
257}
258
259pub(super) struct DisabledSubmissionWaveDispatchTimingSink;
260
261impl DeviceSubmissionTimingSink for DisabledSubmissionWaveDispatchTimingSink {
262 const ENABLED: bool = false;
263
264 fn record_device_submission(&self, _stage: DeviceSubmissionStage, _elapsed: Duration) {
265 unreachable!("disabled device submission timing cannot record")
266 }
267}
268
269impl SubmissionWaveDispatchTimingSink for DisabledSubmissionWaveDispatchTimingSink {
270 fn record(&self, _stage: SubmissionWaveDispatchStage, _elapsed: Duration) {
271 unreachable!("disabled submission timing cannot record")
272 }
273}
274
275pub(super) struct SubmissionWaveDispatchStageTimer<'sink, S>
276where
277 S: SubmissionWaveDispatchTimingSink,
278{
279 sink: &'sink S,
280 stage: SubmissionWaveDispatchStage,
281 started: Option<Instant>,
282}
283
284impl<'sink, S> SubmissionWaveDispatchStageTimer<'sink, S>
285where
286 S: SubmissionWaveDispatchTimingSink,
287{
288 #[inline(always)]
289 pub(super) fn start(sink: &'sink S, stage: SubmissionWaveDispatchStage) -> Self {
290 Self {
291 sink,
292 stage,
293 started: S::ENABLED.then(Instant::now),
294 }
295 }
296}
297
298impl<S> Drop for SubmissionWaveDispatchStageTimer<'_, S>
299where
300 S: SubmissionWaveDispatchTimingSink,
301{
302 fn drop(&mut self) {
303 if let Some(started) = self.started.take() {
304 if !std::thread::panicking() {
305 self.sink.record(self.stage, started.elapsed());
306 }
307 }
308 }
309}
310
311#[cfg(test)]
312mod submission_wave_dispatch_timing_tests {
313 use std::time::Duration;
314
315 use super::{
316 DeviceSubmissionStage, DeviceSubmissionTimingSink, SubmissionWaveDispatchStage,
317 SubmissionWaveDispatchStageTimer, SubmissionWaveDispatchTimingSink,
318 };
319
320 struct DisabledPanicSink;
321
322 impl DeviceSubmissionTimingSink for DisabledPanicSink {
323 const ENABLED: bool = false;
324
325 fn record_device_submission(&self, _stage: DeviceSubmissionStage, _elapsed: Duration) {
326 panic!("disabled device timing sink was called");
327 }
328 }
329
330 impl SubmissionWaveDispatchTimingSink for DisabledPanicSink {
331 fn record(&self, _stage: SubmissionWaveDispatchStage, _elapsed: Duration) {
332 panic!("disabled timing sink was called");
333 }
334 }
335
336 #[test]
337 fn disabled_submission_timing_does_not_record() {
338 let timer = SubmissionWaveDispatchStageTimer::start(
339 &DisabledPanicSink,
340 SubmissionWaveDispatchStage::ProviderNodeEncode,
341 );
342 drop(timer);
343
344 assert!(!DisabledPanicSink::ENABLED);
345 }
346}
347
348#[cfg(test)]
349mod terminal_timing_attribution_tests {
350 use super::*;
351 use crate::vnext::{
352 DeviceBatchingForm, DeviceExecutionInterval, DeviceExecutionIntervalKind,
353 DeviceExecutionPath, DeviceExecutionSpanKind, DeviceNativeOperationId,
354 DeviceNativeWorkAttribution, DeviceSubmissionExecutionSpan, DeviceTimingUnavailableReason,
355 };
356
357 fn attributed_command(command_index: u32) -> DeviceNativeWorkAttribution {
358 DeviceNativeWorkAttribution::new(
359 command_index,
360 None,
361 DeviceCommandPhase::Compute,
362 DeviceNativeOperationId::new("test.compute").unwrap(),
363 DeviceExecutionPath::Eager,
364 DeviceBatchingForm::Scalar,
365 1,
366 1,
367 1,
368 0,
369 None,
370 )
371 .unwrap()
372 }
373
374 fn measured_span(command_index: u32) -> DeviceSubmissionExecutionSpan {
375 DeviceSubmissionExecutionSpan::measured(
376 command_index,
377 command_index + 1,
378 DeviceExecutionSpanKind::EagerCommand,
379 vec![DeviceExecutionInterval::new(
380 DeviceExecutionIntervalKind::Compute,
381 u64::from(command_index) * 10,
382 u64::from(command_index) * 10 + 5,
383 )
384 .unwrap()],
385 )
386 .unwrap()
387 }
388
389 fn unavailable_span(command_index: u32) -> DeviceSubmissionExecutionSpan {
390 DeviceSubmissionExecutionSpan::unavailable(
391 command_index,
392 command_index + 1,
393 DeviceExecutionSpanKind::EagerCommand,
394 DeviceTimingUnavailableReason::BackendMeasurementFailed,
395 )
396 .unwrap()
397 }
398
399 #[test]
400 fn terminal_timing_accepts_unavailable_hole_between_attributed_commands() {
401 let attribution =
402 DeviceSubmissionAttribution::new(vec![attributed_command(0), attributed_command(2)])
403 .unwrap();
404 let timing = DeviceSubmissionExecutionTiming::from_spans(
405 3,
406 vec![measured_span(0), unavailable_span(1), measured_span(2)],
407 )
408 .unwrap();
409
410 assert!(terminal_timing_matches_submission_attribution(
411 &timing,
412 &attribution
413 ));
414 }
415
416 #[test]
417 fn terminal_timing_rejects_measured_unattributed_work_or_missing_attributed_command() {
418 let attribution =
419 DeviceSubmissionAttribution::new(vec![attributed_command(0), attributed_command(2)])
420 .unwrap();
421 let unattributed = DeviceSubmissionExecutionTiming::from_spans(
422 3,
423 vec![measured_span(0), measured_span(1), measured_span(2)],
424 )
425 .unwrap();
426 let truncated = DeviceSubmissionExecutionTiming::from_spans(
427 2,
428 vec![measured_span(0), unavailable_span(1)],
429 )
430 .unwrap();
431
432 assert!(!terminal_timing_matches_submission_attribution(
433 &unattributed,
434 &attribution
435 ));
436 assert!(!terminal_timing_matches_submission_attribution(
437 &truncated,
438 &attribution
439 ));
440 }
441
442 #[test]
443 fn terminal_timing_rejects_reusable_span_covering_unattributed_command() {
444 let attribution = DeviceSubmissionAttribution::new(vec![attributed_command(0)]).unwrap();
445 let timing = DeviceSubmissionExecutionTiming::from_spans(
446 2,
447 vec![DeviceSubmissionExecutionSpan::measured(
448 0,
449 2,
450 DeviceExecutionSpanKind::ReusableExecutable,
451 vec![
452 DeviceExecutionInterval::new(DeviceExecutionIntervalKind::Compute, 0, 5)
453 .unwrap(),
454 ],
455 )
456 .unwrap()],
457 )
458 .unwrap();
459
460 assert!(!terminal_timing_matches_submission_attribution(
461 &timing,
462 &attribution
463 ));
464 }
465}
466
467impl<R, Retry> fmt::Debug for OperationDispatchError<R, Retry>
468where
469 R: DeviceRuntime,
470 Retry: DispatchRetryAuthority,
471{
472 fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
473 match self {
474 Self::Contract(error) => formatter.debug_tuple("Contract").field(error).finish(),
475 Self::Provider(error) => formatter.debug_tuple("Provider").field(error).finish(),
476 Self::Initialization(error) => formatter
477 .debug_tuple("Initialization")
478 .field(error)
479 .finish(),
480 Self::InputUpload(error) => formatter.debug_tuple("InputUpload").field(error).finish(),
481 Self::DefinitelyNotSubmitted { failures, retry } => formatter
482 .debug_struct("DefinitelyNotSubmitted")
483 .field("failures", failures)
484 .field("retry", retry)
485 .finish(),
486 Self::SubmissionIndeterminate { recovery } => formatter
487 .debug_struct("SubmissionIndeterminate")
488 .field("recovery", recovery)
489 .finish(),
490 Self::PostSubmitContract { error, completion } => formatter
491 .debug_struct("PostSubmitContract")
492 .field("error", error)
493 .field("completion", completion)
494 .finish(),
495 }
496 }
497}
498
499impl<R, Retry> fmt::Display for OperationDispatchError<R, Retry>
500where
501 R: DeviceRuntime,
502 Retry: DispatchRetryAuthority,
503{
504 fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
505 match self {
506 Self::Contract(error) => {
507 write!(formatter, "operation dispatch contract failed: {error}")
508 }
509 Self::Provider(error) => write!(
510 formatter,
511 "operation provider failed with {}: {}",
512 error.code(),
513 error.message()
514 ),
515 Self::Initialization(error) => write!(
516 formatter,
517 "operation backing initialization failed with {}: {}",
518 error.failure().code(),
519 error.failure().message()
520 ),
521 Self::InputUpload(error) => write!(
522 formatter,
523 "operation input upload failed with {}: {}",
524 error.failure().code(),
525 error.failure().message()
526 ),
527 Self::DefinitelyNotSubmitted { failures, retry } => write!(
528 formatter,
529 "operation attempt {} with {} participants was definitely not submitted: {}",
530 retry.prior_attempt(),
531 failures.len(),
532 failures
533 .first()
534 .map(|failure| failure.failure().message())
535 .unwrap_or("missing classified participant failure")
536 ),
537 Self::SubmissionIndeterminate { recovery } => write!(
538 formatter,
539 "operation submission may have reached the device; completion slot {} retains ownership",
540 recovery.slot_id().get()
541 ),
542 Self::PostSubmitContract { error, completion } => write!(
543 formatter,
544 "operation submission reached the device but slot {} observed a contract failure: {error}",
545 completion.slot_id().get()
546 ),
547 }
548 }
549}
550
551#[derive(Debug, Clone, Copy, Default, PartialEq, Eq, Serialize, Deserialize)]
556#[serde(rename_all = "snake_case")]
557pub enum SubmissionScratchInitialization {
558 #[default]
559 ProviderContract,
560 FillByte(u8),
561}
562
563#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
569#[serde(deny_unknown_fields)]
570pub struct SubmissionExecutionPolicy {
571 compute_path: DeviceComputePathRequirement,
572 scratch_initialization: SubmissionScratchInitialization,
573}
574
575impl SubmissionExecutionPolicy {
576 pub const fn adaptive() -> Self {
577 Self {
578 compute_path: DeviceComputePathRequirement::Adaptive,
579 scratch_initialization: SubmissionScratchInitialization::ProviderContract,
580 }
581 }
582
583 pub const fn determinism_eager(scratch_fill: u8) -> Self {
584 Self {
585 compute_path: DeviceComputePathRequirement::EagerOnly,
586 scratch_initialization: SubmissionScratchInitialization::FillByte(scratch_fill),
587 }
588 }
589
590 pub const fn determinism_replayed(scratch_fill: u8) -> Self {
591 Self {
592 compute_path: DeviceComputePathRequirement::ReplayedOnly,
593 scratch_initialization: SubmissionScratchInitialization::FillByte(scratch_fill),
594 }
595 }
596
597 pub const fn compute_path(self) -> DeviceComputePathRequirement {
598 self.compute_path
599 }
600
601 pub const fn scratch_initialization(self) -> SubmissionScratchInitialization {
602 self.scratch_initialization
603 }
604}
605
606impl Default for SubmissionExecutionPolicy {
607 fn default() -> Self {
608 Self::adaptive()
609 }
610}
611
612#[derive(Debug, Clone, PartialEq, Eq)]
616pub struct SubmissionWaveInputUpload {
617 node_id: NodeId,
618 participant_index: u32,
619 input_ordinal: u32,
620 logical_offset_bytes: u64,
621 source_layout: HostTransferLayout,
622 bytes: Vec<u8>,
623}
624
625impl SubmissionWaveInputUpload {
626 pub fn new(
627 node_id: NodeId,
628 participant_index: u32,
629 input_ordinal: u32,
630 logical_offset_bytes: u64,
631 source_layout: HostTransferLayout,
632 bytes: Vec<u8>,
633 ) -> Result<Self, VNextError> {
634 source_layout.validate_bytes(bytes.len())?;
635 let byte_len = source_layout.byte_len()?;
636 if logical_offset_bytes.checked_add(byte_len).is_none()
637 || logical_offset_bytes % source_layout.element_type().size_bytes() != 0
638 {
639 return Err(invalid_operation(
640 "submission input upload has an invalid aligned logical range",
641 ));
642 }
643 Ok(Self {
644 node_id,
645 participant_index,
646 input_ordinal,
647 logical_offset_bytes,
648 source_layout,
649 bytes,
650 })
651 }
652
653 pub fn node_id(&self) -> &NodeId {
654 &self.node_id
655 }
656
657 pub const fn participant_index(&self) -> u32 {
658 self.participant_index
659 }
660
661 pub const fn input_ordinal(&self) -> u32 {
662 self.input_ordinal
663 }
664
665 pub const fn logical_offset_bytes(&self) -> u64 {
666 self.logical_offset_bytes
667 }
668
669 pub const fn source_layout(&self) -> HostTransferLayout {
670 self.source_layout
671 }
672
673 pub fn bytes(&self) -> &[u8] {
674 &self.bytes
675 }
676}
677
678impl<R, Retry> std::error::Error for OperationDispatchError<R, Retry>
679where
680 R: DeviceRuntime,
681 Retry: DispatchRetryAuthority,
682{
683}