eredu_runtime/

capture.rs

//! Budget admission at the observation boundary, before native handles are retained.

use eredu_core::capture::*;

mod checkpoint;
#[cfg(test)]
mod tests;
pub use checkpoint::{
    CaptureCheckpoint, CaptureForkRequest, InterventionForkRequest, PreparedCaptureRestore,
};

/// A run owns one ledger and at most one step of host records. Consumers must drain
/// each step before another is started; there is no producer queue.
pub struct CaptureSession {
    // Identity is deliberately not serialized or copied into child sessions.
    owner: std::sync::Arc<()>,
    checkpoint_ready: bool,
    has_step: bool,
    pub(crate) plan: AdmittedCapturePlan,
    pub(crate) ledger: CaptureLedger,
    pub(crate) records: Option<Vec<CaptureRecord>>,
    pub(crate) prediction: u64,
    pub(crate) phase: CapturePhase,
    pub(crate) capture_seconds: f64,
    pub(crate) interventions: Option<crate::intervention::InterventionRun>,
}

impl CaptureSession {
    /// Creates an unstarted capture run owning its admission and ledger.
    pub fn new(plan: AdmittedCapturePlan) -> Self {
        Self {
            owner: std::sync::Arc::new(()),
            checkpoint_ready: true,
            has_step: false,
            ledger: CaptureLedger::new(&plan),
            plan,
            records: None,
            prediction: 0,
            phase: CapturePhase::Prefill,
            capture_seconds: 0.0,
            interventions: None,
        }
    }

    /// Borrows this run's immutable admission.
    pub fn plan(&self) -> &AdmittedCapturePlan {
        &self.plan
    }

    /// Immutable intervention admission currently paired with this shared owner.
    pub fn intervention_plan(&self) -> Option<&eredu_core::intervention::AdmittedInterventionPlan> {
        self.interventions.as_ref().map(|run| &run.plan)
    }

    /// Reserves diagnostic envelopes before execution, including scheduled skips and
    /// missing values. Exhaustion here fails the step: emitting an unaccounted skip
    /// record would itself violate the export limit.
    pub fn begin_step(&mut self, phase: CapturePhase, prediction: u64) -> Result<(), CaptureError> {
        if self.records.is_some() {
            return Err(CaptureError::Invalid(
                "previous capture step has not been consumed".into(),
            ));
        }
        if prediction >= self.plan.request().max_predictions {
            return Err(CaptureError::Invalid(
                "generation exceeds admitted prediction range".into(),
            ));
        }
        // Even a failed reservation can consume cumulative resources. A failed
        // attempt is not a resumable boundary merely because records were drained.
        self.checkpoint_ready = false;
        self.has_step = true;
        self.ledger.begin_step();
        let mut records = Vec::new();
        for (selection, point) in self.plan.plan().selections.iter().zip(self.plan.points()) {
            let charged = metadata_reservation(selection, point)?;
            if let Some(CaptureSkipReason::Limit { budget, cumulative }) =
                self.ledger.reserve(charged)?
            {
                return Err(CaptureError::Limit { budget, cumulative });
            }
            records.push(CaptureRecord {
                schema_version: CAPTURE_SCHEMA_VERSION,
                selection_id: selection.id.clone(),
                path: selection.path.clone(),
                node_id: point.node_id.clone(),
                position: point.position,
                source_shape: None,
                selected_shape: None,
                outcome: if selection.schedule.includes(phase, prediction) {
                    CaptureOutcome::Missing
                } else {
                    CaptureOutcome::Skipped {
                        reason: CaptureSkipReason::Schedule,
                    }
                },
                payload: None,
                charged,
            });
        }
        self.records = Some(records);
        self.phase = phase;
        self.prediction = prediction;
        self.capture_seconds = 0.0;
        if let Some(interventions) = &mut self.interventions {
            interventions.begin_step(&mut self.ledger, phase, prediction)?;
        }
        Ok(())
    }

