car-inference 0.50.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 serde::{Deserialize, Serialize};
use std::path::PathBuf;
use std::sync::{Arc, OnceLock, RwLock};

#[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,
}

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

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;

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