Skip to main content

eredu_runtime/
partitioned_execution.rs

1//! Backend-neutral rank-local execution over opaque communication resources.
2//!
3//! Architecture adapters remain statically dispatched and own tensor equations,
4//! routed providers, and typed boundary encoding. This module owns the common
5//! partition schedule, mechanical wire validation, exact completion, publication,
6//! and the execution strategy installed in [`crate::ReplicatedTextSession`].
7
8use 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
31/// Backend-native tensor metadata used only for mechanical communication validation.
32pub trait CommunicationTensorMetadata<B: NeuralBackend> {
33    /// Returns the portable logical dtype of one native tensor.
34    fn dtype(&self, tensor: &B::Tensor) -> TensorDtype;
35
36    /// Returns the exact logical shape of one native tensor.
37    fn shape(&self, tensor: &B::Tensor) -> Vec<usize>;
38}
39
40/// One backend-native group paired with the opaque identity selected by the manifest.
41pub struct RealizedCommunicationGroup<G> {
42    id: CollectiveGroupId,
43    resource: G,
44}
45
46impl<G> RealizedCommunicationGroup<G> {
47    /// Binds one native resource to its selected opaque group identity.
48    pub const fn new(id: CollectiveGroupId, resource: G) -> Self {
49        Self { id, resource }
50    }
51}
52
53/// One backend-native route paired with the opaque identity selected by the manifest.
54pub struct RealizedCommunicationRoute<R> {
55    id: CommunicationRouteId,
56    resource: R,
57}
58
59impl<R> RealizedCommunicationRoute<R> {
60    /// Binds one native resource to its selected opaque route identity.
61    pub const fn new(id: CommunicationRouteId, resource: R) -> Self {
62        Self { id, resource }
63    }
64}
65
66/// Opaque native communication resources paired in manifest order.
67pub 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/// Cloneable bounded-completion and poison authority for one selected communication session.
88///
89/// Backend-owned operations adjacent to the neutral partition driver must
90/// retain this authority instead of a bare native group. Every clone observes
91/// the first terminal communication failure and fails before later submission.
92#[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    /// Creates the poison and bounded-completion authority selected by one manifest.
107    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    /// Rejects work after any submission, completion, cancellation, or deadline failure.
125    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    /// Returns the bounded completion policy selected by the communication manifest.
138    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    /// Records a native submission failure after preflight admission succeeded.
154    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    /// Records an exact native completion failure in the shared poison domain.
180    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    /// Waits under the selected bound and permanently poisons this session on failure.
206    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    /// Pairs already-created opaque resources with their authoritative manifest order.
369    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    /// Pairs native resources with an authority already retained by adjacent session APIs.
380    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    /// Authoritative local-rank manifest.
425    pub const fn manifest(&self) -> &CommunicationManifest {
426        &self.manifest
427    }
428
429    /// Shares the selected deadline and terminal poison domain with adjacent operations.
430    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    /// Executes one exact sum reduction using only its narrow backend capability.
590    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    /// Submits one globally ordered wave of exact sum reductions before
629    /// waiting for any member of the wave.
630    ///
631    /// This is required for zero-work pipeline participants: an active peer
632    /// may retain a lazy chain containing the entire wave, so waiting after
633    /// the first zero-work submission would prevent later matching
634    /// submissions from ever entering the native executor.
635    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    /// Executes one exact equal-count gather.
682    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    /// Executes one exact unequal-count gather.
736    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    /// Executes the selected expert plan's exact variable-count exchange without exposing EP semantics.
792    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    /// Completes exact local tensor dependencies before the caller advances an
1274    /// architecture-declared distributed execution wave.
1275    ///
1276    /// This is distinct from a collective: it contributes no tensor value and
1277    /// exists only to make lazy predecessors reach the selected bounded
1278    /// completion policy while every rank is still at the matching wave
1279    /// position.
1280    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/// Canonical shared-session phases whose local status must be propagated.
1308#[derive(Debug, Clone, Copy, Eq, PartialEq)]
1309#[non_exhaustive]
1310pub enum DistributedExecutionPhase {
1311    /// Every rank captured the state checkpoint required for transactional rollback.
1312    StateCheckpoint,
1313    /// Every rank accepted the exact prompt-cache load identity and input.
1314    PromptCacheLoadPreflight,
1315    /// Every stateful rank loaded and validated a provisional cache shard.
1316    PromptCacheLoadPreparation,
1317    /// Every rank accepted the exact prompt-cache save identity and input.
1318    PromptCacheSavePreflight,
1319    /// Every stateful rank serialized and validated an unpublished cache shard.
1320    PromptCacheSavePreparation,
1321    /// Every stateful rank reversibly published its prepared cache shard.
1322    PromptCacheSavePublication,
1323    /// Every stateful rank captured an opaque manual-control checkpoint.
1324    SessionCheckpoint,
1325    /// Every stateful rank prepared a replacement selected state.
1326    SessionResetPreparation,
1327    /// Every stateful rank restored a checkpoint into provisional state.
1328    SessionRollbackPreparation,
1329    /// Source tensor dependencies reached exact bounded completion for one route.
1330    BoundarySourceCompletion(CommunicationRouteId),
1331    /// Source execution and both endpoint preparations completed for one route.
1332    BoundarySourceReady(CommunicationRouteId),
1333    /// Rank-local graph execution, including boundary transfer and publication.
1334    Execution,
1335    /// Final output observation and intervention.
1336    OutputObservation,
1337    /// Output-owner hidden capture is present and published for prediction.
1338    PredictionTargetCapture,
1339    /// Every rank completed publication of the exact prediction target capture.
1340    PredictionTargetCapturePublication,
1341    /// Every rank validated a provisional prediction-lane target state swap.
1342    PredictionTargetStatePreparation,
1343    /// Every rank completed one typed prediction-only unit operation.
1344    PredictionExtensionExecution,
1345    /// Publication of the authoritative intervened output to every rank.
1346    OutputPublication,
1347    /// Publication of the authoritative sampled token and stop status.
1348    SamplingSynchronization,
1349    /// Exact output and mutable-state mechanism completion.
1350    MechanismCompletion,
1351    /// Final distributed state-commit decision.
1352    Commit,
1353}
1354
1355/// One exact directed architecture boundary after a graph group completes.
1356#[derive(Debug, Clone, Eq, PartialEq)]
1357pub struct PartitionBoundaryRoute {
1358    /// Producing architecture group.
1359    pub source_group: usize,
1360    /// Consuming architecture group.
1361    pub destination_group: usize,
1362    /// World rank selected by architecture ownership as the producer.
1363    pub source_rank: usize,
1364    /// World rank selected by architecture ownership as the consumer.
1365    pub destination_rank: usize,
1366    /// Opaque route selected in the communication manifest.
1367    pub route: CommunicationRouteId,
1368}
1369
1370/// Authoritative output owner and opaque publication group.
1371#[derive(Debug, Clone, Copy)]
1372pub struct PartitionOutputPublication {
1373    /// Opaque group containing the output owner and every session participant.
1374    pub group: CollectiveGroupId,
1375    /// World rank that owns architecture projection and supplies source data.
1376    pub owner_rank: usize,
1377}
1378
1379/// Exact public-output authority projected into one selected communication group.
1380#[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    /// Architecture-selected world rank supplying authoritative logits.
1389    pub const fn owner_rank(self) -> usize {
1390        self.owner_rank
1391    }
1392
1393    /// Owner's exact local rank within the selected publication group.
1394    pub const fn owner_group_rank(self) -> usize {
1395        self.owner_group_rank
1396    }
1397
1398    /// Whether this manifest rank may expose logits through the public adapter.
1399    pub const fn local_public_output(self) -> bool {
1400        self.local_public_output
1401    }
1402}
1403
1404/// Immutable rank-local execution and communication plan.
1405#[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    /// Validates group slots, local ownership, route dependencies, and output ownership.
1418    #[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    /// Rank-local execution drivers in architecture-group order.
1575    pub fn drivers(&self) -> &[Option<LayeredPartitionDriver>] {
1576        &self.drivers
1577    }
1578
1579    /// Architecture-selected point-to-point boundary routes.
1580    pub fn routes(&self) -> &[PartitionBoundaryRoute] {
1581        &self.routes
1582    }
1583
1584    /// Selected output publication, including its exact opaque session group.
1585    pub const fn publication(&self) -> Option<PartitionOutputPublication> {
1586        self.publication
1587    }
1588
1589    /// Resolves architecture-selected publication authority without backend
1590    /// topology inference.
1591    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    /// Exact opaque session group used for post-publication commit synchronization.
1619    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
1639/// Statically dispatched architecture/provider adapter used by the production driver.
1640///
1641/// Direct, routed-provider, and composite prepared-input implementations share this
1642/// interface without erasing tensors or per-unit execution. Typed architecture boundary
1643/// values are encoded and decoded inside the adapter; only their validated native tensor
1644/// bundle crosses the communication seam.
1645pub 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    /// Per-invocation architecture-owned state, including the prepared model input.
1655    type Pass<'a>;
1656
1657    /// Starts one invocation without traversing unowned groups.
1658    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    /// Reports request activity only for architecture-declared optional roots.
1667    fn request_group_active(&self, pass: &Self::Pass<'_>, group: usize) -> Result<bool, A::Error>;
1668
1669    /// Whether inactive ranks submit collectives during pipeline-stage waves.
1670    ///
1671    /// A failed wave is communication-indeterminate and therefore permanently
1672    /// fences the selected communication authority on every rank after phase
1673    /// agreement. Ordinary pipeline execution leaves this disabled so a
1674    /// deterministic architecture failure remains retryable after rollback.
1675    fn has_cross_stage_collective_waves(&self) -> bool {
1676        false
1677    }
1678
1679    /// Executes exactly one locally owned group through its validated partition driver.
1680    #[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    /// Participates in one globally ordered pipeline-stage execution wave.
1693    ///
1694    /// Ordinary executors perform work only on the active stage. Executors
1695    /// with architecture-selected cross-stage collectives may use the wave
1696    /// ordinal on inactive stages to submit their exact zero-work protocol.
1697    /// The execution plan has already proved that an active rank owns a local
1698    /// driver, so absence of a driver is meaningful only for inactive ranks.
1699    #[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    /// Produces source tensors or destination placeholders for one endpoint route.
1735    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    /// Resolves the exact invocation-dependent architecture boundary schema.
1745    fn boundary_schema(
1746        &self,
1747        pass: &Self::Pass<'_>,
1748        route: &PartitionBoundaryRoute,
1749    ) -> Result<ResolvedBoundaryWireSchema, A::Error>;
1750
1751    /// Installs a validated received typed-boundary bundle before its consumer runs.
1752    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    /// Returns projected output on the owner and matching destination storage elsewhere.
1760    ///
1761    /// The value is source data only on the publication root. On other ranks it is the
1762    /// architecture-selected destination/placeholder tensor validated before submission.
1763    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    /// Resolves an output-owner target capture or a matching rank-local placeholder.
1771    ///
1772    /// The default serves non-pipeline executors, where every participant owns
1773    /// the complete target result. Pipeline executors override this through
1774    /// their retained architecture and allocator.
1775    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    /// Runs one typed prediction-only operation against this rank's target partition.
1784    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
1797/// Additive policy for an optional point-to-point boundary path.
1798pub 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    /// Whether this policy admits a selected boundary route.
1806    const ENABLED: bool;
1807
1808    /// Moves one exact architecture boundary.
1809    #[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/// Point-to-point boundary transport requiring no collective capabilities.
1822#[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/// Proof that this execution path selects no point-to-point boundary routes.
1848#[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
1875/// Additive policy for optional root-owned output publication.
1876pub 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    /// Whether this policy admits selected output publication.
1884    const ENABLED: bool;
1885
1886    /// Publishes source data on root and destination storage on other members.
1887    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/// Broadcast output publication requiring no point-to-point or barrier capability.
1898#[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/// Proof that this execution path selects no root-to-group publication.
1923#[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
1949/// Additive policy for optional distributed state-commit agreement.
1950pub 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    /// Whether this policy admits distributed commit agreement.
1958    const ENABLED: bool;
1959
1960    /// Whether this policy propagates an explicit success status at every
1961    /// canonical shared-session phase.
1962    const PHASE_FAILURE_AGREEMENT: bool = false;
1963
1964    /// Returns the conjunction of every member's local phase status.
1965    ///
1966    /// Barrier-only policies intentionally retain the local status here; they
1967    /// do not claim failure propagation.
1968    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    /// Performs the canonical all-rank agreement after an earlier subgroup
1980    /// communication failure poisoned the local session authority.
1981    ///
1982    /// Implementations may bypass the prior poison only for this bounded
1983    /// recovery agreement. The selected session remains poisoned afterwards.
1984    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    /// Returns this rank's honest observation of the globally identified final decision.
1996    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/// Barrier commit agreement requiring no tensor communication capabilities.
2006#[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/// Explicit all-rank failure propagation and final commit agreement.
2033///
2034/// This policy is intentionally separate from [`OpaqueCommitAgreement`]: a
2035/// barrier cannot reveal that another rank reported a local phase failure.
2036#[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/// Proof that this execution path selects no distributed commit barrier.
2088#[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
2129/// Complete rank-local runtime installed behind the shared replicated session.
2130pub 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/// Failure while a fully local layered traversal runs through a selected
2153/// partitioned runtime.
2154///
2155/// This narrow path serves architectures whose public invocation contains
2156/// more than one tensor and whose architecture-owned traversal hook must run
2157/// between execution groups. Communication ownership and admission remain in
2158/// [`PartitionedTextRuntime`]; only the already selected local layered
2159/// traversal is delegated to its reusable executor.
2160#[derive(Debug, thiserror::Error)]
2161pub enum PartitionedTraversalError<ArchitectureError, PolicyError>
2162where
2163    ArchitectureError: std::fmt::Display,
2164    PolicyError: std::fmt::Display,
2165{
2166    /// The selected partition plan is not a fully local traversal.
2167    #[error("partitioned traversal contract failed: {0}")]
2168    Contract(String),
2169    /// The architecture or residency policy rejected execution.
2170    #[error(transparent)]
2171    Execution(#[from] LayerwiseRuntimeError<ArchitectureError, PolicyError>),
2172}
2173
2174/// Output of one architecture-owned traversal through a neutral partition runtime.
2175pub type PartitionedTraversalResult<Tensor, ForwardContext, ArchitectureError, PolicyError> =
2176    Result<(Tensor, ForwardContext), PartitionedTraversalError<ArchitectureError, PolicyError>>;
2177
2178/// Reusable executor for one architecture-owned, fully local traversal inside
2179/// a neutral partitioned runtime.
2180///
2181/// The parallel context is an opaque backend mechanism selected from the
2182/// manifest. Family input, forward context, and traversal decisions remain
2183/// statically typed by the architecture.
2184pub 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
2196/// Neutral choice between an ordinary local traversal and the same traversal
2197/// installed in a selected partition runtime.
2198///
2199/// Concrete backends construct one of these alternatives but do not retain a
2200/// parallel execution enum or select a separate forward implementation.
2201pub enum LayerwiseTraversalRuntime<Direct, Partitioned> {
2202    /// Replicated local traversal.
2203    Direct(Direct),
2204    /// Architecture-selected partition traversal.
2205    Partitioned(Partitioned),
2206}
2207
2208impl<Direct, Partitioned> LayerwiseTraversalRuntime<Direct, Partitioned> {
2209    /// Wraps an ordinary local layered runtime.
2210    pub const fn direct(runtime: Direct) -> Self {
2211        Self::Direct(runtime)
2212    }
2213
2214    /// Wraps an architecture-selected partition runtime.
2215    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    /// Pairs one local layered runtime with its selected opaque parallel context.
2229    pub const fn new(runtime: LayerwiseRuntime<A, B, S, P>, parallel: B::ParallelContext) -> Self {
2230        Self { runtime, parallel }
2231    }
2232
2233    /// Borrows the installed layered runtime for residency reporting.
2234    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    /// Pairs one immutable plan with its rank-local executable and opaque resources.
2250    #[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    /// Runs one architecture-owned traversal through the selected neutral
2343    /// partitioned base.
2344    ///
2345    /// Every architecture group must be local. Point-to-point routes, output
2346    /// publication, and phase agreement use the ordinary partitioned session
2347    /// entry instead and are rejected here rather than silently bypassed.
2348    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    /// Borrows the installed traversal executor for neutral residency reports.
2381    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    /// Runs the selected local traversal without exposing its realization kind
2423    /// to concrete composition.
2424    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    /// Borrows the selected residency policy independently of realization kind.
2445    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    /// Borrows the selected local architecture independently of realization kind.
2453    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/// Stateless strategy marker selecting [`PartitionedTextRuntime`] in the shared session.
2462#[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    /// Creates the statically dispatched partition strategy.
2472    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                // With explicit failure agreement, advance the rank graph one pipeline wave at
2580                // a time. Every source in a wave executes concurrently so TP peers reach their
2581                // collectives together; only after agreement may matching destinations receive.
2582                // A middle stage becomes a source only after its incoming wave was removed.
2583                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/// Cold-path validation or execution failure in the neutral partition driver.
3571#[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    /// A partition communication object omitted the selected bounded-wait contract.
3579    #[error("partition communication manifest has no bounded completion policy")]
3580    MissingBoundedCompletionPolicy,
3581    /// Exact communication did not complete by its selected deadline.
3582    #[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    /// A submitted communication reported a terminal error while completing.
3592    #[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    /// Native submission failed after entering one communication operation.
3602    #[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    /// A prior timed-out or failed operation left this communicator unsafe for reuse.
3612    #[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    /// A recovery-only agreement was invoked without a local failure or prior poison.
3622    #[error("recovery agreement was requested without a prior failure during {phase:?}")]
3623    RecoveryAgreementWithoutFailure { phase: DistributedExecutionPhase },
3624    /// Native resources were not paired one-for-one with manifest descriptors.
3625    #[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    /// A native resource was paired with a different opaque manifest identity.
3633    #[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}