wowlab_sentinel/scheduler/

runtime_aggregate.rs

//! In-process runtime state for base-config jobs (single and stat-weights).

#![expect(
    clippy::cast_possible_truncation,
    clippy::cast_precision_loss,
    clippy::cast_sign_loss,
    reason = "aggregate sample counts are bounded by job limits and statistical interpolation intentionally uses f64"
)]

use std::collections::VecDeque;

use prost::Message;
use wowlab_analytics as merge;
use wowlab_common::{RuntimeWorkFidelity, RuntimeWorkItem, RuntimeWorkItemKind, RuntimeWorkResult};
use wowlab_types::proto::{self, JobResult, JobTimeline};

use crate::strategy::{FinalOutput, StrategyError, encode_and_check};

const RESULT_BUDGET: usize = 32_768;

#[derive(Clone, Copy, Debug)]
pub(super) struct PlannedChunk {
    iterations: u32,
    seed_offset: u64,
    tag: u64,
    full: bool,
}

#[derive(Clone, Debug)]
struct AggregateRun {
    tag: u64,
    chunks_total: u32,
    completed: u32,
    merge: proto::RunningAggregateStateV1,
}

#[derive(Debug)]
pub(super) struct SingleRuntime {
    queue: VecDeque<PlannedChunk>,
    run: AggregateRun,
}

impl SingleRuntime {
    pub(super) fn new(chunk_iterations: &[u32]) -> Self {
        let chunks_total = chunk_iterations.len() as u32;
        let mut queue = VecDeque::new();
        let mut seed = 0u64;

        for (i, &iters) in chunk_iterations.iter().enumerate() {
            queue.push_back(PlannedChunk {
                iterations: iters,
                seed_offset: seed,
                tag: i as u64,
                full: true,
            });
            seed += u64::from(iters);
        }

        Self {
            queue,
            run: AggregateRun {
                tag: 0,
                chunks_total,
                completed: 0,
                merge: merge::new_running_state(chunks_total),
            },
        }
    }

    pub(super) fn next_work(
        &mut self,
        max_items: usize,
        next_item_id: &mut u64,
    ) -> Vec<RuntimeWorkItem> {
        next_from_queue(&mut self.queue, max_items, next_item_id)
    }

    pub(super) fn ingest(&mut self, results: &[RuntimeWorkResult]) -> Result<(), StrategyError> {
        let mut merged = self.run.merge.clone();

        for r in results {
            let telemetry = r
                .telemetry_pb
                .as_ref()
                .ok_or_else(StrategyError::full_telemetry_required)?;
            let chunk = proto::ChunkTelemetry::decode(telemetry.as_slice())?;

            merge::merge_chunk(&mut merged, &chunk)?;
        }

        self.run.merge = merged;
        self.run.completed += results.len() as u32;

        Ok(())
    }

    pub(super) fn is_complete(&self) -> bool {
        self.run.completed >= self.run.chunks_total
    }

    pub(super) fn has_pending(&self) -> bool {
        !self.queue.is_empty()
    }

    pub(super) fn peek_shape(&self) -> (u32, bool) {
        peek_queue(&self.queue)
    }

    pub(super) fn finalize(&self) -> Result<FinalOutput, StrategyError> {
        finalize_single("single", &self.run.merge)
    }
}

#[derive(Debug)]
pub(super) struct StatWeightsRuntime {
    queue: VecDeque<PlannedChunk>,
    runs: Vec<AggregateRun>,
    runs_completed: u32,
}

impl StatWeightsRuntime {
    /// Build from `(tag, chunk_iterations)` runs; tag 0 is the full-fidelity baseline.
    pub(super) fn new(runs: &[(u64, Vec<u32>)]) -> Self {
        let mut queue = VecDeque::new();
        let mut run_states = Vec::with_capacity(runs.len());

        for (tag, chunk_iters) in runs {
            let mut seed = 0u64;

            for &iters in chunk_iters {
                queue.push_back(PlannedChunk {
                    iterations: iters,
                    seed_offset: seed,
                    tag: *tag,
                    full: *tag == 0,
                });
                seed += u64::from(iters);
            }

            run_states.push(AggregateRun {
                tag: *tag,
                chunks_total: chunk_iters.len() as u32,
                completed: 0,
                merge: merge::new_running_state(chunk_iters.len() as u32),
            });
        }

        Self {
            queue,
            runs: run_states,
            runs_completed: 0,
        }
    }

