diff --git a/.gitignore b/.gitignore index 6166323b8..f674ce706 100644 --- a/.gitignore +++ b/.gitignore @@ -44,3 +44,4 @@ examples/**/uv.lock # Claude Code local state .claude/*.lock +.idea/ diff --git a/Cargo.lock b/Cargo.lock index 3f03d8486..9d23f1ee9 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -2208,6 +2208,7 @@ dependencies = [ "serde_with", "sqlx", "storekey 0.9.0", + "temp-env", "tokio", "tokio-util", "tracing", @@ -11407,6 +11408,15 @@ version = "0.13.5" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "adb6935a6f5c20170eeceb1a3835a49e12e19d792f6dd344ccc76a985ca5a6ca" +[[package]] +name = "temp-env" +version = "0.3.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "96374855068f47402c3121c6eed88d29cb1de8f3ab27090e273e420bdabcf050" +dependencies = [ + "parking_lot", +] + [[package]] name = "tempfile" version = "3.27.0" diff --git a/python/cocoindex/_internal/api.py b/python/cocoindex/_internal/api.py index e6a5b71e2..875d1e1fa 100644 --- a/python/cocoindex/_internal/api.py +++ b/python/cocoindex/_internal/api.py @@ -89,9 +89,10 @@ from .environment import Environment, EnvironmentBuilder, LifespanFn from .environment import lifespan +from .core import GPUPool + from .runner import ( GPU, - GPUPool, GPURunner, Runner, configure_gpu_pool, diff --git a/python/cocoindex/_internal/core.pyi b/python/cocoindex/_internal/core.pyi index 79e3741bc..6ac64418d 100644 --- a/python/cocoindex/_internal/core.pyi +++ b/python/cocoindex/_internal/core.pyi @@ -704,3 +704,39 @@ class RateLimiter: cls, max_rows_per_second: float, burst_window_secs: float = 1.0 ) -> "RateLimiter": ... def acquire(self, n: int = 1) -> Coroutine[Any, Any, None]: ... + +######################################################## +# GPU pool +######################################################## + +class GPUPool: + def __init__(self, num_gpus: int) -> None: ... + + @property + def num_gpus(self) -> int: ... + + async def acquire(self, fraction: float) -> int: + """ + Acquires a fraction of a GPU and returns the GPU ID. + """ + ... + + async def acquire_full(self, gpu_count: int) -> list[int]: + """ + Acquires a given integer number of fully available GPUs (capacity == 1.0) from the GPU pool. + + The gpu_count should be greater or equal to 1 and less than or equal to the number of GPUs in the pool. + """ + ... + + def release(self, gpu_id: int, fraction: float) -> None: + """ + Releases a fraction of capacity back to the specified GPU ID. + """ + ... + + @staticmethod + def default() -> "GPUPool": ... + + @staticmethod + def from_config(config_str: str) -> "GPUPool": ... diff --git a/python/cocoindex/_internal/runner.py b/python/cocoindex/_internal/runner.py index 3544b7a67..565a6585b 100644 --- a/python/cocoindex/_internal/runner.py +++ b/python/cocoindex/_internal/runner.py @@ -213,110 +213,6 @@ def in_subprocess() -> bool: # ============================================================================ -def _detect_num_gpus() -> int: - """Detect the number of GPUs available for the default pool. - - Detection order: - - 1. ``COCOINDEX_NUM_GPUS`` environment variable (explicit override). - 2. ``CUDA_VISIBLE_DEVICES`` environment variable (count of entries). - 3. ``nvidia-smi`` command output (if available). - 4. Default to ``1``. - """ - env_num = os.environ.get("COCOINDEX_NUM_GPUS") - if env_num is not None: - return max(1, int(env_num)) - - cuda_visible = os.environ.get("CUDA_VISIBLE_DEVICES") - if cuda_visible is not None: - devices = [d.strip() for d in cuda_visible.split(",") if d.strip()] - if devices: - return len(devices) - # Empty CUDA_VISIBLE_DEVICES disables CUDA devices; keep a logical - # pool size of 1 so GPU runners still behave predictably. - return 1 - - try: - result = subprocess.run( - ["nvidia-smi", "--query-gpu=count", "--format=csv,noheader"], - capture_output=True, - text=True, - timeout=5, - check=False, - ) - if result.returncode == 0: - output = result.stdout.strip().splitlines() - if output: - count = int(output[0].strip()) - if count > 0: - return count - except Exception: - pass - - return 1 - - -class GPUPool: - """Tracks fractional GPU capacity across multiple GPUs. - - Each GPU starts with capacity 1.0. ``acquire(fraction)`` blocks until a - GPU with enough remaining capacity is available, then returns its id. - ``release(gpu_id, fraction)`` restores capacity and wakes waiters. - - The default pool size is auto-detected from ``COCOINDEX_NUM_GPUS``, - ``CUDA_VISIBLE_DEVICES``, or ``nvidia-smi`` (falling back to 1). - Call ``configure_gpu_pool(N)`` to override programmatically. - """ - - _num_gpus: int - _capacity: list[float] - _cond: asyncio.Condition | None - _bound_loop: asyncio.AbstractEventLoop | None - - def __init__(self, num_gpus: int) -> None: - if num_gpus < 1: - raise ValueError(f"num_gpus must be >= 1, got {num_gpus}") - self._num_gpus = num_gpus - self._capacity = [1.0] * num_gpus - self._cond = None - self._bound_loop = None - - @property - def num_gpus(self) -> int: - return self._num_gpus - - def _get_cond(self) -> asyncio.Condition: - loop = asyncio.get_running_loop() - if self._cond is None or self._bound_loop is not loop: - self._cond = asyncio.Condition() - self._bound_loop = loop - return self._cond - - def _find_available(self, fraction: float) -> int | None: - best_gpu = None - best_cap = -1.0 - for i, cap in enumerate(self._capacity): - if cap >= fraction and cap > best_cap: - best_gpu = i - best_cap = cap - return best_gpu - - async def acquire(self, fraction: float) -> int: - async with self._get_cond(): - while True: - gpu_id = self._find_available(fraction) - if gpu_id is not None: - self._capacity[gpu_id] -= fraction - return gpu_id - await self._get_cond().wait() - - async def release(self, gpu_id: int, fraction: float) -> None: - cond = self._get_cond() - async with cond: - self._capacity[gpu_id] += fraction - cond.notify_all() - - # ============================================================================ # GPU identity propagation # ============================================================================ @@ -359,15 +255,15 @@ def _run_with_gpu_context( # Default GPU pool # ============================================================================ -_default_gpu_pool: GPUPool | None = None +_default_gpu_pool: core.GPUPool | None = None _default_gpu_pool_lock = threading.Lock() -def _get_default_gpu_pool() -> GPUPool: +def _get_default_gpu_pool() -> core.GPUPool: global _default_gpu_pool with _default_gpu_pool_lock: if _default_gpu_pool is None: - _default_gpu_pool = GPUPool(num_gpus=_detect_num_gpus()) + _default_gpu_pool = core.GPUPool.default() return _default_gpu_pool @@ -375,7 +271,7 @@ def configure_gpu_pool(num_gpus: int) -> None: """Override the default GPU pool. Must be called before any GPU function runs.""" global _default_gpu_pool with _default_gpu_pool_lock: - _default_gpu_pool = GPUPool(num_gpus=num_gpus) + _default_gpu_pool = core.GPUPool(num_gpus=num_gpus) # ============================================================================ diff --git a/python/tests/core/test_gpu_pool.py b/python/tests/core/test_gpu_pool.py index 7b4f1c0f9..9a6de9650 100644 --- a/python/tests/core/test_gpu_pool.py +++ b/python/tests/core/test_gpu_pool.py @@ -12,13 +12,11 @@ import cocoindex as coco from cocoindex._internal import runner as _runner_mod from cocoindex._internal.runner import ( - GPUPool, GPURunner, configure_gpu_pool, current_gpu, current_gpus, current_gpu_fraction, - _detect_num_gpus, ) @@ -30,73 +28,6 @@ def _reset_gpu_pool() -> Iterator[None]: _runner_mod._default_gpu_pool = old -@pytest.mark.asyncio -async def test_acquire_returns_gpu_id() -> None: - pool = GPUPool(num_gpus=2) - gpu = await pool.acquire(1.0) - assert 0 <= gpu < 2 - await pool.release(gpu, 1.0) - - -@pytest.mark.asyncio -async def test_acquire_different_gpus() -> None: - pool = GPUPool(num_gpus=2) - gpu0 = await pool.acquire(1.0) - gpu1 = await pool.acquire(1.0) - assert gpu0 != gpu1 - await pool.release(gpu0, 1.0) - await pool.release(gpu1, 1.0) - - -@pytest.mark.asyncio -async def test_acquire_blocks_when_capacity_full() -> None: - pool = GPUPool(num_gpus=1) - gpu = await pool.acquire(1.0) - - task = asyncio.create_task(pool.acquire(1.0)) - await asyncio.sleep(0.02) - assert not task.done() - - await pool.release(gpu, 1.0) - result = await asyncio.wait_for(task, timeout=1.0) - assert result == 0 - await pool.release(result, 1.0) - - -@pytest.mark.asyncio -async def test_fractional_shares_same_gpu() -> None: - pool = GPUPool(num_gpus=1) - gpu0 = await pool.acquire(0.5) - gpu1 = await pool.acquire(0.5) - assert gpu0 == gpu1 - - task = asyncio.create_task(pool.acquire(0.5)) - await asyncio.sleep(0.02) - assert not task.done() - - await pool.release(gpu0, 0.5) - result = await asyncio.wait_for(task, timeout=1.0) - assert result == gpu0 - await pool.release(gpu1, 0.5) - await pool.release(result, 0.5) - - -@pytest.mark.asyncio -async def test_multi_gpu_all_parallel() -> None: - pool = GPUPool(num_gpus=3) - tasks = [asyncio.create_task(pool.acquire(1.0)) for _ in range(3)] - results = await asyncio.gather(*tasks) - assert len(set(results)) == 3 - for g in results: - await pool.release(g, 1.0) - - -@pytest.mark.asyncio -async def test_invalid_num_gpus_raises() -> None: - with pytest.raises(ValueError): - GPUPool(num_gpus=0) - - @pytest.mark.asyncio async def test_runner_sets_current_gpu_sync() -> None: configure_gpu_pool(2) @@ -339,131 +270,3 @@ def _gpu_work(x: int) -> int: assert 0 <= gpus[0] < 2 assert fraction == 0.5 - -def test_detect_num_gpus_explicit_env(monkeypatch: pytest.MonkeyPatch) -> None: - monkeypatch.setenv("COCOINDEX_NUM_GPUS", "4") - monkeypatch.delenv("CUDA_VISIBLE_DEVICES", raising=False) - assert _detect_num_gpus() == 4 - - -def test_detect_num_gpus_cuda_visible_devices(monkeypatch: pytest.MonkeyPatch) -> None: - monkeypatch.delenv("COCOINDEX_NUM_GPUS", raising=False) - monkeypatch.setenv("CUDA_VISIBLE_DEVICES", "0,2,3") - assert _detect_num_gpus() == 3 - - -def test_detect_num_gpus_cuda_visible_empty(monkeypatch: pytest.MonkeyPatch) -> None: - monkeypatch.delenv("COCOINDEX_NUM_GPUS", raising=False) - monkeypatch.setenv("CUDA_VISIBLE_DEVICES", "") - assert _detect_num_gpus() == 1 - - -def test_detect_num_gpus_explicit_env_zero(monkeypatch: pytest.MonkeyPatch) -> None: - monkeypatch.setenv("COCOINDEX_NUM_GPUS", "0") - monkeypatch.delenv("CUDA_VISIBLE_DEVICES", raising=False) - assert _detect_num_gpus() == 1 - - -def test_detect_num_gpus_explicit_env_overrides_cuda_visible( - monkeypatch: pytest.MonkeyPatch, -) -> None: - monkeypatch.setenv("COCOINDEX_NUM_GPUS", "2") - monkeypatch.setenv("CUDA_VISIBLE_DEVICES", "0,1,2,3") - assert _detect_num_gpus() == 2 - - -def test_detect_num_gpus_cuda_visible_single_device( - monkeypatch: pytest.MonkeyPatch, -) -> None: - monkeypatch.delenv("COCOINDEX_NUM_GPUS", raising=False) - monkeypatch.setenv("CUDA_VISIBLE_DEVICES", "0") - assert _detect_num_gpus() == 1 - - -def test_detect_num_gpus_cuda_visible_with_whitespace( - monkeypatch: pytest.MonkeyPatch, -) -> None: - monkeypatch.delenv("COCOINDEX_NUM_GPUS", raising=False) - monkeypatch.setenv("CUDA_VISIBLE_DEVICES", "0, 1 , 2") - assert _detect_num_gpus() == 3 - - -def test_detect_num_gpus_nvidia_smi_returns_count( - monkeypatch: pytest.MonkeyPatch, -) -> None: - monkeypatch.delenv("COCOINDEX_NUM_GPUS", raising=False) - monkeypatch.delenv("CUDA_VISIBLE_DEVICES", raising=False) - - def _mock_run(*args: Any, **kwargs: Any) -> Any: - class _Completed: - returncode = 0 - stdout = "8\n" - - return _Completed() - - monkeypatch.setattr(subprocess, "run", _mock_run) - assert _detect_num_gpus() == 8 - - -def test_detect_num_gpus_nvidia_smi_empty_output( - monkeypatch: pytest.MonkeyPatch, -) -> None: - monkeypatch.delenv("COCOINDEX_NUM_GPUS", raising=False) - monkeypatch.delenv("CUDA_VISIBLE_DEVICES", raising=False) - - def _mock_run(*args: Any, **kwargs: Any) -> Any: - class _Completed: - returncode = 0 - stdout = "" - - return _Completed() - - monkeypatch.setattr(subprocess, "run", _mock_run) - assert _detect_num_gpus() == 1 - - -def test_detect_num_gpus_nvidia_smi_nonzero_exit( - monkeypatch: pytest.MonkeyPatch, -) -> None: - monkeypatch.delenv("COCOINDEX_NUM_GPUS", raising=False) - monkeypatch.delenv("CUDA_VISIBLE_DEVICES", raising=False) - - def _mock_run(*args: Any, **kwargs: Any) -> Any: - class _Completed: - returncode = 1 - stdout = "" - - return _Completed() - - monkeypatch.setattr(subprocess, "run", _mock_run) - assert _detect_num_gpus() == 1 - - -def test_detect_num_gpus_nvidia_smi_not_found( - monkeypatch: pytest.MonkeyPatch, -) -> None: - monkeypatch.delenv("COCOINDEX_NUM_GPUS", raising=False) - monkeypatch.delenv("CUDA_VISIBLE_DEVICES", raising=False) - - def _mock_run(*args: Any, **kwargs: Any) -> Any: - raise FileNotFoundError("nvidia-smi not found") - - monkeypatch.setattr(subprocess, "run", _mock_run) - assert _detect_num_gpus() == 1 - - -def test_detect_num_gpus_all_missing_fallback( - monkeypatch: pytest.MonkeyPatch, -) -> None: - monkeypatch.delenv("COCOINDEX_NUM_GPUS", raising=False) - monkeypatch.delenv("CUDA_VISIBLE_DEVICES", raising=False) - - def _mock_run(*args: Any, **kwargs: Any) -> Any: - class _Completed: - returncode = 1 - stdout = "" - - return _Completed() - - monkeypatch.setattr(subprocess, "run", _mock_run) - assert _detect_num_gpus() == 1 diff --git a/rust/py/src/gpu_pool.rs b/rust/py/src/gpu_pool.rs new file mode 100644 index 000000000..1f9a26257 --- /dev/null +++ b/rust/py/src/gpu_pool.rs @@ -0,0 +1,82 @@ +//! Python bindings for batching infrastructure. +//! +//! Exposes BatchQueue and Batcher to Python for implementing batched function execution. +//! +//! Design: Multiple batchers can share the same queue (e.g., for GPU serialization), +//! and each batcher has its own runner function. When a batcher creates a batch, +//! that batch carries the batcher's runner function. + +use crate::prelude::*; +use cocoindex_utils::gpu_pool::{gpu_capacity::GPUCapacity, GPUPool}; +use pyo3::exceptions::PyValueError; +use pyo3_async_runtimes::tokio::future_into_py; +use std::num::NonZeroUsize; + +#[pyclass(name = "GPUPool")] +#[derive(Clone)] +pub struct PyGPUPool { + inner: Arc, +} + +#[pymethods] +impl PyGPUPool { + #[new] + pub fn new(num_gpus: usize) -> PyResult { + NonZeroUsize::new(num_gpus) + .ok_or_else(|| PyValueError::new_err("num_gpus must be > 0")) + .map(GPUPool::new) + .map(Arc::new) + .map(|gpu_pool| Self { inner: gpu_pool }) + } + + #[staticmethod] + pub fn default() -> Self { + Self { + inner: Arc::new(GPUPool::default()), + } + } + + #[getter] + pub fn num_gpus(&self) -> usize { + self.inner.num_gpus() + } + + pub fn acquire<'py>(&self, py: Python<'py>, fraction: f32) -> PyResult> { + let fraction = + GPUCapacity::try_from(fraction).map_err(|e| PyValueError::new_err(e.to_string()))?; + let gpu_pool = self.inner.clone(); + future_into_py(py, async move { + gpu_pool + .acquire(fraction) + .await + .map_err(|e| PyValueError::new_err(e.to_string())) + }) + } + + pub fn acquire_full<'py>( + &self, + py: Python<'py>, + gpu_count: usize, + ) -> PyResult> { + if gpu_count <= 0 { + return Err(PyValueError::new_err(format!( + "gpu_count must be > 0, got {gpu_count}" + ))); + }; + let gpu_pool = self.inner.clone(); + future_into_py(py, async move { + gpu_pool + .acquire_full(NonZeroUsize::new(gpu_count).unwrap()) + .await + .map_err(|e| PyValueError::new_err(e.to_string())) + }) + } + + pub fn release<'py>(&self, gpu_id: usize, fraction: f32) -> PyResult<()> { + let fraction = + GPUCapacity::try_from(fraction).map_err(|e| PyValueError::new_err(e.to_string()))?; + self.inner + .release(gpu_id, fraction) + .map_err(|e| PyValueError::new_err(e.to_string())) + } +} diff --git a/rust/py/src/lib.rs b/rust/py/src/lib.rs index af7ca85e8..127528427 100644 --- a/rust/py/src/lib.rs +++ b/rust/py/src/lib.rs @@ -7,6 +7,7 @@ mod deadline; mod environment; mod fingerprint; mod function; +mod gpu_pool; mod inspect; pub mod live_component; mod logic_registry; @@ -169,5 +170,8 @@ fn core_module(m: &pyo3::Bound<'_, pyo3::types::PyModule>) -> pyo3::PyResult<()> // Rate limiting m.add_class::()?; + // GPU Pool + m.add_class::()?; + Ok(()) } diff --git a/rust/utils/Cargo.toml b/rust/utils/Cargo.toml index 80215d98a..ddd1fcfb9 100644 --- a/rust/utils/Cargo.toml +++ b/rust/utils/Cargo.toml @@ -45,6 +45,9 @@ yaml-rust2 = { version = "0.10.4", optional = true } serde_with = { workspace = true, features = ["base64"] } storekey = { workspace = true } +[dev-dependencies] +temp-env = "0.3" + [features] default = [] reqwest = ["dep:reqwest"] diff --git a/rust/utils/src/gpu_pool.rs b/rust/utils/src/gpu_pool.rs new file mode 100644 index 000000000..17b0f4dac --- /dev/null +++ b/rust/utils/src/gpu_pool.rs @@ -0,0 +1,1325 @@ +use crate::error::Result; +use crate::{client_bail, internal_error}; +use container::SortedVec; +use gpu_capacity::GPUCapacity; +use std::collections::VecDeque; +use std::num::NonZeroUsize; +use std::sync::Mutex; +use tokio::sync::oneshot; + +/// Tracks fractional GPU capacity across multiple GPUs. +/// +/// Each GPU starts with capacity 1.0. ``acquire(fraction)`` blocks until a +/// GPU with enough remaining capacity is available, then returns its id. +/// ``release(gpu_id, fraction)`` restores capacity and wakes waiters. +/// +/// The default pool size is auto-detected from ``COCOINDEX_NUM_GPUS``, +/// ``CUDA_VISIBLE_DEVICES``, or ``nvidia-smi`` (falling back to 1). +/// Call ``configure_gpu_pool(N)`` to override programmatically. +pub struct GPUPool { + num_gpus: usize, + state: Mutex, +} + +struct PoolState { + capacities: SortedVec, + acquisition_queue: VecDeque, +} + +struct Acquisition { + demand: GPUCapacity, + notifier: oneshot::Sender, +} + +impl GPUPool { + pub fn new(num_gpus: NonZeroUsize) -> Self { + let num_gpus = num_gpus.get(); + let capacities = std::iter::repeat_n(GPUCapacity::MAX, num_gpus).collect(); + let state = PoolState { + capacities, + acquisition_queue: VecDeque::new(), + }; + GPUPool { + num_gpus, + state: Mutex::new(state), + } + } + + pub fn num_gpus(&self) -> usize { + self.num_gpus + } + + /// acquire a fraction of a GPU, if not available the acquisition is thrown into the queue. + /// + /// The function would first attempt to find a GPU with just enough capacity for the demanded + /// fraction, and the GPU cannot be the top n, in the following rules: + /// + /// 1. The `n` is the size of the queue. + /// 2. The top n are defined by capacities. + /// 3. The top n are reserved for acquisitions already in the queue. + /// + /// Reservation is for an abstract concept, + /// e.g. "the GPU with the most available capacity" is reserved for the head acquisition in the queue. + /// + /// The function would try to host the acquisition using the remaining GPUs, + /// or it will send it to the queue which will reserve a GPU now or later. + /// + pub async fn acquire(&self, fraction: GPUCapacity) -> Result { + if fraction == GPUCapacity::ZERO { + client_bail!("Acquired fraction must be between 0.0 and 1.0, got 0"); + } + let receiver = { + let mut pool = self.state.lock().expect("lock poisoned"); + if pool.acquisition_queue.len() < self.num_gpus + && let Some(gpu_id) = pool + .capacities + .find_excluding_top_n(&fraction, pool.acquisition_queue.len()) + { + // excluding top_n, because the acquisitions in the queue have already reserved the top n GPUs. + let updated_capacity = pool.capacities[gpu_id] - fraction; + pool.capacities.update(gpu_id, updated_capacity); + return Ok(gpu_id); + } + Self::send_acquisition_to_queue(&mut pool.acquisition_queue, fraction) + }; + receiver + .await + .map_err(|err| internal_error!("GPUPool dropped while waiting: {err}")) + } + + fn send_acquisition_to_queue( + acquisition_queue: &mut VecDeque, + demand: GPUCapacity, + ) -> oneshot::Receiver { + let (notifier, receiver) = oneshot::channel(); + acquisition_queue.push_back(Acquisition { demand, notifier }); + receiver + } + + /// Acquires a given integer number of fully available GPUs (capacity == 1.0) from the GPU pool. + /// + /// # Error: + /// * When the given gpu_count is larger than the total gpus, it returns an error. + /// + /// # Warning + /// * When unable to acquire all GPUs, the system will be acquired the ones that can be acquired first. + /// For instance, if user attempts to acquire 5 GPUs, + /// the function will partially acquire 4 and wait for the last GPU. + pub async fn acquire_full(&self, gpu_count: NonZeroUsize) -> Result> { + let gpu_count = gpu_count.get(); + if gpu_count > self.num_gpus() { + client_bail!( + "Attempted to acquire {} GPUs but only has {}.", + gpu_count, + self.num_gpus + ); + } + let (mut acquired_gpus, receivers) = { + let mut pool = self.state.lock().expect("lock poisoned"); + let mut acquired_gpus = Vec::with_capacity(gpu_count); + if pool.acquisition_queue.len() < self.num_gpus { + let taken_gpus = pool.capacities.find_many_excluding_top_n( + &GPUCapacity::MAX, + gpu_count, + pool.acquisition_queue.len(), + ); + for gpu_id in taken_gpus { + acquired_gpus.push(gpu_id); + pool.capacities.update(gpu_id, GPUCapacity::ZERO); + } + } + if acquired_gpus.len() == gpu_count { + return Ok(acquired_gpus); + } + let gpus_to_be_acquired = gpu_count - acquired_gpus.len(); + let receivers = std::iter::repeat_with(|| { + Self::send_acquisition_to_queue(&mut pool.acquisition_queue, GPUCapacity::MAX) + }) + .take(gpus_to_be_acquired) + .collect::>(); + (acquired_gpus, receivers) + }; + match futures::future::try_join_all(receivers).await { + Ok(gpu_ids) => { + acquired_gpus.extend(gpu_ids); + Ok(acquired_gpus) + } + Err(err) => { + for gpu_id in acquired_gpus { + let _ = self.release(gpu_id, GPUCapacity::MAX); + } + client_bail!("GPUPool reservation cancelled while waiting: {err}") + } + } + } + + /// release adds back capacities to GPUs, and processes pending acquisitions afterward. + /// + /// # Example + /// Initially: + /// ```text + /// GPUs: G1(capacity=0), G2(capacity=0), G3(capacity=0) + /// Queue: T1(req=0.7, reserved=[G1]) T2(req=0.5, reserved=[G2]) + /// ``` + /// After releasing 0.5 capacity to G1: + /// ```text + /// GPUs: G1(capacity=0.5), G2(capacity=0), G3(capacity=0) + /// Queue: T1(req=0.7, reserved=[G1]), T2(req=0.5, reserved=[G2]) + /// ``` + /// After releasing 0.6 capacity to G2: + /// ```text + /// GPUs: G1(capacity=0.5), G2(capacity=0.6), G3(capacity=0) + /// Queue: T1(req=0.7, reserved=[G2]), T2(req=0.5, reserved=[G1]) + /// ``` + /// After releasing 0.1 capacity to G2, T1 will be hosted by G2, then get popped: + /// ```text + /// GPUs: G1(capacity=0.5), G2(capacity=0), G3(capacity=0) + /// Queue: T2(req=0.5, reserved=[G1]) + /// ``` + pub fn release(&self, gpu_id: usize, fraction: GPUCapacity) -> Result<()> { + if gpu_id >= self.num_gpus() { + client_bail!("Releasing to a gpu_id that does not exist: {gpu_id}",); + } + if fraction == GPUCapacity::ZERO { + client_bail!("Cannot release a zero fraction"); + } + let mut state = self.state.lock().expect("lock poisoned"); + let updated_capacity = state.capacities[gpu_id].checked_add(&fraction)?; + state.capacities.update(gpu_id, updated_capacity); + Self::process_acquisition_queue(&mut state); + Ok(()) + } + + /// processes pending acquisition queue following the rules: + /// + /// 1. The first task always reserves the GPU with the most availability at this moment + /// 2. Processing does not change the order of pending acquisitions + /// + fn process_acquisition_queue(pool: &mut PoolState) { + let length = pool.capacities.len(); + let mut pending_acquisitions = Vec::with_capacity(length); + while pending_acquisitions.len() < length + && let Some(acquisition) = pool.acquisition_queue.pop_front() + { + if let Some(gpu_id) = pool + .capacities + .find_excluding_top_n(&acquisition.demand, pending_acquisitions.len()) + { + if acquisition.notifier.send(gpu_id).is_ok() { + let updated_capacity = pool.capacities[gpu_id] - acquisition.demand; + pool.capacities.update(gpu_id, updated_capacity); + } + } else { + pending_acquisitions.push(acquisition); + } + } + while let Some(acquisition) = pending_acquisitions.pop() { + pool.acquisition_queue.push_front(acquisition); + } + } + + /// detect the number of GPUs available for the default pool. + /// + /// # Returns: + /// * number of GPUs + /// + /// # Errors: + /// * failed to find environment variables + /// * failed to read environment variable values + /// * failed to parse an environment variable value to a number + /// * failed to find given commands + /// + /// # Detection order: + /// + /// 1. ``COCOINDEX_NUM_GPUS`` environment variable (explicit override). + /// 2. ``CUDA_VISIBLE_DEVICES`` environment variable (count of entries). + /// 3. ``nvidia-smi`` command output (if available). + /// 4. Default to ``1``. + /// + fn detect_num_gpus() -> Result { + if let Some(env_num) = std::env::var("COCOINDEX_NUM_GPUS") + .ok() + .and_then(|s| s.parse::().ok()) + { + return Ok(std::cmp::max(1, env_num)); + } + if let Ok(cuda_visible) = std::env::var("CUDA_VISIBLE_DEVICES") { + let count = cuda_visible + .split(',') + .map(str::trim) + .filter(|s| !s.is_empty()) + .count(); + return Ok(std::cmp::max(1, count)); + } + #[cfg(not(test))] + let output = std::process::Command::new("nvidia-smi") + .arg("--query-gpu=count") + .arg("--format=csv,noheader") + .output()?; + #[cfg(test)] + let output = { + if std::env::var("MOCK_NVIDIA_SMI_NOT_FOUND").is_ok() { + return Err(crate::error::Error::internal(std::io::Error::new( + std::io::ErrorKind::NotFound, + "nvidia-smi not found", + ))); + } + let mock_gpu_count = std::env::var("MOCK_NVIDIA_SMI_STDOUT").unwrap_or_default(); + let mock_exit_code = std::env::var("MOCK_NVIDIA_SMI_EXIT_CODE") + .ok() + .and_then(|s| s.parse::().ok()) + .unwrap_or(0); + std::process::Command::new("sh") + .arg("-c") + .arg(format!("echo \"{mock_gpu_count}\"; exit {mock_exit_code}")) + .output() + }?; + + if !output.status.success() { + return Ok(1); + } + let count = String::from_utf8_lossy(&output.stdout) + .lines() + .next() + .unwrap_or_default() + .trim() + .parse::()?; + Ok(std::cmp::max(1, count)) + } +} + +impl Default for GPUPool { + fn default() -> Self { + Self::new(NonZeroUsize::new(Self::detect_num_gpus().unwrap_or(1)).unwrap()) + } +} + +pub mod gpu_capacity { + use crate::client_bail; + use crate::error::{Error, Result}; + use std::ops::{Add, AddAssign, Sub, SubAssign}; + + #[derive(Debug, Copy, Clone, PartialEq, Eq, PartialOrd, Ord)] + pub struct GPUCapacity(u32); + + impl GPUCapacity { + const SCALE: f32 = 1_000_000.0; + pub const ZERO: Self = Self(0); + pub const MAX: Self = Self(Self::SCALE as u32); + + #[cfg(test)] + pub(crate) fn unchecked(value: f32) -> Self { + GPUCapacity::try_from(value).expect("Unchecked value initialization should not fail") + } + + pub fn checked_add(&self, other: &Self) -> Result { + if self.0 + other.0 > GPUCapacity::MAX.0 { + client_bail!( + "The sum of {self} and {other} is greater than the max value {}", + Self::MAX + ); + } else { + Ok(GPUCapacity(self.0 + other.0)) + } + } + } + + impl TryFrom for GPUCapacity { + type Error = Error; + + fn try_from(value: f32) -> Result { + if !(0.0..=1.0).contains(&value) { + client_bail!("Fraction must be between 0.0 and 1.0, got {value}",); + } + Ok(Self((value * Self::SCALE) as u32)) + } + } + + impl Add for GPUCapacity { + type Output = Self; + + fn add(self, other: Self) -> Self { + Self(self.0 + other.0) + } + } + + impl AddAssign for GPUCapacity { + fn add_assign(&mut self, other: Self) { + self.0 += other.0; + } + } + + impl Sub for GPUCapacity { + type Output = Self; + + fn sub(self, other: Self) -> Self { + Self(self.0 - other.0) + } + } + + impl SubAssign for GPUCapacity { + fn sub_assign(&mut self, other: Self) { + self.0 -= other.0; + } + } + + impl std::fmt::Display for GPUCapacity { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + write!(f, "{}", self.0 as f32 / Self::SCALE) + } + } +} + +mod container { + use std::collections::BTreeSet; + + #[derive(Debug, Default, Clone)] + pub struct SortedVec { + values: Vec, + sorted: BTreeSet<(T, usize)>, + } + + impl SortedVec { + pub fn len(&self) -> usize { + self.values.len() + } + } + + impl FromIterator for SortedVec { + fn from_iter>(iter: I) -> Self { + let values = iter.into_iter().collect::>(); + let sorted = BTreeSet::from_iter(values.iter().cloned().zip(0..)); + Self { values, sorted } + } + } + + impl SortedVec { + /// find_excluding_top_n should return the first index which points to minimal value that + /// is greater or equal to `target`. + /// + /// When the target value is greater than all values, return None. + pub fn find_excluding_top_n(&self, target: &T, top_n: usize) -> Option { + self.find_excluding_top_n_iter(target, top_n).next() + } + + /// find_many_excluding_top_n should return the `count` number of indices + /// which points to minimal value that is greater or equal to `target`. + /// + /// When the target value is greater than all values, return empty vec. + pub fn find_many_excluding_top_n( + &self, + target: &T, + count: usize, + top_n: usize, + ) -> Vec { + self.find_excluding_top_n_iter(target, top_n) + .take(count) + .collect() + } + + fn find_excluding_top_n_iter( + &self, + target: &T, + top_n: usize, + ) -> impl Iterator { + let upper_bound = if top_n >= self.sorted.len() { + None + } else if top_n > self.sorted.len() / 2 { + self.sorted.iter().nth(self.sorted.len() - 1 - top_n) + } else { + self.sorted.iter().rev().nth(top_n) + }; + upper_bound + .into_iter() + .filter(move |(upper_bound_value, _)| target <= upper_bound_value) + .flat_map(|upper_bound| self.sorted.range(&(target.clone(), 0)..=upper_bound)) + .map(|(_, index)| *index) + } + + pub fn update(&mut self, index: usize, value: T) { + let Some(old_value) = self.values.get_mut(index) else { + return; + }; + self.sorted.remove(&(old_value.clone(), index)); + *old_value = value.clone(); + self.sorted.insert((value, index)); + } + } + + impl std::ops::Index for SortedVec { + type Output = T; + + fn index(&self, index: usize) -> &Self::Output { + &self.values[index] + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::gpu_pool::container::SortedVec; + use itertools::Itertools; + use rand::Rng; + use std::sync::Arc; + + #[tokio::test] + async fn test_acquire_returns_gpu_id() -> Result<()> { + let pool = GPUPool::new(NonZeroUsize::new(2).unwrap()); + let gpu = pool.acquire(GPUCapacity::MAX).await?; + assert!(gpu < 2); + pool.release(gpu, GPUCapacity::MAX)?; + Ok(()) + } + + #[tokio::test] + async fn test_acquire_different_gpus() -> Result<()> { + let pool = GPUPool::new(NonZeroUsize::new(2).unwrap()); + let gpu0 = pool.acquire(GPUCapacity::MAX).await?; + let gpu1 = pool.acquire(GPUCapacity::MAX).await?; + assert_ne!(gpu0, gpu1); + pool.release(gpu0, GPUCapacity::MAX)?; + pool.release(gpu1, GPUCapacity::MAX)?; + Ok(()) + } + + #[tokio::test] + async fn test_acquire_blocks_when_capacity_full() -> Result<()> { + let pool = Arc::new(GPUPool::new(NonZeroUsize::new(1).unwrap())); + let gpu = pool.acquire(GPUCapacity::MAX).await?; + + let cloned_pool = pool.clone(); + let task = tokio::spawn(async move { cloned_pool.acquire(GPUCapacity::MAX).await }); + tokio::time::sleep(std::time::Duration::from_secs_f32(0.02)).await; + assert!(!task.is_finished()); + + pool.release(gpu, GPUCapacity::MAX)?; + let result = tokio::time::timeout(std::time::Duration::from_secs(1), task) + .await + .expect("task finished")?; + assert!(matches!(result, Ok(0))); + pool.release(result.unwrap(), GPUCapacity::MAX)?; + Ok(()) + } + + #[tokio::test] + async fn test_fractional_shares_same_gpu() -> Result<()> { + let pool = Arc::new(GPUPool::new(NonZeroUsize::new(1).unwrap())); + let half_fraction = GPUCapacity::try_from(0.5).expect("0.5 is a valid fraction"); + let gpu0 = pool.acquire(half_fraction).await?; + let gpu1 = pool.acquire(half_fraction).await?; + assert_eq!(gpu0, gpu1); + + let cloned_pool = pool.clone(); + let task = tokio::spawn(async move { cloned_pool.acquire(half_fraction).await }); + tokio::time::sleep(std::time::Duration::from_secs_f32(0.02)).await; + assert!(!task.is_finished()); + + pool.release(gpu0, half_fraction)?; + let result = tokio::time::timeout(std::time::Duration::from_secs(1), task) + .await + .expect("task finished")?; + assert!(matches!(result, Ok(0))); + pool.release(gpu1, half_fraction)?; + pool.release(result.unwrap(), half_fraction)?; + Ok(()) + } + + #[tokio::test] + async fn test_multi_gpu_all_parallel() -> Result<()> { + let pool = Arc::new(GPUPool::new(NonZeroUsize::new(3).unwrap())); + let mut tasks = Vec::with_capacity(3); + for _ in 0..3 { + let pool = pool.clone(); + tasks.push(tokio::spawn( + async move { pool.acquire(GPUCapacity::MAX).await }, + )); + } + let results = futures::future::try_join_all(tasks).await?; + let gpus = results.into_iter().collect::, _>>()?; + assert_eq!(gpus.len(), 3); + for g in gpus { + pool.release(g, GPUCapacity::MAX)?; + } + Ok(()) + } + + #[tokio::test] + async fn test_acquire_fractions_equals_to_zero() { + let pool = GPUPool::new(NonZeroUsize::new(1).unwrap()); + let result = pool.acquire(GPUCapacity::ZERO).await; + assert!(result.is_err()); + assert!( + result + .unwrap_err() + .to_string() + .contains("Acquired fraction must be between 0.0 and 1.0, got 0") + ); + } + + #[tokio::test] + async fn test_acquire_fractions_not_enough_with_release_not_enough() -> Result<()> { + let pool = Arc::new(GPUPool::new(NonZeroUsize::new(3).unwrap())); + let occupied_gpu_1 = pool.acquire(GPUCapacity::unchecked(0.6)).await?; + assert_eq!(occupied_gpu_1, 0); + let occupied_gpu_2 = pool.acquire(GPUCapacity::unchecked(0.6)).await?; + assert_eq!(occupied_gpu_2, 1); + let cloned_pool = pool.clone(); + let not_enough_task = tokio::spawn(async move { + cloned_pool + .acquire_full(NonZeroUsize::new(3).unwrap()) + .await + }); + tokio::time::sleep(std::time::Duration::from_secs_f32(0.02)).await; + assert!(!not_enough_task.is_finished()); + pool.release(occupied_gpu_2, GPUCapacity::unchecked(0.2))?; + tokio::time::sleep(std::time::Duration::from_secs_f32(0.02)).await; + assert!(!not_enough_task.is_finished()); + pool.release(occupied_gpu_2, GPUCapacity::unchecked(0.4))?; + pool.release(occupied_gpu_1, GPUCapacity::unchecked(0.6))?; + tokio::time::sleep(std::time::Duration::from_secs_f32(0.02)).await; + assert!(not_enough_task.is_finished()); + let gpus = tokio::time::timeout(std::time::Duration::from_secs(1), not_enough_task) + .await + .expect("task finished") + .expect("no timeout")?; + assert_eq!(gpus.len(), 3); + for gpu in gpus { + pool.release(gpu, GPUCapacity::MAX)?; + } + Ok(()) + } + + #[tokio::test] + async fn test_acquire_full_gpus_enough() -> Result<()> { + let pool = GPUPool::new(NonZeroUsize::new(2).unwrap()); + let gpus = pool + .acquire_full(NonZeroUsize::new(2).expect("2 is not zero")) + .await?; + assert_eq!(gpus, vec![0, 1]); + for g in gpus { + pool.release(g, GPUCapacity::MAX)?; + } + Ok(()) + } + + #[tokio::test] + async fn test_acquire_full_gpus_not_enough() -> Result<()> { + let pool = Arc::new(GPUPool::new(NonZeroUsize::new(3).unwrap())); + let partially_used_gpu = pool.acquire(GPUCapacity::unchecked(0.6)).await?; + assert_eq!(partially_used_gpu, 0); + let cloned_pool = pool.clone(); + let task = tokio::spawn(async move { + cloned_pool + .acquire_full(NonZeroUsize::new(3).expect("3 is not zero")) + .await + }); + tokio::time::sleep(std::time::Duration::from_secs_f32(0.02)).await; + assert!(!task.is_finished()); + pool.release(partially_used_gpu, GPUCapacity::unchecked(0.6))?; + let result = tokio::time::timeout(std::time::Duration::from_secs(1), task) + .await + .expect("task finished") + .expect("no timeout")?; + assert_eq!(&result, &[1, 2, 0]); + for gpu in result { + pool.release(gpu, GPUCapacity::MAX)?; + } + Ok(()) + } + + #[tokio::test] + async fn test_acquire_full_gpus_with_partial_acquiring() -> Result<()> { + let pool = Arc::new(GPUPool::new(NonZeroUsize::new(3).unwrap())); + let partially_used_gpu = pool.acquire(GPUCapacity::unchecked(0.6)).await?; + assert_eq!(partially_used_gpu, 0); + let cloned_pool = pool.clone(); + let task = tokio::spawn(async move { + cloned_pool + .acquire_full(NonZeroUsize::new(3).expect("3 is not zero")) + .await + }); + let cloned_pool = pool.clone(); + let second_acquired_gpu = + tokio::spawn(async move { cloned_pool.acquire(GPUCapacity::unchecked(0.2)).await }); + tokio::time::sleep(std::time::Duration::from_secs_f32(0.02)).await; + assert!(!task.is_finished()); + assert!(!second_acquired_gpu.is_finished()); + pool.release(partially_used_gpu, GPUCapacity::unchecked(0.6))?; + let result = tokio::time::timeout(std::time::Duration::from_secs(1), task) + .await + .expect("task finished") + .expect("no timeout")?; + // initial 0.6 occupied index 0, then GPU 1 and 2 are reserved, until 0 is added. + assert_eq!(&result, &[1, 2, 0]); + for gpu in result { + pool.release(gpu, GPUCapacity::MAX)?; + } + Ok(()) + } + + #[tokio::test] + async fn test_acquire_full_cancelled_releases_acquired_gpus() -> Result<()> { + let pool = Arc::new(GPUPool::new(NonZeroUsize::new(3).unwrap())); + let partially_used_gpu = pool.acquire(GPUCapacity::unchecked(0.6)).await?; + assert_eq!(partially_used_gpu, 0); + + let cloned_pool = pool.clone(); + let task = tokio::spawn(async move { + cloned_pool + .acquire_full(NonZeroUsize::new(3).expect("3 is not zero")) + .await + }); + tokio::time::sleep(std::time::Duration::from_secs_f32(0.02)).await; + assert!(!task.is_finished()); + + // Verify GPU 1 and 2 were acquired and 1 request is waiting in queue. + // then cancel + { + let mut state = pool.state.lock().expect("lock poisoned"); + assert_eq!(state.capacities[1], GPUCapacity::ZERO); + assert_eq!(state.capacities[2], GPUCapacity::ZERO); + assert_eq!(state.acquisition_queue.len(), 1); + // Cancel the acquisition by dropping the receiver. + state.acquisition_queue.clear(); + } + + let result = tokio::time::timeout(std::time::Duration::from_secs(1), task) + .await + .expect("task finished") + .expect("task did not panic"); + assert!(result.is_err()); + assert!( + result + .unwrap_err() + .to_string() + .contains("GPUPool reservation cancelled while waiting") + ); + + // Verify all partially acquired GPUs (1 and 2) are released back to full capacity. + { + let state = pool.state.lock().expect("lock poisoned"); + assert_eq!(state.capacities[1], GPUCapacity::MAX); + assert_eq!(state.capacities[2], GPUCapacity::MAX); + assert_eq!(state.capacities[0], GPUCapacity::unchecked(0.4)); + } + + pool.release(partially_used_gpu, GPUCapacity::unchecked(0.6))?; + Ok(()) + } + + #[tokio::test] + async fn test_acquire_more_gpus_than_allowed() { + let pool = GPUPool::new(NonZeroUsize::new(2).unwrap()); + let result = pool + .acquire_full(NonZeroUsize::new(3).expect("3 is not zero")) + .await; + assert!(result.is_err()); + assert!( + result + .unwrap_err() + .to_string() + .contains("Attempted to acquire 3 GPUs but only has 2.") + ); + } + + #[tokio::test] + async fn test_reserve_gpus_then_release() -> Result<()> { + let pool = Arc::new(GPUPool::new(NonZeroUsize::new(2).unwrap())); + let gpu_0 = pool.acquire(GPUCapacity::unchecked(0.5)).await?; + assert_eq!(gpu_0, 0); + let gpu_1 = pool.acquire(GPUCapacity::unchecked(0.6)).await?; + assert_eq!(gpu_1, 1); + let cloned_pool = pool.clone(); + let reserving_task_1 = + tokio::spawn(async move { cloned_pool.acquire(GPUCapacity::unchecked(0.6)).await }); + let cloned_pool = pool.clone(); + let reserving_task_2 = + tokio::spawn(async move { cloned_pool.acquire(GPUCapacity::unchecked(0.7)).await }); + tokio::time::sleep(std::time::Duration::from_secs_f32(0.02)).await; + assert!(!reserving_task_1.is_finished()); + assert!(!reserving_task_2.is_finished()); + + pool.release(gpu_0, GPUCapacity::unchecked(0.1))?; + let reserving_task_1_acquired_gpu = + tokio::time::timeout(std::time::Duration::from_secs(1), reserving_task_1) + .await + .expect("task finished") + .expect("no timeout")?; + assert_eq!(reserving_task_1_acquired_gpu, gpu_0); + assert!(!reserving_task_2.is_finished()); + + pool.release(gpu_1, GPUCapacity::unchecked(0.3))?; + let reserving_task_2_acquired_gpu = + tokio::time::timeout(std::time::Duration::from_secs(1), reserving_task_2) + .await + .expect("task finished") + .expect("no timeout")?; + assert_eq!(reserving_task_2_acquired_gpu, gpu_1); + + pool.release(gpu_0, GPUCapacity::MAX)?; + pool.release(gpu_1, GPUCapacity::MAX)?; + Ok(()) + } + + #[tokio::test] + async fn test_reserve_gpus_without_affecting_unreserved() -> Result<()> { + let pool = Arc::new(GPUPool::new(NonZeroUsize::new(2).unwrap())); + let gpu_0 = pool.acquire(GPUCapacity::unchecked(0.5)).await?; + assert_eq!(gpu_0, 0); + let gpu_1 = pool.acquire(GPUCapacity::unchecked(0.6)).await?; + assert_eq!(gpu_1, 1); + let cloned_pool = pool.clone(); + let reserving_task = + tokio::spawn(async move { cloned_pool.acquire(GPUCapacity::unchecked(0.6)).await }); + let cloned_pool = pool.clone(); + let task_not_blocked = + tokio::spawn(async move { cloned_pool.acquire(GPUCapacity::unchecked(0.2)).await }); + tokio::time::sleep(std::time::Duration::from_secs_f32(0.02)).await; + assert!(!reserving_task.is_finished()); + assert!(task_not_blocked.is_finished()); + pool.release(gpu_0, GPUCapacity::unchecked(0.1))?; + + pool.release(gpu_1, GPUCapacity::unchecked(0.8))?; + let reserving_task_acquired_gpu = + tokio::time::timeout(std::time::Duration::from_secs(1), reserving_task) + .await + .expect("task finished") + .expect("no timeout")?; + assert_eq!(reserving_task_acquired_gpu, gpu_0); + + pool.release(gpu_0, GPUCapacity::MAX)?; + Ok(()) + } + + #[tokio::test] + async fn test_reserve_the_same_gpu_in_a_queue() -> Result<()> { + let pool = Arc::new(GPUPool::new(NonZeroUsize::new(1).unwrap())); + let gpu_0 = pool.acquire(GPUCapacity::unchecked(0.5)).await?; + assert_eq!(gpu_0, 0); + let cloned_pool = pool.clone(); + let reserving_task_1 = + tokio::spawn(async move { cloned_pool.acquire(GPUCapacity::unchecked(0.6)).await }); + let cloned_pool = pool.clone(); + let reserving_task_2 = + tokio::spawn(async move { cloned_pool.acquire(GPUCapacity::unchecked(0.7)).await }); + tokio::time::sleep(std::time::Duration::from_secs_f32(0.02)).await; + assert!(!reserving_task_1.is_finished()); + assert!(!reserving_task_2.is_finished()); + + pool.release(gpu_0, GPUCapacity::unchecked(0.1))?; + let reserving_task_1_acquired_gpu = + tokio::time::timeout(std::time::Duration::from_secs(1), reserving_task_1) + .await + .expect("task finished") + .expect("no timeout")?; + assert_eq!(reserving_task_1_acquired_gpu, gpu_0); + assert!(!reserving_task_2.is_finished()); + + pool.release(gpu_0, GPUCapacity::unchecked(0.7))?; + let reserving_task_2_acquired_gpu = + tokio::time::timeout(std::time::Duration::from_secs(1), reserving_task_2) + .await + .expect("task finished") + .expect("no timeout")?; + assert_eq!(reserving_task_2_acquired_gpu, gpu_0); + + pool.release(gpu_0, GPUCapacity::MAX)?; + Ok(()) + } + + #[tokio::test] + async fn test_reserve_front_queue_not_block_later_items() -> Result<()> { + let pool = Arc::new(GPUPool::new(NonZeroUsize::new(2).unwrap())); + let gpu_1 = pool.acquire(GPUCapacity::unchecked(0.5)).await?; + let gpu_2 = pool.acquire(GPUCapacity::unchecked(0.8)).await?; + let cloned_pool = pool.clone(); + let task_1 = + tokio::spawn(async move { cloned_pool.acquire(GPUCapacity::unchecked(0.6)).await }); + tokio::time::sleep(std::time::Duration::from_secs_f32(0.02)).await; + assert!(!task_1.is_finished()); + let cloned_pool = pool.clone(); + let task_2 = + tokio::spawn(async move { cloned_pool.acquire(GPUCapacity::unchecked(0.4)).await }); + tokio::time::sleep(std::time::Duration::from_secs_f32(0.02)).await; + assert!(!task_2.is_finished()); + + pool.release(gpu_2, GPUCapacity::unchecked(0.2))?; + let task_2_acquired_gpu = tokio::time::timeout(std::time::Duration::from_secs(1), task_2) + .await + .expect("task finished") + .expect("no timeout")?; + assert_eq!(task_2_acquired_gpu, gpu_2); + assert!(!task_1.is_finished()); + + pool.release(gpu_1, GPUCapacity::unchecked(0.5))?; + let task_1_acquired_gpu = tokio::time::timeout(std::time::Duration::from_secs(1), task_1) + .await + .expect("task finished") + .expect("no timeout")?; + assert_eq!(task_1_acquired_gpu, gpu_1); + + pool.release(gpu_2, GPUCapacity::MAX) + } + + #[tokio::test] + async fn test_reserve_queue_assigned_task_not_blocking() -> Result<()> { + let pool = Arc::new(GPUPool::new(NonZeroUsize::new(2).unwrap())); + let gpu_1 = pool.acquire(GPUCapacity::unchecked(0.3)).await?; + let gpu_2 = pool.acquire(GPUCapacity::unchecked(0.8)).await?; + let cloned_pool = pool.clone(); + let task_1 = tokio::spawn(async move { cloned_pool.acquire(GPUCapacity::MAX).await }); + tokio::time::sleep(std::time::Duration::from_secs_f32(0.02)).await; + assert!(!task_1.is_finished()); + let cloned_pool = pool.clone(); + let task_2 = + tokio::spawn(async move { cloned_pool.acquire(GPUCapacity::unchecked(0.4)).await }); + tokio::time::sleep(std::time::Duration::from_secs_f32(0.02)).await; + assert!(!task_2.is_finished()); + let cloned_pool = pool.clone(); + let task_3 = + tokio::spawn(async move { cloned_pool.acquire(GPUCapacity::unchecked(0.2)).await }); + tokio::time::sleep(std::time::Duration::from_secs_f32(0.02)).await; + assert!(!task_3.is_finished()); + + pool.release(gpu_2, GPUCapacity::unchecked(0.4))?; + assert!(!task_1.is_finished()); + let task_2_acquired_gpu = tokio::time::timeout(std::time::Duration::from_secs(1), task_2) + .await + .expect("task finished") + .expect("no timeout")?; + assert_eq!(task_2_acquired_gpu, gpu_2); + let task_3_acquired_gpu = tokio::time::timeout(std::time::Duration::from_secs(1), task_3) + .await + .expect("task finished") + .expect("no timeout")?; + assert_eq!(task_3_acquired_gpu, gpu_2); + + pool.release(gpu_1, GPUCapacity::unchecked(0.3))?; + let task_1_acquired_gpu = tokio::time::timeout(std::time::Duration::from_secs(1), task_1) + .await + .expect("task finished") + .expect("no timeout")?; + assert_eq!(task_1_acquired_gpu, gpu_1); + + pool.release(gpu_2, GPUCapacity::MAX) + } + + #[tokio::test] + async fn test_release_gpus() -> Result<()> { + let pool = GPUPool::new(NonZeroUsize::new(1).unwrap()); + let gpu_0 = pool.acquire(GPUCapacity::unchecked(0.5)).await?; + assert_eq!(gpu_0, 0); + pool.release(gpu_0, GPUCapacity::unchecked(0.5))?; + Ok(()) + } + + #[tokio::test] + async fn test_release_to_wrong_gpu_id() { + let pool = GPUPool::new(NonZeroUsize::new(1).unwrap()); + let release_result = pool.release(1, GPUCapacity::unchecked(0.5)); + assert!(release_result.is_err()); + assert!( + release_result + .unwrap_err() + .to_string() + .contains("Releasing to a gpu_id that does not exist: 1") + ); + } + + #[tokio::test] + async fn test_release_zero_fraction() { + let pool = GPUPool::new(NonZeroUsize::new(1).unwrap()); + let release_result = pool.release(0, GPUCapacity::ZERO); + assert!(release_result.is_err()); + assert!( + release_result + .unwrap_err() + .to_string() + .contains("Cannot release a zero fraction") + ); + } + + #[tokio::test] + async fn test_release_overflown_gpus() -> Result<()> { + let pool = GPUPool::new(NonZeroUsize::new(1).unwrap()); + let gpu_0 = pool.acquire(GPUCapacity::unchecked(0.5)).await?; + assert_eq!(gpu_0, 0); + let release_result = pool.release(gpu_0, GPUCapacity::unchecked(0.6)); + assert!(release_result.is_err()); + assert!( + release_result + .unwrap_err() + .to_string() + .contains("The sum of 0.5 and 0.6 is greater than the max value 1") + ); + Ok(()) + } + + #[test] + fn test_detect_num_gpus_explicit_env() { + temp_env::with_vars( + [ + ("COCOINDEX_NUM_GPUS", Some("4")), + ("CUDA_VISIBLE_DEVICES", None), + ], + || { + let pool = GPUPool::default(); + assert_eq!(pool.num_gpus(), 4); + }, + ); + } + + #[test] + fn test_detect_num_gpus_cuda_visible_devices() { + temp_env::with_vars( + [ + ("CUDA_VISIBLE_DEVICES", Some("0,2,3")), + ("COCOINDEX_NUM_GPUS", None), + ], + || { + let pool = GPUPool::default(); + assert_eq!(pool.num_gpus(), 3); + }, + ); + } + + #[test] + fn test_detect_num_gpus_cuda_visible_empty() { + temp_env::with_vars( + [ + ("CUDA_VISIBLE_DEVICES", Some("")), + ("COCOINDEX_NUM_GPUS", None), + ], + || { + let pool = GPUPool::default(); + assert_eq!(pool.num_gpus(), 1); + }, + ); + } + + #[test] + fn test_detect_num_gpus_explicit_env_zero() { + temp_env::with_vars( + [ + ("CUDA_VISIBLE_DEVICES", None), + ("COCOINDEX_NUM_GPUS", Some("0")), + ], + || { + let pool = GPUPool::default(); + assert_eq!(pool.num_gpus(), 1); + }, + ); + } + + #[test] + fn test_detect_num_gpus_explicit_env_overrides_cuda_visible() { + temp_env::with_vars( + [ + ("CUDA_VISIBLE_DEVICES", Some("0,1,2,3")), + ("COCOINDEX_NUM_GPUS", Some("2")), + ], + || { + let pool = GPUPool::default(); + assert_eq!(pool.num_gpus(), 2); + }, + ); + } + + #[test] + fn test_detect_num_gpus_cuda_visible_single_device() { + temp_env::with_vars( + [ + ("CUDA_VISIBLE_DEVICES", Some("0")), + ("COCOINDEX_NUM_GPUS", None), + ], + || { + let pool = GPUPool::default(); + assert_eq!(pool.num_gpus(), 1); + }, + ); + } + + #[test] + fn test_detect_num_gpus_cuda_visible_with_whitespace() { + temp_env::with_vars( + [ + ("CUDA_VISIBLE_DEVICES", Some("0, 1 , 2")), + ("COCOINDEX_NUM_GPUS", None), + ], + || { + let pool = GPUPool::default(); + assert_eq!(pool.num_gpus(), 3); + }, + ); + } + + #[test] + fn test_detect_num_gpus_nvidia_smi_returns_count() { + temp_env::with_vars( + [ + ("MOCK_NVIDIA_SMI_STDOUT", Some("8")), + ("CUDA_VISIBLE_DEVICES", None), + ("COCOINDEX_NUM_GPUS", None), + ], + || { + let pool = GPUPool::default(); + assert_eq!(pool.num_gpus(), 8); + }, + ); + } + + #[test] + fn test_detect_num_gpus_nvidia_smi_empty_output() { + temp_env::with_vars_unset(["CUDA_VISIBLE_DEVICES", "COCOINDEX_NUM_GPUS"], || { + let pool = GPUPool::default(); + assert_eq!(pool.num_gpus(), 1); + }) + } + + #[test] + fn test_detect_num_gpus_nvidia_smi_nonzero_exit() { + temp_env::with_vars( + [ + ("MOCK_NVIDIA_SMI_STDOUT", Some("8")), + ("MOCK_NVIDIA_SMI_EXIT_CODE", Some("1")), + ("CUDA_VISIBLE_DEVICES", None), + ("COCOINDEX_NUM_GPUS", None), + ], + || { + let pool = GPUPool::default(); + assert_eq!(pool.num_gpus(), 1); + }, + ); + } + + #[test] + fn test_detect_num_gpus_nvidia_smi_not_found() { + temp_env::with_vars( + [ + ("MOCK_NVIDIA_SMI_NOT_FOUND", Some("1")), + ("MOCK_NVIDIA_SMI_STDOUT", Some("8")), + ("CUDA_VISIBLE_DEVICES", None), + ("COCOINDEX_NUM_GPUS", None), + ], + || { + let pool = GPUPool::default(); + assert_eq!(pool.num_gpus(), 1); + }, + ); + } + + #[test] + fn test_detect_num_gpus_all_missing_fallback() { + temp_env::with_vars( + [ + ("MOCK_NVIDIA_SMI_EXIT_CODE", Some("1")), + ("MOCK_NVIDIA_SMI_STDOUT", None), + ("CUDA_VISIBLE_DEVICES", None), + ("COCOINDEX_NUM_GPUS", None), + ], + || { + let pool = GPUPool::default(); + assert_eq!(pool.num_gpus(), 1); + }, + ); + } + + #[test] + fn test_gpu_capacity_larger_than_one() { + let result = GPUCapacity::try_from(1.1); + assert!(result.is_err()); + assert!( + result + .unwrap_err() + .to_string() + .contains("Fraction must be between 0.0 and 1.0, got 1.1") + ); + } + + #[test] + fn test_gpu_capacity_less_than_zero() { + let result = GPUCapacity::try_from(-1.1); + assert!(result.is_err()); + assert!( + result + .unwrap_err() + .to_string() + .contains("Fraction must be between 0.0 and 1.0, got -1.1") + ); + } + + #[test] + fn test_gpu_capacity_zero() -> Result<()> { + let half = GPUCapacity::try_from(0.5)?; + let result = GPUCapacity::MAX - half + half - GPUCapacity::ZERO; + assert_eq!(result, GPUCapacity::MAX); + Ok(()) + } + + #[test] + fn test_gpu_capacity_repeat_acquire_then_release() -> Result<()> { + let mut full = GPUCapacity::MAX; + let mut rng = rand::rng(); + for _ in 0..100_000 { + let random_portion: f32 = rng.random_range(0.0..=1.0); + full -= GPUCapacity::try_from(random_portion)?; + full += GPUCapacity::try_from(random_portion)?; + } + assert_eq!(full, GPUCapacity::MAX); + Ok(()) + } + + #[test] + fn test_gpu_capacity_repeat_acquire_then_release_later() -> Result<()> { + let mut full = GPUCapacity::MAX; + let mut rng = rand::rng(); + let mut random_capacities = vec![]; + for _ in 0..100_000 { + let random_portion: f32 = rng.random_range(0.0..=1.0); + let capacity = GPUCapacity::try_from(random_portion)?; + if full <= capacity { + for cap in &random_capacities { + full += *cap; + } + assert_eq!( + full, + GPUCapacity::MAX, + "full ({full}) + sum({random_capacities:?}) != 1.0 (should be 1.0)", + ); + random_capacities.clear(); + } + random_capacities.push(capacity); + full -= capacity; + } + for cap in &random_capacities { + full += *cap; + } + assert_eq!( + full, + GPUCapacity::MAX, + "full ({full}) + sum({random_capacities:?}) != 1.0 (should be 1.0) (final)" + ); + Ok(()) + } + + #[test] + fn test_sorted_vec_find_lowest_index() { + let original = [1; 10]; + let capacity = SortedVec::from_iter(original); + let index = capacity.find_excluding_top_n(&1, 0); + assert_eq!(index, Some(0)); + } + + #[test] + fn test_sorted_vec_find_missing() { + let original = [4, 3, 0]; + let capacity = SortedVec::from_iter(original); // [0, 3, 4] + let index = capacity.find_excluding_top_n(&1, 0); + let expected = original.iter().position(|x| *x == 3); + assert_eq!(index, expected); + } + + #[test] + fn test_sorted_vec_find_exact() { + let original = [4, 3, 0]; + let capacity = SortedVec::from_iter(original); + let index = capacity.find_excluding_top_n(&3, 0); + let expected = original.iter().position(|x| *x == 3); + assert_eq!(index, expected); + } + + #[test] + fn test_sorted_vec_find_over_max() { + let capacity = SortedVec::from_iter([0, 3, 4].into_iter().rev()); + let index = capacity.find_excluding_top_n(&i32::MAX, 0); + assert_eq!(index, None); + } + + #[test] + fn test_sorted_vec_find_empty() { + let capacity = SortedVec::::from_iter([]); + let index = capacity.find_excluding_top_n(&3, 0); + assert_eq!(index, None); + } + + #[test] + fn test_sorted_vec_first_excluding_top_n_found() { + let original = (0..10).rev().collect::>(); + let capacity = SortedVec::from_iter(original.clone()); + let index = capacity.find_excluding_top_n(&5, 3); + let expected = original.iter().position(|x| *x == 5); + assert_eq!(index, expected); + } + + #[test] + fn test_sorted_vec_first_excluding_top_n_found_repeated() { + let capacity = SortedVec::from_iter([1; 10]); + let index = capacity.find_excluding_top_n(&1, 3); + assert_eq!(index, Some(0)); + } + + #[test] + fn test_sorted_vec_first_excluding_top_n_excluded() { + let capacity = SortedVec::from_iter((0..10).rev()); + let index = capacity.find_excluding_top_n(&5, 6); + assert_eq!(index, None); + } + + #[test] + fn test_sorted_vec_first_excluding_top_n_missing_excluded() { + let capacity = SortedVec::from_iter([0, 4, 3, 5]); // [0, 3, 4, 5] + let index = capacity.find_excluding_top_n(&2, 3); + assert_eq!(index, None); + } + + #[test] + fn test_sorted_vec_first_excluding_top_n_missing_found() { + let original = [0, 19, 15, 9, 20]; + let capacity = SortedVec::from_iter(original); // [0, 9, 15, 19, 20] + let index = capacity.find_excluding_top_n(&9, 2); + let expected = original.iter().position(|x| *x == 9); + assert_eq!(index, expected); + } + + #[test] + fn test_sorted_vec_take_excluding_top_n_success() { + let original = [20, 19, 15, 9, 20, 20, 11, 20]; + let capacity = SortedVec::from_iter(original); + let indices = capacity.find_many_excluding_top_n(&20, usize::MAX, 0); + let expected = original.iter().positions(|x| *x == 20).collect::>(); + assert_eq!(indices, expected); + } + + #[test] + fn test_sorted_vec_take_excluding_top_n_excluded() { + let original = [20, 19, 15, 9, 20, 20, 11, 20]; + let capacity = SortedVec::from_iter(original); + let indices = capacity.find_many_excluding_top_n(&20, usize::MAX, 1); + let mut expected = original.iter().positions(|x| *x == 20).collect::>(); + expected.pop(); + assert_eq!(indices, expected); + } + + #[test] + fn test_sorted_vec_take_excluding_top_n_take_few() { + let original = [20, 19, 15, 9, 20, 20, 11, 20]; + let capacity = SortedVec::from_iter(original); + let indices = capacity.find_many_excluding_top_n(&20, 2, 0); + let expected = original.iter().positions(|x| *x == 20).collect::>(); + assert_eq!(indices, expected[..2]); + } + + #[test] + fn test_sorted_vec_take_excluding_top_n_take_close_values() { + let original = [20, 19, 15, 9, 20, 20, 11, 20]; + let capacity = SortedVec::from_iter(original); + let indices = capacity.find_many_excluding_top_n(&19, usize::MAX, 0); + let expected = original + .iter() + .positions(|x| *x == 19) + .chain(original.iter().positions(|x| *x == 20)) + .collect::>(); + assert_eq!(indices, expected); + } +} diff --git a/rust/utils/src/lib.rs b/rust/utils/src/lib.rs index 46a2ccee3..3bb6f2a15 100644 --- a/rust/utils/src/lib.rs +++ b/rust/utils/src/lib.rs @@ -4,6 +4,7 @@ pub mod db; pub mod deser; pub mod error; pub mod fingerprint; +pub mod gpu_pool; pub mod immutable; pub mod ratelimit; pub mod retryable;