car_inference/

offload.rs

//! On-device generation offload — process isolation for local inference.
//!
//! Closes the second half of [Parslee-ai/car-releases#74]. A single heavy
//! on-device generation (a large MLX Qwen prompt at a high `max_tokens`) can
//! abort the process from the **Metal/MLX C++ side** — an allocation/OOM abort
//! or a C++ exception crossing the FFI boundary *below* every Rust
//! `catch_unwind` guard ([`crate::InferenceEngine`]'s `catch_mlx`,
//! `handle_infer`'s `catch_unwind`). In the shared `car-server` daemon that
//! takes down inference for **every** connected client at once.
//!
//! The fix is to run on-device generation in a **separate worker process** the
//! daemon owns. When a `LocalGenerationOffload` is installed (via
//! [`set_local_offload`]), [`crate::InferenceEngine`]'s on-device branch hands
//! the fully-resolved request to it instead of running the Metal decode loop
//! in-process. A Metal abort then kills only the worker: the offload call
//! returns an [`InferenceError`], the daemon fails that one RPC gracefully and
//! stays up, and the next call respawns the worker.
//!
//! The slot is process-wide and defaults to `None`, so in-process consumers
//! (NAPI/PyO3/CLI, the worker itself) run on-device generation directly and
//! are completely unaffected — only a host that explicitly installs an
//! offloader (the daemon) gets the subprocess boundary.

use crate::stream::StreamEvent;
use crate::tasks::generate::GenerateRequest;
use crate::{InferenceError, InferenceResult};
use serde::{Deserialize, Serialize};
use std::future::Future;
use std::path::PathBuf;
use std::sync::atomic::{AtomicBool, Ordering};
use std::sync::{Arc, OnceLock, RwLock};

#[derive(Clone, Debug)]
struct InferenceControlScope {
    inference_id: String,
    termination: ControlledTerminationToken,
}

tokio::task_local! {
    static INFERENCE_CONTROL_SCOPE: InferenceControlScope;
}

/// Request-scoped evidence that the exact isolated backend was killed and
/// reaped. It is deliberately monotonic: only a confirmed kill + wait may set
/// it, and every adaptive/retry boundary can cheaply observe it afterward.
#[derive(Clone, Debug, Default)]
pub struct ControlledTerminationToken(Arc<AtomicBool>);

impl ControlledTerminationToken {
    /// Publish exact backend termination evidence. Callers must only invoke
    /// this after the isolated process has both accepted termination and been
    /// waited/reaped; body drops and async kill requests do not qualify.
    pub fn confirm_exact_backend_termination(&self) {
        self.0.store(true, Ordering::Release);
    }

    pub fn is_confirmed(&self) -> bool {
        self.0.load(Ordering::Acquire)
    }

    pub fn error_if_confirmed(&self) -> Result<(), InferenceError> {
        if self.is_confirmed() {
            Err(InferenceError::ControlledTermination)
        } else {
            Ok(())
        }
    }
}

/// Scope a server-minted lifecycle ID to one inference future without adding a
/// serializable field to `GenerateRequest`. Direct local-offload awaits inherit
/// this value; spawned tasks do not inherit it.
pub async fn scope_inference_control_id<F>(inference_id: String, future: F) -> F::Output
where
    F: Future,
{
    INFERENCE_CONTROL_SCOPE
        .scope(
            InferenceControlScope {
                inference_id,
                termination: ControlledTerminationToken::default(),
            },
            future,
        )
        .await
}

/// Snapshot the lifecycle ID in the current inference task, if one was scoped.
pub fn current_inference_control_id() -> Option<String> {
    INFERENCE_CONTROL_SCOPE
        .try_with(|scope| scope.inference_id.clone())
        .ok()
}

/// Caller deadlines for in-flight inferences, keyed by the server-minted
/// lifecycle ID (`scope_inference_control_id`). The daemon's `infer.deadline`
/// handler arms an entry AFTER the request has started — that is the wire
/// protocol's shape — so this is a mutable registry rather than a field on the
/// request, and the remote retry loop re-reads it at every phase boundary so a
/// late-armed deadline still takes effect mid-attempt (car-eyj: a 600 s
/// deadline used to be abandoned at 549 s by transport ceilings that had never
/// heard of it).
static REMOTE_DEADLINES: OnceLock<RwLock<std::collections::HashMap<String, RemoteDeadline>>> =
    OnceLock::new();

/// One armed caller deadline: when it expires, and the caller's own number
/// (`timeout_ms`) so a termination can NAME the deadline that was applied
/// instead of reporting an anonymous transport ceiling.
#[derive(Clone, Copy, Debug)]
pub struct RemoteDeadline {
    pub expires_at: std::time::Instant,
    pub applied_ms: u64,
}

impl RemoteDeadline {
    pub fn remaining(&self) -> std::time::Duration {
        self.expires_at
            .saturating_duration_since(std::time::Instant::now())
    }
}

