car-inference 0.36.0

//! 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 std::sync::{Arc, OnceLock, RwLock};

/// 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>;
}

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)]
mod tests {
    use super::*;

    struct NoopOffload;

    #[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()))
        }
    }

    #[test]
    fn slot_set_and_clear() {
        // 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());
    }
}