rivet/multigpu/
gpu_policy.rs

1//! GPU-pool construction helpers derived from [`crate::spec::EncodePolicy`].
2
3use std::sync::Arc;
4
5use crate::gpu_pool::GpuPool;
6use crate::spec::{EncodePolicy, GpuFamily};
7
8/// Build a [`GpuPool`] from the host's detected GPU inventory.
9pub fn detect_gpu_pool() -> Arc<GpuPool> {
10    Arc::new(GpuPool::new(&codec::gpu::detect_gpus()))
11}
12
13fn policy_vendor(fam: GpuFamily) -> codec::gpu::GpuVendor {
14    match fam {
15        GpuFamily::Nvidia => codec::gpu::GpuVendor::Nvidia,
16        GpuFamily::Amd => codec::gpu::GpuVendor::Amd,
17        GpuFamily::Intel => codec::gpu::GpuVendor::Intel,
18    }
19}
20
21/// The host GPUs selected by an [`EncodePolicy`]: all of them for `AllGpus`,
22/// the first / pinned index for `SingleGpu`, every device of one vendor for
23/// `Family`.
24fn select_gpus_for_policy(policy: EncodePolicy) -> Vec<codec::gpu::GpuDevice> {
25    let gpus = codec::gpu::detect_gpus();
26    match policy {
27        EncodePolicy::AllGpus => gpus,
28        EncodePolicy::SingleGpu(None) => gpus.into_iter().take(1).collect(),
29        EncodePolicy::SingleGpu(Some(idx)) => gpus.into_iter().filter(|g| g.index == idx).collect(),
30        EncodePolicy::Family(fam) => {
31            let v = policy_vendor(fam);
32            gpus.into_iter().filter(|g| g.vendor == v).collect()
33        }
34    }
35}
36
37/// Build a [`GpuPool`] constrained to the given [`EncodePolicy`]. An empty pool
38/// (e.g. a pinned index or vendor family that isn't present) yields capacity 0,
39/// so the orchestrator's pre-flight probe / lease claim surfaces a clear error.
40///
41/// When more than one GPU is selected, cards that can't actually encode the
42/// REQUESTED `codec` (e.g. a pre-Ada NVIDIA that decodes via NVDEC but has no
43/// AV1 encode silicon — yet can still encode H.264/H.265) are dropped from the
44/// **encode** pool, so a worker never leases an incapable card and hard-fails
45/// the run; the capable cards do the encoding. A single selected GPU is left
46/// as-is, since the serial path's non-pinned encoder dispatch already falls
47/// through vendors. Dropped cards stay available for the decode pump
48/// ([`policy_gpu_indices`] is intentionally NOT filtered).
49pub fn gpu_pool_for_policy(policy: EncodePolicy, codec: codec::frame::VideoCodec) -> Arc<GpuPool> {
50    let selected = select_gpus_for_policy(policy);
51    let pool_gpus = if selected.len() > 1 {
52        selected.into_iter().filter(|g| codec::encode::encode_capable(g, codec)).collect()
53    } else {
54        selected
55    };
56    Arc::new(GpuPool::new(&pool_gpus))
57}
58
59/// The GPU indices an [`EncodePolicy`] selects, in detection order. Used to pin
60/// the decode pump to a device consistent with the policy (so decode honors a
61/// `Family` / `SingleGpu` constraint, not just encode).
62pub fn policy_gpu_indices(policy: EncodePolicy) -> Vec<u32> {
63    select_gpus_for_policy(policy).into_iter().map(|g| g.index).collect()
64}
65
66/// The GPU index to pin a *serial* (single-GPU) encode/decode to under a
67/// policy: `None` (auto/first-available) for `AllGpus`, the pinned index for
68/// `SingleGpu`, the first device of the vendor for `Family`.
69pub fn serial_gpu_for_policy(policy: EncodePolicy) -> Option<u32> {
70    match policy {
71        EncodePolicy::AllGpus => None,
72        EncodePolicy::SingleGpu(idx) => idx,
73        EncodePolicy::Family(_) => select_gpus_for_policy(policy).first().map(|g| g.index),
74    }
75}
rivet/multigpu/gpu_policy.rs

rivet/multigpu/
gpu_policy.rs