    /// Invoked while the architecture borrows a tensor. A skipped point never calls
    /// the native transform and never clones a native tensor handle.
    pub fn observe<B: CaptureBackend>(
        &mut self,
        backend: &mut B,
        path: &str,
        tensor: &B::Tensor,
    ) -> Result<(), CaptureExecutionError<B::Error>> {
        let Some(records) = self.records.as_mut() else {
            return Err(CaptureError::Invalid("capture step not started".into()).into());
        };
        for ((selection, point), record) in self
            .plan
            .plan()
            .selections
            .iter()
            .zip(self.plan.points())
            .zip(records)
        {
            if selection.path != path || matches!(record.outcome, CaptureOutcome::Skipped { .. }) {
                continue;
            }
            if !matches!(record.outcome, CaptureOutcome::Missing) {
                return Err(
                    CaptureError::Invalid(format!("observation emitted twice: {path}")).into(),
                );
            }
            let started = std::time::Instant::now();
            let result = capture_value(
                backend,
                tensor,
                selection,
                point,
                record,
                self.plan.request(),
                self.phase,
                self.prediction,
                &mut self.ledger,
            );
            self.capture_seconds += started.elapsed().as_secs_f64();
            if let Err(error) = result {
                let reason = match &error {
                    CaptureExecutionError::Admission(CaptureError::Limit {
                        budget,
                        cumulative,
                    }) => CaptureFailureReason::Limit {
                        budget: *budget,
                        cumulative: *cumulative,
                    },
                    CaptureExecutionError::Admission(CaptureError::Unsupported(_)) => {
                        CaptureFailureReason::Unsupported
                    }
                    CaptureExecutionError::Admission(_) => CaptureFailureReason::Invalid,
                    CaptureExecutionError::Backend(_) => CaptureFailureReason::Native,
                };
                record.payload = None;
                record.outcome = CaptureOutcome::Failed {
                    reason,
                    message: bounded_diagnostic(&error),
                };
                return Err(error);
            }
        }
        Ok(())
    }

    /// Moves the current bounded record batch to the consumer.
    pub fn take_step(&mut self) -> Option<CapturedStep> {
        if let Some(records) = &self.records {
            self.checkpoint_ready = self.finish_interventions().is_ok()
                && !records
                    .iter()
                    .any(|record| matches!(record.outcome, CaptureOutcome::Failed { .. }));
        }
        self.records.take().map(|records| CapturedStep {
            phase: self.phase,
            prediction_index: self.prediction,
            records,
            interventions: self
                .interventions
                .as_mut()
                .map_or_else(Vec::new, |run| run.take_records()),
            step_usage: self.ledger.step(),
            cumulative_usage: self.ledger.total(),
            capture_seconds: self.capture_seconds,
        })
    }
}

/// Shared transformation path for ordinary captures and intervention evidence.
#[allow(clippy::too_many_arguments)]
pub(crate) fn capture_value<B: CaptureBackend>(
    backend: &mut B,
    tensor: &B::Tensor,
    selection: &CaptureSelection,
    point: &eredu_core::ObservationPoint,
    record: &mut CaptureRecord,
    request: CaptureRequestShape,
    phase: CapturePhase,
    prediction: u64,
    ledger: &mut CaptureLedger,
) -> Result<(), CaptureExecutionError<B::Error>> {
    let path = &selection.path;

    let shape = backend
        .shape(tensor)
        .map_err(CaptureExecutionError::Backend)?;
    request.validate_actual(point, phase, prediction, &shape)?;
    if let Some(expected) = request.resolve(point, phase, prediction)? {
        if expected != shape {
            return Err(CaptureError::Invalid(format!(
                "runtime shape for {path}: expected {expected:?}, got {shape:?}"
            ))
            .into());
        }
    }
    let slice = resolve_slice(point, selection, &shape)?;
    let usage = backend.estimate(tensor, selection, &slice)?;
    record.source_shape = Some(shape);
    record.selected_shape = Some(slice.shape.clone());
    if let Some(reason) = ledger.reserve(usage)? {
        record.outcome = CaptureOutcome::Skipped { reason };
        return Ok(());
    }
    record.charged = record.charged.checked_add(usage)?;
    let mut payload = backend
        .transform(tensor, selection, &slice)
        .map_err(CaptureExecutionError::Backend)?;
    if let CapturePayload::Candidates(candidates) = &mut payload {
        candidates.source = if record.position == eredu_core::ObservationPosition::AfterIntervention
        {
            CandidateLogitsSource::Effective
        } else {
            CandidateLogitsSource::Original
        };
    }
    let available = elements(&slice.shape)?;
    record.outcome = match selection.transform {
        CaptureTransform::Preview { max_elements } if max_elements < available => {
            CaptureOutcome::Truncated {
                available_elements: available,
                emitted_elements: max_elements,
            }
        }
        _ => CaptureOutcome::Captured,
    };
    record.payload = Some(payload);
    // Count through a bounded sink, without allocating a second JSON buffer.
    // This is a backend-contract check, not a substitute for the pre-copy estimate.
    let mut sink = CountingWriter {
        written: 0,
        limit: record.charged.encoded_bytes,
    };
    serde_json::to_writer(&mut sink, record)
        .map_err(|_| CaptureError::Invalid("backend underestimated encoded capture size".into()))?;
    Ok(())
}

