atlas_core/
cuda_host.rs

1// SPDX-License-Identifier: AGPL-3.0-only
2
3//! The **process-scoped** half of the CUDA runtime: the context, the stream,
4//! and the teardown check that guards a model swap.
5//!
6//! Split from [`crate::registry`] because the two have different lifetimes and
7//! conflating them is what made in-process model swapping impossible. The host
8//! is derived from the *device* and lives for the process; the module registry
9//! is derived from the *checkpoint* and lives for one model.
10
11use std::sync::{Arc, OnceLock};
12
13use cudarc::driver::{CudaContext, CudaStream};
14
15use crate::error::{AtlasError, Result};
16use crate::registry::AtlasRegistry;
17
18/// The CUDA context and stream — **process-scoped, on purpose**.
19///
20/// Creating and destroying a CUDA context is precisely what in-process model
21/// swapping exists to avoid: it is slow, it invalidates every handle in the
22/// process, and on GB10 UVM it is the operation with the worst failure modes.
23/// The context is derived from the *device*, not from the checkpoint, so one
24/// model's context serves the next one unchanged.
25pub struct CudaHost {
26    pub ctx: Arc<CudaContext>,
27    pub stream: Arc<CudaStream>,
28    ordinal: usize,
29}
30
31impl CudaHost {
32    pub fn ordinal(&self) -> usize {
33        self.ordinal
34    }
35}
36
37// SAFETY: as for `AtlasRegistry` below — these are handles into a context that
38// outlives every reader, and CUDA serialises the work itself via streams.
39unsafe impl Send for CudaHost {}
40unsafe impl Sync for CudaHost {}
41
42/// The process's CUDA host.
43///
44/// **STATIC, AND THIS IS THE CASE WHERE IT MUST BE.** The argument, in full,
45/// because the standing rule is that every surviving static earns one:
46///
47/// 1. **It is derived from the process and the device, not from any model.** A
48///    CUDA context belongs to a (process, GPU) pair. Nothing about a checkpoint
49///    reaches it, so it cannot go stale when the model changes — the property
50///    that makes every other static here a hazard simply does not apply.
51/// 2. **The driver enforces the singleton anyway.** There is one primary
52///    context per device per process. Holding two handles would not give two
53///    contexts; it would give two names for the same one, with the ownership
54///    question moved from the type system to a convention.
55/// 3. **Its whole purpose is to outlive every model.** Not destroying and
56///    recreating the context across a swap IS the feature — that operation is
57///    slow, invalidates every handle in the process, and has the worst failure
58///    modes on GB10 UVM. A context propagated per-model would have to be
59///    recreated per-model, which is precisely what is being avoided.
60/// 4. **Propagating it changes nothing it does not already guarantee.** Every
61///    consumer reaches it through the `AtlasRegistry` it already holds
62///    (`registry.host()`), so the *usage* is propagated. This static is only
63///    the place the one context is created and found.
64///
65/// Rebinding to a different ordinal is an error rather than a silent no-op —
66/// that was the bug in the previous `get_or_init(ordinal, kernel_blobs)`.
67static HOST: OnceLock<std::result::Result<Arc<CudaHost>, String>> = OnceLock::new();
68
69/// Get (or create) the process CUDA host on `ordinal`.
70///
71/// The ordinal is fixed by the first call — a process serves one GPU, and a
72/// later call asking for a different one is a bug worth reporting rather than
73/// silently ignoring, which is exactly what the old `get_or_init` did with its
74/// `kernel_blobs` argument.
75pub fn host(ordinal: usize) -> Result<Arc<CudaHost>> {
76    let result = HOST.get_or_init(|| {
77        let ctx = CudaContext::new(ordinal).map_err(|e| format!("{e}"))?;
78        let stream = ctx.new_stream().map_err(|e| format!("{e}"))?;
79        Ok(Arc::new(CudaHost {
80            ctx,
81            stream,
82            ordinal,
83        }))
84    });
85    match result {
86        Ok(h) if h.ordinal == ordinal => Ok(h.clone()),
87        Ok(h) => Err(AtlasError::ModuleLoad(format!(
88            "CUDA host already bound to GPU {} — cannot rebind to {ordinal}",
89            h.ordinal
90        ))),
91        Err(msg) => Err(AtlasError::ModuleLoad(msg.clone())),
92    }
93}
94
95/// Release a registry, refusing to pretend if anything still holds a handle.
96///
97/// This is the **missed-propagation detector**. The whole hazard of in-process
98/// swapping is a reference to the old model surviving somewhere nobody
99/// remembered to scope; `Arc::strong_count` turns that from silent wrong output
100/// into a named error at teardown, before the next model loads.
101pub fn release(registry: Arc<AtlasRegistry>) -> Result<()> {
102    let outstanding = Arc::strong_count(&registry) - 1;
103    if outstanding > 0 {
104        return Err(AtlasError::ModuleLoad(format!(
105            "cannot release the kernel modules: {outstanding} handle(s) are still live. \
106             Something is holding the previous model's registry — find it before swapping, \
107             or the next model will run against unloaded modules."
108        )));
109    }
110    let mut owned = Arc::try_unwrap(registry).map_err(|_| {
111        AtlasError::ModuleLoad("registry handle count changed during release".to_string())
112    })?;
113    let failures = owned.unload_raw();
114    if !failures.is_empty() {
115        return Err(AtlasError::ModuleLoad(format!(
116            "{} module(s) failed to unload: {}",
117            failures.len(),
118            failures.join("; ")
119        )));
120    }
121    Ok(())
122}
atlas_core/cuda_host.rs

atlas_core/
cuda_host.rs
