metrale_gpu_runtime/
cuda_host.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! 2026-09-25: The process-scoped CUDA context and stream (`CudaHost`), and
4//! `release`, which unloads one model's [`crate::registry`] modules.
5//!
6//! Owner: gpu-runtime.
7//! Invariants:
8//! - The process has at most one `CudaHost`, on the ordinal of the first `host`
9//!   call; a later call with another ordinal is an error, and a failed first
10//!   creation is returned as the same error to every later call.
11//! - `release` unloads nothing while another `Arc` to the registry is live.
12
13use std::sync::{Arc, OnceLock};
14
15use cudarc::driver::{CudaContext, CudaStream};
16
17use crate::registry::MetraleRegistry;
18use metrale_core::error::{MetraleError, Result};
19
20/// 2026-09-25: The CUDA context and stream for the process. Each model's
21/// `MetraleRegistry` loads its modules into this one context.
22pub struct CudaHost {
23    pub ctx: Arc<CudaContext>,
24    pub stream: Arc<CudaStream>,
25    ordinal: usize,
26}
27
28impl CudaHost {
29    pub fn ordinal(&self) -> usize {
30        self.ordinal
31    }
32}
33
34// 2026-09-25: SAFETY: the same argument as for `MetraleRegistry` in
35// `registry.rs`: these are handles into a context that is never destroyed.
36unsafe impl Send for CudaHost {}
37unsafe impl Sync for CudaHost {}
38
39/// 2026-09-25: The process's CUDA host. Static because it belongs to the
40/// (process, device) pair, not to a model: it is created once and kept across
41/// model loads.
42static HOST: OnceLock<std::result::Result<Arc<CudaHost>, String>> = OnceLock::new();
43
44/// 2026-09-25: Get, or on the first call create, the process CUDA host on
45/// `ordinal`. The first call fixes the ordinal; asking for another is an error.
46pub fn host(ordinal: usize) -> Result<Arc<CudaHost>> {
47    let result = HOST.get_or_init(|| {
48        let ctx = CudaContext::new(ordinal).map_err(|e| format!("{e}"))?;
49        let stream = ctx.new_stream().map_err(|e| format!("{e}"))?;
50        Ok(Arc::new(CudaHost {
51            ctx,
52            stream,
53            ordinal,
54        }))
55    });
56    match result {
57        Ok(h) if h.ordinal == ordinal => Ok(h.clone()),
58        Ok(h) => Err(MetraleError::ModuleLoad(format!(
59            "CUDA host already bound to GPU {} — cannot rebind to {ordinal}",
60            h.ordinal
61        ))),
62        Err(msg) => Err(MetraleError::ModuleLoad(msg.clone())),
63    }
64}
65
66/// 2026-09-25: Unload a model's kernel modules. Errors, unloading nothing, while
67/// another strong `Arc` to the registry is live; otherwise drops every module
68/// the registry holds.
69pub fn release(registry: Arc<MetraleRegistry>) -> Result<()> {
70    let outstanding = Arc::strong_count(&registry) - 1;
71    if outstanding > 0 {
72        return Err(MetraleError::ModuleLoad(format!(
73            "cannot release the kernel modules: {outstanding} handle(s) are still live. \
74             Something is holding the previous model's registry — find it before swapping, \
75             or the next model will run against unloaded modules."
76        )));
77    }
78    let mut owned = Arc::try_unwrap(registry).map_err(|_| {
79        MetraleError::ModuleLoad("registry handle count changed during release".to_string())
80    })?;
81    let failures = owned.unload_raw();
82    if !failures.is_empty() {
83        return Err(MetraleError::ModuleLoad(format!(
84            "{} module(s) failed to unload: {}",
85            failures.len(),
86            failures.join("; ")
87        )));
88    }
89    Ok(())
90}
metrale_gpu_runtime/cuda_host.rs

metrale_gpu_runtime/
cuda_host.rs