    pub(super) fn next_work(
        &mut self,
        max_items: usize,
        next_item_id: &mut u64,
    ) -> Vec<RuntimeWorkItem> {
        next_from_queue(&mut self.queue, max_items, next_item_id)
    }

    pub(super) fn ingest(&mut self, results: &[RuntimeWorkResult]) -> Result<(), StrategyError> {
        let mut runs = self.runs.clone();
        let mut runs_completed = self.runs_completed;

        for r in results {
            let run = runs
                .iter_mut()
                .find(|run| run.tag == r.tag)
                .ok_or_else(|| StrategyError::unknown_tag(r.tag))?;

            if let Some(telemetry) = &r.telemetry_pb {
                let chunk = proto::ChunkTelemetry::decode(telemetry.as_slice())?;

                merge::merge_chunk(&mut run.merge, &chunk)?;
            } else {
                run.merge.chunks_completed += 1;
                run.merge.iterations_total += u64::from(r.iterations);
                let n = run.merge.iterations_total.max(1);
                let old_mean = run.merge.weighted_mean_num_x10 as f64;
                let new_mean = old_mean
                    + (f64::from(r.mean_dps_x10) - old_mean) * f64::from(r.iterations) / n as f64;

                run.merge.weighted_mean_num_x10 = new_mean as u64;
            }

            run.completed += 1;

            if run.completed == run.chunks_total {
                runs_completed += 1;
            }
        }

        self.runs = runs;
        self.runs_completed = runs_completed;

        Ok(())
    }

    pub(super) fn is_complete(&self) -> bool {
        self.runs_completed >= self.runs.len() as u32
    }

    pub(super) fn has_pending(&self) -> bool {
        !self.queue.is_empty()
    }

    pub(super) fn peek_shape(&self) -> (u32, bool) {
        peek_queue(&self.queue)
    }

    pub(super) fn finalize(&self) -> Result<FinalOutput, StrategyError> {
        let baseline = self
            .runs
            .iter()
            .find(|r| r.tag == 0)
            .ok_or_else(StrategyError::missing_baseline)?;

        finalize_single("stat_weights", &baseline.merge)
    }
}

fn next_from_queue(
    queue: &mut VecDeque<PlannedChunk>,
    max_items: usize,
    next_item_id: &mut u64,
) -> Vec<RuntimeWorkItem> {
    let mut items = Vec::new();

    while items.len() < max_items {
        let Some(chunk) = queue.pop_front() else {
            break;
        };
        let item_id = *next_item_id;

        *next_item_id += 1;
        items.push(work_item_from_chunk(item_id, &chunk));
    }

    items
}

fn peek_queue(queue: &VecDeque<PlannedChunk>) -> (u32, bool) {
    queue.front().map_or((0, false), |c| (c.iterations, c.full))
}

const fn work_item_from_chunk(item_id: u64, chunk: &PlannedChunk) -> RuntimeWorkItem {
    RuntimeWorkItem {
        item_id,
        kind: RuntimeWorkItemKind::Base,
        tag: chunk.tag,
        iterations: chunk.iterations,
        seed_offset: chunk.seed_offset,
        fidelity: if chunk.full {
            RuntimeWorkFidelity::Full
        } else {
            RuntimeWorkFidelity::DpsOnly
        },
    }
}

fn finalize_single(
    strategy: &str,
    merge_state: &proto::RunningAggregateStateV1,
) -> Result<FinalOutput, StrategyError> {
    let result_view = merge::emit_result_snapshot(merge_state);
    let timeline_view = merge::emit_timeline_snapshot(merge_state);
    let result_pb = encode_and_check(
        &JobResult {
            strategy: strategy.into(),
            total_iterations: merge_state.iterations_total,
            duration_ms: 0,
            result: Some(proto::job_result::Result::Single(proto::SingleResult {
                result: Some(result_view),
            })),
        },
        RESULT_BUDGET,
    )?;
    let timeline_pb = JobTimeline {
        strategy: strategy.into(),
        timeline: Some(proto::job_timeline::Timeline::Single(
            proto::SingleTimeline {
                timeline: Some(timeline_view),
            },
        )),
    }
    .encode_to_vec();

    Ok(FinalOutput {
        result_pb,
        timeline_pb: Some(timeline_pb),
    })
}