fn remote_deadlines() -> &'static RwLock<std::collections::HashMap<String, RemoteDeadline>> {
    REMOTE_DEADLINES.get_or_init(|| RwLock::new(std::collections::HashMap::new()))
}

/// Arm (or re-arm) the caller's deadline for one inference. `timeout` is
/// measured from now, matching `infer.deadline`'s `timeout_ms` semantics.
pub fn set_remote_deadline(inference_id: &str, timeout: std::time::Duration) {
    if let Ok(mut map) = remote_deadlines().write() {
        map.insert(
            inference_id.to_string(),
            RemoteDeadline {
                expires_at: std::time::Instant::now() + timeout,
                applied_ms: timeout.as_millis() as u64,
            },
        );
    }
}

/// Drop the deadline entry once the inference reaches ANY terminal state.
/// Idempotent; the registry must never outlive its inference or the next
/// request to reuse an ID would inherit a stale deadline.
pub fn clear_remote_deadline(inference_id: &str) {
    if let Ok(mut map) = remote_deadlines().write() {
        map.remove(inference_id);
    }
}

/// The armed deadline for the CURRENT inference task, if any. Reads the
/// task-local lifecycle ID and then the registry, so a call outside a scoped
/// inference (embeddings, health probes) sees `None` and keeps its defaults.
pub fn current_remote_deadline() -> Option<RemoteDeadline> {
    let id = current_inference_control_id()?;
    remote_deadlines().read().ok()?.get(&id).copied()
}

/// Snapshot the request's controlled-termination token. Spawned backend tasks
/// must clone this explicitly because Tokio task-local values are not inherited.
pub fn current_controlled_termination_token() -> Option<ControlledTerminationToken> {
    INFERENCE_CONTROL_SCOPE
        .try_with(|scope| scope.termination.clone())
        .ok()
}

/// Stop an adaptive retry/fallback boundary after exact backend termination.
pub fn ensure_not_controlled_terminated() -> Result<(), InferenceError> {
    match current_controlled_termination_token() {
        Some(token) => token.error_if_confirmed(),
        None => Ok(()),
    }
}

#[derive(Clone, Debug, Serialize, Deserialize, PartialEq, Eq)]
pub struct LocalWorkerAdmission {
    pub policy: crate::resource_policy::ResourcePolicy,
    pub policy_generation: u64,
    pub state_root: PathBuf,
    pub measured_weights_bytes: u64,
}

#[derive(Clone, Debug, Serialize, Deserialize, PartialEq, Eq)]
pub struct LocalWorkerResidency {
    pub model_id: String,
    pub measured_weights_bytes: u64,
}

pub struct LocalOffloadResult {
    pub result: InferenceResult,
    pub residency: LocalWorkerResidency,
    pub retention: crate::backend_cache::BackendRetention,
}

pub struct LocalOffloadStream {
    pub events: tokio::sync::mpsc::Receiver<StreamEvent>,
    pub residency: LocalWorkerResidency,
    pub retention: crate::backend_cache::BackendRetention,
}

/// Evidence returned by an isolated generation backend after a termination
/// request. Dropping a body/stream/future is not evidence and must remain
/// `Unconfirmed`.
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum InferenceTerminationAck {
    Confirmed,
    Unconfirmed,
}

/// Handler that runs a fully-resolved on-device [`GenerateRequest`] somewhere
/// isolated from the calling process (the daemon's `car-server --mlx-worker`
/// child). Implementations must run the request against a real in-process
/// engine and return the same [`InferenceResult`] / [`StreamEvent`] stream the
/// caller would have produced itself — the boundary is transparent except that
/// a native abort surfaces as an `Err`/dropped stream instead of a crash.
#[async_trait::async_trait]
pub trait LocalGenerationOffload: Send + Sync {
    /// Run a non-streaming generation to completion in the worker.
    async fn generate(&self, request: GenerateRequest) -> Result<InferenceResult, InferenceError>;

    /// Run a streaming generation in the worker, returning a receiver of the
    /// same [`StreamEvent`]s the in-process path emits. The channel closes when
    /// the worker finishes (or dies); a mid-stream worker death is observed as
    /// the sender dropping, exactly like any other stream end.
    async fn stream(
        &self,
        request: GenerateRequest,
    ) -> Result<tokio::sync::mpsc::Receiver<StreamEvent>, InferenceError>;

    /// Admission-aware extension. The default preserves source compatibility
    /// for external legacy implementers and treats their result as transient;
    /// only implementations that explicitly override this method may report
    /// retained worker/process weights.
    async fn generate_admitted(
        &self,
        request: GenerateRequest,
        _admission: LocalWorkerAdmission,
    ) -> Result<LocalOffloadResult, InferenceError> {
        let model_id = request.model.clone().unwrap_or_default();
        let result = self.generate(request).await?;
        Ok(LocalOffloadResult {
            result,
            residency: LocalWorkerResidency {
                model_id,
                measured_weights_bytes: 0,
            },
            retention: crate::backend_cache::BackendRetention::Transient,
        })
    }

