Skip to main content

eredu_core/
scheduler.rs

1//! Fair transactional scheduling independent of event and tensor runtimes.
2
3use crate::consensus::{
4    resolve_completions, validate_disposition, validate_schedule, CompletionObservation,
5    CompletionResolution, ConsensusTransport, ScheduledWork,
6};
7use serde::{Deserialize, Serialize};
8use std::{
9    collections::{BTreeMap, BTreeSet, VecDeque},
10    time::{Duration, Instant},
11};
12
13/// Stable caller-assigned request identity.
14#[derive(Debug, Clone, Copy, Eq, Ord, PartialEq, PartialOrd, Hash, Serialize, Deserialize)]
15pub struct RequestId(u64);
16
17impl RequestId {
18    /// Creates an identity. Zero is valid.
19    pub const fn new(value: u64) -> Self {
20        Self(value)
21    }
22
23    /// Returns its numeric value.
24    pub const fn value(self) -> u64 {
25        self.0
26    }
27}
28
29/// Scheduler-assigned ordered transition identity.
30#[derive(Debug, Clone, Copy, Eq, Ord, PartialEq, PartialOrd, Hash, Serialize, Deserialize)]
31pub struct WorkId {
32    request: RequestId,
33    sequence: u64,
34}
35
36impl WorkId {
37    /// Creates an ordered transition identity.
38    pub const fn new(request: RequestId, sequence: u64) -> Self {
39        Self { request, sequence }
40    }
41
42    /// Request updated by this work.
43    pub const fn request(self) -> RequestId {
44        self.request
45    }
46
47    /// Zero-based transition number.
48    pub const fn sequence(self) -> u64 {
49        self.sequence
50    }
51}
52
53/// Authoritative lifecycle of one transition.
54#[derive(Debug, Clone, Copy, Eq, PartialEq, Serialize, Deserialize)]
55#[serde(rename_all = "snake_case")]
56pub enum WorkLifecycle {
57    /// Accepted but not branched.
58    Queued,
59    /// Branch exists; nothing submitted.
60    Prepared,
61    /// Exact completion exists and is incomplete.
62    Submitted,
63    /// Completion succeeded and commit is resolving.
64    Completing,
65    /// Branch was published.
66    Committed,
67    /// Submitted work was cancelled and is retained until exact completion.
68    Abandoned,
69    /// Preparation, execution, completion, or commit failed.
70    Failed,
71}
72
73/// Observable lifecycle of a request.
74#[derive(Debug, Clone, Copy, Eq, PartialEq, Serialize, Deserialize)]
75#[serde(rename_all = "snake_case")]
76pub enum RequestStatus {
77    /// Owns canonical state.
78    Active,
79    /// Finished normally.
80    Finished,
81    /// Explicitly cancelled.
82    Cancelled,
83    /// Deadline expired.
84    DeadlineExceeded,
85    /// Request-local failure.
86    Failed,
87}
88
89/// Cause attached to cancellation.
90#[derive(Debug, Clone, Copy, Eq, PartialEq, Serialize, Deserialize)]
91#[serde(rename_all = "snake_case")]
92pub enum CancellationCause {
93    /// Caller cancellation.
94    Explicit,
95    /// Deadline expiry.
96    Deadline,
97}
98
99/// Capacity and cooperative-preemption controls.
100#[derive(Debug, Clone, Copy, Eq, PartialEq, Serialize, Deserialize)]
101pub struct SchedulerLimits {
102    /// Maximum active requests.
103    pub max_active_requests: usize,
104    /// Maximum accepted nonterminal work.
105    pub max_queued_work: usize,
106    /// Maximum submissions per scheduling turn.
107    pub max_new_submissions_per_turn: usize,
108    /// Maximum global in-flight work.
109    pub max_in_flight_global: usize,
110    /// Maximum per-request in-flight work.
111    pub max_in_flight_per_request: usize,
112    /// Program-defined operations per transition.
113    pub execution_slice: usize,
114}
115
116impl SchedulerLimits {
117    /// Creates conservative bounds.
118    pub fn new(max_active_requests: usize, max_queued_work: usize) -> Result<Self, SchedulerError> {
119        Self::with_execution_bounds(
120            max_active_requests,
121            max_queued_work,
122            1,
123            max_active_requests,
124            1,
125            usize::MAX,
126        )
127    }
128
129    /// Creates explicit bounds.
130    pub fn with_execution_bounds(
131        max_active_requests: usize,
132        max_queued_work: usize,
133        max_new_submissions_per_turn: usize,
134        max_in_flight_global: usize,
135        max_in_flight_per_request: usize,
136        execution_slice: usize,
137    ) -> Result<Self, SchedulerError> {
138        let values = [
139            max_active_requests,
140            max_queued_work,
141            max_new_submissions_per_turn,
142            max_in_flight_global,
143            max_in_flight_per_request,
144            execution_slice,
145        ];
146        if values.contains(&0) {
147            return Err(SchedulerError::InvalidLimits(values));
148        }
149        Ok(Self {
150            max_active_requests,
151            max_queued_work,
152            max_new_submissions_per_turn,
153            max_in_flight_global,
154            max_in_flight_per_request,
155            execution_slice,
156        })
157    }
158}
159
160impl Default for SchedulerLimits {
161    fn default() -> Self {
162        Self {
163            max_active_requests: 64,
164            max_queued_work: 256,
165            max_new_submissions_per_turn: 1,
166            max_in_flight_global: 64,
167            max_in_flight_per_request: 1,
168            execution_slice: usize::MAX,
169        }
170    }
171}
172
173/// Stable semantic description of one work item.
174pub trait WorkDescriptor {
175    /// Descriptor construction error.
176    type Error: std::error::Error;
177
178    /// Appends every semantic field in stable wire order.
179    fn encode_descriptor(&self, output: &mut Vec<u32>) -> Result<(), Self::Error>;
180
181    /// Number of program-defined non-preemptible operations in this transition.
182    fn execution_slice_size(&self) -> usize {
183        1
184    }
185}
186
187/// Transaction over canonical semantic state.
188pub trait SemanticStateTransaction {
189    /// Transition-local branch.
190    type Branch;
191    /// State error.
192    type Error: std::error::Error;
193
194    /// Creates an unpublished branch.
195    fn branch(&self) -> Result<Self::Branch, Self::Error>;
196
197    /// Publishes a completed branch.
198    fn commit_branch(&mut self, branch: Self::Branch) -> Result<(), Self::Error>;
199
200    /// Discards an unpublished branch.
201    fn discard_branch(branch: Self::Branch) -> Result<(), Self::Error> {
202        drop(branch);
203        Ok(())
204    }
205
206    /// Whether independent branches may share the current canonical base.
207    fn permits_parallel_branches(&self) -> bool {
208        false
209    }
210}
211
212/// Exact backend completion retaining submitted resources.
213pub trait TransitionOutput {
214    /// Completion observation error.
215    type Error: std::error::Error;
216
217    /// Nonblocking exact completion observation.
218    fn is_complete(&self) -> Result<bool, Self::Error>;
219
220    /// Stable backend name used only for capability telemetry.
221    fn backend_name(&self) -> Option<String> {
222        None
223    }
224
225    /// Whether already executing work can be physically interrupted.
226    fn physically_preemptible(&self) -> bool {
227        false
228    }
229
230    /// Explicitly retained resource count.
231    fn retained_resources(&self) -> usize;
232}
233
234#[derive(Debug)]
235struct Request<W, S> {
236    state: S,
237    next: u64,
238    pending: VecDeque<Queued<W>>,
239}
240
241#[derive(Debug)]
242struct Queued<W> {
243    id: WorkId,
244    work: W,
245    deadline: Option<Instant>,
246}
247
248#[derive(Debug)]
249struct Prepared<W, B> {
250    id: WorkId,
251    work: W,
252    descriptor: Vec<u32>,
253    branch: B,
254    deadline: Option<Instant>,
255}
256
257#[derive(Debug, Clone, Copy)]
258enum Disposition {
259    Publish,
260    Abandon { cancelled_at: Instant },
261    Fail,
262}
263
264#[derive(Debug)]
265struct Submitted<W, B, O> {
266    id: WorkId,
267    work: W,
268    branch: B,
269    output: O,
270    disposition: Disposition,
271}
272
273/// Result of one scheduling turn.
274#[derive(Debug)]
275pub struct SchedulerProgress<W, O> {
276    /// Newly submitted transitions.
277    pub newly_submitted: usize,
278    /// Successfully committed work and backend outputs.
279    pub committed: Vec<(WorkId, W, O)>,
280    /// Failed work with structured scheduler context.
281    pub failed: Vec<(WorkId, SchedulerError)>,
282}
283
284impl<W, O> Default for SchedulerProgress<W, O> {
285    fn default() -> Self {
286        Self {
287            newly_submitted: 0,
288            committed: Vec::new(),
289            failed: Vec::new(),
290        }
291    }
292}
293
294/// Static scheduler capabilities and configured bounds.
295#[derive(Debug, Clone, Eq, PartialEq, Serialize, Deserialize)]
296pub struct SchedulerCapabilities {
297    /// Configured scheduler limits.
298    pub limits: SchedulerLimits,
299    /// Backends observed on submitted transitions.
300    pub observed_backends: Vec<String>,
301    /// Whether every observed backend output supports physical preemption.
302    pub executing_work_physically_preemptible: bool,
303    /// Exact description of the unavoidable cancellation interval.
304    pub non_preemptible_interval: String,
305}
306
307/// Snapshot of scheduler occupancy and cumulative telemetry.
308#[derive(Debug, Clone, Copy, Default, Eq, PartialEq, Serialize, Deserialize)]
309pub struct SchedulerReport {
310    /// Active requests.
311    pub active_requests: usize,
312    /// Queued work.
313    pub queued_work: usize,
314    /// Prepared work.
315    pub prepared_work: usize,
316    /// Submitted work still eligible to publish.
317    pub submitted_in_flight_work: usize,
318    /// Work currently resolving publication.
319    pub completing_work: usize,
320    /// Cancelled submitted work awaiting exact completion.
321    pub abandoned_in_flight_work: usize,
322    /// Globally failed work retained until every rank reaches completion.
323    pub failed_in_flight_work: usize,
324    /// All submitted work, including abandoned work.
325    pub current_in_flight_work: usize,
326    /// Peak submitted work.
327    pub peak_in_flight_work: usize,
328    /// Peak accepted nonterminal work.
329    pub peak_queued_work: usize,
330    /// Total accepted work.
331    pub submitted_work: u64,
332    /// Total committed work.
333    pub completed_work: u64,
334    /// Total failed work.
335    pub failed_work: u64,
336    /// Work discarded before submission.
337    pub discarded_work: u64,
338    /// Work cancelled before submission.
339    pub cancellation_before_submission: u64,
340    /// Work cancelled after submission.
341    pub cancellation_after_submission: u64,
342    /// Abandoned work released after exact completion.
343    pub abandoned_released_work: u64,
344    /// Resources retained by abandoned work.
345    pub abandoned_retained_resources: usize,
346    /// Maximum resources retained by abandoned work.
347    pub peak_abandoned_retained_resources: usize,
348    /// Last cancellation-to-release latency in nanoseconds.
349    pub last_cancellation_to_release_ns: Option<u128>,
350    /// Maximum cancellation-to-release latency in nanoseconds.
351    pub max_cancellation_to_release_ns: Option<u128>,
352    /// Requests ended normally.
353    pub finished_requests: u64,
354    /// Requests cancelled explicitly.
355    pub cancelled_requests: u64,
356    /// Requests ended by deadline.
357    pub deadline_expired_requests: u64,
358    /// Turns which submitted at least one transition.
359    pub drain_cycles: u64,
360    /// Configured per-turn submission bound.
361    pub configured_submission_bound: usize,
362    /// Configured execution-slice bound.
363    pub configured_slice_bound: usize,
364    /// Whether unsafe distributed ordering poisoned the scheduler.
365    pub poisoned: bool,
366}
367
368/// Fair scheduler whose state transitions are independent of backend objects.
369#[derive(Debug)]
370pub struct Scheduler<W, S: SemanticStateTransaction, O: TransitionOutput> {
371    limits: SchedulerLimits,
372    requests: BTreeMap<RequestId, Request<W, S>>,
373    terminal: BTreeMap<RequestId, RequestStatus>,
374    ready: VecDeque<RequestId>,
375    prepared: VecDeque<Prepared<W, S::Branch>>,
376    submitted: Vec<Submitted<W, S::Branch, O>>,
377    lifecycle: BTreeMap<WorkId, WorkLifecycle>,
378    accepted_work: usize,
379    peak_accepted_work: usize,
380    peak_in_flight_work: usize,
381    submitted_work: u64,
382    completed_work: u64,
383    failed_work: u64,
384    discarded_work: u64,
385    cancellation_before_submission: u64,
386    cancellation_after_submission: u64,
387    abandoned_released_work: u64,
388    peak_abandoned_retained_resources: usize,
389    last_cancellation_to_release: Option<Duration>,
390    max_cancellation_to_release: Option<Duration>,
391    finished_requests: u64,
392    cancelled_requests: u64,
393    deadline_expired_requests: u64,
394    drain_cycles: u64,
395    observed_backends: BTreeSet<String>,
396    all_outputs_preemptible: bool,
397    poisoned: Option<String>,
398}
399
400impl<W, S: SemanticStateTransaction, O: TransitionOutput> Scheduler<W, S, O> {
401    /// Creates an empty scheduler under validated limits.
402    pub fn new(limits: SchedulerLimits) -> Result<Self, SchedulerError> {
403        let limits = SchedulerLimits::with_execution_bounds(
404            limits.max_active_requests,
405            limits.max_queued_work,
406            limits.max_new_submissions_per_turn,
407            limits.max_in_flight_global,
408            limits.max_in_flight_per_request,
409            limits.execution_slice,
410        )?;
411        Ok(Self {
412            limits,
413            requests: BTreeMap::new(),
414            terminal: BTreeMap::new(),
415            ready: VecDeque::new(),
416            prepared: VecDeque::new(),
417            submitted: Vec::new(),
418            lifecycle: BTreeMap::new(),
419            accepted_work: 0,
420            peak_accepted_work: 0,
421            peak_in_flight_work: 0,
422            submitted_work: 0,
423            completed_work: 0,
424            failed_work: 0,
425            discarded_work: 0,
426            cancellation_before_submission: 0,
427            cancellation_after_submission: 0,
428            abandoned_released_work: 0,
429            peak_abandoned_retained_resources: 0,
430            last_cancellation_to_release: None,
431            max_cancellation_to_release: None,
432            finished_requests: 0,
433            cancelled_requests: 0,
434            deadline_expired_requests: 0,
435            drain_cycles: 0,
436            observed_backends: BTreeSet::new(),
437            all_outputs_preemptible: true,
438            poisoned: None,
439        })
440    }
441
442    /// Validates identity and active-request capacity before allocating state.
443    pub fn validate_registration(&self, id: RequestId) -> Result<(), SchedulerError> {
444        self.ensure_ready()?;
445        if self.requests.contains_key(&id) || self.terminal.contains_key(&id) {
446            return Err(SchedulerError::DuplicateRequest(id));
447        }
448        if self.requests.len() >= self.limits.max_active_requests {
449            return Err(SchedulerError::Capacity(format!(
450                "scheduler active-request capacity {} is exhausted",
451                self.limits.max_active_requests
452            )));
453        }
454        Ok(())
455    }
456
457    /// Registers canonical request state.
458    pub fn register(&mut self, id: RequestId, state: S) -> Result<(), SchedulerError> {
459        self.validate_registration(id)?;
460        self.requests.insert(
461            id,
462            Request {
463                state,
464                next: 0,
465                pending: VecDeque::new(),
466            },
467        );
468        Ok(())
469    }
470
471    /// Returns immutable canonical state for an active request.
472    pub fn request_state(&self, id: RequestId) -> Option<&S> {
473        self.requests.get(&id).map(|entry| &entry.state)
474    }
475
476    /// Returns mutable canonical state only while no branch exists.
477    pub fn request_state_mut(&mut self, id: RequestId) -> Result<&mut S, SchedulerError> {
478        self.ensure_ready()?;
479        if self.branch_count(id) != 0 {
480            return Err(SchedulerError::State(format!(
481                "request {} has prepared or submitted state branches",
482                id.value()
483            )));
484        }
485        self.requests
486            .get_mut(&id)
487            .map(|entry| &mut entry.state)
488            .ok_or(SchedulerError::UnknownRequest(id))
489    }
490
491    /// Enqueues one transition without a deadline.
492    pub fn enqueue(&mut self, request: RequestId, work: W) -> Result<WorkId, SchedulerError> {
493        self.enqueue_with_deadline(request, work, None)
494    }
495
496    /// Enqueues one transition with an optional absolute deadline.
497    pub fn enqueue_with_deadline(
498        &mut self,
499        request: RequestId,
500        work: W,
501        deadline: Option<Instant>,
502    ) -> Result<WorkId, SchedulerError> {
503        Ok(self
504            .enqueue_batch_with_deadlines(request, vec![(work, deadline)])?
505            .pop()
506            .expect("one work item was supplied"))
507    }
508
509    /// Atomically enqueues an ordered batch.
510    pub fn enqueue_batch(
511        &mut self,
512        request: RequestId,
513        work: Vec<W>,
514    ) -> Result<Vec<WorkId>, SchedulerError> {
515        self.enqueue_batch_with_deadlines(
516            request,
517            work.into_iter().map(|work| (work, None)).collect(),
518        )
519    }
520
521    fn enqueue_batch_with_deadlines(
522        &mut self,
523        request: RequestId,
524        work: Vec<(W, Option<Instant>)>,
525    ) -> Result<Vec<WorkId>, SchedulerError> {
526        self.ensure_ready()?;
527        let requested = work.len();
528        let accepted_after = self
529            .accepted_work
530            .checked_add(requested)
531            .ok_or_else(|| SchedulerError::Capacity("scheduler occupancy overflow".into()))?;
532        if accepted_after > self.limits.max_queued_work {
533            return Err(SchedulerError::Capacity(format!(
534                "scheduler queue capacity {} cannot accept {requested} items with {} outstanding",
535                self.limits.max_queued_work, self.accepted_work
536            )));
537        }
538        let entry = self
539            .requests
540            .get_mut(&request)
541            .ok_or(SchedulerError::UnknownRequest(request))?;
542        let count = u64::try_from(requested)
543            .map_err(|_| SchedulerError::Capacity("work batch length exceeds u64".into()))?;
544        let next = entry
545            .next
546            .checked_add(count)
547            .ok_or_else(|| SchedulerError::Capacity("work identity space exhausted".into()))?;
548        let was_empty = entry.pending.is_empty();
549        let mut ids = Vec::with_capacity(requested);
550        for (offset, (work, deadline)) in work.into_iter().enumerate() {
551            let id = WorkId::new(request, entry.next + offset as u64);
552            entry.pending.push_back(Queued { id, work, deadline });
553            self.lifecycle.insert(id, WorkLifecycle::Queued);
554            ids.push(id);
555        }
556        entry.next = next;
557        if was_empty && requested != 0 {
558            self.ready.push_back(request);
559        }
560        self.accepted_work = accepted_after;
561        self.peak_accepted_work = self.peak_accepted_work.max(accepted_after);
562        self.submitted_work = self.submitted_work.saturating_add(count);
563        Ok(ids)
564    }
565
566    /// Prepares bounded branches in round-robin request order.
567    pub fn prepare_bounded(&mut self, limit: usize, now: Instant) -> Result<usize, SchedulerError>
568    where
569        W: WorkDescriptor,
570    {
571        self.ensure_ready()?;
572        if limit == 0 {
573            return Err(SchedulerError::Capacity(
574                "scheduler preparation bound must be positive".into(),
575            ));
576        }
577        self.expire_deadlines(now)?;
578        let mut count = 0;
579        let mut stalled = 0;
580        while count < limit && !self.ready.is_empty() {
581            let request = self.ready.pop_front().expect("ready queue is nonempty");
582            let branches = self.branch_count(request);
583            let can_branch = self
584                .requests
585                .get(&request)
586                .is_some_and(|entry| branches == 0 || entry.state.permits_parallel_branches());
587            if !can_branch || branches >= self.limits.max_in_flight_per_request {
588                self.ready.push_back(request);
589                stalled += 1;
590                if stalled >= self.ready.len() {
591                    break;
592                }
593                continue;
594            }
595            stalled = 0;
596            let queued = self
597                .requests
598                .get_mut(&request)
599                .and_then(|entry| entry.pending.pop_front())
600                .expect("ready request owns queued work");
601            if self
602                .requests
603                .get(&request)
604                .is_some_and(|entry| !entry.pending.is_empty())
605            {
606                self.ready.push_back(request);
607            }
608            let slice = queued.work.execution_slice_size();
609            if slice == 0 || slice > self.limits.execution_slice {
610                self.fail_before_submission(queued.id);
611                return Err(SchedulerError::Descriptor(format!(
612                    "work {:?} execution slice {slice} exceeds configured bound {}",
613                    queued.id, self.limits.execution_slice
614                )));
615            }
616            let mut descriptor = Vec::new();
617            if let Err(error) = queued.work.encode_descriptor(&mut descriptor) {
618                self.fail_before_submission(queued.id);
619                return Err(SchedulerError::Descriptor(error.to_string()));
620            }
621            let branch = match self
622                .requests
623                .get(&request)
624                .expect("active request exists")
625                .state
626                .branch()
627            {
628                Ok(branch) => branch,
629                Err(error) => {
630                    self.fail_before_submission(queued.id);
631                    return Err(SchedulerError::State(error.to_string()));
632                }
633            };
634            self.lifecycle.insert(queued.id, WorkLifecycle::Prepared);
635            self.prepared.push_back(Prepared {
636                id: queued.id,
637                work: queued.work,
638                descriptor,
639                branch,
640                deadline: queued.deadline,
641            });
642            count += 1;
643        }
644        Ok(count)
645    }
646
647    /// Submits a bounded prepared prefix through the injected backend adapter.
648    pub fn submit_prepared<E>(
649        &mut self,
650        now: Instant,
651        mut execute: impl FnMut(WorkId, &W, &mut S::Branch) -> Result<O, E>,
652    ) -> Result<usize, SchedulerError>
653    where
654        E: std::error::Error,
655    {
656        self.ensure_ready()?;
657        self.expire_deadlines(now)?;
658        let capacity = self
659            .limits
660            .max_in_flight_global
661            .saturating_sub(self.submitted.len())
662            .min(self.limits.max_new_submissions_per_turn);
663        let mut count = 0;
664        while count < capacity {
665            let Some(mut prepared) = self.prepared.pop_front() else {
666                break;
667            };
668            if prepared.deadline.is_some_and(|deadline| deadline <= now) {
669                let request = prepared.id.request();
670                self.prepared.push_front(prepared);
671                self.cancel_internal(request, CancellationCause::Deadline, now)?;
672                continue;
673            }
674            let output = match execute(prepared.id, &prepared.work, &mut prepared.branch) {
675                Ok(output) => output,
676                Err(error) => {
677                    let id = prepared.id;
678                    let discard = S::discard_branch(prepared.branch).err();
679                    self.lifecycle.insert(id, WorkLifecycle::Failed);
680                    self.failed_work = self.failed_work.saturating_add(1);
681                    self.accepted_work = self.accepted_work.saturating_sub(1);
682                    self.fail_request(id.request());
683                    let message = discard.map_or_else(
684                        || error.to_string(),
685                        |discard| format!("{error}; branch discard also failed: {discard}"),
686                    );
687                    return Err(SchedulerError::Submission(message));
688                }
689            };
690            if let Some(backend) = output.backend_name() {
691                self.observed_backends.insert(backend);
692            }
693            self.all_outputs_preemptible &= output.physically_preemptible();
694            self.lifecycle.insert(prepared.id, WorkLifecycle::Submitted);
695            self.submitted.push(Submitted {
696                id: prepared.id,
697                work: prepared.work,
698                branch: prepared.branch,
699                output,
700                disposition: Disposition::Publish,
701            });
702            count += 1;
703            self.peak_in_flight_work = self.peak_in_flight_work.max(self.submitted.len());
704        }
705        if count != 0 {
706            self.drain_cycles = self.drain_cycles.saturating_add(1);
707        }
708        self.update_abandoned_resource_peak();
709        Ok(count)
710    }
711
712    /// Polls exact completions and publishes only successful work.
713    pub fn poll_completions(&mut self, now: Instant) -> SchedulerProgress<W, O> {
714        let mut progress = SchedulerProgress::default();
715        let mut retained = Vec::with_capacity(self.submitted.len());
716        for submitted in std::mem::take(&mut self.submitted) {
717            match submitted.output.is_complete() {
718                Ok(false) => retained.push(submitted),
719                Ok(true) => self.resolve_completed(submitted, now, &mut progress),
720                Err(error) => {
721                    let id = submitted.id;
722                    let already_failed = matches!(submitted.disposition, Disposition::Fail);
723                    let discard = S::discard_branch(submitted.branch).err();
724                    self.lifecycle.insert(id, WorkLifecycle::Failed);
725                    if !already_failed {
726                        self.failed_work = self.failed_work.saturating_add(1);
727                    }
728                    self.accepted_work = self.accepted_work.saturating_sub(1);
729                    self.fail_request(id.request());
730                    progress.failed.push((
731                        id,
732                        SchedulerError::Completion(discard.map_or_else(
733                            || error.to_string(),
734                            |discard| format!("{error}; branch discard also failed: {discard}"),
735                        )),
736                    ));
737                }
738            }
739        }
740        // A completion failure can make its whole request terminal while sibling
741        // submissions from the same poll batch are still incomplete. Those
742        // siblings remain backend-owned until their exact completions resolve,
743        // but can no longer publish into the removed canonical state.
744        for submitted in &mut retained {
745            if !self.requests.contains_key(&submitted.id.request())
746                && matches!(submitted.disposition, Disposition::Publish)
747            {
748                submitted.disposition = Disposition::Abandon { cancelled_at: now };
749                self.lifecycle
750                    .insert(submitted.id, WorkLifecycle::Abandoned);
751            }
752        }
753        self.submitted = retained;
754        self.update_abandoned_resource_peak();
755        progress
756    }
757
758    /// Runs one bounded local scheduling turn.
759    pub fn run_local_turn<E>(
760        &mut self,
761        now: Instant,
762        execute: impl FnMut(WorkId, &W, &mut S::Branch) -> Result<O, E>,
763    ) -> Result<SchedulerProgress<W, O>, SchedulerError>
764    where
765        W: WorkDescriptor,
766        E: std::error::Error,
767    {
768        self.ensure_ready()?;
769        let mut progress = self.poll_completions(now);
770        self.prepare_bounded(self.limits.max_new_submissions_per_turn, now)?;
771        progress.newly_submitted = self.submit_prepared(now, execute)?;
772        let after_submit = self.poll_completions(now);
773        progress.committed.extend(after_submit.committed);
774        progress.failed.extend(after_submit.failed);
775        Ok(progress)
776    }
777
778    /// Runs one bounded turn with topology-wide schedule and completion consensus.
779    pub fn run_distributed_turn<T, E>(
780        &mut self,
781        protocol: u64,
782        transport: &T,
783        now: Instant,
784        execute: impl FnMut(WorkId, &W, &mut S::Branch) -> Result<O, E>,
785    ) -> Result<SchedulerProgress<W, O>, SchedulerError>
786    where
787        W: WorkDescriptor,
788        T: ConsensusTransport,
789        E: std::error::Error,
790    {
791        self.ensure_ready()?;
792        self.expire_deadlines_distributed(protocol, transport, now)?;
793        let mut progress = self.poll_distributed(protocol, transport, now)?;
794        self.prepare_bounded(self.limits.max_new_submissions_per_turn, now)?;
795        let plan = self
796            .prepared
797            .iter()
798            .take(
799                self.limits
800                    .max_in_flight_global
801                    .saturating_sub(self.submitted.len())
802                    .min(self.limits.max_new_submissions_per_turn),
803            )
804            .map(|work| ScheduledWork {
805                id: work.id,
806                descriptor: &work.descriptor,
807            })
808            .collect::<Vec<_>>();
809        if let Err(error) = validate_schedule(transport, &plan, self.drain_cycles, protocol) {
810            self.poison(error.to_string(), now);
811            return Err(SchedulerError::Consensus(error.to_string()));
812        }
813        progress.newly_submitted = match self.submit_prepared(now, execute) {
814            Ok(count) => count,
815            Err(error) => {
816                self.poison(error.to_string(), now);
817                return Err(error);
818            }
819        };
820        Ok(progress)
821    }
822
823    /// Reaches topology-wide consensus before cancelling a request.
824    pub fn cancel_distributed<T: ConsensusTransport>(
825        &mut self,
826        protocol: u64,
827        request: RequestId,
828        transport: &T,
829        now: Instant,
830    ) -> Result<(), SchedulerError> {
831        self.ensure_ready()?;
832        if let Err(error) =
833            validate_disposition(transport, protocol, request, CancellationCause::Explicit)
834        {
835            self.poison(error.to_string(), now);
836            return Err(SchedulerError::Consensus(error.to_string()));
837        }
838        self.cancel_internal(request, CancellationCause::Explicit, now)
839    }
840
841    /// Marks a request finished and discards its queued work.
842    pub fn finish(&mut self, request: RequestId) -> Result<(), SchedulerError> {
843        self.ensure_ready()?;
844        if self
845            .submitted
846            .iter()
847            .any(|work| work.id.request() == request)
848        {
849            return Err(SchedulerError::State(format!(
850                "request {} still has submitted work",
851                request.value()
852            )));
853        }
854        let entry = self
855            .requests
856            .remove(&request)
857            .ok_or(SchedulerError::UnknownRequest(request))?;
858        let queued = entry.pending.len();
859        for work in entry.pending {
860            self.lifecycle.insert(work.id, WorkLifecycle::Abandoned);
861        }
862        self.ready.retain(|candidate| *candidate != request);
863        let (prepared, discard_error) = self.discard_prepared_for_request(request);
864        let discarded = queued + prepared;
865        self.accepted_work = self.accepted_work.saturating_sub(discarded);
866        self.discarded_work = self.discarded_work.saturating_add(discarded as u64);
867        self.terminal.insert(
868            request,
869            if discard_error.is_some() {
870                RequestStatus::Failed
871            } else {
872                RequestStatus::Finished
873            },
874        );
875        self.finished_requests = self.finished_requests.saturating_add(1);
876        discard_error.map_or(Ok(()), |error| Err(SchedulerError::State(error)))
877    }
878
879    /// Cancels a request, retaining submitted resources until exact completion.
880    pub fn cancel(&mut self, request: RequestId) -> Result<(), SchedulerError> {
881        self.cancel_internal(request, CancellationCause::Explicit, Instant::now())
882    }
883
884    /// Releases an idle request and returns its canonical state.
885    pub fn release(&mut self, request: RequestId) -> Result<S, SchedulerError> {
886        self.ensure_ready()?;
887        let entry = self
888            .requests
889            .get(&request)
890            .ok_or(SchedulerError::UnknownRequest(request))?;
891        if !entry.pending.is_empty() || self.branch_count(request) != 0 {
892            return Err(SchedulerError::State(format!(
893                "request {} still owns unpublished work",
894                request.value()
895            )));
896        }
897        Ok(self
898            .requests
899            .remove(&request)
900            .expect("checked active request")
901            .state)
902    }
903
904    /// Removes a terminal identity so it may be reused explicitly.
905    pub fn forget_terminal(&mut self, request: RequestId) -> Result<RequestStatus, SchedulerError> {
906        self.ensure_ready()?;
907        if self
908            .submitted
909            .iter()
910            .any(|work| work.id.request() == request)
911        {
912            return Err(SchedulerError::State(format!(
913                "request {} still has retained abandoned work",
914                request.value()
915            )));
916        }
917        self.terminal
918            .remove(&request)
919            .ok_or(SchedulerError::UnknownRequest(request))
920    }
921
922    /// Returns active or terminal request status.
923    pub fn request_status(&self, id: RequestId) -> Option<RequestStatus> {
924        self.requests
925            .contains_key(&id)
926            .then_some(RequestStatus::Active)
927            .or_else(|| self.terminal.get(&id).copied())
928    }
929
930    /// Returns a work lifecycle state.
931    pub fn work_lifecycle(&self, id: WorkId) -> Option<WorkLifecycle> {
932        self.lifecycle.get(&id).copied()
933    }
934
935    /// Returns queued transitions for one active request.
936    pub fn queued_for_request(&self, request: RequestId) -> usize {
937        self.requests
938            .get(&request)
939            .map_or(0, |entry| entry.pending.len())
940    }
941
942    /// Returns configured bounds and observed backend capabilities.
943    pub fn capabilities(&self) -> SchedulerCapabilities {
944        SchedulerCapabilities {
945            limits: self.limits,
946            observed_backends: self.observed_backends.iter().cloned().collect(),
947            executing_work_physically_preemptible: !self.observed_backends.is_empty()
948                && self.all_outputs_preemptible,
949            non_preemptible_interval:
950                "from exact backend submission until that transition's completion resolves".into(),
951        }
952    }
953
954    /// Returns current occupancy and cumulative telemetry.
955    pub fn report(&self) -> SchedulerReport {
956        let abandoned = self
957            .submitted
958            .iter()
959            .filter(|work| matches!(work.disposition, Disposition::Abandon { .. }))
960            .count();
961        let failed = self
962            .submitted
963            .iter()
964            .filter(|work| matches!(work.disposition, Disposition::Fail))
965            .count();
966        SchedulerReport {
967            active_requests: self.requests.len(),
968            queued_work: self
969                .requests
970                .values()
971                .map(|entry| entry.pending.len())
972                .sum(),
973            prepared_work: self.prepared.len(),
974            submitted_in_flight_work: self.submitted.len() - abandoned - failed,
975            completing_work: 0,
976            abandoned_in_flight_work: abandoned,
977            failed_in_flight_work: failed,
978            current_in_flight_work: self.submitted.len(),
979            peak_in_flight_work: self.peak_in_flight_work,
980            peak_queued_work: self.peak_accepted_work,
981            submitted_work: self.submitted_work,
982            completed_work: self.completed_work,
983            failed_work: self.failed_work,
984            discarded_work: self.discarded_work,
985            cancellation_before_submission: self.cancellation_before_submission,
986            cancellation_after_submission: self.cancellation_after_submission,
987            abandoned_released_work: self.abandoned_released_work,
988            abandoned_retained_resources: self.abandoned_retained_resources(),
989            peak_abandoned_retained_resources: self.peak_abandoned_retained_resources,
990            last_cancellation_to_release_ns: self
991                .last_cancellation_to_release
992                .map(|duration| duration.as_nanos()),
993            max_cancellation_to_release_ns: self
994                .max_cancellation_to_release
995                .map(|duration| duration.as_nanos()),
996            finished_requests: self.finished_requests,
997            cancelled_requests: self.cancelled_requests,
998            deadline_expired_requests: self.deadline_expired_requests,
999            drain_cycles: self.drain_cycles,
1000            configured_submission_bound: self.limits.max_new_submissions_per_turn,
1001            configured_slice_bound: self.limits.execution_slice,
1002            poisoned: self.poisoned.is_some(),
1003        }
1004    }
1005
1006    /// Returns the unsafe distributed-ordering failure, if one occurred.
1007    pub fn poison_reason(&self) -> Option<&str> {
1008        self.poisoned.as_deref()
1009    }
1010
1011    fn resolve_completed(
1012        &mut self,
1013        submitted: Submitted<W, S::Branch, O>,
1014        now: Instant,
1015        progress: &mut SchedulerProgress<W, O>,
1016    ) {
1017        match submitted.disposition {
1018            Disposition::Abandon { cancelled_at } => {
1019                if let Err(error) = S::discard_branch(submitted.branch) {
1020                    self.lifecycle.insert(submitted.id, WorkLifecycle::Failed);
1021                    self.failed_work = self.failed_work.saturating_add(1);
1022                    progress.failed.push((
1023                        submitted.id,
1024                        SchedulerError::State(format!(
1025                            "failed to discard abandoned state branch: {error}"
1026                        )),
1027                    ));
1028                }
1029                self.abandoned_released_work = self.abandoned_released_work.saturating_add(1);
1030                self.accepted_work = self.accepted_work.saturating_sub(1);
1031                let latency = now.saturating_duration_since(cancelled_at);
1032                self.last_cancellation_to_release = Some(latency);
1033                self.max_cancellation_to_release = Some(
1034                    self.max_cancellation_to_release
1035                        .map_or(latency, |previous| previous.max(latency)),
1036                );
1037            }
1038            Disposition::Publish => {
1039                self.lifecycle
1040                    .insert(submitted.id, WorkLifecycle::Completing);
1041                let Some(request) = self.requests.get_mut(&submitted.id.request()) else {
1042                    if let Err(error) = S::discard_branch(submitted.branch) {
1043                        self.lifecycle.insert(submitted.id, WorkLifecycle::Failed);
1044                        self.failed_work = self.failed_work.saturating_add(1);
1045                        progress.failed.push((
1046                            submitted.id,
1047                            SchedulerError::State(format!(
1048                                "failed to discard unpublished state branch: {error}"
1049                            )),
1050                        ));
1051                    } else {
1052                        self.lifecycle
1053                            .insert(submitted.id, WorkLifecycle::Abandoned);
1054                    }
1055                    self.accepted_work = self.accepted_work.saturating_sub(1);
1056                    return;
1057                };
1058                if let Err(error) = request.state.commit_branch(submitted.branch) {
1059                    self.lifecycle.insert(submitted.id, WorkLifecycle::Failed);
1060                    self.failed_work = self.failed_work.saturating_add(1);
1061                    self.accepted_work = self.accepted_work.saturating_sub(1);
1062                    self.fail_request(submitted.id.request());
1063                    progress
1064                        .failed
1065                        .push((submitted.id, SchedulerError::State(error.to_string())));
1066                    return;
1067                }
1068                self.lifecycle
1069                    .insert(submitted.id, WorkLifecycle::Committed);
1070                self.completed_work = self.completed_work.saturating_add(1);
1071                self.accepted_work = self.accepted_work.saturating_sub(1);
1072                progress
1073                    .committed
1074                    .push((submitted.id, submitted.work, submitted.output));
1075            }
1076            Disposition::Fail => {
1077                let discard = S::discard_branch(submitted.branch).err();
1078                self.lifecycle.insert(submitted.id, WorkLifecycle::Failed);
1079                self.accepted_work = self.accepted_work.saturating_sub(1);
1080                self.fail_request(submitted.id.request());
1081                progress.failed.push((
1082                    submitted.id,
1083                    SchedulerError::DistributedCompletion(discard.map_or_else(
1084                        || "backend completion failed on at least one rank".into(),
1085                        |error| {
1086                            format!(
1087                                "backend completion failed on at least one rank; branch discard also failed: {error}"
1088                            )
1089                        },
1090                    )),
1091                ));
1092            }
1093        }
1094    }
1095
1096    fn poll_distributed<T: ConsensusTransport>(
1097        &mut self,
1098        protocol: u64,
1099        transport: &T,
1100        now: Instant,
1101    ) -> Result<SchedulerProgress<W, O>, SchedulerError> {
1102        let local = self
1103            .submitted
1104            .iter()
1105            .map(|work| {
1106                let status = match work.output.is_complete() {
1107                    Ok(false) => CompletionObservation::Incomplete,
1108                    Ok(true) => CompletionObservation::Complete,
1109                    Err(_) => CompletionObservation::Failed,
1110                };
1111                (work.id, status)
1112            })
1113            .collect::<Vec<_>>();
1114        let global = match resolve_completions(transport, protocol, &local) {
1115            Ok(global) => global,
1116            Err(error) => {
1117                self.poison(error.to_string(), now);
1118                return Err(SchedulerError::Consensus(error.to_string()));
1119            }
1120        };
1121
1122        let mut progress = SchedulerProgress::default();
1123        let mut retained = Vec::with_capacity(self.submitted.len());
1124        for (mut work, status) in std::mem::take(&mut self.submitted).into_iter().zip(global) {
1125            match status {
1126                CompletionResolution::Incomplete => retained.push(work),
1127                CompletionResolution::Complete => {
1128                    self.resolve_completed(work, now, &mut progress);
1129                }
1130                CompletionResolution::FailedPending => {
1131                    if !matches!(work.disposition, Disposition::Fail) {
1132                        work.disposition = Disposition::Fail;
1133                        self.lifecycle.insert(work.id, WorkLifecycle::Failed);
1134                        self.failed_work = self.failed_work.saturating_add(1);
1135                    }
1136                    retained.push(work);
1137                }
1138                CompletionResolution::FailedComplete => {
1139                    if !matches!(work.disposition, Disposition::Fail) {
1140                        work.disposition = Disposition::Fail;
1141                        self.failed_work = self.failed_work.saturating_add(1);
1142                    }
1143                    self.resolve_completed(work, now, &mut progress);
1144                }
1145            }
1146        }
1147        for work in &mut retained {
1148            if !self.requests.contains_key(&work.id.request())
1149                && matches!(work.disposition, Disposition::Publish)
1150            {
1151                work.disposition = Disposition::Abandon { cancelled_at: now };
1152                self.lifecycle.insert(work.id, WorkLifecycle::Abandoned);
1153            }
1154        }
1155        self.submitted = retained;
1156        self.update_abandoned_resource_peak();
1157        Ok(progress)
1158    }
1159
1160    fn expire_deadlines(&mut self, now: Instant) -> Result<(), SchedulerError> {
1161        let mut expired = BTreeSet::new();
1162        for (request, entry) in &self.requests {
1163            if entry
1164                .pending
1165                .iter()
1166                .any(|work| work.deadline.is_some_and(|deadline| deadline <= now))
1167            {
1168                expired.insert(*request);
1169            }
1170        }
1171        for work in &self.prepared {
1172            if work.deadline.is_some_and(|deadline| deadline <= now) {
1173                expired.insert(work.id.request());
1174            }
1175        }
1176        for request in expired {
1177            self.cancel_internal(request, CancellationCause::Deadline, now)?;
1178        }
1179        Ok(())
1180    }
1181
1182    fn expire_deadlines_distributed<T: ConsensusTransport>(
1183        &mut self,
1184        protocol: u64,
1185        transport: &T,
1186        now: Instant,
1187    ) -> Result<(), SchedulerError> {
1188        let mut expired = BTreeSet::new();
1189        for (request, entry) in &self.requests {
1190            if entry
1191                .pending
1192                .iter()
1193                .any(|work| work.deadline.is_some_and(|deadline| deadline <= now))
1194            {
1195                expired.insert(*request);
1196            }
1197        }
1198        for work in &self.prepared {
1199            if work.deadline.is_some_and(|deadline| deadline <= now) {
1200                expired.insert(work.id.request());
1201            }
1202        }
1203        for request in expired {
1204            if let Err(error) =
1205                validate_disposition(transport, protocol, request, CancellationCause::Deadline)
1206            {
1207                self.poison(error.to_string(), now);
1208                return Err(SchedulerError::Consensus(error.to_string()));
1209            }
1210            self.cancel_internal(request, CancellationCause::Deadline, now)?;
1211        }
1212        Ok(())
1213    }
1214
1215    fn cancel_internal(
1216        &mut self,
1217        request: RequestId,
1218        cause: CancellationCause,
1219        now: Instant,
1220    ) -> Result<(), SchedulerError> {
1221        self.ensure_ready()?;
1222        if self.terminal.contains_key(&request) {
1223            return Err(SchedulerError::State(format!(
1224                "request {} is already terminal",
1225                request.value()
1226            )));
1227        }
1228        let entry = self
1229            .requests
1230            .remove(&request)
1231            .ok_or(SchedulerError::UnknownRequest(request))?;
1232        let queued = entry.pending.len();
1233        for work in entry.pending {
1234            self.lifecycle.insert(work.id, WorkLifecycle::Abandoned);
1235        }
1236        self.ready.retain(|candidate| *candidate != request);
1237        let (prepared, discard_error) = self.discard_prepared_for_request(request);
1238        let before_submission = queued + prepared;
1239        self.accepted_work = self.accepted_work.saturating_sub(before_submission);
1240        self.discarded_work = self.discarded_work.saturating_add(before_submission as u64);
1241        self.cancellation_before_submission = self
1242            .cancellation_before_submission
1243            .saturating_add(before_submission as u64);
1244        let mut after_submission = 0u64;
1245        for work in &mut self.submitted {
1246            if work.id.request() == request && matches!(work.disposition, Disposition::Publish) {
1247                work.disposition = Disposition::Abandon { cancelled_at: now };
1248                self.lifecycle.insert(work.id, WorkLifecycle::Abandoned);
1249                after_submission += 1;
1250            }
1251        }
1252        self.cancellation_after_submission = self
1253            .cancellation_after_submission
1254            .saturating_add(after_submission);
1255        let status = match cause {
1256            CancellationCause::Explicit => {
1257                self.cancelled_requests = self.cancelled_requests.saturating_add(1);
1258                RequestStatus::Cancelled
1259            }
1260            CancellationCause::Deadline => {
1261                self.deadline_expired_requests = self.deadline_expired_requests.saturating_add(1);
1262                RequestStatus::DeadlineExceeded
1263            }
1264        };
1265        self.terminal.insert(
1266            request,
1267            if discard_error.is_some() {
1268                RequestStatus::Failed
1269            } else {
1270                status
1271            },
1272        );
1273        self.update_abandoned_resource_peak();
1274        discard_error.map_or(Ok(()), |error| Err(SchedulerError::State(error)))
1275    }
1276
1277    fn discard_prepared_for_request(&mut self, request: RequestId) -> (usize, Option<String>) {
1278        let mut retained = VecDeque::with_capacity(self.prepared.len());
1279        let mut discarded = 0;
1280        let mut errors = Vec::new();
1281        for work in std::mem::take(&mut self.prepared) {
1282            if work.id.request() == request {
1283                discarded += 1;
1284                match S::discard_branch(work.branch) {
1285                    Ok(()) => {
1286                        self.lifecycle.insert(work.id, WorkLifecycle::Abandoned);
1287                    }
1288                    Err(error) => {
1289                        self.lifecycle.insert(work.id, WorkLifecycle::Failed);
1290                        self.failed_work = self.failed_work.saturating_add(1);
1291                        errors.push(format!("work {:?}: {error}", work.id));
1292                    }
1293                }
1294            } else {
1295                retained.push_back(work);
1296            }
1297        }
1298        self.prepared = retained;
1299        let error = (!errors.is_empty()).then(|| errors.join("; "));
1300        (discarded, error)
1301    }
1302
1303    fn branch_count(&self, request: RequestId) -> usize {
1304        self.prepared
1305            .iter()
1306            .filter(|work| work.id.request() == request)
1307            .count()
1308            + self
1309                .submitted
1310                .iter()
1311                .filter(|work| work.id.request() == request)
1312                .count()
1313    }
1314
1315    fn fail_before_submission(&mut self, id: WorkId) {
1316        self.lifecycle.insert(id, WorkLifecycle::Failed);
1317        self.failed_work = self.failed_work.saturating_add(1);
1318        self.accepted_work = self.accepted_work.saturating_sub(1);
1319        self.fail_request(id.request());
1320    }
1321
1322    fn fail_request(&mut self, request: RequestId) {
1323        let Some(entry) = self.requests.remove(&request) else {
1324            return;
1325        };
1326        let queued = entry.pending.len();
1327        for work in entry.pending {
1328            self.lifecycle.insert(work.id, WorkLifecycle::Failed);
1329        }
1330        self.ready.retain(|candidate| *candidate != request);
1331        let (prepared, _) = self.discard_prepared_for_request(request);
1332        let discarded = queued + prepared;
1333        self.accepted_work = self.accepted_work.saturating_sub(discarded);
1334        self.discarded_work = self.discarded_work.saturating_add(discarded as u64);
1335        for work in &mut self.submitted {
1336            if work.id.request() == request && matches!(work.disposition, Disposition::Publish) {
1337                work.disposition = Disposition::Abandon {
1338                    cancelled_at: Instant::now(),
1339                };
1340                self.lifecycle.insert(work.id, WorkLifecycle::Abandoned);
1341            }
1342        }
1343        self.terminal.insert(request, RequestStatus::Failed);
1344    }
1345
1346    fn ensure_ready(&self) -> Result<(), SchedulerError> {
1347        self.poisoned.as_ref().map_or(Ok(()), |reason| {
1348            Err(SchedulerError::Poisoned(reason.clone()))
1349        })
1350    }
1351
1352    fn poison(&mut self, reason: String, now: Instant) {
1353        if self.poisoned.is_some() {
1354            return;
1355        }
1356        let mut discarded = 0usize;
1357        for (request, entry) in std::mem::take(&mut self.requests) {
1358            self.terminal.insert(request, RequestStatus::Failed);
1359            for work in entry.pending {
1360                self.lifecycle.insert(work.id, WorkLifecycle::Failed);
1361                discarded += 1;
1362            }
1363        }
1364        self.ready.clear();
1365        let mut cleanup_errors = Vec::new();
1366        for work in std::mem::take(&mut self.prepared) {
1367            let id = work.id;
1368            if let Err(error) = S::discard_branch(work.branch) {
1369                cleanup_errors.push(format!("work {id:?}: {error}"));
1370                self.failed_work = self.failed_work.saturating_add(1);
1371            }
1372            self.lifecycle.insert(id, WorkLifecycle::Failed);
1373            discarded += 1;
1374        }
1375        self.accepted_work = self.accepted_work.saturating_sub(discarded);
1376        self.discarded_work = self.discarded_work.saturating_add(discarded as u64);
1377        for work in &mut self.submitted {
1378            if matches!(work.disposition, Disposition::Publish) {
1379                work.disposition = Disposition::Abandon { cancelled_at: now };
1380                self.lifecycle.insert(work.id, WorkLifecycle::Abandoned);
1381            }
1382        }
1383        self.poisoned = Some(if cleanup_errors.is_empty() {
1384            reason
1385        } else {
1386            format!(
1387                "{reason}; prepared branch cleanup failed: {}",
1388                cleanup_errors.join("; ")
1389            )
1390        });
1391        self.update_abandoned_resource_peak();
1392    }
1393
1394    fn abandoned_retained_resources(&self) -> usize {
1395        self.submitted
1396            .iter()
1397            .filter(|work| matches!(work.disposition, Disposition::Abandon { .. }))
1398            .map(|work| work.output.retained_resources())
1399            .sum()
1400    }
1401
1402    fn update_abandoned_resource_peak(&mut self) {
1403        self.peak_abandoned_retained_resources = self
1404            .peak_abandoned_retained_resources
1405            .max(self.abandoned_retained_resources());
1406    }
1407}
1408
1409/// Structured neutral scheduler error.
1410#[derive(Debug, Clone, Eq, PartialEq, thiserror::Error)]
1411pub enum SchedulerError {
1412    /// One or more limits were zero.
1413    #[error("scheduler limits must be positive, got {0:?}")]
1414    InvalidLimits([usize; 6]),
1415    /// Duplicate request identity.
1416    #[error("request {} is already registered", .0.value())]
1417    DuplicateRequest(RequestId),
1418    /// Unknown request identity.
1419    #[error("request {} is not active", .0.value())]
1420    UnknownRequest(RequestId),
1421    /// Capacity exhausted with contextual detail.
1422    #[error("{0}")]
1423    Capacity(String),
1424    /// Work descriptor validation failed.
1425    #[error("work descriptor failed: {0}")]
1426    Descriptor(String),
1427    /// Semantic state operation failed.
1428    #[error("semantic state transaction failed: {0}")]
1429    State(String),
1430    /// Backend submission failed.
1431    #[error("backend submission failed: {0}")]
1432    Submission(String),
1433    /// Exact completion observation failed.
1434    #[error("exact completion observation failed: {0}")]
1435    Completion(String),
1436    /// Topology-wide scheduler agreement failed.
1437    #[error("distributed scheduler consensus failed: {0}")]
1438    Consensus(String),
1439    /// A distributed backend failed and every rank reached a terminal completion.
1440    #[error("distributed {0}")]
1441    DistributedCompletion(String),
1442    /// Unsafe distributed ordering prevents further scheduler mutation.
1443    #[error("scheduler is poisoned after unsafe distributed ordering: {0}")]
1444    Poisoned(String),
1445}
1446
1447#[cfg(test)]
1448mod tests {
1449    use super::*;
1450    use std::{
1451        cell::{Cell, RefCell},
1452        collections::VecDeque,
1453        convert::Infallible,
1454        rc::Rc,
1455    };
1456
1457    #[derive(Default)]
1458    struct State(u32);
1459
1460    impl SemanticStateTransaction for State {
1461        type Branch = u32;
1462        type Error = Infallible;
1463
1464        fn branch(&self) -> Result<u32, Infallible> {
1465            Ok(self.0 + 1)
1466        }
1467
1468        fn commit_branch(&mut self, branch: u32) -> Result<(), Infallible> {
1469            self.0 = branch;
1470            Ok(())
1471        }
1472
1473        fn permits_parallel_branches(&self) -> bool {
1474            true
1475        }
1476    }
1477
1478    impl WorkDescriptor for u32 {
1479        type Error = Infallible;
1480
1481        fn encode_descriptor(&self, output: &mut Vec<u32>) -> Result<(), Self::Error> {
1482            output.push(*self);
1483            Ok(())
1484        }
1485    }
1486
1487    #[derive(Debug)]
1488    struct Output {
1489        complete: Rc<Cell<bool>>,
1490        fail: bool,
1491    }
1492
1493    impl TransitionOutput for Output {
1494        type Error = std::io::Error;
1495
1496        fn is_complete(&self) -> Result<bool, Self::Error> {
1497            if self.fail {
1498                Err(std::io::Error::other("mock failure"))
1499            } else {
1500                Ok(self.complete.get())
1501            }
1502        }
1503
1504        fn backend_name(&self) -> Option<String> {
1505            Some("mock".into())
1506        }
1507
1508        fn retained_resources(&self) -> usize {
1509            2
1510        }
1511    }
1512
1513    #[derive(Default)]
1514    struct GatherStep {
1515        replacements: Vec<(usize, usize, u32)>,
1516    }
1517
1518    struct ScriptedTransport {
1519        participants: usize,
1520        steps: RefCell<VecDeque<GatherStep>>,
1521    }
1522
1523    impl ScriptedTransport {
1524        fn new(participants: usize, steps: Vec<GatherStep>) -> Self {
1525            Self {
1526                participants,
1527                steps: RefCell::new(steps.into()),
1528            }
1529        }
1530    }
1531
1532    impl ConsensusTransport for ScriptedTransport {
1533        type Error = Infallible;
1534
1535        fn participant_count(&self) -> usize {
1536            self.participants
1537        }
1538
1539        fn all_gather_words(&self, local: &[u32]) -> Result<Vec<u32>, Self::Error> {
1540            let mut gathered = local.repeat(self.participants);
1541            let step = self.steps.borrow_mut().pop_front().unwrap_or_default();
1542            for (rank, offset, value) in step.replacements {
1543                gathered[rank * local.len() + offset] = value;
1544            }
1545            Ok(gathered)
1546        }
1547    }
1548
1549    fn scheduler() -> Scheduler<u32, State, Output> {
1550        Scheduler::new(SchedulerLimits::default()).unwrap()
1551    }
1552
1553    #[test]
1554    fn scheduler_construction_revalidates_public_limit_fields() {
1555        let invalid = SchedulerLimits {
1556            max_active_requests: 0,
1557            ..SchedulerLimits::default()
1558        };
1559        assert!(matches!(
1560            Scheduler::<u32, State, Output>::new(invalid),
1561            Err(SchedulerError::InvalidLimits(_))
1562        ));
1563    }
1564
1565    #[test]
1566    fn queued_prepared_submitted_committed_exactly() {
1567        let done = Rc::new(Cell::new(false));
1568        let mut scheduler = scheduler();
1569        let request = RequestId::new(1);
1570        scheduler.register(request, State::default()).unwrap();
1571        let id = scheduler.enqueue(request, 7).unwrap();
1572        assert_eq!(scheduler.work_lifecycle(id), Some(WorkLifecycle::Queued));
1573        scheduler.prepare_bounded(1, Instant::now()).unwrap();
1574        assert_eq!(scheduler.work_lifecycle(id), Some(WorkLifecycle::Prepared));
1575        scheduler
1576            .submit_prepared(Instant::now(), |_, _, _| {
1577                Ok::<_, Infallible>(Output {
1578                    complete: done.clone(),
1579                    fail: false,
1580                })
1581            })
1582            .unwrap();
1583        assert!(scheduler
1584            .poll_completions(Instant::now())
1585            .committed
1586            .is_empty());
1587        done.set(true);
1588        assert_eq!(
1589            scheduler.poll_completions(Instant::now()).committed.len(),
1590            1
1591        );
1592        assert_eq!(scheduler.request_state(request).unwrap().0, 1);
1593    }
1594
1595    #[test]
1596    fn cancellation_retains_submitted_resources_until_completion() {
1597        let done = Rc::new(Cell::new(false));
1598        let mut scheduler = scheduler();
1599        let request = RequestId::new(2);
1600        scheduler.register(request, State::default()).unwrap();
1601        let id = scheduler.enqueue(request, 1).unwrap();
1602        scheduler.prepare_bounded(1, Instant::now()).unwrap();
1603        scheduler
1604            .submit_prepared(Instant::now(), |_, _, _| {
1605                Ok::<_, Infallible>(Output {
1606                    complete: done.clone(),
1607                    fail: false,
1608                })
1609            })
1610            .unwrap();
1611        scheduler.cancel(request).unwrap();
1612        assert_eq!(scheduler.work_lifecycle(id), Some(WorkLifecycle::Abandoned));
1613        assert_eq!(scheduler.report().abandoned_retained_resources, 2);
1614        assert_eq!(
1615            scheduler.request_status(request),
1616            Some(RequestStatus::Cancelled)
1617        );
1618        done.set(true);
1619        scheduler.poll_completions(Instant::now());
1620        assert_eq!(scheduler.report().current_in_flight_work, 0);
1621        assert_eq!(scheduler.report().abandoned_released_work, 1);
1622    }
1623
1624    #[test]
1625    fn exact_completion_failure_does_not_commit() {
1626        let mut scheduler = scheduler();
1627        let request = RequestId::new(3);
1628        scheduler.register(request, State::default()).unwrap();
1629        let id = scheduler.enqueue(request, 1).unwrap();
1630        scheduler.prepare_bounded(1, Instant::now()).unwrap();
1631        scheduler
1632            .submit_prepared(Instant::now(), |_, _, _| {
1633                Ok::<_, Infallible>(Output {
1634                    complete: Rc::new(Cell::new(true)),
1635                    fail: true,
1636                })
1637            })
1638            .unwrap();
1639        assert_eq!(scheduler.poll_completions(Instant::now()).failed.len(), 1);
1640        assert_eq!(scheduler.work_lifecycle(id), Some(WorkLifecycle::Failed));
1641        assert_eq!(
1642            scheduler.request_status(request),
1643            Some(RequestStatus::Failed)
1644        );
1645    }
1646
1647    #[test]
1648    fn request_failure_abandons_incomplete_sibling_until_exact_completion() {
1649        let waiting = Rc::new(Cell::new(false));
1650        let limits = SchedulerLimits::with_execution_bounds(1, 2, 2, 2, 2, 1).unwrap();
1651        let mut scheduler = Scheduler::new(limits).unwrap();
1652        let request = RequestId::new(31);
1653        scheduler.register(request, State::default()).unwrap();
1654        let ids = scheduler.enqueue_batch(request, vec![1, 2]).unwrap();
1655        scheduler.prepare_bounded(2, Instant::now()).unwrap();
1656        scheduler
1657            .submit_prepared(Instant::now(), |id, _, _| {
1658                Ok::<_, Infallible>(Output {
1659                    complete: if id.sequence() == 0 {
1660                        Rc::new(Cell::new(true))
1661                    } else {
1662                        waiting.clone()
1663                    },
1664                    fail: id.sequence() == 0,
1665                })
1666            })
1667            .unwrap();
1668
1669        let progress = scheduler.poll_completions(Instant::now());
1670        assert_eq!(progress.failed.len(), 1);
1671        assert_eq!(
1672            scheduler.work_lifecycle(ids[0]),
1673            Some(WorkLifecycle::Failed)
1674        );
1675        assert_eq!(
1676            scheduler.work_lifecycle(ids[1]),
1677            Some(WorkLifecycle::Abandoned)
1678        );
1679        assert_eq!(scheduler.report().abandoned_retained_resources, 2);
1680
1681        waiting.set(true);
1682        scheduler.poll_completions(Instant::now());
1683        assert_eq!(scheduler.report().current_in_flight_work, 0);
1684        assert_eq!(scheduler.report().abandoned_released_work, 1);
1685    }
1686
1687    #[test]
1688    fn submission_failure_marks_work_and_request_failed() {
1689        let mut scheduler = scheduler();
1690        let request = RequestId::new(4);
1691        scheduler.register(request, State::default()).unwrap();
1692        let id = scheduler.enqueue(request, 1).unwrap();
1693        scheduler.prepare_bounded(1, Instant::now()).unwrap();
1694        let error = scheduler
1695            .submit_prepared(Instant::now(), |_, _, _| {
1696                Err::<Output, _>(std::io::Error::other("mock submit"))
1697            })
1698            .unwrap_err();
1699        assert!(matches!(error, SchedulerError::Submission(_)));
1700        assert_eq!(scheduler.work_lifecycle(id), Some(WorkLifecycle::Failed));
1701        assert_eq!(
1702            scheduler.request_status(request),
1703            Some(RequestStatus::Failed)
1704        );
1705    }
1706
1707    #[test]
1708    fn deadlines_batches_and_fairness_are_backend_neutral() {
1709        let mut scheduler = scheduler();
1710        let first = RequestId::new(10);
1711        let second = RequestId::new(20);
1712        scheduler.register(first, State::default()).unwrap();
1713        scheduler.register(second, State::default()).unwrap();
1714        scheduler.enqueue_batch(first, vec![1, 2]).unwrap();
1715        scheduler.enqueue(second, 3).unwrap();
1716        scheduler.prepare_bounded(3, Instant::now()).unwrap();
1717        let mut order = Vec::new();
1718        scheduler
1719            .submit_prepared(Instant::now(), |id, _, _| {
1720                order.push(id.request());
1721                Ok::<_, Infallible>(Output {
1722                    complete: Rc::new(Cell::new(true)),
1723                    fail: false,
1724                })
1725            })
1726            .unwrap();
1727        assert_eq!(order, vec![first]);
1728
1729        let expired = RequestId::new(30);
1730        scheduler.register(expired, State::default()).unwrap();
1731        scheduler
1732            .enqueue_with_deadline(
1733                expired,
1734                4,
1735                Some(Instant::now().checked_sub(Duration::from_secs(1)).unwrap()),
1736            )
1737            .unwrap();
1738        scheduler.prepare_bounded(1, Instant::now()).unwrap();
1739        assert_eq!(
1740            scheduler.request_status(expired),
1741            Some(RequestStatus::DeadlineExceeded)
1742        );
1743    }
1744
1745    #[test]
1746    fn distributed_schedule_mismatch_poisons_the_canonical_machine() {
1747        let transport = ScriptedTransport::new(
1748            2,
1749            vec![
1750                GatherStep::default(),
1751                GatherStep {
1752                    replacements: vec![(1, 10, 99)],
1753                },
1754            ],
1755        );
1756        let mut scheduler = scheduler();
1757        let request = RequestId::new(50);
1758        scheduler.register(request, State::default()).unwrap();
1759        let work = scheduler.enqueue(request, 7).unwrap();
1760
1761        let error = scheduler
1762            .run_distributed_turn(0xAA, &transport, Instant::now(), |_, _, _| {
1763                Ok::<_, Infallible>(Output {
1764                    complete: Rc::new(Cell::new(false)),
1765                    fail: false,
1766                })
1767            })
1768            .unwrap_err();
1769        assert!(matches!(error, SchedulerError::Consensus(_)));
1770        assert!(scheduler.poison_reason().is_some());
1771        assert!(scheduler.report().poisoned);
1772        assert_eq!(
1773            scheduler.request_status(request),
1774            Some(RequestStatus::Failed)
1775        );
1776        assert_eq!(scheduler.work_lifecycle(work), Some(WorkLifecycle::Failed));
1777        assert!(matches!(
1778            scheduler.enqueue(request, 8),
1779            Err(SchedulerError::Poisoned(_))
1780        ));
1781    }
1782
1783    #[test]
1784    fn poisoned_submissions_remain_retained_until_local_exact_completion() {
1785        let transport = ScriptedTransport::new(
1786            2,
1787            vec![
1788                GatherStep::default(),
1789                GatherStep::default(),
1790                GatherStep::default(),
1791                GatherStep {
1792                    replacements: vec![(1, 4, 99)],
1793                },
1794            ],
1795        );
1796        let done = Rc::new(Cell::new(false));
1797        let mut scheduler = scheduler();
1798        let request = RequestId::new(54);
1799        scheduler.register(request, State::default()).unwrap();
1800        scheduler.enqueue(request, 1).unwrap();
1801        scheduler
1802            .run_distributed_turn(0xEE, &transport, Instant::now(), |_, _, _| {
1803                Ok::<_, Infallible>(Output {
1804                    complete: done.clone(),
1805                    fail: false,
1806                })
1807            })
1808            .unwrap();
1809
1810        let error = scheduler
1811            .run_distributed_turn(0xEE, &transport, Instant::now(), |_, _, _| {
1812                Ok::<_, Infallible>(Output {
1813                    complete: done.clone(),
1814                    fail: false,
1815                })
1816            })
1817            .unwrap_err();
1818        assert!(matches!(error, SchedulerError::Consensus(_)));
1819        assert_eq!(scheduler.report().abandoned_in_flight_work, 1);
1820        assert_eq!(scheduler.report().abandoned_retained_resources, 2);
1821
1822        assert!(scheduler
1823            .poll_completions(Instant::now())
1824            .committed
1825            .is_empty());
1826        assert_eq!(scheduler.report().current_in_flight_work, 1);
1827        done.set(true);
1828        assert!(scheduler
1829            .poll_completions(Instant::now())
1830            .committed
1831            .is_empty());
1832        assert_eq!(scheduler.report().current_in_flight_work, 0);
1833        assert_eq!(scheduler.report().abandoned_released_work, 1);
1834    }
1835
1836    #[test]
1837    fn distributed_failure_retains_output_until_every_rank_is_terminal() {
1838        let transport = ScriptedTransport::new(
1839            3,
1840            vec![
1841                GatherStep::default(),
1842                GatherStep::default(),
1843                GatherStep {
1844                    replacements: vec![(1, 7, 2), (2, 7, 0)],
1845                },
1846                GatherStep::default(),
1847                GatherStep {
1848                    replacements: vec![(1, 7, 2), (2, 7, 1)],
1849                },
1850                GatherStep::default(),
1851            ],
1852        );
1853        let mut scheduler = scheduler();
1854        let request = RequestId::new(51);
1855        scheduler.register(request, State::default()).unwrap();
1856        scheduler.enqueue(request, 1).unwrap();
1857        let output = || Output {
1858            complete: Rc::new(Cell::new(true)),
1859            fail: false,
1860        };
1861
1862        scheduler
1863            .run_distributed_turn(0xBB, &transport, Instant::now(), |_, _, _| {
1864                Ok::<_, Infallible>(output())
1865            })
1866            .unwrap();
1867        let pending = scheduler
1868            .run_distributed_turn(0xBB, &transport, Instant::now(), |_, _, _| {
1869                Ok::<_, Infallible>(output())
1870            })
1871            .unwrap();
1872        assert!(pending.failed.is_empty());
1873        assert_eq!(scheduler.report().failed_in_flight_work, 1);
1874        assert_eq!(scheduler.report().current_in_flight_work, 1);
1875
1876        let terminal = scheduler
1877            .run_distributed_turn(0xBB, &transport, Instant::now(), |_, _, _| {
1878                Ok::<_, Infallible>(output())
1879            })
1880            .unwrap();
1881        assert_eq!(terminal.failed.len(), 1);
1882        assert_eq!(scheduler.report().current_in_flight_work, 0);
1883        assert_eq!(
1884            scheduler.request_status(request),
1885            Some(RequestStatus::Failed)
1886        );
1887    }
1888
1889    #[test]
1890    fn distributed_cancellation_and_deadline_use_the_same_core_lifecycle() {
1891        let transport = ScriptedTransport::new(2, vec![GatherStep::default()]);
1892        let mut cancelled = scheduler();
1893        let request = RequestId::new(52);
1894        cancelled.register(request, State::default()).unwrap();
1895        cancelled.enqueue(request, 1).unwrap();
1896        cancelled
1897            .cancel_distributed(0xCC, request, &transport, Instant::now())
1898            .unwrap();
1899        assert_eq!(
1900            cancelled.request_status(request),
1901            Some(RequestStatus::Cancelled)
1902        );
1903
1904        let transport = ScriptedTransport::new(
1905            2,
1906            vec![
1907                GatherStep::default(),
1908                GatherStep::default(),
1909                GatherStep::default(),
1910            ],
1911        );
1912        let mut scheduler = scheduler();
1913        let request = RequestId::new(53);
1914        scheduler.register(request, State::default()).unwrap();
1915        scheduler
1916            .enqueue_with_deadline(
1917                request,
1918                1,
1919                Some(Instant::now().checked_sub(Duration::from_secs(1)).unwrap()),
1920            )
1921            .unwrap();
1922        scheduler
1923            .run_distributed_turn(0xDD, &transport, Instant::now(), |_, _, _| {
1924                Ok::<_, Infallible>(Output {
1925                    complete: Rc::new(Cell::new(true)),
1926                    fail: false,
1927                })
1928            })
1929            .unwrap();
1930        assert_eq!(
1931            scheduler.request_status(request),
1932            Some(RequestStatus::DeadlineExceeded)
1933        );
1934    }
1935
1936    #[test]
1937    fn production_telemetry_round_trips_without_backend_types() {
1938        let mut scheduler = scheduler();
1939        let request = RequestId::new(40);
1940        scheduler.register(request, State::default()).unwrap();
1941        scheduler.enqueue(request, 1).unwrap();
1942        scheduler.prepare_bounded(1, Instant::now()).unwrap();
1943        scheduler
1944            .submit_prepared(Instant::now(), |_, _, _| {
1945                Ok::<_, Infallible>(Output {
1946                    complete: Rc::new(Cell::new(false)),
1947                    fail: false,
1948                })
1949            })
1950            .unwrap();
1951
1952        let report = scheduler.report();
1953        let report_json = serde_json::to_string(&report).unwrap();
1954        assert_eq!(
1955            serde_json::from_str::<SchedulerReport>(&report_json).unwrap(),
1956            report
1957        );
1958
1959        let capabilities = scheduler.capabilities();
1960        assert_eq!(capabilities.observed_backends, ["mock"]);
1961        let capabilities_json = serde_json::to_string(&capabilities).unwrap();
1962        assert_eq!(
1963            serde_json::from_str::<SchedulerCapabilities>(&capabilities_json).unwrap(),
1964            capabilities
1965        );
1966    }
1967}