pub(crate) fn bounded_diagnostic(error: &impl std::fmt::Display) -> String {
    use std::fmt::Write;
    struct Message(String);
    impl std::fmt::Write for Message {
        fn write_str(&mut self, text: &str) -> std::fmt::Result {
            let mut end = text.len().min(256 - self.0.len());
            while !text.is_char_boundary(end) {
                end -= 1;
            }
            self.0.push_str(&text[..end]);
            if end < text.len() {
                Err(std::fmt::Error)
            } else {
                Ok(())
            }
        }
    }
    let mut message = Message(String::with_capacity(256));
    let _ = write!(&mut message, "{error}");
    message.0
}

/// Conservative JSON envelope reservation: every UTF-8 byte can escape to at most
/// six bytes; each dimension can use twenty decimal digits. Numeric payloads are
/// reserved by the native estimator. Selection count and metadata are bounded before
/// generation begins, so missing/skipped points cannot form an unbounded queue.
pub fn metadata_reservation(
    selection: &CaptureSelection,
    point: &eredu_core::ObservationPoint,
) -> Result<CaptureUsage, CaptureError> {
    let strings = add(
        add(selection.id.len() as u64, selection.path.len() as u64)?,
        point.node_id.len() as u64,
    )?;
    let rank = point.axes.as_ref().map_or(32, |axes| axes.len() as u64);
    Ok(CaptureUsage {
        captures: 0,
        retained_bytes: 0,
        host_bytes: add(512, add(strings, mul(rank, 128)?)?)?,
        encoded_bytes: add(2048, add(mul(strings, 6)?, mul(rank, 64)?)?)?,
    })
}

/// Checks conservative known per-step and run-total costs before execution.
/// Backends supply only their side-effect-free storage/transfer estimator. Unknown
/// dimensions remain subject to exact observation-time reservation. Step estimates
/// conservatively allow every phase-enabled selection to coincide.
pub fn preflight(
    plan: &AdmittedCapturePlan,
    estimate: impl FnMut(
        &[u64],
        &CaptureSelection,
        &ResolvedCaptureSlice,
    ) -> Result<CaptureUsage, CaptureError>,
) -> Result<(), CaptureError> {
    preflight_with_extra(plan, &[], CaptureUsage::default(), &[], estimate)
}

/// Revalidates admission against the loaded catalog, then checks known geometry
/// and budgets with backend estimates. Native execution-mode checks stay with the
/// caller; this helper never accesses a device or starts a submission.
pub fn validate_session(
    plan: &AdmittedCapturePlan,
    discovery: &CaptureDiscovery,
    estimate: impl FnMut(
        &[u64],
        &CaptureSelection,
        &ResolvedCaptureSlice,
    ) -> Result<CaptureUsage, CaptureError>,
) -> Result<(), CaptureError> {
    validate_continuation(plan, discovery, 0, CaptureUsage::default(), estimate)
}

pub(crate) fn validate_continuation(
    plan: &AdmittedCapturePlan,
    discovery: &CaptureDiscovery,
    next_prediction: u64,
    inherited: CaptureUsage,
    estimate: impl FnMut(
        &[u64],
        &CaptureSelection,
        &ResolvedCaptureSlice,
    ) -> Result<CaptureUsage, CaptureError>,
) -> Result<(), CaptureError> {
    let checked = plan.plan().clone().admit(
        &discovery.catalog,
        &discovery.support,
        &discovery.support.capture,
        plan.request(),
    )?;
    if checked.identity() != plan.identity() {
        return Err(CaptureError::Invalid(
            "capture admission does not match this session's catalog".into(),
        ));
    }
    preflight_continuation(
        &checked,
        &[],
        CaptureUsage::default(),
        &[],
        next_prediction,
        inherited,
        estimate,
    )
}