    async fn stream_admitted(
        &self,
        request: GenerateRequest,
        _admission: LocalWorkerAdmission,
    ) -> Result<LocalOffloadStream, InferenceError> {
        let model_id = request.model.clone().unwrap_or_default();
        let events = self.stream(request).await?;
        Ok(LocalOffloadStream {
            events,
            residency: LocalWorkerResidency {
                model_id,
                measured_weights_bytes: 0,
            },
            retention: crate::backend_cache::BackendRetention::Transient,
        })
    }

    fn refresh_resource_policy(&self, _generation: u64) {}

    async fn resident_models(&self) -> Vec<String> {
        Vec::new()
    }

    /// Exact process-generation allocation owner for parent-side residency
    /// publication. Legacy implementations may omit it and remain transient.
    fn resident_allocation_id(&self, _model_id: &str) -> Option<String> {
        None
    }

    /// Release worker-owned residency for one model. Success is acknowledged
    /// only after the owning worker has exited.
    async fn release_model(&self, _model_id: &str) -> Result<bool, InferenceError> {
        Ok(false)
    }

    /// Terminate the exact isolated process currently executing this ID.
    /// The compatibility default is intentionally conservative: only an exact
    /// kill-and-wait implementation may report confirmed termination.
    async fn terminate_inference(&self, _inference_id: &str) -> InferenceTerminationAck {
        InferenceTerminationAck::Unconfirmed
    }
}

fn offload_slot() -> &'static RwLock<Option<Arc<dyn LocalGenerationOffload>>> {
    static SLOT: OnceLock<RwLock<Option<Arc<dyn LocalGenerationOffload>>>> = OnceLock::new();
    SLOT.get_or_init(|| RwLock::new(None))
}

/// Install a process-wide on-device generation offloader. Pass `None` to
/// clear the slot. Re-registering overwrites any previous offloader.
pub fn set_local_offload(offload: Option<Arc<dyn LocalGenerationOffload>>) {
    let mut guard = offload_slot().write().expect("local offload slot poisoned");
    *guard = offload;
}

/// Snapshot the currently installed offloader (if any). Cheap `Arc` clone.
///
/// Returns `None` inside a worker process (`CAR_INFERENCE_WORKER` set) even if
/// a slot were somehow installed — a hard guard against a worker offloading to
/// itself and spawning an infinite chain of workers.
pub fn current_local_offload() -> Option<Arc<dyn LocalGenerationOffload>> {
    if is_offload_worker() {
        return None;
    }
    offload_slot()
        .read()
        .expect("local offload slot poisoned")
        .clone()
}

/// Whether this process is an on-device inference worker. Set by the daemon on
/// the `car-server --mlx-worker` child it spawns. A worker runs generation
/// in-process (the real Metal path), never offloads, and disables outcome
/// persistence so it can't race the parent daemon writing the shared
/// `~/.car` profile/ledger files.
pub fn is_offload_worker() -> bool {
    std::env::var_os("CAR_INFERENCE_WORKER").is_some()
}

#[cfg(test)]
pub(crate) fn test_offload_lock() -> &'static tokio::sync::Mutex<()> {
    static LOCK: OnceLock<tokio::sync::Mutex<()>> = OnceLock::new();
    LOCK.get_or_init(|| tokio::sync::Mutex::new(()))
}

#[cfg(test)]
mod tests {
    use super::*;

    struct NoopOffload;

    #[tokio::test]
    async fn legacy_or_nonisolated_offload_never_confirms_termination() {
        let offload = NoopOffload;
        assert_eq!(
            offload.terminate_inference("opaque").await,
            InferenceTerminationAck::Unconfirmed
        );
    }

    #[async_trait::async_trait]
    impl LocalGenerationOffload for NoopOffload {
        async fn generate(
            &self,
            _request: GenerateRequest,
        ) -> Result<InferenceResult, InferenceError> {
            Err(InferenceError::InferenceFailed("noop".into()))
        }
        async fn stream(
            &self,
            _request: GenerateRequest,
        ) -> Result<tokio::sync::mpsc::Receiver<StreamEvent>, InferenceError> {
            Err(InferenceError::InferenceFailed("noop".into()))
        }
    }

    #[allow(dead_code)]
    fn external_legacy_trait_contract_still_compiles(
        implementation: Arc<dyn LocalGenerationOffload>,
    ) -> Arc<dyn LocalGenerationOffload> {
        implementation
    }

    #[tokio::test]
    async fn slot_set_and_clear() {
        let _guard = test_offload_lock().lock().await;
        // Serialize against other tests touching the global slot is not
        // needed — this is the only test that writes it, and it restores None.
        assert!(current_local_offload().is_none());
        set_local_offload(Some(Arc::new(NoopOffload)));
        // Note: current_local_offload() would still return None here if
        // CAR_INFERENCE_WORKER is set in the test env; it isn't by default.
        assert!(current_local_offload().is_some());
        set_local_offload(None);
        assert!(current_local_offload().is_none());
    }
}