1use std::{borrow::Borrow, marker::PhantomData};
9
10use eredu_core::{
11 checkpoint::TensorDtype, BoundedCompletion, BoundedCompletionOutcome, BoundedSubmissionOutcome,
12 CollectiveGroupId, CompletionCancellationMode, DistributedCommitEpoch,
13 DistributedCommitOutcome, DistributedCommitPhase,
14};
15use eredu_nn::{NeuralBackend, Tensor};
16
17use crate::{
18 ActivationObserver, ArchitectureGroupKind, BarrierBackend, BroadcastBackend,
19 CommunicationBackend, CommunicationGroupDescriptor, CommunicationManifest,
20 CommunicationOperation, CommunicationOperationRequirement, CommunicationPeerCounts,
21 CommunicationRouteDescriptor, CommunicationRouteId, EvenGatherBackend, ExecutionGraph,
22 ExecutionResidency, ExpertPass, FailureAgreementBackend, LayeredArchitecture,
23 LayeredPartitionDriver, LayeredPipelineSchedule, LayeredPipelineScheduleError,
24 LayeredTraversalHook, LayerwisePolicy, LayerwiseRuntime, LayerwiseRuntimeError,
25 ParallelLayeredArchitecture, PipelineActivationDtype, PipelineWireContract,
26 PointToPointBackend, ReplicatedTextExecutionStrategy, ReplicatedTextSessionError,
27 ResolvedBoundaryTensorSpec, ResolvedBoundaryWireSchema, RuntimeState, SubmissionBackend,
28 SumReductionBackend, UnevenGatherBackend, VariableAllToAllBackend,
29};
30
31pub trait CommunicationTensorMetadata<B: NeuralBackend> {
33 fn dtype(&self, tensor: &B::Tensor) -> TensorDtype;
35
36 fn shape(&self, tensor: &B::Tensor) -> Vec<usize>;
38}
39
40pub struct RealizedCommunicationGroup<G> {
42 id: CollectiveGroupId,
43 resource: G,
44}
45
46impl<G> RealizedCommunicationGroup<G> {
47 pub const fn new(id: CollectiveGroupId, resource: G) -> Self {
49 Self { id, resource }
50 }
51}
52
53pub struct RealizedCommunicationRoute<R> {
55 id: CommunicationRouteId,
56 resource: R,
57}
58
59impl<R> RealizedCommunicationRoute<R> {
60 pub const fn new(id: CommunicationRouteId, resource: R) -> Self {
62 Self { id, resource }
63 }
64}
65
66pub struct PartitionCommunication<B, G, R, I>
68where
69 B: CommunicationBackend,
70{
71 manifest: CommunicationManifest,
72 groups: Vec<RealizedCommunicationGroup<G>>,
73 routes: Vec<RealizedCommunicationRoute<R>>,
74 inspector: I,
75 authority: PartitionCommunicationAuthority,
76 backend: PhantomData<fn() -> B>,
77}
78
79#[derive(Debug, Clone, Copy, Eq, PartialEq)]
80struct CommunicationPoison {
81 operation: CommunicationOperation,
82 phase: DistributedExecutionPhase,
83 route: Option<CommunicationRouteId>,
84 cancellation: CompletionCancellationMode,
85}
86
87#[derive(Debug, Clone)]
93pub struct PartitionCommunicationAuthority {
94 policy: Option<crate::CommunicationCompletionPolicy>,
95 poison: std::sync::Arc<std::sync::Mutex<Option<CommunicationPoison>>>,
96}
97
98impl PartitionCommunicationAuthority {
99 fn new(policy: Option<crate::CommunicationCompletionPolicy>) -> Self {
100 Self {
101 policy,
102 poison: std::sync::Arc::new(std::sync::Mutex::new(None)),
103 }
104 }
105
106 pub fn from_manifest(
108 manifest: &CommunicationManifest,
109 ) -> Result<Self, PartitionExecutionError> {
110 if manifest.completion_policy().is_none()
111 && (!manifest.groups().is_empty() || !manifest.routes().is_empty())
112 {
113 return Err(PartitionExecutionError::MissingBoundedCompletionPolicy);
114 }
115 Ok(Self::new(manifest.completion_policy()))
116 }
117
118 fn poison_guard(&self) -> std::sync::MutexGuard<'_, Option<CommunicationPoison>> {
119 self.poison
120 .lock()
121 .unwrap_or_else(std::sync::PoisonError::into_inner)
122 }
123
124 pub fn ensure_active(&self) -> Result<(), PartitionExecutionError> {
126 match *self.poison_guard() {
127 Some(poison) => Err(PartitionExecutionError::CommunicationPoisoned {
128 operation: poison.operation,
129 phase: poison.phase,
130 route: poison.route,
131 cancellation: poison.cancellation,
132 }),
133 None => Ok(()),
134 }
135 }
136
137 pub const fn completion_policy(&self) -> Option<crate::CommunicationCompletionPolicy> {
139 self.policy
140 }
141
142 fn mark_poisoned(&self, poison: CommunicationPoison) {
143 let mut current = self.poison_guard();
144 if current.is_none() {
145 *current = Some(poison);
146 }
147 }
148
149 fn is_poisoned(&self) -> bool {
150 self.poison_guard().is_some()
151 }
152
153 pub fn submission_error(
155 &self,
156 error: impl std::fmt::Display,
157 operation: CommunicationOperation,
158 phase: DistributedExecutionPhase,
159 route: Option<CommunicationRouteId>,
160 ) -> PartitionExecutionError {
161 let cancellation = self
162 .policy
163 .expect("communication submission requires a selected completion policy")
164 .cancellation();
165 self.mark_poisoned(CommunicationPoison {
166 operation,
167 phase,
168 route,
169 cancellation,
170 });
171 PartitionExecutionError::CommunicationSubmissionFailed {
172 operation,
173 phase,
174 route,
175 error: error.to_string(),
176 }
177 }
178
179 pub fn completion_error(
181 &self,
182 error: impl std::fmt::Display,
183 operation: CommunicationOperation,
184 phase: DistributedExecutionPhase,
185 route: Option<CommunicationRouteId>,
186 ) -> PartitionExecutionError {
187 let cancellation = self
188 .policy
189 .expect("communication completion requires a selected completion policy")
190 .cancellation();
191 self.mark_poisoned(CommunicationPoison {
192 operation,
193 phase,
194 route,
195 cancellation,
196 });
197 PartitionExecutionError::CommunicationCompletionFailed {
198 operation,
199 phase,
200 route,
201 error: error.to_string(),
202 }
203 }
204
205 pub fn wait<T, C>(
207 &self,
208 submission: eredu_core::Submission<T, C>,
209 operation: CommunicationOperation,
210 phase: DistributedExecutionPhase,
211 route: Option<CommunicationRouteId>,
212 ) -> Result<T, PartitionExecutionError>
213 where
214 C: BoundedCompletion,
215 C::Error: std::fmt::Display,
216 {
217 self.ensure_active()?;
218 let policy = self
219 .policy
220 .ok_or(PartitionExecutionError::MissingBoundedCompletionPolicy)?
221 .bounded_wait();
222 let outcome = match submission.wait_bounded(policy) {
223 Ok(outcome) => outcome,
224 Err(error) => {
225 self.mark_poisoned(CommunicationPoison {
226 operation,
227 phase,
228 route,
229 cancellation: policy.cancellation(),
230 });
231 return Err(PartitionExecutionError::CommunicationCompletionFailed {
232 operation,
233 phase,
234 route,
235 error: error.to_string(),
236 });
237 }
238 };
239 match outcome {
240 BoundedSubmissionOutcome::Completed(output) => Ok(output),
241 BoundedSubmissionOutcome::DeadlineExceeded { cancellation } => {
242 self.mark_poisoned(CommunicationPoison {
243 operation,
244 phase,
245 route,
246 cancellation,
247 });
248 Err(PartitionExecutionError::CommunicationDeadlineExceeded {
249 operation,
250 phase,
251 route,
252 cancellation,
253 })
254 }
255 }
256 }
257
258 fn wait_after_prior_failure<T, C>(
259 &self,
260 submission: eredu_core::Submission<T, C>,
261 operation: CommunicationOperation,
262 phase: DistributedExecutionPhase,
263 route: Option<CommunicationRouteId>,
264 ) -> Result<T, PartitionExecutionError>
265 where
266 C: BoundedCompletion,
267 C::Error: std::fmt::Display,
268 {
269 let policy = self
270 .policy
271 .ok_or(PartitionExecutionError::MissingBoundedCompletionPolicy)?
272 .bounded_wait();
273 let outcome = submission.wait_bounded(policy).map_err(|error| {
274 self.mark_poisoned(CommunicationPoison {
275 operation,
276 phase,
277 route,
278 cancellation: policy.cancellation(),
279 });
280 PartitionExecutionError::CommunicationCompletionFailed {
281 operation,
282 phase,
283 route,
284 error: error.to_string(),
285 }
286 })?;
287 match outcome {
288 BoundedSubmissionOutcome::Completed(output) => Ok(output),
289 BoundedSubmissionOutcome::DeadlineExceeded { cancellation } => {
290 self.mark_poisoned(CommunicationPoison {
291 operation,
292 phase,
293 route,
294 cancellation,
295 });
296 Err(PartitionExecutionError::CommunicationDeadlineExceeded {
297 operation,
298 phase,
299 route,
300 cancellation,
301 })
302 }
303 }
304 }
305
306 fn wait_before_failure_agreement<T, C>(
307 &self,
308 submission: eredu_core::Submission<T, C>,
309 operation: CommunicationOperation,
310 phase: DistributedExecutionPhase,
311 route: Option<CommunicationRouteId>,
312 ) -> Result<T, PartitionExecutionError>
313 where
314 C: BoundedCompletion,
315 C::Error: std::fmt::Display,
316 {
317 self.ensure_active()?;
318 let policy = self
319 .policy
320 .ok_or(PartitionExecutionError::MissingBoundedCompletionPolicy)?
321 .bounded_wait();
322 match submission.wait_bounded(policy).map_err(|error| {
323 PartitionExecutionError::CommunicationCompletionFailed {
324 operation,
325 phase,
326 route,
327 error: error.to_string(),
328 }
329 })? {
330 BoundedSubmissionOutcome::Completed(output) => Ok(output),
331 BoundedSubmissionOutcome::DeadlineExceeded { cancellation } => {
332 Err(PartitionExecutionError::CommunicationDeadlineExceeded {
333 operation,
334 phase,
335 route,
336 cancellation,
337 })
338 }
339 }
340 }
341
342 fn fence_protocol_failure(
343 &self,
344 operation: CommunicationOperation,
345 phase: DistributedExecutionPhase,
346 route: Option<CommunicationRouteId>,
347 ) {
348 let cancellation = self
349 .policy
350 .expect("communication protocol failure requires a selected completion policy")
351 .cancellation();
352 self.mark_poisoned(CommunicationPoison {
353 operation,
354 phase,
355 route,
356 cancellation,
357 });
358 }
359}
360
361impl<B, G, R, I> PartitionCommunication<B, G, R, I>
362where
363 B: CommunicationBackend,
364 G: Borrow<B::CommunicationGroup>,
365 R: Borrow<B::CommunicationRoute>,
366 I: CommunicationTensorMetadata<B>,
367{
368 pub fn new(
370 manifest: CommunicationManifest,
371 groups: Vec<RealizedCommunicationGroup<G>>,
372 routes: Vec<RealizedCommunicationRoute<R>>,
373 inspector: I,
374 ) -> Result<Self, PartitionExecutionError> {
375 let authority = PartitionCommunicationAuthority::from_manifest(&manifest)?;
376 Self::new_with_authority(manifest, groups, routes, inspector, authority)
377 }
378
379 pub fn new_with_authority(
381 manifest: CommunicationManifest,
382 groups: Vec<RealizedCommunicationGroup<G>>,
383 routes: Vec<RealizedCommunicationRoute<R>>,
384 inspector: I,
385 authority: PartitionCommunicationAuthority,
386 ) -> Result<Self, PartitionExecutionError> {
387 if authority.policy != manifest.completion_policy() {
388 return Err(PartitionExecutionError::MissingBoundedCompletionPolicy);
389 }
390 if groups.len() != manifest.groups().len() || routes.len() != manifest.routes().len() {
391 return Err(PartitionExecutionError::ResourceCount {
392 expected_groups: manifest.groups().len(),
393 actual_groups: groups.len(),
394 expected_routes: manifest.routes().len(),
395 actual_routes: routes.len(),
396 });
397 }
398 for (descriptor, resource) in manifest.groups().iter().zip(&groups) {
399 if descriptor.id() != resource.id {
400 return Err(PartitionExecutionError::ResourceIdentity {
401 expected: u64::from(descriptor.id().value()),
402 actual: u64::from(resource.id.value()),
403 });
404 }
405 }
406 for (descriptor, resource) in manifest.routes().iter().zip(&routes) {
407 if descriptor.id() != resource.id {
408 return Err(PartitionExecutionError::ResourceIdentity {
409 expected: descriptor.id().value(),
410 actual: resource.id.value(),
411 });
412 }
413 }
414 Ok(Self {
415 manifest,
416 groups,
417 routes,
418 inspector,
419 authority,
420 backend: PhantomData,
421 })
422 }
423
424 pub const fn manifest(&self) -> &CommunicationManifest {
426 &self.manifest
427 }
428
429 pub fn authority(&self) -> PartitionCommunicationAuthority {
431 self.authority.clone()
432 }
433
434 fn ensure_active(&self) -> Result<(), PartitionExecutionError> {
435 self.authority.ensure_active()
436 }
437
438 fn submission_error(
439 &self,
440 error: impl std::fmt::Display,
441 operation: CommunicationOperation,
442 phase: DistributedExecutionPhase,
443 route: Option<CommunicationRouteId>,
444 ) -> PartitionExecutionError {
445 self.authority
446 .submission_error(error, operation, phase, route)
447 }
448
449 fn group(
450 &self,
451 id: CollectiveGroupId,
452 operation: CommunicationOperation,
453 ) -> Result<(&CommunicationGroupDescriptor, &B::CommunicationGroup), PartitionExecutionError>
454 {
455 self.ensure_active()?;
456 let index = self
457 .manifest
458 .groups()
459 .iter()
460 .position(|candidate| candidate.id() == id)
461 .ok_or(PartitionExecutionError::UnknownGroup(id))?;
462 let descriptor = &self.manifest.groups()[index];
463 if descriptor.local_index().is_none() {
464 return Err(PartitionExecutionError::NotGroupMember(id));
465 }
466 if !descriptor
467 .requirements()
468 .operations()
469 .iter()
470 .any(|requirement| requirement.operation() == operation)
471 {
472 return Err(PartitionExecutionError::OperationNotSelected {
473 resource: format!("group {}", id.value()),
474 operation,
475 });
476 }
477 Ok((descriptor, self.groups[index].resource.borrow()))
478 }
479
480 fn route(
481 &self,
482 id: CommunicationRouteId,
483 ) -> Result<(&CommunicationRouteDescriptor, &B::CommunicationRoute), PartitionExecutionError>
484 {
485 self.ensure_active()?;
486 let index = self
487 .manifest
488 .routes()
489 .iter()
490 .position(|candidate| candidate.id() == id)
491 .ok_or(PartitionExecutionError::UnknownRoute(id))?;
492 Ok((
493 &self.manifest.routes()[index],
494 self.routes[index].resource.borrow(),
495 ))
496 }
497
498 fn group_requirement(
499 descriptor: &CommunicationGroupDescriptor,
500 operation: CommunicationOperation,
501 ) -> &CommunicationOperationRequirement {
502 descriptor
503 .requirements()
504 .operations()
505 .iter()
506 .find(|requirement| requirement.operation() == operation)
507 .expect("group() established the selected operation")
508 }
509
510 fn validate_tensor(
511 &self,
512 value: &B::Tensor,
513 requirement: &CommunicationOperationRequirement,
514 completed: bool,
515 ) -> Result<(), PartitionExecutionError> {
516 let dtype = self.inspector.dtype(value);
517 let shape = self.inspector.shape(value);
518 if !requirement.dtypes().contains(&dtype) {
519 return Err(PartitionExecutionError::TensorDtype { dtype });
520 }
521 let limits = requirement
522 .limits()
523 .ok_or(PartitionExecutionError::MissingTensorLimits)?;
524 let elements = shape
525 .iter()
526 .try_fold(1usize, |product, dimension| product.checked_mul(*dimension));
527 let maximum = if completed {
528 limits.max_output_tensor_elements()
529 } else {
530 limits.max_tensor_elements()
531 };
532 if shape.len() > limits.max_tensor_rank()
533 || elements.is_none_or(|elements| elements > maximum)
534 {
535 return Err(PartitionExecutionError::TensorLimits { shape });
536 }
537 Ok(())
538 }
539
540 fn wait<T>(
541 &self,
542 submission: eredu_core::Submission<T, B::CommunicationCompletion>,
543 operation: CommunicationOperation,
544 phase: DistributedExecutionPhase,
545 route: Option<CommunicationRouteId>,
546 ) -> Result<T, PartitionExecutionError> {
547 self.authority.wait(submission, operation, phase, route)
548 }
549
550 fn expected_axis_output_shape(
551 &self,
552 value: &B::Tensor,
553 axis: usize,
554 output_width: usize,
555 ) -> Result<Vec<usize>, PartitionExecutionError> {
556 let mut shape = self.inspector.shape(value);
557 let rank = shape.len();
558 let dimension = shape
559 .get_mut(axis)
560 .ok_or(PartitionExecutionError::CommunicationAxis { axis, rank })?;
561 *dimension = output_width;
562 Ok(shape)
563 }
564
565 fn validate_output_shape(
566 &self,
567 output: &B::Tensor,
568 expected: Vec<usize>,
569 ) -> Result<(), PartitionExecutionError> {
570 let actual = self.inspector.shape(output);
571 if actual != expected {
572 return Err(PartitionExecutionError::CommunicationOutputShape { expected, actual });
573 }
574 Ok(())
575 }
576
577 fn output_contract_error(
578 &self,
579 error: PartitionExecutionError,
580 operation: CommunicationOperation,
581 phase: DistributedExecutionPhase,
582 route: Option<CommunicationRouteId>,
583 ) -> PartitionExecutionError {
584 self.authority
585 .fence_protocol_failure(operation, phase, route);
586 error
587 }
588
589 pub fn all_reduce_sum(
591 &self,
592 value: B::Tensor,
593 group: CollectiveGroupId,
594 executor: &B::Executor,
595 ) -> Result<B::Tensor, PartitionExecutionError>
596 where
597 B: SumReductionBackend,
598 {
599 let (descriptor, native) = self.group(group, CommunicationOperation::AllReduceSum)?;
600 let requirement = Self::group_requirement(descriptor, CommunicationOperation::AllReduceSum);
601 self.validate_tensor(&value, requirement, false)?;
602 let output = B::all_reduce_sum(value, native, executor).map_err(|error| {
603 self.submission_error(
604 error,
605 CommunicationOperation::AllReduceSum,
606 DistributedExecutionPhase::Execution,
607 None,
608 )
609 })?;
610 let output = self.wait(
611 output,
612 CommunicationOperation::AllReduceSum,
613 DistributedExecutionPhase::Execution,
614 None,
615 )?;
616 self.validate_tensor(&output, requirement, true)
617 .map_err(|error| {
618 self.output_contract_error(
619 error,
620 CommunicationOperation::AllReduceSum,
621 DistributedExecutionPhase::Execution,
622 None,
623 )
624 })?;
625 Ok(output)
626 }
627
628 pub fn all_reduce_sum_wave(
636 &self,
637 values: impl IntoIterator<Item = B::Tensor>,
638 group: CollectiveGroupId,
639 executor: &B::Executor,
640 ) -> Result<Vec<B::Tensor>, PartitionExecutionError>
641 where
642 B: SumReductionBackend,
643 {
644 let (descriptor, native) = self.group(group, CommunicationOperation::AllReduceSum)?;
645 let requirement = Self::group_requirement(descriptor, CommunicationOperation::AllReduceSum);
646 let mut submissions = Vec::new();
647 for value in values {
648 self.validate_tensor(&value, requirement, false)?;
649 submissions.push(B::all_reduce_sum(value, native, executor).map_err(|error| {
650 self.submission_error(
651 error,
652 CommunicationOperation::AllReduceSum,
653 DistributedExecutionPhase::Execution,
654 None,
655 )
656 })?);
657 }
658 submissions
659 .into_iter()
660 .map(|submission| {
661 let output = self.wait(
662 submission,
663 CommunicationOperation::AllReduceSum,
664 DistributedExecutionPhase::Execution,
665 None,
666 )?;
667 self.validate_tensor(&output, requirement, true)
668 .map_err(|error| {
669 self.output_contract_error(
670 error,
671 CommunicationOperation::AllReduceSum,
672 DistributedExecutionPhase::Execution,
673 None,
674 )
675 })?;
676 Ok(output)
677 })
678 .collect()
679 }
680
681 pub fn all_gather_even(
683 &self,
684 value: B::Tensor,
685 axis: usize,
686 group: CollectiveGroupId,
687 executor: &B::Executor,
688 ) -> Result<B::Tensor, PartitionExecutionError>
689 where
690 B: EvenGatherBackend,
691 {
692 let (descriptor, native) = self.group(group, CommunicationOperation::AllGatherEven)?;
693 let requirement =
694 Self::group_requirement(descriptor, CommunicationOperation::AllGatherEven);
695 self.validate_tensor(&value, requirement, false)?;
696 let input_shape = self.inspector.shape(&value);
697 let input_width =
698 *input_shape
699 .get(axis)
700 .ok_or(PartitionExecutionError::CommunicationAxis {
701 axis,
702 rank: input_shape.len(),
703 })?;
704 let output_width = input_width
705 .checked_mul(descriptor.members().len())
706 .ok_or(PartitionExecutionError::CommunicationShapeOverflow)?;
707 let expected = self.expected_axis_output_shape(&value, axis, output_width)?;
708 let output = B::all_gather_even(value, axis, native, executor).map_err(|error| {
709 self.submission_error(
710 error,
711 CommunicationOperation::AllGatherEven,
712 DistributedExecutionPhase::Execution,
713 None,
714 )
715 })?;
716 let output = self.wait(
717 output,
718 CommunicationOperation::AllGatherEven,
719 DistributedExecutionPhase::Execution,
720 None,
721 )?;
722 self.validate_tensor(&output, requirement, true)
723 .and_then(|()| self.validate_output_shape(&output, expected))
724 .map_err(|error| {
725 self.output_contract_error(
726 error,
727 CommunicationOperation::AllGatherEven,
728 DistributedExecutionPhase::Execution,
729 None,
730 )
731 })?;
732 Ok(output)
733 }
734
735 pub fn all_gather_uneven(
737 &self,
738 value: B::Tensor,
739 counts: &[usize],
740 axis: usize,
741 group: CollectiveGroupId,
742 executor: &B::Executor,
743 ) -> Result<B::Tensor, PartitionExecutionError>
744 where
745 B: UnevenGatherBackend,
746 {
747 let (descriptor, native) = self.group(group, CommunicationOperation::AllGatherUneven)?;
748 if counts.len() != descriptor.members().len() {
749 return Err(PartitionExecutionError::PeerCount {
750 expected: descriptor.members().len(),
751 actual: counts.len(),
752 });
753 }
754 let requirement =
755 Self::group_requirement(descriptor, CommunicationOperation::AllGatherUneven);
756 self.validate_tensor(&value, requirement, false)?;
757 let output_width = counts.iter().try_fold(0usize, |total, count| {
758 total
759 .checked_add(*count)
760 .ok_or(PartitionExecutionError::CommunicationShapeOverflow)
761 })?;
762 let expected = self.expected_axis_output_shape(&value, axis, output_width)?;
763 let output =
764 B::all_gather_uneven(value, counts, axis, native, executor).map_err(|error| {
765 self.submission_error(
766 error,
767 CommunicationOperation::AllGatherUneven,
768 DistributedExecutionPhase::Execution,
769 None,
770 )
771 })?;
772 let output = self.wait(
773 output,
774 CommunicationOperation::AllGatherUneven,
775 DistributedExecutionPhase::Execution,
776 None,
777 )?;
778 self.validate_tensor(&output, requirement, true)
779 .and_then(|()| self.validate_output_shape(&output, expected))
780 .map_err(|error| {
781 self.output_contract_error(
782 error,
783 CommunicationOperation::AllGatherUneven,
784 DistributedExecutionPhase::Execution,
785 None,
786 )
787 })?;
788 Ok(output)
789 }
790
791 pub fn variable_all_to_all(
793 &self,
794 value: B::Tensor,
795 counts: &CommunicationPeerCounts,
796 axis: usize,
797 group: CollectiveGroupId,
798 executor: &B::Executor,
799 ) -> Result<B::Tensor, PartitionExecutionError>
800 where
801 B: VariableAllToAllBackend,
802 {
803 let (descriptor, native) = self.group(group, CommunicationOperation::VariableAllToAll)?;
804 if counts.group_size() != descriptor.members().len() {
805 return Err(PartitionExecutionError::PeerCount {
806 expected: descriptor.members().len(),
807 actual: counts.group_size(),
808 });
809 }
810 let requirement =
811 Self::group_requirement(descriptor, CommunicationOperation::VariableAllToAll);
812 self.validate_tensor(&value, requirement, false)?;
813 let max = requirement
814 .limits()
815 .and_then(|limits| limits.max_count_per_peer())
816 .ok_or(PartitionExecutionError::MissingTensorLimits)?;
817 if counts
818 .send()
819 .iter()
820 .chain(counts.receive())
821 .any(|count| *count > max)
822 {
823 return Err(PartitionExecutionError::PeerCountLimit { maximum: max });
824 }
825 let output_width = counts.receive().iter().try_fold(0usize, |total, count| {
826 total
827 .checked_add(*count)
828 .ok_or(PartitionExecutionError::CommunicationShapeOverflow)
829 })?;
830 let expected = self.expected_axis_output_shape(&value, axis, output_width)?;
831 let output =
832 B::variable_all_to_all(value, counts, axis, native, executor).map_err(|error| {
833 self.submission_error(
834 error,
835 CommunicationOperation::VariableAllToAll,
836 DistributedExecutionPhase::Execution,
837 None,
838 )
839 })?;
840 let output = self.wait(
841 output,
842 CommunicationOperation::VariableAllToAll,
843 DistributedExecutionPhase::Execution,
844 None,
845 )?;
846 self.validate_tensor(&output, requirement, true)
847 .and_then(|()| self.validate_output_shape(&output, expected))
848 .map_err(|error| {
849 self.output_contract_error(
850 error,
851 CommunicationOperation::VariableAllToAll,
852 DistributedExecutionPhase::Execution,
853 None,
854 )
855 })?;
856 Ok(output)
857 }
858
859 fn transfer_boundary(
860 &self,
861 route: CommunicationRouteId,
862 values: Vec<crate::ArchitectureBoundaryValue<B::Tensor>>,
863 schema: &ResolvedBoundaryWireSchema,
864 wire: PipelineWireContract,
865 executor: &B::Executor,
866 ) -> Result<Vec<B::Tensor>, PartitionExecutionError>
867 where
868 B: PointToPointBackend,
869 {
870 let (descriptor, native) = self.route(route)?;
871 let rank = self.manifest.rank();
872 if rank != descriptor.source() && rank != descriptor.destination() {
873 return Err(PartitionExecutionError::NotRouteEndpoint(route));
874 }
875 validate_tagged_boundary_bundle::<B, I>(&self.inspector, &values, schema, wire)?;
876 let tensors = values
877 .iter()
878 .map(crate::ArchitectureBoundaryValue::tensor)
879 .cloned()
880 .collect::<Vec<_>>();
881 validate_bundle_requirement::<B, I>(
882 &self.inspector,
883 &tensors,
884 descriptor.requirement(),
885 false,
886 )?;
887 let contract = descriptor.boundary_contract().ok_or_else(|| {
888 PartitionExecutionError::BoundaryFraming(
889 "point-to-point boundary route has no role-exact framing contract".into(),
890 )
891 })?;
892 if contract.schema() != schema.identity() {
893 return Err(PartitionExecutionError::BoundaryFraming(format!(
894 "route schema {:?} differs from execution schema {:?}",
895 contract.schema(),
896 schema.identity()
897 )));
898 }
899 let actual_roles = resolved_tagged_boundary_roles(&values, schema, wire)?;
900 let values = contract
901 .frame_values(
902 route,
903 &actual_roles,
904 values
905 .into_iter()
906 .map(crate::ArchitectureBoundaryValue::into_parts)
907 .map(|(_, tensor)| tensor)
908 .collect(),
909 )
910 .map_err(|error| PartitionExecutionError::BoundaryFraming(error.to_string()))?;
911 let submission = B::send_receive(values, native, executor).map_err(|error| {
912 self.submission_error(
913 error,
914 CommunicationOperation::SendReceive,
915 DistributedExecutionPhase::Execution,
916 Some(route),
917 )
918 })?;
919 let output = self.wait(
920 submission,
921 CommunicationOperation::SendReceive,
922 DistributedExecutionPhase::Execution,
923 Some(route),
924 )?;
925 validate_boundary_bundle::<B, I>(&self.inspector, &output, schema, wire)
926 .and_then(|()| {
927 validate_bundle_requirement::<B, I>(
928 &self.inspector,
929 &output,
930 descriptor.requirement(),
931 true,
932 )
933 })
934 .map_err(|error| {
935 self.output_contract_error(
936 error,
937 CommunicationOperation::SendReceive,
938 DistributedExecutionPhase::Execution,
939 Some(route),
940 )
941 })?;
942 Ok(output)
943 }
944
945 fn boundary_endpoint_is_source(
946 &self,
947 route: CommunicationRouteId,
948 ) -> Result<bool, PartitionExecutionError> {
949 let descriptor = self
950 .manifest
951 .routes()
952 .iter()
953 .find(|candidate| candidate.id() == route)
954 .ok_or(PartitionExecutionError::UnknownRoute(route))?;
955 let rank = self.manifest.rank();
956 if rank != descriptor.source() && rank != descriptor.destination() {
957 return Err(PartitionExecutionError::NotRouteEndpoint(route));
958 }
959 Ok(rank == descriptor.source())
960 }
961
962 fn validate_prepared_boundary(
963 &self,
964 route: CommunicationRouteId,
965 values: &[crate::ArchitectureBoundaryValue<B::Tensor>],
966 schema: &ResolvedBoundaryWireSchema,
967 wire: PipelineWireContract,
968 ) -> Result<(), PartitionExecutionError> {
969 let descriptor = self
970 .manifest
971 .routes()
972 .iter()
973 .find(|candidate| candidate.id() == route)
974 .ok_or(PartitionExecutionError::UnknownRoute(route))?;
975 let rank = self.manifest.rank();
976 if rank != descriptor.source() && rank != descriptor.destination() {
977 return Err(PartitionExecutionError::NotRouteEndpoint(route));
978 }
979 validate_tagged_boundary_bundle::<B, I>(&self.inspector, values, schema, wire)?;
980 let contract = descriptor.boundary_contract().ok_or_else(|| {
981 PartitionExecutionError::BoundaryFraming(
982 "point-to-point boundary route has no role-exact framing contract".into(),
983 )
984 })?;
985 if contract.schema() != schema.identity() {
986 return Err(PartitionExecutionError::BoundaryFraming(format!(
987 "route schema {:?} differs from prepared schema {:?}",
988 contract.schema(),
989 schema.identity(),
990 )));
991 }
992 contract
993 .validate_invocation(&resolved_tagged_boundary_roles(values, schema, wire)?)
994 .map_err(|error| PartitionExecutionError::BoundaryFraming(error.to_string()))?;
995 let tensors = values
996 .iter()
997 .map(crate::ArchitectureBoundaryValue::tensor)
998 .cloned()
999 .collect::<Vec<_>>();
1000 validate_bundle_requirement::<B, I>(
1001 &self.inspector,
1002 &tensors,
1003 descriptor.requirement(),
1004 false,
1005 )
1006 }
1007
1008 fn broadcast_output(
1009 &self,
1010 value: B::Tensor,
1011 publication: PartitionOutputPublication,
1012 phase: DistributedExecutionPhase,
1013 executor: &B::Executor,
1014 ) -> Result<B::Tensor, PartitionExecutionError>
1015 where
1016 B: BroadcastBackend,
1017 {
1018 let (descriptor, native) =
1019 self.group(publication.group, CommunicationOperation::Broadcast)?;
1020 let root = descriptor
1021 .members()
1022 .iter()
1023 .position(|rank| *rank == publication.owner_rank)
1024 .ok_or(PartitionExecutionError::OutputOwnerNotMember {
1025 rank: publication.owner_rank,
1026 group: publication.group,
1027 })?;
1028 let requirement = Self::group_requirement(descriptor, CommunicationOperation::Broadcast);
1029 self.validate_tensor(&value, requirement, false)?;
1030 let submission = B::broadcast(value, root, native, executor).map_err(|error| {
1031 self.submission_error(error, CommunicationOperation::Broadcast, phase, None)
1032 })?;
1033 let output = self.wait(submission, CommunicationOperation::Broadcast, phase, None)?;
1034 self.validate_tensor(&output, requirement, true)
1035 .map_err(|error| {
1036 self.output_contract_error(error, CommunicationOperation::Broadcast, phase, None)
1037 })?;
1038 Ok(output)
1039 }
1040
1041 fn barrier(
1042 &self,
1043 group: CollectiveGroupId,
1044 executor: &B::Executor,
1045 ) -> Result<(), PartitionExecutionError>
1046 where
1047 B: BarrierBackend,
1048 {
1049 let (_, native) = self.group(group, CommunicationOperation::Barrier)?;
1050 let completion = B::barrier(native, executor).map_err(|error| {
1051 self.submission_error(
1052 error,
1053 CommunicationOperation::Barrier,
1054 DistributedExecutionPhase::Commit,
1055 None,
1056 )
1057 })?;
1058 let policy = self
1059 .manifest
1060 .completion_policy()
1061 .expect("partition communication requires bounded completion")
1062 .bounded_wait();
1063 let outcome = match completion.wait_bounded(policy) {
1064 Ok(outcome) => outcome,
1065 Err(error) => {
1066 self.authority.mark_poisoned(CommunicationPoison {
1067 operation: CommunicationOperation::Barrier,
1068 phase: DistributedExecutionPhase::Commit,
1069 route: None,
1070 cancellation: policy.cancellation(),
1071 });
1072 return Err(PartitionExecutionError::CommunicationCompletionFailed {
1073 operation: CommunicationOperation::Barrier,
1074 phase: DistributedExecutionPhase::Commit,
1075 route: None,
1076 error: error.to_string(),
1077 });
1078 }
1079 };
1080 match outcome {
1081 BoundedCompletionOutcome::Completed => Ok(()),
1082 BoundedCompletionOutcome::DeadlineExceeded { cancellation } => {
1083 self.authority.mark_poisoned(CommunicationPoison {
1084 operation: CommunicationOperation::Barrier,
1085 phase: DistributedExecutionPhase::Commit,
1086 route: None,
1087 cancellation,
1088 });
1089 Err(PartitionExecutionError::CommunicationDeadlineExceeded {
1090 operation: CommunicationOperation::Barrier,
1091 phase: DistributedExecutionPhase::Commit,
1092 route: None,
1093 cancellation,
1094 })
1095 }
1096 }
1097 }
1098
1099 fn agree_success(
1100 &self,
1101 local_success: bool,
1102 group: CollectiveGroupId,
1103 phase: DistributedExecutionPhase,
1104 executor: &B::Executor,
1105 ) -> Result<bool, PartitionExecutionError>
1106 where
1107 B: FailureAgreementBackend,
1108 {
1109 let (_, native) = self.group(group, CommunicationOperation::FailureAgreement)?;
1110 let output = self.wait(
1111 B::agree_success(local_success, native, executor).map_err(|error| {
1112 self.submission_error(error, CommunicationOperation::FailureAgreement, phase, None)
1113 })?,
1114 CommunicationOperation::FailureAgreement,
1115 phase,
1116 None,
1117 )?;
1118 B::resolve_failure_agreement(output).map_err(|error| {
1119 self.authority.mark_poisoned(CommunicationPoison {
1120 operation: CommunicationOperation::FailureAgreement,
1121 phase,
1122 route: None,
1123 cancellation: self
1124 .manifest
1125 .completion_policy()
1126 .expect("partition communication requires bounded completion")
1127 .cancellation(),
1128 });
1129 PartitionExecutionError::CommunicationCompletionFailed {
1130 operation: CommunicationOperation::FailureAgreement,
1131 phase,
1132 route: None,
1133 error: error.to_string(),
1134 }
1135 })
1136 }
1137
1138 fn agree_success_after_prior_failure(
1139 &self,
1140 local_success: bool,
1141 group: CollectiveGroupId,
1142 phase: DistributedExecutionPhase,
1143 executor: &B::Executor,
1144 ) -> Result<bool, PartitionExecutionError>
1145 where
1146 B: FailureAgreementBackend,
1147 {
1148 if local_success && !self.authority.is_poisoned() {
1149 return Err(PartitionExecutionError::RecoveryAgreementWithoutFailure { phase });
1150 }
1151 let index = self
1152 .manifest
1153 .groups()
1154 .iter()
1155 .position(|candidate| candidate.id() == group)
1156 .ok_or(PartitionExecutionError::UnknownGroup(group))?;
1157 let descriptor = &self.manifest.groups()[index];
1158 if descriptor.local_index().is_none() {
1159 return Err(PartitionExecutionError::NotGroupMember(group));
1160 }
1161 if !descriptor
1162 .requirements()
1163 .operations()
1164 .iter()
1165 .any(|requirement| requirement.operation() == CommunicationOperation::FailureAgreement)
1166 {
1167 return Err(PartitionExecutionError::OperationNotSelected {
1168 resource: format!("group {}", group.value()),
1169 operation: CommunicationOperation::FailureAgreement,
1170 });
1171 }
1172 let native = self.groups[index].resource.borrow();
1173 let submission = B::agree_success(local_success, native, executor).map_err(|error| {
1174 self.authority.mark_poisoned(CommunicationPoison {
1175 operation: CommunicationOperation::FailureAgreement,
1176 phase,
1177 route: None,
1178 cancellation: self
1179 .manifest
1180 .completion_policy()
1181 .expect("partition communication requires bounded completion")
1182 .cancellation(),
1183 });
1184 PartitionExecutionError::CommunicationSubmissionFailed {
1185 operation: CommunicationOperation::FailureAgreement,
1186 phase,
1187 route: None,
1188 error: error.to_string(),
1189 }
1190 })?;
1191 let output = self.authority.wait_after_prior_failure(
1192 submission,
1193 CommunicationOperation::FailureAgreement,
1194 phase,
1195 None,
1196 )?;
1197 B::resolve_failure_agreement(output).map_err(|error| {
1198 self.authority.mark_poisoned(CommunicationPoison {
1199 operation: CommunicationOperation::FailureAgreement,
1200 phase,
1201 route: None,
1202 cancellation: self
1203 .manifest
1204 .completion_policy()
1205 .expect("partition communication requires bounded completion")
1206 .cancellation(),
1207 });
1208 PartitionExecutionError::CommunicationCompletionFailed {
1209 operation: CommunicationOperation::FailureAgreement,
1210 phase,
1211 route: None,
1212 error: error.to_string(),
1213 }
1214 })
1215 }
1216
1217 fn complete_local_dependencies(
1218 &self,
1219 values: &[crate::ArchitectureBoundaryValue<B::Tensor>],
1220 route: CommunicationRouteId,
1221 executor: &B::Executor,
1222 before_failure_agreement: bool,
1223 ) -> Result<(), PartitionExecutionError> {
1224 self.ensure_active()?;
1225 let descriptor = self
1226 .manifest
1227 .routes()
1228 .iter()
1229 .find(|candidate| candidate.id() == route)
1230 .ok_or(PartitionExecutionError::UnknownRoute(route))?;
1231 if descriptor.source() != self.manifest.rank() {
1232 return Err(PartitionExecutionError::NotRouteEndpoint(route));
1233 }
1234 let phase = DistributedExecutionPhase::BoundarySourceCompletion(route);
1235 let submission = B::submit_local_dependencies(
1236 values.iter().map(crate::ArchitectureBoundaryValue::tensor),
1237 executor,
1238 )
1239 .map_err(|error| {
1240 if before_failure_agreement {
1241 PartitionExecutionError::CommunicationSubmissionFailed {
1242 operation: CommunicationOperation::SendReceive,
1243 phase,
1244 route: Some(route),
1245 error: error.to_string(),
1246 }
1247 } else {
1248 self.submission_error(
1249 error,
1250 CommunicationOperation::SendReceive,
1251 phase,
1252 Some(route),
1253 )
1254 }
1255 })?;
1256 if before_failure_agreement {
1257 self.authority.wait_before_failure_agreement(
1258 submission,
1259 CommunicationOperation::SendReceive,
1260 phase,
1261 Some(route),
1262 )
1263 } else {
1264 self.wait(
1265 submission,
1266 CommunicationOperation::SendReceive,
1267 phase,
1268 Some(route),
1269 )
1270 }
1271 }
1272
1273 pub fn complete_execution_dependencies<'a, V>(
1281 &self,
1282 values: V,
1283 executor: &B::Executor,
1284 ) -> Result<(), PartitionExecutionError>
1285 where
1286 V: IntoIterator<Item = &'a B::Tensor>,
1287 B::Tensor: 'a,
1288 {
1289 self.ensure_active()?;
1290 let submission = B::submit_local_dependencies(values, executor).map_err(|error| {
1291 self.submission_error(
1292 error,
1293 CommunicationOperation::SendReceive,
1294 DistributedExecutionPhase::Execution,
1295 None,
1296 )
1297 })?;
1298 self.wait(
1299 submission,
1300 CommunicationOperation::SendReceive,
1301 DistributedExecutionPhase::Execution,
1302 None,
1303 )
1304 }
1305}
1306
1307#[derive(Debug, Clone, Copy, Eq, PartialEq)]
1309#[non_exhaustive]
1310pub enum DistributedExecutionPhase {
1311 StateCheckpoint,
1313 PromptCacheLoadPreflight,
1315 PromptCacheLoadPreparation,
1317 PromptCacheSavePreflight,
1319 PromptCacheSavePreparation,
1321 PromptCacheSavePublication,
1323 SessionCheckpoint,
1325 SessionResetPreparation,
1327 SessionRollbackPreparation,
1329 BoundarySourceCompletion(CommunicationRouteId),
1331 BoundarySourceReady(CommunicationRouteId),
1333 Execution,
1335 OutputObservation,
1337 PredictionTargetCapture,
1339 PredictionTargetCapturePublication,
1341 PredictionTargetStatePreparation,
1343 PredictionExtensionExecution,
1345 OutputPublication,
1347 SamplingSynchronization,
1349 MechanismCompletion,
1351 Commit,
1353}
1354
1355#[derive(Debug, Clone, Eq, PartialEq)]
1357pub struct PartitionBoundaryRoute {
1358 pub source_group: usize,
1360 pub destination_group: usize,
1362 pub source_rank: usize,
1364 pub destination_rank: usize,
1366 pub route: CommunicationRouteId,
1368}
1369
1370#[derive(Debug, Clone, Copy)]
1372pub struct PartitionOutputPublication {
1373 pub group: CollectiveGroupId,
1375 pub owner_rank: usize,
1377}
1378
1379#[derive(Debug, Clone, Copy, Eq, PartialEq)]
1381pub struct PartitionOutputAuthority {
1382 owner_rank: usize,
1383 owner_group_rank: usize,
1384 local_public_output: bool,
1385}
1386
1387impl PartitionOutputAuthority {
1388 pub const fn owner_rank(self) -> usize {
1390 self.owner_rank
1391 }
1392
1393 pub const fn owner_group_rank(self) -> usize {
1395 self.owner_group_rank
1396 }
1397
1398 pub const fn local_public_output(self) -> bool {
1400 self.local_public_output
1401 }
1402}
1403
1404#[derive(Debug, Clone)]
1406pub struct PartitionedExecutionPlan {
1407 graph: ExecutionGraph,
1408 group_contracts: Vec<(ArchitectureGroupKind, bool)>,
1409 drivers: Vec<Option<LayeredPartitionDriver>>,
1410 routes: Vec<PartitionBoundaryRoute>,
1411 publication: Option<PartitionOutputPublication>,
1412 commit_barrier: Option<CollectiveGroupId>,
1413 wire: PipelineWireContract,
1414}
1415
1416impl PartitionedExecutionPlan {
1417 #[allow(clippy::too_many_arguments)]
1419 pub fn new(
1420 graph: ExecutionGraph,
1421 group_contracts: Vec<(ArchitectureGroupKind, bool)>,
1422 drivers: Vec<Option<LayeredPartitionDriver>>,
1423 routes: Vec<PartitionBoundaryRoute>,
1424 publication: Option<PartitionOutputPublication>,
1425 commit_barrier: Option<CollectiveGroupId>,
1426 wire: PipelineWireContract,
1427 ) -> Result<Self, PartitionExecutionError> {
1428 if group_contracts.len() != graph.groups().len() || drivers.len() != graph.groups().len() {
1429 return Err(PartitionExecutionError::GroupCount {
1430 graph: graph.groups().len(),
1431 contracts: group_contracts.len(),
1432 drivers: drivers.len(),
1433 });
1434 }
1435 for (group, driver) in drivers.iter().enumerate() {
1436 if driver
1437 .as_ref()
1438 .is_some_and(|driver| driver.group_index() != group)
1439 {
1440 return Err(PartitionExecutionError::DriverGroup { group });
1441 }
1442 }
1443 for route in &routes {
1444 if route.source_group >= graph.groups().len()
1445 || route.destination_group >= graph.groups().len()
1446 {
1447 return Err(PartitionExecutionError::RouteDependency(route.route));
1448 }
1449 if route.source_group != route.destination_group {
1450 let dependencies = graph
1451 .dependencies(route.destination_group)
1452 .ok_or(PartitionExecutionError::RouteGroup(route.route))?;
1453 if !dependencies.contains(&route.source_group) {
1454 return Err(PartitionExecutionError::RouteDependency(route.route));
1455 }
1456 }
1457 }
1458 Ok(Self {
1459 graph,
1460 group_contracts,
1461 drivers,
1462 routes,
1463 publication,
1464 commit_barrier,
1465 wire,
1466 })
1467 }
1468
1469 fn validate_manifest(
1470 &self,
1471 manifest: &CommunicationManifest,
1472 ) -> Result<(), PartitionExecutionError> {
1473 if self.routes.len() != manifest.routes().len() {
1474 return Err(PartitionExecutionError::RouteCount {
1475 plan: self.routes.len(),
1476 manifest: manifest.routes().len(),
1477 });
1478 }
1479 for (planned, descriptor) in self.routes.iter().zip(manifest.routes()) {
1480 if planned.route != descriptor.id()
1481 || planned.source_rank != descriptor.source()
1482 || planned.destination_rank != descriptor.destination()
1483 {
1484 return Err(PartitionExecutionError::RouteDescriptorMismatch {
1485 route: planned.route,
1486 planned_source: planned.source_rank,
1487 planned_destination: planned.destination_rank,
1488 manifest_source: descriptor.source(),
1489 manifest_destination: descriptor.destination(),
1490 });
1491 }
1492 let rank = manifest.rank();
1493 if rank == planned.source_rank && self.drivers[planned.source_group].is_none() {
1494 return Err(PartitionExecutionError::RouteOwnerMissing {
1495 route: planned.route,
1496 rank,
1497 group: planned.source_group,
1498 });
1499 }
1500 if rank == planned.destination_rank && self.drivers[planned.destination_group].is_none()
1501 {
1502 return Err(PartitionExecutionError::RouteOwnerMissing {
1503 route: planned.route,
1504 rank,
1505 group: planned.destination_group,
1506 });
1507 }
1508 }
1509 if let Some(publication) = self.publication {
1510 let descriptor = manifest
1511 .groups()
1512 .iter()
1513 .find(|descriptor| descriptor.id() == publication.group)
1514 .ok_or(PartitionExecutionError::UnknownGroup(publication.group))?;
1515 if !descriptor.members().contains(&publication.owner_rank) {
1516 return Err(PartitionExecutionError::OutputOwnerNotMember {
1517 rank: publication.owner_rank,
1518 group: publication.group,
1519 });
1520 }
1521 let local_owns_output = self
1522 .drivers
1523 .iter()
1524 .flatten()
1525 .any(LayeredPartitionDriver::owns_output);
1526 if manifest.rank() == publication.owner_rank && !local_owns_output {
1527 return Err(PartitionExecutionError::OutputOwnership {
1528 rank: manifest.rank(),
1529 owner: publication.owner_rank,
1530 });
1531 }
1532 Self::validate_exact_group_operation(descriptor, CommunicationOperation::Broadcast)?;
1533 }
1534 Ok(())
1535 }
1536
1537 fn validate_exact_group_operation(
1538 descriptor: &CommunicationGroupDescriptor,
1539 operation: CommunicationOperation,
1540 ) -> Result<(), PartitionExecutionError> {
1541 let selected = descriptor
1542 .requirements()
1543 .operations()
1544 .iter()
1545 .find(|requirement| requirement.operation() == operation)
1546 .ok_or_else(|| PartitionExecutionError::OperationNotSelected {
1547 resource: format!("group {}", descriptor.id().value()),
1548 operation,
1549 })?;
1550 if !selected.exact_completion() {
1551 return Err(PartitionExecutionError::InexactOperationRequirement {
1552 group: descriptor.id(),
1553 operation,
1554 });
1555 }
1556 Ok(())
1557 }
1558
1559 fn validate_exact_commit_agreement(
1560 &self,
1561 manifest: &CommunicationManifest,
1562 ) -> Result<(), PartitionExecutionError> {
1563 let group = self
1564 .commit_barrier
1565 .ok_or(PartitionExecutionError::CommunicationPolicyMismatch)?;
1566 let descriptor = manifest
1567 .groups()
1568 .iter()
1569 .find(|descriptor| descriptor.id() == group)
1570 .ok_or(PartitionExecutionError::UnknownGroup(group))?;
1571 Self::validate_exact_group_operation(descriptor, CommunicationOperation::FailureAgreement)
1572 }
1573
1574 pub fn drivers(&self) -> &[Option<LayeredPartitionDriver>] {
1576 &self.drivers
1577 }
1578
1579 pub fn routes(&self) -> &[PartitionBoundaryRoute] {
1581 &self.routes
1582 }
1583
1584 pub const fn publication(&self) -> Option<PartitionOutputPublication> {
1586 self.publication
1587 }
1588
1589 pub fn publication_authority(
1592 &self,
1593 manifest: &CommunicationManifest,
1594 ) -> Result<Option<PartitionOutputAuthority>, PartitionExecutionError> {
1595 let Some(publication) = self.publication else {
1596 return Ok(None);
1597 };
1598 let descriptor = manifest
1599 .groups()
1600 .iter()
1601 .find(|descriptor| descriptor.id() == publication.group)
1602 .ok_or(PartitionExecutionError::UnknownGroup(publication.group))?;
1603 let owner_group_rank = descriptor
1604 .members()
1605 .iter()
1606 .position(|rank| *rank == publication.owner_rank)
1607 .ok_or(PartitionExecutionError::OutputOwnerNotMember {
1608 rank: publication.owner_rank,
1609 group: publication.group,
1610 })?;
1611 Ok(Some(PartitionOutputAuthority {
1612 owner_rank: publication.owner_rank,
1613 owner_group_rank,
1614 local_public_output: manifest.rank() == publication.owner_rank,
1615 }))
1616 }
1617
1618 pub const fn commit_barrier(&self) -> Option<CollectiveGroupId> {
1620 self.commit_barrier
1621 }
1622
1623 fn selects_phase_failure_agreement(&self, manifest: &CommunicationManifest) -> bool {
1624 self.commit_barrier.is_some_and(|group| {
1625 manifest.groups().iter().any(|descriptor| {
1626 descriptor.id() == group
1627 && descriptor
1628 .requirements()
1629 .operations()
1630 .iter()
1631 .any(|requirement| {
1632 requirement.operation() == CommunicationOperation::FailureAgreement
1633 })
1634 })
1635 })
1636 }
1637}
1638
1639pub trait PartitionedGroupExecutor<A, B, S, G, R, I>
1646where
1647 B: CommunicationBackend,
1648 S: RuntimeState<B>,
1649 A: LayeredArchitecture<B, S>,
1650 G: Borrow<B::CommunicationGroup>,
1651 R: Borrow<B::CommunicationRoute>,
1652 I: CommunicationTensorMetadata<B>,
1653{
1654 type Pass<'a>;
1656
1657 fn begin<'a>(
1659 &mut self,
1660 input: A::Input<'a>,
1661 state: &mut S,
1662 pass: ExpertPass,
1663 context: &<B::Tensor as Tensor>::Context,
1664 ) -> Result<Self::Pass<'a>, A::Error>;
1665
1666 fn request_group_active(&self, pass: &Self::Pass<'_>, group: usize) -> Result<bool, A::Error>;
1668
1669 fn has_cross_stage_collective_waves(&self) -> bool {
1676 false
1677 }
1678
1679 #[allow(clippy::too_many_arguments)]
1681 fn execute_group<O: ActivationObserver<B::Tensor, A::Error> + ?Sized>(
1682 &mut self,
1683 pass: &mut Self::Pass<'_>,
1684 driver: &LayeredPartitionDriver,
1685 state: &mut S,
1686 communication: &PartitionCommunication<B, G, R, I>,
1687 communication_executor: &B::Executor,
1688 context: &<B::Tensor as Tensor>::Context,
1689 observer: &mut O,
1690 ) -> Result<(), A::Error>;
1691
1692 #[allow(clippy::too_many_arguments)]
1700 fn execute_pipeline_wave<O: ActivationObserver<B::Tensor, A::Error> + ?Sized>(
1701 &mut self,
1702 pass: &mut Self::Pass<'_>,
1703 group: usize,
1704 driver: Option<&LayeredPartitionDriver>,
1705 active: bool,
1706 _wave: usize,
1707 state: &mut S,
1708 communication: &PartitionCommunication<B, G, R, I>,
1709 communication_executor: &B::Executor,
1710 context: &<B::Tensor as Tensor>::Context,
1711 observer: &mut O,
1712 ) -> Result<(), A::Error> {
1713 if active {
1714 let driver = driver.expect("validated active pipeline rank owns a partition driver");
1715 assert_eq!(
1716 driver.group_index(),
1717 group,
1718 "active pipeline driver differs from the scheduled architecture group"
1719 );
1720 self.execute_group(
1721 pass,
1722 driver,
1723 state,
1724 communication,
1725 communication_executor,
1726 context,
1727 observer,
1728 )
1729 } else {
1730 Ok(())
1731 }
1732 }
1733
1734 fn boundary_values(
1736 &mut self,
1737 pass: &mut Self::Pass<'_>,
1738 route: &PartitionBoundaryRoute,
1739 schema: &ResolvedBoundaryWireSchema,
1740 source: bool,
1741 context: &<B::Tensor as Tensor>::Context,
1742 ) -> Result<Vec<crate::ArchitectureBoundaryValue<B::Tensor>>, A::Error>;
1743
1744 fn boundary_schema(
1746 &self,
1747 pass: &Self::Pass<'_>,
1748 route: &PartitionBoundaryRoute,
1749 ) -> Result<ResolvedBoundaryWireSchema, A::Error>;
1750
1751 fn accept_boundary(
1753 &mut self,
1754 pass: &mut Self::Pass<'_>,
1755 route: &PartitionBoundaryRoute,
1756 values: Vec<B::Tensor>,
1757 ) -> Result<(), A::Error>;
1758
1759 fn finish(
1764 &mut self,
1765 pass: Self::Pass<'_>,
1766 state: &mut S,
1767 context: &<B::Tensor as Tensor>::Context,
1768 ) -> Result<(B::Tensor, A::ForwardContext), A::Error>;
1769
1770 fn prediction_target_capture(
1776 &mut self,
1777 forward: &A::ForwardContext,
1778 _context: &<B::Tensor as Tensor>::Context,
1779 ) -> Result<Option<B::Tensor>, A::Error> {
1780 Ok(<A as crate::LayeredArchitecture<B, S>>::prediction_target_capture(forward).cloned())
1781 }
1782
1783 fn apply_prediction_target_operation<O>(
1785 &mut self,
1786 _state: &mut S,
1787 _operation: O,
1788 _context: &<B::Tensor as Tensor>::Context,
1789 ) -> Result<Option<O::Output>, A::Error>
1790 where
1791 O: crate::PredictionTargetOperation<A, B, S>,
1792 {
1793 Ok(None)
1794 }
1795}
1796
1797pub trait PartitionBoundaryTransport<B, G, R, I>
1799where
1800 B: CommunicationBackend,
1801 G: Borrow<B::CommunicationGroup>,
1802 R: Borrow<B::CommunicationRoute>,
1803 I: CommunicationTensorMetadata<B>,
1804{
1805 const ENABLED: bool;
1807
1808 #[allow(clippy::too_many_arguments)]
1810 fn transfer(
1811 &mut self,
1812 communication: &PartitionCommunication<B, G, R, I>,
1813 route: CommunicationRouteId,
1814 values: Vec<crate::ArchitectureBoundaryValue<B::Tensor>>,
1815 schema: &ResolvedBoundaryWireSchema,
1816 wire: PipelineWireContract,
1817 executor: &B::Executor,
1818 ) -> Result<Vec<B::Tensor>, PartitionExecutionError>;
1819}
1820
1821#[derive(Debug, Default, Clone, Copy)]
1823pub struct OpaqueBoundaryTransport;
1824
1825impl<B, G, R, I> PartitionBoundaryTransport<B, G, R, I> for OpaqueBoundaryTransport
1826where
1827 B: CommunicationBackend + PointToPointBackend,
1828 G: Borrow<B::CommunicationGroup>,
1829 R: Borrow<B::CommunicationRoute>,
1830 I: CommunicationTensorMetadata<B>,
1831{
1832 const ENABLED: bool = true;
1833
1834 fn transfer(
1835 &mut self,
1836 communication: &PartitionCommunication<B, G, R, I>,
1837 route: CommunicationRouteId,
1838 values: Vec<crate::ArchitectureBoundaryValue<B::Tensor>>,
1839 schema: &ResolvedBoundaryWireSchema,
1840 wire: PipelineWireContract,
1841 executor: &B::Executor,
1842 ) -> Result<Vec<B::Tensor>, PartitionExecutionError> {
1843 communication.transfer_boundary(route, values, schema, wire, executor)
1844 }
1845}
1846
1847#[derive(Debug, Default, Clone, Copy)]
1849pub struct NoBoundaryTransport;
1850
1851impl<B, G, R, I> PartitionBoundaryTransport<B, G, R, I> for NoBoundaryTransport
1852where
1853 B: CommunicationBackend,
1854 G: Borrow<B::CommunicationGroup>,
1855 R: Borrow<B::CommunicationRoute>,
1856 I: CommunicationTensorMetadata<B>,
1857{
1858 const ENABLED: bool = false;
1859
1860 fn transfer(
1861 &mut self,
1862 _communication: &PartitionCommunication<B, G, R, I>,
1863 _route: CommunicationRouteId,
1864 _values: Vec<crate::ArchitectureBoundaryValue<B::Tensor>>,
1865 _schema: &ResolvedBoundaryWireSchema,
1866 _wire: PipelineWireContract,
1867 _executor: &B::Executor,
1868 ) -> Result<Vec<B::Tensor>, PartitionExecutionError> {
1869 Err(PartitionExecutionError::OperationPolicyUnavailable(
1870 CommunicationOperation::SendReceive,
1871 ))
1872 }
1873}
1874
1875pub trait PartitionOutputPublisher<B, G, R, I>
1877where
1878 B: CommunicationBackend,
1879 G: Borrow<B::CommunicationGroup>,
1880 R: Borrow<B::CommunicationRoute>,
1881 I: CommunicationTensorMetadata<B>,
1882{
1883 const ENABLED: bool;
1885
1886 fn publish(
1888 &mut self,
1889 communication: &PartitionCommunication<B, G, R, I>,
1890 value: B::Tensor,
1891 publication: PartitionOutputPublication,
1892 phase: DistributedExecutionPhase,
1893 executor: &B::Executor,
1894 ) -> Result<B::Tensor, PartitionExecutionError>;
1895}
1896
1897#[derive(Debug, Default, Clone, Copy)]
1899pub struct OpaqueOutputPublisher;
1900
1901impl<B, G, R, I> PartitionOutputPublisher<B, G, R, I> for OpaqueOutputPublisher
1902where
1903 B: CommunicationBackend + BroadcastBackend,
1904 G: Borrow<B::CommunicationGroup>,
1905 R: Borrow<B::CommunicationRoute>,
1906 I: CommunicationTensorMetadata<B>,
1907{
1908 const ENABLED: bool = true;
1909
1910 fn publish(
1911 &mut self,
1912 communication: &PartitionCommunication<B, G, R, I>,
1913 value: B::Tensor,
1914 publication: PartitionOutputPublication,
1915 phase: DistributedExecutionPhase,
1916 executor: &B::Executor,
1917 ) -> Result<B::Tensor, PartitionExecutionError> {
1918 communication.broadcast_output(value, publication, phase, executor)
1919 }
1920}
1921
1922#[derive(Debug, Default, Clone, Copy)]
1924pub struct NoOutputPublisher;
1925
1926impl<B, G, R, I> PartitionOutputPublisher<B, G, R, I> for NoOutputPublisher
1927where
1928 B: CommunicationBackend,
1929 G: Borrow<B::CommunicationGroup>,
1930 R: Borrow<B::CommunicationRoute>,
1931 I: CommunicationTensorMetadata<B>,
1932{
1933 const ENABLED: bool = false;
1934
1935 fn publish(
1936 &mut self,
1937 _communication: &PartitionCommunication<B, G, R, I>,
1938 _value: B::Tensor,
1939 _publication: PartitionOutputPublication,
1940 _phase: DistributedExecutionPhase,
1941 _executor: &B::Executor,
1942 ) -> Result<B::Tensor, PartitionExecutionError> {
1943 Err(PartitionExecutionError::OperationPolicyUnavailable(
1944 CommunicationOperation::Broadcast,
1945 ))
1946 }
1947}
1948
1949pub trait PartitionCommitAgreement<B, G, R, I>
1951where
1952 B: CommunicationBackend,
1953 G: Borrow<B::CommunicationGroup>,
1954 R: Borrow<B::CommunicationRoute>,
1955 I: CommunicationTensorMetadata<B>,
1956{
1957 const ENABLED: bool;
1959
1960 const PHASE_FAILURE_AGREEMENT: bool = false;
1963
1964 fn agree_phase(
1969 &mut self,
1970 _communication: &PartitionCommunication<B, G, R, I>,
1971 _group: CollectiveGroupId,
1972 _phase: DistributedExecutionPhase,
1973 local_success: bool,
1974 _executor: &B::Executor,
1975 ) -> Result<bool, PartitionExecutionError> {
1976 Ok(local_success)
1977 }
1978
1979 fn agree_phase_after_prior_failure(
1985 &mut self,
1986 communication: &PartitionCommunication<B, G, R, I>,
1987 group: CollectiveGroupId,
1988 phase: DistributedExecutionPhase,
1989 local_success: bool,
1990 executor: &B::Executor,
1991 ) -> Result<bool, PartitionExecutionError> {
1992 self.agree_phase(communication, group, phase, local_success, executor)
1993 }
1994
1995 fn commit(
1997 &mut self,
1998 communication: &PartitionCommunication<B, G, R, I>,
1999 group: CollectiveGroupId,
2000 epoch: DistributedCommitEpoch,
2001 executor: &B::Executor,
2002 ) -> DistributedCommitOutcome;
2003}
2004
2005#[derive(Debug, Default, Clone, Copy)]
2007pub struct OpaqueCommitAgreement;
2008
2009impl<B, G, R, I> PartitionCommitAgreement<B, G, R, I> for OpaqueCommitAgreement
2010where
2011 B: CommunicationBackend + BarrierBackend,
2012 G: Borrow<B::CommunicationGroup>,
2013 R: Borrow<B::CommunicationRoute>,
2014 I: CommunicationTensorMetadata<B>,
2015{
2016 const ENABLED: bool = true;
2017
2018 fn commit(
2019 &mut self,
2020 communication: &PartitionCommunication<B, G, R, I>,
2021 group: CollectiveGroupId,
2022 epoch: DistributedCommitEpoch,
2023 executor: &B::Executor,
2024 ) -> DistributedCommitOutcome {
2025 match communication.barrier(group, executor) {
2026 Ok(()) => DistributedCommitOutcome::Committed(epoch),
2027 Err(error) => indeterminate_commit(epoch, &error),
2028 }
2029 }
2030}
2031
2032#[derive(Debug, Default, Clone, Copy)]
2037pub struct OpaqueFailureAgreement;
2038
2039impl<B, G, R, I> PartitionCommitAgreement<B, G, R, I> for OpaqueFailureAgreement
2040where
2041 B: CommunicationBackend + FailureAgreementBackend,
2042 G: Borrow<B::CommunicationGroup>,
2043 R: Borrow<B::CommunicationRoute>,
2044 I: CommunicationTensorMetadata<B>,
2045{
2046 const ENABLED: bool = true;
2047 const PHASE_FAILURE_AGREEMENT: bool = true;
2048
2049 fn agree_phase(
2050 &mut self,
2051 communication: &PartitionCommunication<B, G, R, I>,
2052 group: CollectiveGroupId,
2053 phase: DistributedExecutionPhase,
2054 local_success: bool,
2055 executor: &B::Executor,
2056 ) -> Result<bool, PartitionExecutionError> {
2057 communication.agree_success(local_success, group, phase, executor)
2058 }
2059
2060 fn agree_phase_after_prior_failure(
2061 &mut self,
2062 communication: &PartitionCommunication<B, G, R, I>,
2063 group: CollectiveGroupId,
2064 phase: DistributedExecutionPhase,
2065 local_success: bool,
2066 executor: &B::Executor,
2067 ) -> Result<bool, PartitionExecutionError> {
2068 communication.agree_success_after_prior_failure(local_success, group, phase, executor)
2069 }
2070
2071 fn commit(
2072 &mut self,
2073 communication: &PartitionCommunication<B, G, R, I>,
2074 group: CollectiveGroupId,
2075 epoch: DistributedCommitEpoch,
2076 executor: &B::Executor,
2077 ) -> DistributedCommitOutcome {
2078 match communication.agree_success(true, group, DistributedExecutionPhase::Commit, executor)
2079 {
2080 Ok(true) => DistributedCommitOutcome::Committed(epoch),
2081 Ok(false) => DistributedCommitOutcome::Aborted(epoch),
2082 Err(error) => indeterminate_commit(epoch, &error),
2083 }
2084 }
2085}
2086
2087#[derive(Debug, Default, Clone, Copy)]
2089pub struct NoCommitAgreement;
2090
2091impl<B, G, R, I> PartitionCommitAgreement<B, G, R, I> for NoCommitAgreement
2092where
2093 B: CommunicationBackend,
2094 G: Borrow<B::CommunicationGroup>,
2095 R: Borrow<B::CommunicationRoute>,
2096 I: CommunicationTensorMetadata<B>,
2097{
2098 const ENABLED: bool = false;
2099
2100 fn commit(
2101 &mut self,
2102 _communication: &PartitionCommunication<B, G, R, I>,
2103 _group: CollectiveGroupId,
2104 epoch: DistributedCommitEpoch,
2105 _executor: &B::Executor,
2106 ) -> DistributedCommitOutcome {
2107 DistributedCommitOutcome::Aborted(epoch)
2108 }
2109}
2110
2111fn indeterminate_commit(
2112 epoch: DistributedCommitEpoch,
2113 error: &PartitionExecutionError,
2114) -> DistributedCommitOutcome {
2115 let phase = match error {
2116 PartitionExecutionError::CommunicationSubmissionFailed { .. } => {
2117 DistributedCommitPhase::DecisionSubmission
2118 }
2119 PartitionExecutionError::CommunicationCompletionFailed { .. }
2120 | PartitionExecutionError::CommunicationDeadlineExceeded { .. }
2121 | PartitionExecutionError::CommunicationPoisoned { .. } => {
2122 DistributedCommitPhase::DecisionCompletion
2123 }
2124 _ => DistributedCommitPhase::DecisionObservation,
2125 };
2126 DistributedCommitOutcome::Indeterminate { epoch, phase }
2127}
2128
2129pub struct PartitionedTextRuntime<A, B, S, P, E, G, R, I, T, U, V>
2131where
2132 B: CommunicationBackend,
2133 G: Borrow<B::CommunicationGroup>,
2134 R: Borrow<B::CommunicationRoute>,
2135 I: CommunicationTensorMetadata<B>,
2136 T: PartitionBoundaryTransport<B, G, R, I>,
2137 U: PartitionOutputPublisher<B, G, R, I>,
2138 V: PartitionCommitAgreement<B, G, R, I>,
2139{
2140 plan: PartitionedExecutionPlan,
2141 executor: E,
2142 communication: PartitionCommunication<B, G, R, I>,
2143 communication_executor: B::OwnedExecutor,
2144 boundary_transport: T,
2145 output_publisher: U,
2146 commit_agreement: V,
2147 residency: ExecutionResidency,
2148 bounded_policy: Option<P>,
2149 marker: PhantomData<fn(A, S)>,
2150}
2151
2152#[derive(Debug, thiserror::Error)]
2161pub enum PartitionedTraversalError<ArchitectureError, PolicyError>
2162where
2163 ArchitectureError: std::fmt::Display,
2164 PolicyError: std::fmt::Display,
2165{
2166 #[error("partitioned traversal contract failed: {0}")]
2168 Contract(String),
2169 #[error(transparent)]
2171 Execution(#[from] LayerwiseRuntimeError<ArchitectureError, PolicyError>),
2172}
2173
2174pub type PartitionedTraversalResult<Tensor, ForwardContext, ArchitectureError, PolicyError> =
2176 Result<(Tensor, ForwardContext), PartitionedTraversalError<ArchitectureError, PolicyError>>;
2177
2178pub struct LayerwiseTraversalPartitionExecutor<A, B, S, P>
2185where
2186 B: SubmissionBackend<Executor = <<B as NeuralBackend>::Tensor as Tensor>::Context>,
2187 S: RuntimeState<B>,
2188 A: LayeredArchitecture<B, S>,
2189 P: LayerwisePolicy<B, A::Unit>,
2190 B::ParallelContext: Sized,
2191{
2192 runtime: LayerwiseRuntime<A, B, S, P>,
2193 parallel: B::ParallelContext,
2194}
2195
2196pub enum LayerwiseTraversalRuntime<Direct, Partitioned> {
2202 Direct(Direct),
2204 Partitioned(Partitioned),
2206}
2207
2208impl<Direct, Partitioned> LayerwiseTraversalRuntime<Direct, Partitioned> {
2209 pub const fn direct(runtime: Direct) -> Self {
2211 Self::Direct(runtime)
2212 }
2213
2214 pub const fn partitioned(runtime: Partitioned) -> Self {
2216 Self::Partitioned(runtime)
2217 }
2218}
2219
2220impl<A, B, S, P> LayerwiseTraversalPartitionExecutor<A, B, S, P>
2221where
2222 B: SubmissionBackend<Executor = <<B as NeuralBackend>::Tensor as Tensor>::Context>,
2223 S: RuntimeState<B>,
2224 A: LayeredArchitecture<B, S>,
2225 P: LayerwisePolicy<B, A::Unit>,
2226 B::ParallelContext: Sized,
2227{
2228 pub const fn new(runtime: LayerwiseRuntime<A, B, S, P>, parallel: B::ParallelContext) -> Self {
2230 Self { runtime, parallel }
2231 }
2232
2233 pub const fn runtime(&self) -> &LayerwiseRuntime<A, B, S, P> {
2235 &self.runtime
2236 }
2237}
2238
2239impl<A, B, S, P, E, G, R, I, T, U, V> PartitionedTextRuntime<A, B, S, P, E, G, R, I, T, U, V>
2240where
2241 B: CommunicationBackend,
2242 G: Borrow<B::CommunicationGroup>,
2243 R: Borrow<B::CommunicationRoute>,
2244 I: CommunicationTensorMetadata<B>,
2245 T: PartitionBoundaryTransport<B, G, R, I>,
2246 U: PartitionOutputPublisher<B, G, R, I>,
2247 V: PartitionCommitAgreement<B, G, R, I>,
2248{
2249 #[allow(
2251 clippy::too_many_arguments,
2252 reason = "constructor pairs independently selected runtime and communication policies"
2253 )]
2254 pub fn new(
2255 plan: PartitionedExecutionPlan,
2256 executor: E,
2257 communication: PartitionCommunication<B, G, R, I>,
2258 communication_executor: B::OwnedExecutor,
2259 boundary_transport: T,
2260 output_publisher: U,
2261 commit_agreement: V,
2262 residency: ExecutionResidency,
2263 bounded_policy: Option<P>,
2264 ) -> Result<Self, PartitionExecutionError> {
2265 plan.validate_manifest(communication.manifest())?;
2266 if V::PHASE_FAILURE_AGREEMENT {
2267 plan.validate_exact_commit_agreement(communication.manifest())?;
2268 }
2269 let bounded = !matches!(residency, ExecutionResidency::FullyResident);
2270 if bounded != bounded_policy.is_some() {
2271 return Err(PartitionExecutionError::ResidencyPolicyMismatch);
2272 }
2273 if (!T::ENABLED && !plan.routes.is_empty())
2274 || (!U::ENABLED && plan.publication.is_some())
2275 || (!V::ENABLED && plan.commit_barrier.is_some())
2276 || (V::PHASE_FAILURE_AGREEMENT
2277 != plan.selects_phase_failure_agreement(communication.manifest()))
2278 {
2279 return Err(PartitionExecutionError::CommunicationPolicyMismatch);
2280 }
2281 Ok(Self {
2282 plan,
2283 executor,
2284 communication,
2285 communication_executor,
2286 boundary_transport,
2287 output_publisher,
2288 commit_agreement,
2289 residency,
2290 bounded_policy,
2291 marker: PhantomData,
2292 })
2293 }
2294
2295 fn agree_phase(
2296 &mut self,
2297 phase: DistributedExecutionPhase,
2298 local_success: bool,
2299 ) -> Result<bool, PartitionExecutionError> {
2300 let Some(group) = self.plan.commit_barrier else {
2301 return Ok(local_success);
2302 };
2303 self.commit_agreement.agree_phase(
2304 &self.communication,
2305 group,
2306 phase,
2307 local_success,
2308 self.communication_executor.borrow(),
2309 )
2310 }
2311}
2312
2313impl<A, B, S, P, ExecutionPolicy, G, R, I, T, U, V>
2314 PartitionedTextRuntime<
2315 A,
2316 B,
2317 S,
2318 P,
2319 LayerwiseTraversalPartitionExecutor<A, B, S, ExecutionPolicy>,
2320 G,
2321 R,
2322 I,
2323 T,
2324 U,
2325 V,
2326 >
2327where
2328 B: CommunicationBackend<Executor = <<B as NeuralBackend>::Tensor as Tensor>::Context>,
2329 S: RuntimeState<B>,
2330 A: ParallelLayeredArchitecture<B, S>,
2331 ExecutionPolicy: LayerwisePolicy<B, A::Unit>,
2332 G: Borrow<B::CommunicationGroup>,
2333 R: Borrow<B::CommunicationRoute>,
2334 I: CommunicationTensorMetadata<B>,
2335 T: PartitionBoundaryTransport<B, G, R, I>,
2336 U: PartitionOutputPublisher<B, G, R, I>,
2337 V: PartitionCommitAgreement<B, G, R, I>,
2338 B::ParallelContext: Sized,
2339 A::Error: std::fmt::Display,
2340 ExecutionPolicy::Error: std::fmt::Display,
2341{
2342 pub fn forward_with_traversal_hook<'a, H>(
2349 &mut self,
2350 input: A::Input<'a>,
2351 state: &mut S,
2352 context: &<B::Tensor as Tensor>::Context,
2353 hook: &mut H,
2354 ) -> PartitionedTraversalResult<B::Tensor, A::ForwardContext, A::Error, ExecutionPolicy::Error>
2355 where
2356 H: LayeredTraversalHook<B, A::ForwardContext, A::Error> + ?Sized,
2357 {
2358 if !self.plan.routes.is_empty()
2359 || self.plan.publication.is_some()
2360 || self.plan.commit_barrier.is_some()
2361 || self.plan.drivers.iter().any(Option::is_none)
2362 {
2363 return Err(PartitionedTraversalError::Contract(
2364 "fully local traversal requires every group locally owned and no route, publication, or agreement policy"
2365 .into(),
2366 ));
2367 }
2368 self.executor
2369 .runtime
2370 .forward_parallel_with_traversal_hook(
2371 input,
2372 state,
2373 &self.executor.parallel,
2374 context,
2375 hook,
2376 )
2377 .map_err(PartitionedTraversalError::Execution)
2378 }
2379
2380 pub const fn traversal_executor(
2382 &self,
2383 ) -> &LayerwiseTraversalPartitionExecutor<A, B, S, ExecutionPolicy> {
2384 &self.executor
2385 }
2386}
2387
2388impl<A, B, S, ExecutionPolicy, G, R, I, T, U, V>
2389 LayerwiseTraversalRuntime<
2390 LayerwiseRuntime<A, B, S, ExecutionPolicy>,
2391 Box<
2392 PartitionedTextRuntime<
2393 A,
2394 B,
2395 S,
2396 (),
2397 LayerwiseTraversalPartitionExecutor<A, B, S, ExecutionPolicy>,
2398 G,
2399 R,
2400 I,
2401 T,
2402 U,
2403 V,
2404 >,
2405 >,
2406 >
2407where
2408 B: CommunicationBackend<Executor = <<B as NeuralBackend>::Tensor as Tensor>::Context>,
2409 S: RuntimeState<B>,
2410 A: ParallelLayeredArchitecture<B, S>,
2411 ExecutionPolicy: LayerwisePolicy<B, A::Unit>,
2412 G: Borrow<B::CommunicationGroup>,
2413 R: Borrow<B::CommunicationRoute>,
2414 I: CommunicationTensorMetadata<B>,
2415 T: PartitionBoundaryTransport<B, G, R, I>,
2416 U: PartitionOutputPublisher<B, G, R, I>,
2417 V: PartitionCommitAgreement<B, G, R, I>,
2418 B::ParallelContext: Sized,
2419 A::Error: std::fmt::Display,
2420 ExecutionPolicy::Error: std::fmt::Display,
2421{
2422 pub fn forward_with_traversal_hook<'a, H>(
2425 &mut self,
2426 input: A::Input<'a>,
2427 state: &mut S,
2428 context: &<B::Tensor as Tensor>::Context,
2429 hook: &mut H,
2430 ) -> PartitionedTraversalResult<B::Tensor, A::ForwardContext, A::Error, ExecutionPolicy::Error>
2431 where
2432 H: LayeredTraversalHook<B, A::ForwardContext, A::Error> + ?Sized,
2433 {
2434 match self {
2435 Self::Direct(runtime) => runtime
2436 .forward_with_traversal_hook(input, state, context, hook)
2437 .map_err(PartitionedTraversalError::Execution),
2438 Self::Partitioned(runtime) => {
2439 runtime.forward_with_traversal_hook(input, state, context, hook)
2440 }
2441 }
2442 }
2443
2444 pub fn policy(&self) -> &ExecutionPolicy {
2446 match self {
2447 Self::Direct(runtime) => runtime.policy(),
2448 Self::Partitioned(runtime) => runtime.traversal_executor().runtime().policy(),
2449 }
2450 }
2451
2452 pub fn architecture(&self) -> &A {
2454 match self {
2455 Self::Direct(runtime) => runtime.architecture(),
2456 Self::Partitioned(runtime) => runtime.traversal_executor().runtime().architecture(),
2457 }
2458 }
2459}
2460
2461#[allow(
2463 clippy::type_complexity,
2464 reason = "marker preserves static dispatch across independent additive policies"
2465)]
2466pub struct PartitionedTextExecution<E, G, R, I, T, U, V>(
2467 PhantomData<fn() -> (E, G, R, I, T, U, V)>,
2468);
2469
2470impl<E, G, R, I, T, U, V> PartitionedTextExecution<E, G, R, I, T, U, V> {
2471 pub const fn new() -> Self {
2473 Self(PhantomData)
2474 }
2475}
2476
2477impl<E, G, R, I, T, U, V> Default for PartitionedTextExecution<E, G, R, I, T, U, V> {
2478 fn default() -> Self {
2479 Self::new()
2480 }
2481}
2482
2483impl<A, B, S, Resident, Bounded, E, G, R, I, T, U, V>
2484 ReplicatedTextExecutionStrategy<A, B, S, Resident, Bounded>
2485 for PartitionedTextExecution<E, G, R, I, T, U, V>
2486where
2487 B: CommunicationBackend,
2488 S: RuntimeState<B>,
2489 A: LayeredArchitecture<B, S>,
2490 Resident: LayerwisePolicy<B, A::Unit>,
2491 Bounded: LayerwisePolicy<B, A::Unit, Error = Resident::Error>,
2492 E: PartitionedGroupExecutor<A, B, S, G, R, I>,
2493 G: Borrow<B::CommunicationGroup>,
2494 R: Borrow<B::CommunicationRoute>,
2495 I: CommunicationTensorMetadata<B>,
2496 T: PartitionBoundaryTransport<B, G, R, I>,
2497 U: PartitionOutputPublisher<B, G, R, I>,
2498 V: PartitionCommitAgreement<B, G, R, I>,
2499 A::Error: std::fmt::Display,
2500 Resident::Error: std::fmt::Display,
2501{
2502 const PARTITIONED_SESSION: bool = true;
2503 const DISTRIBUTED_PHASE_AGREEMENT: bool = V::PHASE_FAILURE_AGREEMENT;
2504
2505 type Runtime = PartitionedTextRuntime<A, B, S, Bounded, E, G, R, I, T, U, V>;
2506
2507 fn bounded_policy(runtime: &Self::Runtime) -> Option<&Bounded> {
2508 runtime.bounded_policy.as_ref()
2509 }
2510
2511 fn execution_residency(
2512 runtime: &Self::Runtime,
2513 _selected: &crate::SelectedReplicatedTextRealization,
2514 ) -> ExecutionResidency {
2515 runtime.residency
2516 }
2517
2518 fn forward_with_observer<'a, O>(
2519 &mut self,
2520 runtime: &mut Self::Runtime,
2521 input: A::Input<'a>,
2522 state: &mut S,
2523 pass_kind: ExpertPass,
2524 context: &<B::Tensor as Tensor>::Context,
2525 observer: &mut O,
2526 ) -> Result<
2527 (B::Tensor, A::ForwardContext),
2528 ReplicatedTextSessionError<A::Error, Resident::Error, std::convert::Infallible>,
2529 >
2530 where
2531 O: ActivationObserver<B::Tensor, A::Error> + ?Sized,
2532 {
2533 let mut pass = runtime
2534 .executor
2535 .begin(input, state, pass_kind, context)
2536 .map_err(ReplicatedTextSessionError::Architecture)?;
2537 let phase_agreement_group = runtime.plan.commit_barrier;
2538 let mut schedule = LayeredPipelineSchedule::try_new(
2539 &runtime.plan.graph,
2540 runtime.plan.group_contracts.iter().copied(),
2541 |group| {
2542 runtime
2543 .executor
2544 .request_group_active(&pass, group)
2545 .map_err(PartitionScheduleSetupError::Architecture)
2546 },
2547 )
2548 .map_err(|error| match error {
2549 PartitionScheduleSetupError::Architecture(error) => {
2550 ReplicatedTextSessionError::Architecture(error)
2551 }
2552 PartitionScheduleSetupError::Schedule(error) => {
2553 ReplicatedTextSessionError::Contract(error.to_string())
2554 }
2555 })?;
2556
2557 while !schedule.is_complete() {
2558 let ready = schedule.ready_groups().collect::<Vec<_>>();
2559 let Some(&group) = ready.first() else {
2560 return Err(ReplicatedTextSessionError::Contract(
2561 "partitioned graph schedule made no progress".into(),
2562 ));
2563 };
2564 schedule
2565 .started(group)
2566 .map_err(|error| ReplicatedTextSessionError::Contract(error.to_string()))?;
2567 if schedule.is_active(group) == Some(true) {
2568 let same_group_routes = runtime
2569 .plan
2570 .routes
2571 .iter()
2572 .filter(|route| route.source_group == group && route.destination_group == group)
2573 .cloned()
2574 .collect::<Vec<_>>();
2575 let rank = runtime.communication.manifest().rank();
2576 let mut executed = false;
2577 let mut pipeline_wave = 0usize;
2578
2579 if V::PHASE_FAILURE_AGREEMENT && !same_group_routes.is_empty() {
2584 let mut remaining = vec![true; same_group_routes.len()];
2585 while remaining.iter().any(|remaining| *remaining) {
2586 let wave = same_group_routes
2587 .iter()
2588 .enumerate()
2589 .filter(|(index, route)| {
2590 remaining[*index]
2591 && !same_group_routes.iter().enumerate().any(
2592 |(predecessor, candidate)| {
2593 remaining[predecessor]
2594 && candidate.destination_rank == route.source_rank
2595 },
2596 )
2597 })
2598 .map(|(index, _)| index)
2599 .collect::<Vec<_>>();
2600 if wave.is_empty() {
2601 return Err(ReplicatedTextSessionError::Contract(
2602 "same-group pipeline routes contain a rank cycle".into(),
2603 ));
2604 }
2605 let active = wave
2606 .iter()
2607 .any(|index| same_group_routes[*index].source_rank == rank)
2608 && !executed;
2609 if active {
2610 executed = true;
2611 }
2612 let local_execution = Some(runtime.executor.execute_pipeline_wave(
2613 &mut pass,
2614 group,
2615 runtime.plan.drivers[group].as_ref(),
2616 active,
2617 pipeline_wave,
2618 state,
2619 &runtime.communication,
2620 runtime.communication_executor.borrow(),
2621 context,
2622 observer,
2623 ));
2624 let mut local_error = match local_execution {
2625 Some(Err(error)) => {
2626 Some(PartitionRouteTransferError::Architecture(error))
2627 }
2628 _ => None,
2629 };
2630 let mut prepared_wave = Vec::with_capacity(wave.len());
2631 for index in wave.iter().copied() {
2632 let route = &same_group_routes[index];
2633 let prepared = if local_error.is_none()
2634 && (rank == route.source_rank || rank == route.destination_rank)
2635 {
2636 match prepare_partition_boundary::<A, B, S, E, G, R, I>(
2637 &mut runtime.executor,
2638 &mut pass,
2639 &runtime.communication,
2640 route,
2641 runtime.plan.wire,
2642 context,
2643 ) {
2644 Ok(prepared) => Some(prepared),
2645 Err(error) => {
2646 local_error = Some(error);
2647 None
2648 }
2649 }
2650 } else {
2651 None
2652 };
2653 prepared_wave.push((index, prepared));
2654 }
2655 if local_error.is_none() {
2656 for (index, prepared) in &prepared_wave {
2657 let Some(prepared) = prepared.as_ref().filter(|value| value.source)
2658 else {
2659 continue;
2660 };
2661 let route = &same_group_routes[*index];
2662 if let Err(error) =
2663 runtime.communication.complete_local_dependencies(
2664 &prepared.values,
2665 route.route,
2666 runtime.communication_executor.borrow(),
2667 true,
2668 )
2669 {
2670 local_error =
2671 Some(PartitionRouteTransferError::Contract(error));
2672 break;
2673 }
2674 }
2675 }
2676 let local_success = local_error.is_none();
2677 let mut remote_completion_failure = None;
2678 for index in wave.iter().copied() {
2679 let route = &same_group_routes[index];
2680 let completed = match runtime.commit_agreement.agree_phase(
2681 &runtime.communication,
2682 phase_agreement_group.expect(
2683 "failure-agreement policy requires a selected session group",
2684 ),
2685 DistributedExecutionPhase::BoundarySourceCompletion(route.route),
2686 local_success,
2687 runtime.communication_executor.borrow(),
2688 ) {
2689 Ok(completed) => completed,
2690 Err(error) => {
2691 if let Some(local) = local_error.take() {
2692 return Err(map_partition_route_error(local));
2693 }
2694 return Err(ReplicatedTextSessionError::Contract(
2695 error.to_string(),
2696 ));
2697 }
2698 };
2699 if !completed && remote_completion_failure.is_none() {
2700 remote_completion_failure = Some(route.route);
2701 }
2702 }
2703 if let Some(error) = local_error {
2704 let route = wave.first().map(|index| same_group_routes[*index].route);
2705 runtime.communication.authority.fence_protocol_failure(
2706 CommunicationOperation::SendReceive,
2707 route.map_or(
2708 DistributedExecutionPhase::Execution,
2709 DistributedExecutionPhase::BoundarySourceCompletion,
2710 ),
2711 route,
2712 );
2713 return Err(map_partition_route_error(error));
2714 }
2715 if let Some(route) = remote_completion_failure {
2716 runtime.communication.authority.fence_protocol_failure(
2717 CommunicationOperation::SendReceive,
2718 DistributedExecutionPhase::BoundarySourceCompletion(route),
2719 Some(route),
2720 );
2721 return Err(ReplicatedTextSessionError::Contract(
2722 PartitionExecutionError::RemotePhaseFailure(
2723 DistributedExecutionPhase::BoundarySourceCompletion(route),
2724 )
2725 .to_string(),
2726 ));
2727 }
2728 let mut remote_failure = None;
2729 for index in wave.iter().copied() {
2730 let route = &same_group_routes[index];
2731 let ready = runtime
2732 .commit_agreement
2733 .agree_phase(
2734 &runtime.communication,
2735 phase_agreement_group.expect(
2736 "failure-agreement policy requires a selected session group",
2737 ),
2738 DistributedExecutionPhase::BoundarySourceReady(route.route),
2739 true,
2740 runtime.communication_executor.borrow(),
2741 )
2742 .map_err(|error| {
2743 ReplicatedTextSessionError::Contract(error.to_string())
2744 })?;
2745 if !ready && remote_failure.is_none() {
2746 remote_failure = Some(route.route);
2747 }
2748 }
2749 if let Some(route) = remote_failure {
2750 return Err(ReplicatedTextSessionError::Contract(
2751 PartitionExecutionError::RemotePhaseFailure(
2752 DistributedExecutionPhase::BoundarySourceReady(route),
2753 )
2754 .to_string(),
2755 ));
2756 }
2757 for (index, prepared) in prepared_wave {
2758 let route = &same_group_routes[index];
2759 if let Some(prepared) = prepared {
2760 transfer_prepared_partition_boundary::<A, B, S, E, G, R, I, T>(
2761 &mut runtime.executor,
2762 &mut pass,
2763 &runtime.communication,
2764 &mut runtime.boundary_transport,
2765 route,
2766 runtime.plan.wire,
2767 runtime.communication_executor.borrow(),
2768 prepared,
2769 )
2770 .map_err(map_partition_route_error)?;
2771 }
2772 remaining[index] = false;
2773 }
2774 pipeline_wave = pipeline_wave.checked_add(1).ok_or_else(|| {
2775 ReplicatedTextSessionError::Contract(
2776 "pipeline execution wave ordinal overflowed".into(),
2777 )
2778 })?;
2779 }
2780 } else {
2781 for route in &same_group_routes {
2782 if rank == route.destination_rank {
2783 prepare_and_transfer_partition_boundary::<A, B, S, E, G, R, I, T>(
2784 &mut runtime.executor,
2785 &mut pass,
2786 &runtime.communication,
2787 &mut runtime.boundary_transport,
2788 route,
2789 runtime.plan.wire,
2790 runtime.communication_executor.borrow(),
2791 context,
2792 )
2793 .map_err(map_partition_route_error)?;
2794 }
2795 }
2796 }
2797
2798 let mut local_execution =
2799 if V::PHASE_FAILURE_AGREEMENT && !same_group_routes.is_empty() {
2800 let active = !executed && runtime.plan.drivers[group].is_some();
2801 Some(runtime.executor.execute_pipeline_wave(
2802 &mut pass,
2803 group,
2804 runtime.plan.drivers[group].as_ref(),
2805 active,
2806 pipeline_wave,
2807 state,
2808 &runtime.communication,
2809 runtime.communication_executor.borrow(),
2810 context,
2811 observer,
2812 ))
2813 } else if executed {
2814 None
2815 } else {
2816 runtime.plan.drivers[group].as_ref().map(|driver| {
2817 runtime.executor.execute_group(
2818 &mut pass,
2819 driver,
2820 state,
2821 &runtime.communication,
2822 runtime.communication_executor.borrow(),
2823 context,
2824 observer,
2825 )
2826 })
2827 };
2828
2829 let outgoing_routes = runtime
2830 .plan
2831 .routes
2832 .iter()
2833 .filter(|route| {
2834 route.source_group == group
2835 && (!V::PHASE_FAILURE_AGREEMENT
2836 || route.source_group != route.destination_group)
2837 })
2838 .cloned()
2839 .collect::<Vec<_>>();
2840 if V::PHASE_FAILURE_AGREEMENT {
2841 let manifest = runtime.communication.manifest();
2842 let mut waves = Vec::new();
2843 for descriptor_range in manifest.route_submission_waves() {
2844 let wave = descriptor_range
2845 .clone()
2846 .filter_map(|index| {
2847 let id = manifest.routes()[index].id();
2848 outgoing_routes
2849 .iter()
2850 .find(|route| route.route == id)
2851 .cloned()
2852 })
2853 .collect::<Vec<_>>();
2854 if wave.is_empty() {
2855 continue;
2856 }
2857 if wave.len() != descriptor_range.len() {
2858 return Err(ReplicatedTextSessionError::Contract(
2859 PartitionExecutionError::RouteSubmissionWave {
2860 first: manifest.routes()[descriptor_range.start].id(),
2861 expected: descriptor_range.len(),
2862 actual: wave.len(),
2863 }
2864 .to_string(),
2865 ));
2866 }
2867 waves.push(wave);
2868 }
2869 if waves.iter().map(Vec::len).sum::<usize>() != outgoing_routes.len() {
2870 return Err(ReplicatedTextSessionError::Contract(
2871 "outgoing routes were omitted from manifest submission waves".into(),
2872 ));
2873 }
2874
2875 for wave in waves {
2876 let phase_route = wave[0].route;
2877 let endpoints = wave
2878 .iter()
2879 .filter(|route| {
2880 rank == route.source_rank || rank == route.destination_rank
2881 })
2882 .collect::<Vec<_>>();
2883 if wave.len() > 1 && endpoints.len() != 1 {
2884 return Err(ReplicatedTextSessionError::Contract(
2885 PartitionExecutionError::RouteSubmissionWave {
2886 first: phase_route,
2887 expected: 1,
2888 actual: endpoints.len(),
2889 }
2890 .to_string(),
2891 ));
2892 }
2893 let local_route = endpoints.first().copied();
2894 let mut local_error = if local_route
2895 .is_some_and(|route| rank == route.source_rank)
2896 && local_execution.as_ref().is_some_and(Result::is_err)
2897 {
2898 let Some(Err(error)) = local_execution.take() else {
2899 unreachable!("the local execution result was checked as an error")
2900 };
2901 Some(PartitionRouteTransferError::Architecture(error))
2902 } else {
2903 None
2904 };
2905 let prepared = if local_error.is_none() {
2906 local_route.and_then(|route| {
2907 match prepare_partition_boundary::<A, B, S, E, G, R, I>(
2908 &mut runtime.executor,
2909 &mut pass,
2910 &runtime.communication,
2911 route,
2912 runtime.plan.wire,
2913 context,
2914 ) {
2915 Ok(prepared) => Some(prepared),
2916 Err(error) => {
2917 local_error = Some(error);
2918 None
2919 }
2920 }
2921 })
2922 } else {
2923 None
2924 };
2925 if local_error.is_none() {
2926 if let (Some(route), Some(prepared)) =
2927 (local_route, prepared.as_ref().filter(|value| value.source))
2928 {
2929 if let Err(error) =
2930 runtime.communication.complete_local_dependencies(
2931 &prepared.values,
2932 route.route,
2933 runtime.communication_executor.borrow(),
2934 true,
2935 )
2936 {
2937 local_error =
2938 Some(PartitionRouteTransferError::Contract(error));
2939 }
2940 }
2941 }
2942 let completed = runtime
2943 .commit_agreement
2944 .agree_phase(
2945 &runtime.communication,
2946 phase_agreement_group.expect(
2947 "failure-agreement policy requires a selected session group",
2948 ),
2949 DistributedExecutionPhase::BoundarySourceCompletion(phase_route),
2950 local_error.is_none(),
2951 runtime.communication_executor.borrow(),
2952 )
2953 .map_err(|error| {
2954 ReplicatedTextSessionError::Contract(error.to_string())
2955 })?;
2956 if let Some(error) = local_error {
2957 runtime.communication.authority.fence_protocol_failure(
2958 CommunicationOperation::SendReceive,
2959 DistributedExecutionPhase::BoundarySourceCompletion(phase_route),
2960 Some(phase_route),
2961 );
2962 return Err(map_partition_route_error(error));
2963 }
2964 if !completed {
2965 runtime.communication.authority.fence_protocol_failure(
2966 CommunicationOperation::SendReceive,
2967 DistributedExecutionPhase::BoundarySourceCompletion(phase_route),
2968 Some(phase_route),
2969 );
2970 return Err(ReplicatedTextSessionError::Contract(
2971 PartitionExecutionError::RemotePhaseFailure(
2972 DistributedExecutionPhase::BoundarySourceCompletion(
2973 phase_route,
2974 ),
2975 )
2976 .to_string(),
2977 ));
2978 }
2979 let ready = runtime
2980 .commit_agreement
2981 .agree_phase(
2982 &runtime.communication,
2983 phase_agreement_group.expect(
2984 "failure-agreement policy requires a selected session group",
2985 ),
2986 DistributedExecutionPhase::BoundarySourceReady(phase_route),
2987 true,
2988 runtime.communication_executor.borrow(),
2989 )
2990 .map_err(|error| {
2991 ReplicatedTextSessionError::Contract(error.to_string())
2992 })?;
2993 if !ready {
2994 return Err(ReplicatedTextSessionError::Contract(
2995 PartitionExecutionError::RemotePhaseFailure(
2996 DistributedExecutionPhase::BoundarySourceReady(phase_route),
2997 )
2998 .to_string(),
2999 ));
3000 }
3001 if let (Some(route), Some(prepared)) = (local_route, prepared) {
3002 transfer_prepared_partition_boundary::<A, B, S, E, G, R, I, T>(
3003 &mut runtime.executor,
3004 &mut pass,
3005 &runtime.communication,
3006 &mut runtime.boundary_transport,
3007 route,
3008 runtime.plan.wire,
3009 runtime.communication_executor.borrow(),
3010 prepared,
3011 )
3012 .map_err(map_partition_route_error)?;
3013 }
3014 }
3015 } else {
3016 for route in &outgoing_routes {
3017 let descriptor = runtime
3018 .communication
3019 .manifest()
3020 .routes()
3021 .iter()
3022 .find(|candidate| candidate.id() == route.route)
3023 .ok_or_else(|| {
3024 ReplicatedTextSessionError::Contract(
3025 PartitionExecutionError::UnknownRoute(route.route).to_string(),
3026 )
3027 })?;
3028 let same_group = route.source_group == route.destination_group;
3029 if rank != descriptor.source()
3030 && (same_group || rank != descriptor.destination())
3031 {
3032 continue;
3033 }
3034 prepare_and_transfer_partition_boundary::<A, B, S, E, G, R, I, T>(
3035 &mut runtime.executor,
3036 &mut pass,
3037 &runtime.communication,
3038 &mut runtime.boundary_transport,
3039 route,
3040 runtime.plan.wire,
3041 runtime.communication_executor.borrow(),
3042 context,
3043 )
3044 .map_err(map_partition_route_error)?;
3045 }
3046 }
3047 if let Some(Err(error)) = local_execution {
3048 return Err(ReplicatedTextSessionError::Architecture(error));
3049 }
3050 }
3051 schedule
3052 .ordered(group)
3053 .map_err(|error| ReplicatedTextSessionError::Contract(error.to_string()))?;
3054 }
3055
3056 let (output, forward) = runtime
3057 .executor
3058 .finish(pass, state, context)
3059 .map_err(ReplicatedTextSessionError::Architecture)?;
3060 Ok((output, forward))
3061 }
3062
3063 fn observe_output<O>(
3064 runtime: &mut Self::Runtime,
3065 output: &B::Tensor,
3066 observer: &mut O,
3067 _context: &<B::Tensor as Tensor>::Context,
3068 ) -> Result<
3069 B::Tensor,
3070 ReplicatedTextSessionError<A::Error, Resident::Error, std::convert::Infallible>,
3071 >
3072 where
3073 O: ActivationObserver<B::Tensor, A::Error> + ?Sized,
3074 {
3075 if let Some(publication) = runtime.plan.publication {
3076 let rank = runtime.communication.manifest().rank();
3077 let local_output = runtime
3078 .plan
3079 .drivers
3080 .iter()
3081 .flatten()
3082 .any(LayeredPartitionDriver::owns_output);
3083 if rank == publication.owner_rank && !local_output {
3084 return Err(ReplicatedTextSessionError::Contract(
3085 PartitionExecutionError::OutputOwnership {
3086 rank,
3087 owner: publication.owner_rank,
3088 }
3089 .to_string(),
3090 ));
3091 }
3092 if rank != publication.owner_rank {
3093 return Ok(output.clone());
3094 }
3095 }
3096 crate::observe_model_logits(observer, output)
3097 .map_err(ReplicatedTextSessionError::Architecture)
3098 }
3099
3100 fn publish_observed_output(
3101 runtime: &mut Self::Runtime,
3102 output: B::Tensor,
3103 _context: &<B::Tensor as Tensor>::Context,
3104 ) -> Result<
3105 B::Tensor,
3106 ReplicatedTextSessionError<A::Error, Resident::Error, std::convert::Infallible>,
3107 > {
3108 match runtime.plan.publication {
3109 Some(publication) => runtime
3110 .output_publisher
3111 .publish(
3112 &runtime.communication,
3113 output,
3114 publication,
3115 DistributedExecutionPhase::OutputPublication,
3116 runtime.communication_executor.borrow(),
3117 )
3118 .map_err(|error| ReplicatedTextSessionError::Contract(error.to_string())),
3119 None => Ok(output),
3120 }
3121 }
3122
3123 fn prediction_target_capture(
3124 runtime: &mut Self::Runtime,
3125 forward: &A::ForwardContext,
3126 context: &<B::Tensor as Tensor>::Context,
3127 ) -> Result<
3128 Option<B::Tensor>,
3129 ReplicatedTextSessionError<A::Error, Resident::Error, std::convert::Infallible>,
3130 > {
3131 runtime
3132 .executor
3133 .prediction_target_capture(forward, context)
3134 .map_err(ReplicatedTextSessionError::Architecture)
3135 }
3136
3137 fn publish_prediction_target_capture(
3138 runtime: &mut Self::Runtime,
3139 capture: B::Tensor,
3140 _context: &<B::Tensor as Tensor>::Context,
3141 ) -> Result<
3142 B::Tensor,
3143 ReplicatedTextSessionError<A::Error, Resident::Error, std::convert::Infallible>,
3144 > {
3145 match runtime.plan.publication {
3146 Some(publication) => runtime
3147 .output_publisher
3148 .publish(
3149 &runtime.communication,
3150 capture,
3151 publication,
3152 DistributedExecutionPhase::PredictionTargetCapture,
3153 runtime.communication_executor.borrow(),
3154 )
3155 .map_err(|error| ReplicatedTextSessionError::Contract(error.to_string())),
3156 None => Ok(capture),
3157 }
3158 }
3159
3160 fn apply_prediction_target_operation<O>(
3161 runtime: &mut Self::Runtime,
3162 state: &mut S,
3163 operation: O,
3164 context: &<B::Tensor as Tensor>::Context,
3165 ) -> Result<
3166 Option<O::Output>,
3167 ReplicatedTextSessionError<A::Error, Resident::Error, std::convert::Infallible>,
3168 >
3169 where
3170 O: crate::PredictionTargetOperation<A, B, S>,
3171 {
3172 runtime
3173 .executor
3174 .apply_prediction_target_operation(state, operation, context)
3175 .map_err(ReplicatedTextSessionError::Architecture)
3176 }
3177
3178 fn commit_after_completion(
3179 runtime: &mut Self::Runtime,
3180 epoch: DistributedCommitEpoch,
3181 _context: &<B::Tensor as Tensor>::Context,
3182 ) -> DistributedCommitOutcome {
3183 if let Some(group) = runtime.plan.commit_barrier {
3184 return runtime.commit_agreement.commit(
3185 &runtime.communication,
3186 group,
3187 epoch,
3188 runtime.communication_executor.borrow(),
3189 );
3190 }
3191 DistributedCommitOutcome::Committed(epoch)
3192 }
3193
3194 fn agree_distributed_phase(
3195 runtime: &mut Self::Runtime,
3196 phase: DistributedExecutionPhase,
3197 local_success: bool,
3198 _context: &<B::Tensor as Tensor>::Context,
3199 ) -> Result<bool, ReplicatedTextSessionError<A::Error, Resident::Error, std::convert::Infallible>>
3200 {
3201 let cross_stage_failure = phase == DistributedExecutionPhase::Execution
3202 && runtime.executor.has_cross_stage_collective_waves();
3203 let needs_recovery_agreement = cross_stage_failure
3204 && (!local_success || runtime.communication.authority.is_poisoned());
3205 let agreed = if needs_recovery_agreement {
3206 match runtime.plan.commit_barrier {
3207 Some(group) => runtime.commit_agreement.agree_phase_after_prior_failure(
3208 &runtime.communication,
3209 group,
3210 phase,
3211 local_success,
3212 runtime.communication_executor.borrow(),
3213 ),
3214 None => Ok(local_success),
3215 }
3216 } else {
3217 runtime.agree_phase(phase, local_success)
3218 }
3219 .map_err(|error| ReplicatedTextSessionError::Contract(error.to_string()))?;
3220 if cross_stage_failure && !agreed {
3221 let _ = runtime.communication.authority.submission_error(
3222 "routed pipeline collective wave failed on at least one rank",
3223 CommunicationOperation::AllGatherEven,
3224 DistributedExecutionPhase::Execution,
3225 None,
3226 );
3227 }
3228 Ok(agreed)
3229 }
3230}
3231
3232enum PartitionRouteTransferError<E> {
3233 Architecture(E),
3234 Contract(PartitionExecutionError),
3235}
3236
3237struct PreparedPartitionBoundary<T> {
3238 source: bool,
3239 schema: ResolvedBoundaryWireSchema,
3240 values: Vec<crate::ArchitectureBoundaryValue<T>>,
3241}
3242
3243fn map_partition_route_error<E: std::fmt::Display, P: std::fmt::Display>(
3244 error: PartitionRouteTransferError<E>,
3245) -> ReplicatedTextSessionError<E, P, std::convert::Infallible> {
3246 match error {
3247 PartitionRouteTransferError::Architecture(error) => {
3248 ReplicatedTextSessionError::Architecture(error)
3249 }
3250 PartitionRouteTransferError::Contract(error) => {
3251 ReplicatedTextSessionError::Contract(error.to_string())
3252 }
3253 }
3254}
3255
3256#[allow(clippy::too_many_arguments)]
3257fn prepare_partition_boundary<'a, A, B, S, E, G, R, I>(
3258 executor: &mut E,
3259 pass: &mut E::Pass<'a>,
3260 communication: &PartitionCommunication<B, G, R, I>,
3261 route: &PartitionBoundaryRoute,
3262 wire: PipelineWireContract,
3263 context: &<B::Tensor as Tensor>::Context,
3264) -> Result<PreparedPartitionBoundary<B::Tensor>, PartitionRouteTransferError<A::Error>>
3265where
3266 B: CommunicationBackend,
3267 S: RuntimeState<B>,
3268 A: LayeredArchitecture<B, S>,
3269 E: PartitionedGroupExecutor<A, B, S, G, R, I>,
3270 G: Borrow<B::CommunicationGroup>,
3271 R: Borrow<B::CommunicationRoute>,
3272 I: CommunicationTensorMetadata<B>,
3273{
3274 let source = communication
3275 .boundary_endpoint_is_source(route.route)
3276 .map_err(PartitionRouteTransferError::Contract)?;
3277 let schema = executor
3278 .boundary_schema(pass, route)
3279 .map_err(PartitionRouteTransferError::Architecture)?;
3280 let values = executor
3281 .boundary_values(pass, route, &schema, source, context)
3282 .map_err(PartitionRouteTransferError::Architecture)?;
3283 communication
3284 .validate_prepared_boundary(route.route, &values, &schema, wire)
3285 .map_err(PartitionRouteTransferError::Contract)?;
3286 Ok(PreparedPartitionBoundary {
3287 source,
3288 schema,
3289 values,
3290 })
3291}
3292
3293#[allow(clippy::too_many_arguments)]
3294fn transfer_prepared_partition_boundary<'a, A, B, S, E, G, R, I, T>(
3295 executor: &mut E,
3296 pass: &mut E::Pass<'a>,
3297 communication: &PartitionCommunication<B, G, R, I>,
3298 transport: &mut T,
3299 route: &PartitionBoundaryRoute,
3300 wire: PipelineWireContract,
3301 communication_executor: &B::Executor,
3302 prepared: PreparedPartitionBoundary<B::Tensor>,
3303) -> Result<(), PartitionRouteTransferError<A::Error>>
3304where
3305 B: CommunicationBackend,
3306 S: RuntimeState<B>,
3307 A: LayeredArchitecture<B, S>,
3308 E: PartitionedGroupExecutor<A, B, S, G, R, I>,
3309 G: Borrow<B::CommunicationGroup>,
3310 R: Borrow<B::CommunicationRoute>,
3311 I: CommunicationTensorMetadata<B>,
3312 T: PartitionBoundaryTransport<B, G, R, I>,
3313{
3314 let values = transport
3315 .transfer(
3316 communication,
3317 route.route,
3318 prepared.values,
3319 &prepared.schema,
3320 wire,
3321 communication_executor,
3322 )
3323 .map_err(PartitionRouteTransferError::Contract)?;
3324 if !prepared.source {
3325 executor
3326 .accept_boundary(pass, route, values)
3327 .map_err(PartitionRouteTransferError::Architecture)?;
3328 }
3329 Ok(())
3330}
3331
3332#[allow(clippy::too_many_arguments)]
3333fn prepare_and_transfer_partition_boundary<'a, A, B, S, E, G, R, I, T>(
3334 executor: &mut E,
3335 pass: &mut E::Pass<'a>,
3336 communication: &PartitionCommunication<B, G, R, I>,
3337 transport: &mut T,
3338 route: &PartitionBoundaryRoute,
3339 wire: PipelineWireContract,
3340 communication_executor: &B::Executor,
3341 context: &<B::Tensor as Tensor>::Context,
3342) -> Result<(), PartitionRouteTransferError<A::Error>>
3343where
3344 B: CommunicationBackend,
3345 S: RuntimeState<B>,
3346 A: LayeredArchitecture<B, S>,
3347 E: PartitionedGroupExecutor<A, B, S, G, R, I>,
3348 G: Borrow<B::CommunicationGroup>,
3349 R: Borrow<B::CommunicationRoute>,
3350 I: CommunicationTensorMetadata<B>,
3351 T: PartitionBoundaryTransport<B, G, R, I>,
3352{
3353 let prepared = prepare_partition_boundary::<A, B, S, E, G, R, I>(
3354 executor,
3355 pass,
3356 communication,
3357 route,
3358 wire,
3359 context,
3360 )?;
3361 if prepared.source {
3362 communication
3363 .complete_local_dependencies(
3364 &prepared.values,
3365 route.route,
3366 communication_executor,
3367 false,
3368 )
3369 .map_err(PartitionRouteTransferError::Contract)?;
3370 }
3371 transfer_prepared_partition_boundary::<A, B, S, E, G, R, I, T>(
3372 executor,
3373 pass,
3374 communication,
3375 transport,
3376 route,
3377 wire,
3378 communication_executor,
3379 prepared,
3380 )
3381}
3382
3383fn validate_bundle_requirement<B, I>(
3384 inspector: &I,
3385 values: &[B::Tensor],
3386 requirement: &CommunicationOperationRequirement,
3387 completed: bool,
3388) -> Result<(), PartitionExecutionError>
3389where
3390 B: NeuralBackend,
3391 I: CommunicationTensorMetadata<B>,
3392{
3393 let limits = requirement
3394 .limits()
3395 .ok_or(PartitionExecutionError::MissingTensorLimits)?;
3396 if values.len() > limits.max_tensors() {
3397 return Err(PartitionExecutionError::TensorCount {
3398 expected_at_most: limits.max_tensors(),
3399 actual: values.len(),
3400 });
3401 }
3402 for value in values {
3403 let dtype = inspector.dtype(value);
3404 let shape = inspector.shape(value);
3405 let elements = shape
3406 .iter()
3407 .try_fold(1usize, |product, dimension| product.checked_mul(*dimension));
3408 if !requirement.dtypes().contains(&dtype) {
3409 return Err(PartitionExecutionError::TensorDtype { dtype });
3410 }
3411 let maximum = if completed {
3412 limits.max_output_tensor_elements()
3413 } else {
3414 limits.max_tensor_elements()
3415 };
3416 if shape.len() > limits.max_tensor_rank()
3417 || elements.is_none_or(|elements| elements > maximum)
3418 {
3419 return Err(PartitionExecutionError::TensorLimits { shape });
3420 }
3421 }
3422 Ok(())
3423}
3424
3425fn validate_boundary_bundle<B, I>(
3426 inspector: &I,
3427 values: &[B::Tensor],
3428 schema: &ResolvedBoundaryWireSchema,
3429 wire: PipelineWireContract,
3430) -> Result<(), PartitionExecutionError>
3431where
3432 B: NeuralBackend,
3433 I: CommunicationTensorMetadata<B>,
3434{
3435 let specs = std::iter::once(schema.primary()).chain(schema.auxiliary());
3436 let expected = 1 + schema.auxiliary().len();
3437 if values.len() != expected {
3438 return Err(PartitionExecutionError::BoundaryTensorCount {
3439 boundary: schema.identity(),
3440 expected,
3441 actual: values.len(),
3442 });
3443 }
3444 for (value, spec) in values.iter().zip(specs) {
3445 validate_boundary_tensor::<B, I>(inspector, value, spec, wire)?;
3446 }
3447 Ok(())
3448}
3449
3450fn validate_tagged_boundary_bundle<B, I>(
3451 inspector: &I,
3452 values: &[crate::ArchitectureBoundaryValue<B::Tensor>],
3453 schema: &ResolvedBoundaryWireSchema,
3454 wire: PipelineWireContract,
3455) -> Result<(), PartitionExecutionError>
3456where
3457 B: NeuralBackend,
3458 I: CommunicationTensorMetadata<B>,
3459{
3460 let specs = std::iter::once(schema.primary()).chain(schema.auxiliary());
3461 let expected = 1 + schema.auxiliary().len();
3462 if values.len() != expected {
3463 return Err(PartitionExecutionError::BoundaryTensorCount {
3464 boundary: schema.identity(),
3465 expected,
3466 actual: values.len(),
3467 });
3468 }
3469 for (value, spec) in values.iter().zip(specs) {
3470 if value.role() != spec.role() {
3471 return Err(PartitionExecutionError::BoundaryFraming(format!(
3472 "architecture boundary role {:?} differs from selected role {:?}",
3473 value.role(),
3474 spec.role(),
3475 )));
3476 }
3477 validate_boundary_tensor::<B, I>(inspector, value.tensor(), spec, wire)?;
3478 }
3479 Ok(())
3480}
3481
3482fn resolved_tagged_boundary_roles<T>(
3483 values: &[crate::ArchitectureBoundaryValue<T>],
3484 schema: &ResolvedBoundaryWireSchema,
3485 wire: PipelineWireContract,
3486) -> Result<Vec<crate::BoundaryRoleContract>, PartitionExecutionError> {
3487 values
3488 .iter()
3489 .zip(std::iter::once(schema.primary()).chain(schema.auxiliary()))
3490 .map(|(value, spec)| {
3491 let dtype = match spec.dtype() {
3492 crate::BoundaryTensorDtype::Activation => match wire.activation_dtype() {
3493 PipelineActivationDtype::Float16 => TensorDtype::F16,
3494 PipelineActivationDtype::Bfloat16 => TensorDtype::Bf16,
3495 PipelineActivationDtype::Float32 => TensorDtype::F32,
3496 },
3497 crate::BoundaryTensorDtype::Uint32 => TensorDtype::U32,
3498 crate::BoundaryTensorDtype::Int32 => TensorDtype::I32,
3499 };
3500 let shape = spec
3501 .shape()
3502 .iter()
3503 .map(|dimension| usize::try_from(*dimension))
3504 .collect::<Result<Vec<_>, _>>()
3505 .map_err(|_| {
3506 PartitionExecutionError::BoundaryFraming(
3507 "resolved boundary shape is not representable".into(),
3508 )
3509 })?;
3510 crate::BoundaryRoleContract::new(value.role(), dtype, shape)
3511 .map_err(|error| PartitionExecutionError::BoundaryFraming(error.to_string()))
3512 })
3513 .collect()
3514}
3515
3516fn validate_boundary_tensor<B, I>(
3517 inspector: &I,
3518 value: &B::Tensor,
3519 spec: &ResolvedBoundaryTensorSpec,
3520 wire: PipelineWireContract,
3521) -> Result<(), PartitionExecutionError>
3522where
3523 B: NeuralBackend,
3524 I: CommunicationTensorMetadata<B>,
3525{
3526 let shape = inspector.shape(value);
3527 let expected_shape = spec
3528 .shape()
3529 .iter()
3530 .map(|dimension| usize::try_from(*dimension).unwrap_or(usize::MAX))
3531 .collect::<Vec<_>>();
3532 if shape != expected_shape {
3533 return Err(PartitionExecutionError::BoundaryShape {
3534 role: spec.role().to_owned(),
3535 expected: expected_shape,
3536 actual: shape,
3537 });
3538 }
3539 let expected_dtype = match spec.dtype() {
3540 crate::BoundaryTensorDtype::Activation => match wire.activation_dtype() {
3541 PipelineActivationDtype::Float16 => TensorDtype::F16,
3542 PipelineActivationDtype::Bfloat16 => TensorDtype::Bf16,
3543 PipelineActivationDtype::Float32 => TensorDtype::F32,
3544 },
3545 crate::BoundaryTensorDtype::Uint32 => TensorDtype::U32,
3546 crate::BoundaryTensorDtype::Int32 => TensorDtype::I32,
3547 };
3548 let actual = inspector.dtype(value);
3549 if actual != expected_dtype {
3550 return Err(PartitionExecutionError::BoundaryDtype {
3551 role: spec.role().to_owned(),
3552 expected: expected_dtype,
3553 actual,
3554 });
3555 }
3556 Ok(())
3557}
3558
3559enum PartitionScheduleSetupError<E> {
3560 Architecture(E),
3561 Schedule(LayeredPipelineScheduleError),
3562}
3563
3564impl<E> From<LayeredPipelineScheduleError> for PartitionScheduleSetupError<E> {
3565 fn from(error: LayeredPipelineScheduleError) -> Self {
3566 Self::Schedule(error)
3567 }
3568}
3569
3570#[derive(Debug, thiserror::Error)]
3572#[non_exhaustive]
3573#[allow(
3574 missing_docs,
3575 reason = "field names and error messages document mechanical validation diagnostics"
3576)]
3577pub enum PartitionExecutionError {
3578 #[error("partition communication manifest has no bounded completion policy")]
3580 MissingBoundedCompletionPolicy,
3581 #[error(
3583 "communication {operation:?} exceeded its selected deadline during {phase:?} (route {route:?}, disposition {cancellation:?})"
3584 )]
3585 CommunicationDeadlineExceeded {
3586 operation: CommunicationOperation,
3587 phase: DistributedExecutionPhase,
3588 route: Option<CommunicationRouteId>,
3589 cancellation: CompletionCancellationMode,
3590 },
3591 #[error(
3593 "communication {operation:?} failed during exact completion in {phase:?} (route {route:?}): {error}"
3594 )]
3595 CommunicationCompletionFailed {
3596 operation: CommunicationOperation,
3597 phase: DistributedExecutionPhase,
3598 route: Option<CommunicationRouteId>,
3599 error: String,
3600 },
3601 #[error(
3603 "communication {operation:?} submission failed in {phase:?} (route {route:?}): {error}"
3604 )]
3605 CommunicationSubmissionFailed {
3606 operation: CommunicationOperation,
3607 phase: DistributedExecutionPhase,
3608 route: Option<CommunicationRouteId>,
3609 error: String,
3610 },
3611 #[error(
3613 "communication is poisoned by prior {operation:?} during {phase:?} (route {route:?}, disposition {cancellation:?})"
3614 )]
3615 CommunicationPoisoned {
3616 operation: CommunicationOperation,
3617 phase: DistributedExecutionPhase,
3618 route: Option<CommunicationRouteId>,
3619 cancellation: CompletionCancellationMode,
3620 },
3621 #[error("recovery agreement was requested without a prior failure during {phase:?}")]
3623 RecoveryAgreementWithoutFailure { phase: DistributedExecutionPhase },
3624 #[error("communication resource count mismatch (groups {actual_groups}/{expected_groups}, routes {actual_routes}/{expected_routes})")]
3626 ResourceCount {
3627 expected_groups: usize,
3628 actual_groups: usize,
3629 expected_routes: usize,
3630 actual_routes: usize,
3631 },
3632 #[error("communication resource identity mismatch (actual {actual}, expected {expected})")]
3634 ResourceIdentity { expected: u64, actual: u64 },
3635 #[error("unknown opaque communication group {0:?}")]
3636 UnknownGroup(CollectiveGroupId),
3637 #[error("local rank is not a member of opaque communication group {0:?}")]
3638 NotGroupMember(CollectiveGroupId),
3639 #[error("unknown opaque communication route {0:?}")]
3640 UnknownRoute(CommunicationRouteId),
3641 #[error(
3642 "route submission wave beginning at {first:?} has {actual} local routes/endpoints, expected {expected}"
3643 )]
3644 RouteSubmissionWave {
3645 first: CommunicationRouteId,
3646 expected: usize,
3647 actual: usize,
3648 },
3649 #[error("local rank is not an endpoint of opaque communication route {0:?}")]
3650 NotRouteEndpoint(CommunicationRouteId),
3651 #[error("{resource} did not select operation {operation:?}")]
3652 OperationNotSelected {
3653 resource: String,
3654 operation: CommunicationOperation,
3655 },
3656 #[error("communication tensor has unselected dtype {dtype:?}")]
3657 TensorDtype { dtype: TensorDtype },
3658 #[error("communication tensor shape exceeds selected limits: {shape:?}")]
3659 TensorLimits { shape: Vec<usize> },
3660 #[error("communication axis {axis} is outside tensor rank {rank}")]
3661 CommunicationAxis { axis: usize, rank: usize },
3662 #[error("communication result shape arithmetic overflowed")]
3663 CommunicationShapeOverflow,
3664 #[error("communication result shape {actual:?} differs from expected {expected:?}")]
3665 CommunicationOutputShape {
3666 expected: Vec<usize>,
3667 actual: Vec<usize>,
3668 },
3669 #[error("communication tensor requirement has no payload limits")]
3670 MissingTensorLimits,
3671 #[error("communication bundle contains {actual} tensors, maximum is {expected_at_most}")]
3672 TensorCount {
3673 expected_at_most: usize,
3674 actual: usize,
3675 },
3676 #[error("peer-count cardinality is {actual}, expected {expected}")]
3677 PeerCount { expected: usize, actual: usize },
3678 #[error("a peer count exceeds the selected maximum {maximum}")]
3679 PeerCountLimit { maximum: usize },
3680 #[error("architecture boundary {boundary:?} contains {actual} tensors, expected {expected}")]
3681 BoundaryTensorCount {
3682 boundary: &'static str,
3683 expected: usize,
3684 actual: usize,
3685 },
3686 #[error("architecture boundary role {role:?} has shape {actual:?}, expected {expected:?}")]
3687 BoundaryShape {
3688 role: String,
3689 expected: Vec<usize>,
3690 actual: Vec<usize>,
3691 },
3692 #[error("architecture boundary role {role:?} has dtype {actual:?}, expected {expected:?}")]
3693 BoundaryDtype {
3694 role: String,
3695 expected: TensorDtype,
3696 actual: TensorDtype,
3697 },
3698 #[error("role-exact boundary framing failed: {0}")]
3699 BoundaryFraming(String),
3700 #[error("opaque output group {group:?} does not contain output owner rank {rank}")]
3701 OutputOwnerNotMember {
3702 rank: usize,
3703 group: CollectiveGroupId,
3704 },
3705 #[error("rank {rank} output ownership disagrees with selected owner {owner}")]
3706 OutputOwnership { rank: usize, owner: usize },
3707 #[error("a local output owner has no publication operation")]
3708 MissingOutputPublication,
3709 #[error(
3710 "partition plan declares {contracts} contracts and {drivers} drivers for {graph} groups"
3711 )]
3712 GroupCount {
3713 graph: usize,
3714 contracts: usize,
3715 drivers: usize,
3716 },
3717 #[error("rank-local driver is stored under the wrong group slot {group}")]
3718 DriverGroup { group: usize },
3719 #[error("route {0:?} names a missing architecture group")]
3720 RouteGroup(CommunicationRouteId),
3721 #[error("route {0:?} does not connect an architecture dependency edge")]
3722 RouteDependency(CommunicationRouteId),
3723 #[error("partition plan contains {plan} routes but its manifest contains {manifest}")]
3724 RouteCount { plan: usize, manifest: usize },
3725 #[error(
3726 "route {route:?} ownership {planned_source} -> {planned_destination} differs from manifest endpoints {manifest_source} -> {manifest_destination}"
3727 )]
3728 RouteDescriptorMismatch {
3729 route: CommunicationRouteId,
3730 planned_source: usize,
3731 planned_destination: usize,
3732 manifest_source: usize,
3733 manifest_destination: usize,
3734 },
3735 #[error("route {route:?} endpoint rank {rank} does not own architecture group {group}")]
3736 RouteOwnerMissing {
3737 route: CommunicationRouteId,
3738 rank: usize,
3739 group: usize,
3740 },
3741 #[error("communication mechanism failed: {0}")]
3742 Communication(String),
3743 #[error("another rank reported failure during distributed phase {0:?}")]
3744 RemotePhaseFailure(DistributedExecutionPhase),
3745 #[error("rank-local execution residency and bounded policy disagree")]
3746 ResidencyPolicyMismatch,
3747 #[error("communication operation {0:?} has no selected execution policy")]
3748 OperationPolicyUnavailable(CommunicationOperation),
3749 #[error("opaque communication group {group:?} selected an inexact {operation:?} requirement")]
3750 InexactOperationRequirement {
3751 group: CollectiveGroupId,
3752 operation: CommunicationOperation,
3753 },
3754 #[error("partition communication plan and additive operation policies disagree")]
3755 CommunicationPolicyMismatch,
3756}
3757
3758#[cfg(test)]
3759mod plan_tests {
3760 use super::*;
3761 use crate::{CommunicationGroupRequirements, CommunicationTensorLimits};
3762
3763 #[test]
3764 fn plan_exposes_the_exact_selected_publication_and_commit_group() {
3765 let graph =
3766 ExecutionGraph::new(vec![crate::ExecutionGroupSpec::root("decoder")], "decoder")
3767 .unwrap();
3768 let group = CollectiveGroupId::new(17);
3769 let plan = PartitionedExecutionPlan::new(
3770 graph,
3771 vec![(ArchitectureGroupKind::Decoder, false)],
3772 vec![None],
3773 Vec::new(),
3774 Some(PartitionOutputPublication {
3775 group,
3776 owner_rank: 3,
3777 }),
3778 Some(group),
3779 PipelineWireContract::new(PipelineActivationDtype::Float32),
3780 )
3781 .unwrap();
3782
3783 assert_eq!(plan.publication().unwrap().group, group);
3784 assert_eq!(plan.publication().unwrap().owner_rank, 3);
3785 assert_eq!(plan.commit_barrier(), Some(group));
3786 }
3787
3788 #[test]
3789 fn publication_authority_preserves_selected_owner_and_group_local_rank() {
3790 let graph =
3791 ExecutionGraph::new(vec![crate::ExecutionGroupSpec::root("decoder")], "decoder")
3792 .unwrap();
3793 let group = CollectiveGroupId::new(17);
3794 let plan = PartitionedExecutionPlan::new(
3795 graph,
3796 vec![(ArchitectureGroupKind::Decoder, false)],
3797 vec![None],
3798 Vec::new(),
3799 Some(PartitionOutputPublication {
3800 group,
3801 owner_rank: 3,
3802 }),
3803 Some(group),
3804 PipelineWireContract::new(PipelineActivationDtype::Float32),
3805 )
3806 .unwrap();
3807 let requirements = CommunicationGroupRequirements::new([
3808 CommunicationOperationRequirement::tensors(
3809 CommunicationOperation::Broadcast,
3810 [TensorDtype::F32],
3811 CommunicationTensorLimits::new(1, 3, 32, None).unwrap(),
3812 true,
3813 )
3814 .unwrap(),
3815 CommunicationOperationRequirement::failure_agreement(true),
3816 ])
3817 .unwrap();
3818 let manifest = CommunicationManifest::new(
3819 8,
3820 7,
3821 vec![CommunicationGroupDescriptor::new(
3822 group,
3823 0,
3824 vec![3, 7],
3825 Some(1),
3826 requirements.clone(),
3827 )
3828 .unwrap()],
3829 Vec::new(),
3830 )
3831 .unwrap();
3832 let authority = plan.publication_authority(&manifest).unwrap().unwrap();
3833 assert_eq!(authority.owner_rank(), 3);
3834 assert_eq!(authority.owner_group_rank(), 0);
3835 assert!(!authority.local_public_output());
3836
3837 let substituted = CommunicationManifest::new(
3838 8,
3839 7,
3840 vec![
3841 CommunicationGroupDescriptor::new(group, 0, vec![4, 7], Some(1), requirements)
3842 .unwrap(),
3843 ],
3844 Vec::new(),
3845 )
3846 .unwrap();
3847 assert!(matches!(
3848 plan.publication_authority(&substituted),
3849 Err(PartitionExecutionError::OutputOwnerNotMember { rank: 3, .. })
3850 ));
3851 }
3852
3853 #[test]
3854 fn plan_rejects_manifest_route_endpoints_that_differ_from_architecture_ownership() {
3855 let graph =
3856 ExecutionGraph::new(vec![crate::ExecutionGroupSpec::root("decoder")], "decoder")
3857 .unwrap();
3858 let route = CommunicationRouteId::new(9);
3859 let plan = PartitionedExecutionPlan::new(
3860 graph,
3861 vec![(ArchitectureGroupKind::Decoder, false)],
3862 vec![None],
3863 vec![PartitionBoundaryRoute {
3864 source_group: 0,
3865 destination_group: 0,
3866 source_rank: 0,
3867 destination_rank: 1,
3868 route,
3869 }],
3870 None,
3871 None,
3872 PipelineWireContract::new(PipelineActivationDtype::Float32),
3873 )
3874 .unwrap();
3875 let requirement = CommunicationOperationRequirement::tensors(
3876 CommunicationOperation::SendReceive,
3877 [TensorDtype::F32],
3878 crate::CommunicationTensorLimits::new(1, 3, 8, None).unwrap(),
3879 true,
3880 )
3881 .unwrap();
3882 let manifest = CommunicationManifest::new(
3883 3,
3884 2,
3885 Vec::new(),
3886 vec![CommunicationRouteDescriptor::new(route, 0, 1, 0, requirement).unwrap()],
3887 )
3888 .unwrap();
3889
3890 assert!(matches!(
3891 plan.validate_manifest(&manifest),
3892 Err(PartitionExecutionError::RouteDescriptorMismatch {
3893 planned_source: 0,
3894 planned_destination: 1,
3895 manifest_source: 1,
3896 manifest_destination: 0,
3897 ..
3898 })
3899 ));
3900 }
3901
3902 fn publication_plan(group: CollectiveGroupId) -> PartitionedExecutionPlan {
3903 PartitionedExecutionPlan::new(
3904 ExecutionGraph::new(vec![crate::ExecutionGroupSpec::root("decoder")], "decoder")
3905 .unwrap(),
3906 vec![(ArchitectureGroupKind::Decoder, false)],
3907 vec![None],
3908 Vec::new(),
3909 Some(PartitionOutputPublication {
3910 group,
3911 owner_rank: 0,
3912 }),
3913 Some(group),
3914 PipelineWireContract::new(PipelineActivationDtype::Float32),
3915 )
3916 .unwrap()
3917 }
3918
3919 fn single_group_manifest(
3920 group: CollectiveGroupId,
3921 requirements: CommunicationGroupRequirements,
3922 ) -> CommunicationManifest {
3923 CommunicationManifest::new(
3924 2,
3925 1,
3926 vec![
3927 CommunicationGroupDescriptor::new(group, 0, vec![0, 1], Some(1), requirements)
3928 .unwrap(),
3929 ],
3930 Vec::new(),
3931 )
3932 .unwrap()
3933 }
3934
3935 #[test]
3936 fn plan_rejects_publication_group_without_exact_broadcast_before_execution() {
3937 let group = CollectiveGroupId::new(29);
3938 let plan = publication_plan(group);
3939 let missing = single_group_manifest(
3940 group,
3941 CommunicationGroupRequirements::new([
3942 CommunicationOperationRequirement::failure_agreement(true),
3943 ])
3944 .unwrap(),
3945 );
3946 assert!(matches!(
3947 plan.validate_manifest(&missing),
3948 Err(PartitionExecutionError::OperationNotSelected {
3949 operation: CommunicationOperation::Broadcast,
3950 ..
3951 })
3952 ));
3953
3954 let inexact = single_group_manifest(
3955 group,
3956 CommunicationGroupRequirements::new([
3957 CommunicationOperationRequirement::tensors(
3958 CommunicationOperation::Broadcast,
3959 [TensorDtype::F32],
3960 CommunicationTensorLimits::new(1, 3, 32, None).unwrap(),
3961 false,
3962 )
3963 .unwrap(),
3964 CommunicationOperationRequirement::failure_agreement(true),
3965 ])
3966 .unwrap(),
3967 );
3968 assert!(matches!(
3969 plan.validate_manifest(&inexact),
3970 Err(PartitionExecutionError::InexactOperationRequirement {
3971 operation: CommunicationOperation::Broadcast,
3972 ..
3973 })
3974 ));
3975 }
3976
3977 #[test]
3978 fn plan_rejects_inexact_failure_agreement_before_execution() {
3979 let group = CollectiveGroupId::new(31);
3980 let plan = publication_plan(group);
3981 let manifest = single_group_manifest(
3982 group,
3983 CommunicationGroupRequirements::new([
3984 CommunicationOperationRequirement::tensors(
3985 CommunicationOperation::Broadcast,
3986 [TensorDtype::F32],
3987 CommunicationTensorLimits::new(1, 3, 32, None).unwrap(),
3988 true,
3989 )
3990 .unwrap(),
3991 CommunicationOperationRequirement::failure_agreement(false),
3992 ])
3993 .unwrap(),
3994 );
3995 plan.validate_manifest(&manifest).unwrap();
3996 assert!(matches!(
3997 plan.validate_exact_commit_agreement(&manifest),
3998 Err(PartitionExecutionError::InexactOperationRequirement {
3999 operation: CommunicationOperation::FailureAgreement,
4000 ..
4001 })
4002 ));
4003 }
4004}