pub(crate) fn preflight_with_extra(
    plan: &AdmittedCapturePlan,
    extra: &[(CaptureSelection, eredu_core::ObservationPoint)],
    base: CaptureUsage,
    scheduled_costs: &[(CaptureSchedule, [CaptureUsage; 2])],
    estimate: impl FnMut(
        &[u64],
        &CaptureSelection,
        &ResolvedCaptureSlice,
    ) -> Result<CaptureUsage, CaptureError>,
) -> Result<(), CaptureError> {
    preflight_continuation(
        plan,
        extra,
        base,
        scheduled_costs,
        0,
        CaptureUsage::default(),
        estimate,
    )
}

#[allow(clippy::too_many_arguments)]
pub(crate) fn preflight_continuation(
    plan: &AdmittedCapturePlan,
    extra: &[(CaptureSelection, eredu_core::ObservationPoint)],
    mut base: CaptureUsage,
    scheduled_costs: &[(CaptureSchedule, [CaptureUsage; 2])],
    next_prediction: u64,
    inherited: CaptureUsage,
    mut estimate: impl FnMut(
        &[u64],
        &CaptureSelection,
        &ResolvedCaptureSlice,
    ) -> Result<CaptureUsage, CaptureError>,
) -> Result<(), CaptureError> {
    let remaining = plan
        .request()
        .max_predictions
        .checked_sub(next_prediction)
        .ok_or_else(|| {
            CaptureError::Invalid("continuation exceeds admitted prediction range".into())
        })?;
    let entries: Vec<_> = plan
        .plan()
        .selections
        .iter()
        .zip(plan.points())
        .chain(extra.iter().map(|(selection, point)| (selection, point)))
        .collect();
    for &(selection, point) in &entries {
        base = base.checked_add(metadata_reservation(selection, point)?)?;
    }
    if let Some(budget) = base.exceeded(plan.plan().limits.per_step) {
        return Err(CaptureError::Limit {
            budget,
            cumulative: false,
        });
    }
    let mut total = inherited.checked_add(base.checked_mul(remaining)?)?;
    for phase in [CapturePhase::Prefill, CapturePhase::Decode] {
        if remaining == 0 || (phase == CapturePhase::Prefill && next_prediction > 0) {
            continue;
        }
        if phase == CapturePhase::Decode && plan.request().max_predictions <= 1 {
            continue;
        }
        let mut step = base;
        for (schedule, costs) in scheduled_costs {
            if let Some((count, _)) = schedule.count_and_last_from(
                phase,
                next_prediction,
                plan.request().max_predictions,
            )? {
                let cost = costs[if phase == CapturePhase::Prefill { 0 } else { 1 }];
                step = step.checked_add(cost)?;
                total = total.checked_add(cost.checked_mul(count)?)?;
            }
        }
        for &(selection, point) in &entries {
            let Some((count, last)) = selection.schedule.count_and_last_from(
                phase,
                next_prediction,
                plan.request().max_predictions,
            )?
            else {
                continue;
            };
            if let Some(shape) = plan.request().resolve(point, phase, last)? {
                let slice = resolve_slice(point, selection, &shape)?;
                let cost = estimate(&shape, selection, &slice)?;
                if plan.plan().limits.on_limit == CaptureLimitPolicy::Fail {
                    step = step.checked_add(cost)?;
                    total = total.checked_add(cost.checked_mul(count)?)?;
                }
            }
        }
        if let Some(budget) = step.exceeded(plan.plan().limits.per_step) {
            return Err(CaptureError::Limit {
                budget,
                cumulative: false,
            });
        }
    }
    if let Some(budget) = total.exceeded(plan.plan().limits.cumulative) {
        return Err(CaptureError::Limit {
            budget,
            cumulative: true,
        });
    }
    Ok(())
}

struct CountingWriter {
    written: u64,
    limit: u64,
}
impl std::io::Write for CountingWriter {
    fn write(&mut self, bytes: &[u8]) -> std::io::Result<usize> {
        let next = self
            .written
            .checked_add(bytes.len() as u64)
            .filter(|next| *next <= self.limit)
            .ok_or_else(|| std::io::Error::other("capture JSON budget exceeded"))?;
        self.written = next;
        Ok(bytes.len())
    }
    fn flush(&mut self) -> std::io::Result<()> {
        Ok(())
    }
}

/// Failure before reservation or during the reserved native transformation.
#[derive(Debug, thiserror::Error)]
pub enum CaptureExecutionError<E: std::error::Error + 'static> {
    /// Invalid geometry or rejected budget/capability.
    #[error(transparent)]
    Admission(#[from] CaptureError),
    /// Native transformation failed under the backend's recovery owner.
    #[error("native capture failed: {0}")]
    Backend(E),
}