wowlab_sentinel/scheduler/runtime/
job.rs

1#![expect(
2    clippy::cast_possible_truncation,
3    reason = "runtime collection lengths are bounded by the u32 job permutation contract"
4)]
5
6use std::collections::VecDeque;
7
8use roaring::RoaringBitmap;
9use uuid::Uuid;
10use wowlab_common::{
11    ClaimToken, RuntimeChunkId, RuntimeWorkItem, RuntimeWorkResult, WorkContextHash,
12};
13use wowlab_types::sim::FastMap;
14
15use super::{
16    strategy::RuntimeStrategyState,
17    tournament::{Phase, TournamentRuntime},
18};
19use crate::{
20    scheduler::{
21        runtime_aggregate::{SingleRuntime, StatWeightsRuntime},
22        runtime_breakdown::RuntimeMemoryBreakdown,
23        runtime_completion_log::{CompletionLog, completion_hash},
24        runtime_factorial::FactorialRuntime,
25    },
26    strategy::{FinalOutput, StrategyError},
27};
28
29const BATCH_FULL_FIDELITY: usize = 1;
30const BATCH_TIER_SMALL_MAX_ITERS: u32 = 2_000;
31const BATCH_TIER_MEDIUM_MAX_ITERS: u32 = 10_000;
32const BATCH_TIER_SMALL: usize = 500;
33const BATCH_TIER_MEDIUM: usize = 120;
34const BATCH_TIER_LARGE: usize = 10;
35
36#[derive(Clone, Copy, Debug, Eq, PartialEq)]
37#[non_exhaustive]
38pub enum JobRuntimeStatus {
39    Pending,
40    Running,
41    Completed,
42    Finalizing,
43    Failed,
44}
45
46#[derive(Clone, Debug, Eq, PartialEq)]
47#[non_exhaustive]
48pub(crate) enum CompletionOutcome {
49    Accepted { job_complete: bool },
50    Idempotent,
51    Conflict(String),
52    Stale(String),
53}
54
55// docref:start hosted-compute-claim-record
56#[derive(Clone, Debug)]
57pub struct ClaimRecord {
58    pub chunk_id: RuntimeChunkId,
59    pub node_public_key: String,
60    pub work_context_hash: WorkContextHash,
61    pub work_items: Vec<RuntimeWorkItem>,
62    pub claimed_at_ms: u64,
63    pub reclaim_count: u32,
64}
65// docref:end hosted-compute-claim-record
66
67#[derive(Clone, Copy, Debug)]
68// #t(rust_similar_structs) runtime ranking values are distinct from the serialized scheduler progress contract
69pub struct TopPermutation {
70    pub perm_idx: u32,
71    pub mean_dps_x10: u32,
72    pub iterations: u32,
73}
74
75#[derive(Clone, Debug)]
76pub struct ProgressSnapshot {
77    pub phase: u32,
78    pub phase_count: u32,
79    pub permutations_active: u32,
80    pub top_mean_dps_x10: u32,
81    pub top_k: Vec<TopPermutation>,
82}
83
84#[derive(Debug)]
85// #t(rust_similar_structs) assigned scheduler batches retain claim state separate from the common wire payload
86pub struct AssignedBatch {
87    pub chunk_id: RuntimeChunkId,
88    pub claim_token: ClaimToken,
89    pub work_context_hash: WorkContextHash,
90    pub work_items: Vec<RuntimeWorkItem>,
91}
92
93// docref:start database-job-runtime
94#[derive(Debug)]
95pub struct JobRuntime {
96    pub job_id: Uuid,
97    pub user_id: Uuid,
98    pub strategy_name: String,
99    pub base_sim_config: String,
100    pub sentinel_config: String,
101    /// Encoded `TournamentPayload`, empty for non-tournament jobs.
102    pub payload_bytes: Vec<u8>,
103    pub work_context_hash: WorkContextHash,
104    pub status: JobRuntimeStatus,
105    pub priority: i32,
106    strategy: RuntimeStrategyState,
107    in_flight: FastMap<ClaimToken, ClaimRecord>,
108    pending_reclaims: VecDeque<RuntimeWorkItem>,
109    recent_completions: CompletionLog,
110    failed_items: u32,
111    next_chunk_id: u64,
112    next_item_id: u64,
113}
114// docref:end database-job-runtime
115
116#[derive(Debug, Default)]
117pub struct ReclaimOutcome {
118    pub reclaimed: Vec<ClaimRecord>,
119    /// Dropped permanently after exceeding the retry limit.
120    pub failed: Vec<ClaimRecord>,
121}
122
123#[derive(Debug)]
124pub struct JobRuntimeInit {
125    pub job_id: Uuid,
126    pub user_id: Uuid,
127    pub base_sim_config: String,
128    pub sentinel_config: String,
129    /// Encoded `TournamentPayload`, empty for non-tournament jobs.
130    pub payload_bytes: Vec<u8>,
131    pub priority: i32,
132}
133
134impl JobRuntime {
135    fn new(init: JobRuntimeInit, strategy_name: &str, strategy: RuntimeStrategyState) -> Self {
136        let work_context_hash = WorkContextHash::compute(
137            &init.base_sim_config,
138            &init.payload_bytes,
139            &init.sentinel_config,
140        );
141
142        Self {
143            job_id: init.job_id,
144            user_id: init.user_id,
145            strategy_name: strategy_name.to_owned(),
146            base_sim_config: init.base_sim_config,
147            sentinel_config: init.sentinel_config,
148            payload_bytes: init.payload_bytes,
149            work_context_hash,
150            status: JobRuntimeStatus::Pending,
151            priority: init.priority,
152            strategy,
153            in_flight: FastMap::default(),
154            pending_reclaims: VecDeque::new(),
155            recent_completions: CompletionLog::new(),
156            failed_items: 0,
157            next_chunk_id: 0,
158            next_item_id: 0,
159        }
160    }
161
162    pub fn tournament(init: JobRuntimeInit, phases: Vec<Phase>, to_test: RoaringBitmap) -> Self {
163        let payload = init.payload_bytes.clone();
164        let strategy = RuntimeStrategyState::Tournament(Box::new(TournamentRuntime::new(
165            phases, to_test, payload,
166        )));
167
168        Self::new(init, "tournament", strategy)
169    }
170
171    pub fn tournament_with_factorial(
172        init: JobRuntimeInit,
173        phases: Vec<Phase>,
174        factorial: FactorialRuntime,
175    ) -> Self {
176        let payload = init.payload_bytes.clone();
177        let strategy = RuntimeStrategyState::Tournament(Box::new(
178            TournamentRuntime::with_factorial(phases, payload, factorial),
179        ));
180
181        Self::new(init, "tournament", strategy)
182    }
183
184    pub fn single(init: JobRuntimeInit, chunk_iterations: &[u32]) -> Self {
185        let strategy = RuntimeStrategyState::Single(Box::new(SingleRuntime::new(chunk_iterations)));
186
187        Self::new(init, "single", strategy)
188    }
189
190    /// Tag 0 is the full-fidelity baseline.
191    pub fn stat_weights(init: JobRuntimeInit, runs: &[(u64, Vec<u32>)]) -> Self {
192        let strategy = RuntimeStrategyState::StatWeights(StatWeightsRuntime::new(runs));
193
194        Self::new(init, "stat_weights", strategy)
195    }
196
197    pub fn in_flight_len(&self) -> usize {
198        self.in_flight.len()
199    }
200
201    pub fn in_flight_claims(&self) -> impl Iterator<Item = &ClaimRecord> {
202        self.in_flight.values()
203    }
204
205    pub fn suggested_batch_size(&self) -> usize {
206        let (iterations, full) = self.strategy.peek_next_work_shape();
207
208        if full {
209            return BATCH_FULL_FIDELITY;
210        }
211
212        match iterations {
213            i if i <= BATCH_TIER_SMALL_MAX_ITERS => BATCH_TIER_SMALL,
214            i if i <= BATCH_TIER_MEDIUM_MAX_ITERS => BATCH_TIER_MEDIUM,
215            _ => BATCH_TIER_LARGE,
216        }
217    }
218
219    /// `None` for non-tournament jobs.
220    pub fn progress_snapshot(&self) -> Option<ProgressSnapshot> {
221        match &self.strategy {
222            RuntimeStrategyState::Tournament(t) => Some(t.progress()),
223            _ => None,
224        }
225    }
226
227    /// Reassigns reclaimed work before fresh work.
228    pub fn claim_batch(
229        &mut self,
230        node_public_key: &str,
231        max_items: usize,
232        now_ms: u64,
233    ) -> Option<AssignedBatch> {
234        self.claim_batch_with(node_public_key, max_items, now_ms, || {
235            ClaimToken::new(Uuid::new_v4().to_string())
236        })
237    }
238
239    pub fn authorize_context(
240        &self,
241        claim_token: &ClaimToken,
242        hash: &WorkContextHash,
243        node_public_key: &str,
244    ) -> bool {
245        self.in_flight.get(claim_token).is_some_and(|claim| {
246            claim.node_public_key == node_public_key
247                && &claim.work_context_hash == hash
248                && &self.work_context_hash == hash
249        })
250    }
251
252    pub fn reclaim_stale(
253        &mut self,
254        now_ms: u64,
255        timeout_ms: u64,
256        max_attempts: u32,
257    ) -> ReclaimOutcome {
258        let stale: Vec<ClaimToken> = self
259            .in_flight
260            .iter()
261            .filter(|(_, c)| now_ms.saturating_sub(c.claimed_at_ms) >= timeout_ms)
262            .map(|(t, _)| t.clone())
263            .collect();
264
265        let mut outcome = ReclaimOutcome::default();
266
267        for token in stale {
268            if let Some(mut claim) = self.in_flight.remove(&token) {
269                claim.reclaim_count += 1;
270
271                if claim.reclaim_count > max_attempts {
272                    self.failed_items += claim.work_items.len() as u32;
273                    outcome.failed.push(claim);
274                } else {
275                    self.pending_reclaims
276                        .extend(claim.work_items.iter().cloned());
277                    outcome.reclaimed.push(claim);
278                }
279            }
280        }
281
282        outcome
283    }
284
285    pub fn has_failures(&self) -> bool {
286        self.failed_items > 0
287    }
288
289    pub fn has_outstanding_work(&self) -> bool {
290        !self.in_flight.is_empty()
291            || !self.pending_reclaims.is_empty()
292            || self.strategy.has_pending()
293    }
294
295    pub fn is_complete(&self) -> bool {
296        self.strategy.is_complete()
297    }
298
299    pub fn memory_breakdown(&self) -> RuntimeMemoryBreakdown {
300        let mut breakdown = self.strategy.memory_breakdown();
301
302        let claim_record = size_of::<ClaimRecord>();
303        let work_item = size_of::<RuntimeWorkItem>();
304
305        let mut in_flight_bytes = 0usize;
306        let mut in_flight_items = 0usize;
307
308        for claim in self.in_flight.values() {
309            in_flight_bytes += claim_record
310                + claim.work_items.capacity() * work_item
311                + claim.node_public_key.capacity();
312            in_flight_items += claim.work_items.len();
313        }
314
315        in_flight_bytes += self.in_flight.capacity() * (size_of::<ClaimToken>() + claim_record);
316        breakdown.in_flight_bytes = in_flight_bytes;
317        breakdown.in_flight_items = in_flight_items;
318
319        breakdown.in_flight_bytes += self.pending_reclaims.capacity() * work_item;
320
321        breakdown.completion_log_bytes = self.recent_completions.memory_bytes();
322
323        breakdown
324    }
325
326    pub(in crate::scheduler) fn claim_batch_with<F>(
327        &mut self,
328        node_public_key: &str,
329        max_items: usize,
330        now_ms: u64,
331        claim_token: F,
332    ) -> Option<AssignedBatch>
333    where
334        F: FnOnce() -> ClaimToken,
335    {
336        let mut items = Vec::with_capacity(max_items);
337        let mut reclaim_count = 0u32;
338
339        while items.len() < max_items {
340            if let Some(item) = self.pending_reclaims.pop_front() {
341                reclaim_count = reclaim_count.max(1);
342                items.push(item);
343            } else {
344                break;
345            }
346        }
347
348        if items.len() < max_items {
349            let fresh = self
350                .strategy
351                .next_work(max_items - items.len(), &mut self.next_item_id);
352
353            items.extend(fresh);
354        }
355
356        if items.is_empty() {
357            return None;
358        }
359
360        let chunk_id = RuntimeChunkId::new(self.next_chunk_id);
361
362        self.next_chunk_id += 1;
363        let claim_token = claim_token();
364
365        self.status = JobRuntimeStatus::Running;
366        self.in_flight.insert(
367            claim_token.clone(),
368            ClaimRecord {
369                chunk_id,
370                node_public_key: node_public_key.to_owned(),
371                work_context_hash: self.work_context_hash,
372                work_items: items.clone(),
373                claimed_at_ms: now_ms,
374                reclaim_count,
375            },
376        );
377
378        Some(AssignedBatch {
379            chunk_id,
380            claim_token,
381            work_context_hash: self.work_context_hash,
382            work_items: items,
383        })
384    }
385
386    pub(crate) fn complete(
387        &mut self,
388        chunk_id: RuntimeChunkId,
389        claim_token: &ClaimToken,
390        work_context_hash: &WorkContextHash,
391        node_public_key: &str,
392        results: &[RuntimeWorkResult],
393    ) -> Result<CompletionOutcome, StrategyError> {
394        let hash = completion_hash(chunk_id, work_context_hash, node_public_key, results);
395
396        if let Some(prev) = self.recent_completions.get(claim_token) {
397            return Ok(if *prev == hash {
398                CompletionOutcome::Idempotent
399            } else {
400                CompletionOutcome::Conflict("claim already completed with different results".into())
401            });
402        }
403
404        let Some(claim) = self.in_flight.get(claim_token) else {
405            return Ok(CompletionOutcome::Stale("no active claim for token".into()));
406        };
407
408        if claim.chunk_id != chunk_id
409            || &claim.work_context_hash != work_context_hash
410            || claim.node_public_key != node_public_key
411        {
412            return Ok(CompletionOutcome::Conflict(
413                "claim metadata mismatch".into(),
414            ));
415        }
416
417        if !results_match_claim(claim, results) {
418            return Ok(CompletionOutcome::Conflict(
419                "completion results do not match assigned work items".into(),
420            ));
421        }
422
423        self.strategy.ingest(results)?;
424        self.in_flight.remove(claim_token);
425        self.recent_completions.record(claim_token.clone(), hash);
426
427        let job_complete = self.strategy.is_complete();
428
429        if job_complete {
430            self.status = JobRuntimeStatus::Completed;
431        }
432
433        Ok(CompletionOutcome::Accepted { job_complete })
434    }
435
436    pub(crate) fn finalize(&self) -> Result<FinalOutput, StrategyError> {
437        self.strategy.finalize()
438    }
439}
440
441fn results_match_claim(claim: &ClaimRecord, results: &[RuntimeWorkResult]) -> bool {
442    if claim.work_items.len() != results.len() {
443        return false;
444    }
445
446    let mut by_item_id = FastMap::default();
447
448    for result in results {
449        if by_item_id.insert(result.item_id, result).is_some() {
450            return false;
451        }
452    }
453
454    claim.work_items.iter().all(|item| {
455        by_item_id.get(&item.item_id).is_some_and(|result| {
456            result.tag == item.tag
457                && result.iterations == item.iterations
458                && (!matches!(item.fidelity, wowlab_common::RuntimeWorkFidelity::Full)
459                    || result.telemetry_pb.is_some())
460        })
461    })
462}
463
464#[derive(Debug, Default)]
465pub struct JobRuntimeStore {
466    jobs: FastMap<Uuid, JobRuntime>,
467}
468
469impl JobRuntimeStore {
470    pub fn new() -> Self {
471        Self::default()
472    }
473
474    pub fn insert(&mut self, runtime: JobRuntime) {
475        self.jobs.insert(runtime.job_id, runtime);
476    }
477
478    pub fn contains(&self, job_id: &Uuid) -> bool {
479        self.jobs.contains_key(job_id)
480    }
481
482    pub fn get(&self, job_id: &Uuid) -> Option<&JobRuntime> {
483        self.jobs.get(job_id)
484    }
485
486    pub fn get_mut(&mut self, job_id: &Uuid) -> Option<&mut JobRuntime> {
487        self.jobs.get_mut(job_id)
488    }
489
490    pub fn remove(&mut self, job_id: &Uuid) -> Option<JobRuntime> {
491        self.jobs.remove(job_id)
492    }
493
494    pub fn len(&self) -> usize {
495        self.jobs.len()
496    }
497
498    pub fn is_empty(&self) -> bool {
499        self.jobs.is_empty()
500    }
501
502    pub fn iter(&self) -> impl Iterator<Item = &JobRuntime> {
503        self.jobs.values()
504    }
505
506    pub fn iter_mut(&mut self) -> impl Iterator<Item = &mut JobRuntime> {
507        self.jobs.values_mut()
508    }
509}
510
511impl<'a> IntoIterator for &'a JobRuntimeStore {
512    type Item = &'a JobRuntime;
513    type IntoIter = std::collections::hash_map::Values<'a, Uuid, JobRuntime>;
514
515    fn into_iter(self) -> Self::IntoIter {
516        self.jobs.values()
517    }
518}
519
520impl<'a> IntoIterator for &'a mut JobRuntimeStore {
521    type Item = &'a mut JobRuntime;
522    type IntoIter = std::collections::hash_map::ValuesMut<'a, Uuid, JobRuntime>;
523
524    fn into_iter(self) -> Self::IntoIter {
525        self.jobs.values_mut()
526    }
527}
wowlab_sentinel/scheduler/runtime/job.rs

wowlab_sentinel/scheduler/runtime/
job.rs
