use std::fmt;
use query_planning::{
admission::ExecutionCapability,
error::{QueryError, QueryErrorCode, Result, bytes_with_gib},
};
const OPERATION: &str = "query_engine.gpu.resource_profile";
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
pub struct DeviceCapability {
capability: ExecutionCapability,
}
impl DeviceCapability {
pub const fn new(name: &'static str) -> Self {
Self { capability: ExecutionCapability::new(name) }
}
const fn from_requirement(requirement: ExecutionCapability) -> Self {
Self { capability: requirement }
}
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct DeviceCapabilities {
support: DeviceCapabilitySupport,
}
#[derive(Debug, Clone, PartialEq, Eq)]
enum DeviceCapabilitySupport {
All,
Exact(Vec<DeviceCapability>),
}
impl Default for DeviceCapabilities {
fn default() -> Self {
Self::all()
}
}
impl DeviceCapabilities {
pub(crate) const fn all() -> Self {
Self { support: DeviceCapabilitySupport::All }
}
pub fn only(capabilities: impl IntoIterator<Item = DeviceCapability>) -> Self {
let mut exact = Vec::new();
for capability in capabilities {
if !exact.contains(&capability) {
exact.push(capability);
}
}
Self { support: DeviceCapabilitySupport::Exact(exact) }
}
pub(crate) fn supports(&self, requirement: &ExecutionCapability) -> bool {
match &self.support {
DeviceCapabilitySupport::All => true,
DeviceCapabilitySupport::Exact(capabilities) => {
capabilities.contains(&DeviceCapability::from_requirement(*requirement))
},
}
}
pub(crate) fn supports_all(&self, requirements: &[ExecutionCapability]) -> bool {
requirements.iter().all(|requirement| self.supports(requirement))
}
pub(crate) const fn requires_explicit_attempt_requirement(&self) -> bool {
matches!(self.support, DeviceCapabilitySupport::Exact(_))
}
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct DeviceResourceProfile {
pub device_ordinal: usize,
pub device_capabilities: DeviceCapabilities,
pub(crate) configured_limit_bytes: u64,
pub(crate) headroom_bytes: u64,
pub(crate) backend_reserve_bytes: u64,
pub(crate) usable_capacity_bytes: u64,
pub managed_capacity_bytes: u64,
pub default_query_limit_bytes: u64,
pub max_active_attempts: usize,
pub root_lane_count: usize,
pub root_stream_count: usize,
pub source_work_per_attempt: usize,
pub(crate) aux_source_work_capacity: usize,
pub stream_pool_size: usize,
pub worker_job_capacity: usize,
pub cache_cap_bytes: u64,
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub(crate) struct DeviceResourceEnvelope {
pub(crate) device_ordinal: usize,
pub(crate) device_capabilities: DeviceCapabilities,
pub(crate) configured_limit_bytes: u64,
pub(crate) headroom_bytes: u64,
pub(crate) backend_reserve_bytes: u64,
pub(crate) default_query_limit_bytes: u64,
pub(crate) max_active_attempts: usize,
pub(crate) source_work_per_attempt: usize,
pub(crate) stream_pool_size: usize,
pub(crate) worker_job_capacity: usize,
pub(crate) cache_cap_bytes: u64,
}
impl DeviceResourceProfile {
pub fn query_share_bytes(&self) -> u64 {
self.managed_capacity_bytes.saturating_sub(self.cache_cap_bytes)
}
pub(crate) fn derive_default_query_limit_bytes(
device_ordinal: usize,
managed_capacity_bytes: u64,
cache_cap_bytes: u64,
max_active_attempts: usize,
) -> Result<u64> {
if max_active_attempts == 0 {
return Err(
invalid(device_ordinal, "max_active_attempts", "must be greater than zero").with_fact("reason", "zero")
);
}
let query_share_bytes =
managed_capacity_bytes.checked_sub(cache_cap_bytes).filter(|share| *share > 0).ok_or_else(|| {
invalid(
device_ordinal,
"cache_cap_bytes",
format!(
"is {}, at or above the managed capacity of {}; the largest accepted value is {}",
bytes_with_gib(cache_cap_bytes),
bytes_with_gib(managed_capacity_bytes),
bytes_with_gib(managed_capacity_bytes.saturating_sub(1)),
),
)
.with_fact("reason", "cache_cap_exceeds_managed_capacity")
.with_fact("cache_cap_bytes", cache_cap_bytes)
.with_fact("managed_capacity_bytes", managed_capacity_bytes)
})?;
let divisor = u64::try_from(max_active_attempts).map_err(|_| {
invalid(device_ordinal, "max_active_attempts", "parallelism target exceeds the capacity integer range")
.with_fact("reason", "max_active_attempts_overflow")
})?;
let derived = query_share_bytes.checked_div(divisor).filter(|cap| *cap > 0).ok_or_else(|| {
invalid(device_ordinal, "default_query_limit_bytes", "derived default query limit is zero")
.with_fact("reason", "derived_default_query_limit_is_zero")
.with_fact("query_share_bytes", query_share_bytes)
.with_fact("max_active_attempts", max_active_attempts)
})?;
Ok(derived)
}
pub(crate) fn resolve(envelope: DeviceResourceEnvelope) -> Result<Self> {
let DeviceResourceEnvelope {
device_ordinal,
device_capabilities,
configured_limit_bytes,
headroom_bytes,
backend_reserve_bytes,
default_query_limit_bytes,
max_active_attempts,
source_work_per_attempt,
stream_pool_size,
worker_job_capacity,
cache_cap_bytes,
} = envelope;
let usable_capacity_bytes = configured_limit_bytes.checked_sub(headroom_bytes).ok_or_else(|| {
invalid(
device_ordinal,
"headroom_bytes",
format!(
"is {}, above the configured limit of {}",
bytes_with_gib(headroom_bytes),
bytes_with_gib(configured_limit_bytes),
),
)
.with_fact("reason", "configured_limit_below_headroom")
.with_fact("configured_limit_bytes", configured_limit_bytes)
.with_fact("headroom_bytes", headroom_bytes)
})?;
let managed_capacity_bytes = usable_capacity_bytes.checked_sub(backend_reserve_bytes).ok_or_else(|| {
invalid(
device_ordinal,
"backend_reserve_bytes",
format!(
"is {}, above the usable capacity of {}; the largest accepted value is {}",
bytes_with_gib(backend_reserve_bytes),
bytes_with_gib(usable_capacity_bytes),
bytes_with_gib(usable_capacity_bytes.saturating_sub(1)),
),
)
.with_fact("reason", "backend_reserve_exceeds_usable_capacity")
.with_fact("usable_capacity_bytes", usable_capacity_bytes)
.with_fact("backend_reserve_bytes", backend_reserve_bytes)
})?;
if managed_capacity_bytes == 0 {
return Err(invalid(
device_ordinal,
"managed_capacity_bytes",
format!(
"is zero; backend_reserve_bytes of {} takes the whole usable capacity of {}",
bytes_with_gib(backend_reserve_bytes),
bytes_with_gib(usable_capacity_bytes),
),
)
.with_fact("reason", "zero_managed_capacity"));
}
if cache_cap_bytes >= managed_capacity_bytes {
return Err(invalid(
device_ordinal,
"cache_cap_bytes",
format!(
"is {}, at or above the managed capacity of {}; the largest accepted value is {}",
bytes_with_gib(cache_cap_bytes),
bytes_with_gib(managed_capacity_bytes),
bytes_with_gib(managed_capacity_bytes - 1),
),
)
.with_fact("reason", "cache_cap_exceeds_managed_capacity")
.with_fact("cache_cap_bytes", cache_cap_bytes)
.with_fact("managed_capacity_bytes", managed_capacity_bytes));
}
if max_active_attempts == 0 {
return Err(
invalid(device_ordinal, "max_active_attempts", "must be greater than zero").with_fact("reason", "zero")
);
}
if default_query_limit_bytes == 0 {
return Err(invalid(device_ordinal, "default_query_limit_bytes", "must be greater than zero")
.with_fact("reason", "zero"));
}
let query_share_bytes = managed_capacity_bytes.checked_sub(cache_cap_bytes).ok_or_else(|| {
invalid(device_ordinal, "cache_cap_bytes", "cache cap and query share overflow capacity arithmetic")
.with_fact("reason", "cache_cap_overflow")
})?;
if default_query_limit_bytes > query_share_bytes {
return Err(invalid(
device_ordinal,
"default_query_limit_bytes",
format!(
"is {}, above the query share of {} (managed capacity {} minus cache_cap_bytes {}); the largest accepted value is {}",
bytes_with_gib(default_query_limit_bytes),
bytes_with_gib(query_share_bytes),
bytes_with_gib(managed_capacity_bytes),
bytes_with_gib(cache_cap_bytes),
bytes_with_gib(query_share_bytes),
),
)
.with_fact("reason", "default_query_limit_exceeds_managed_capacity")
.with_fact("default_query_limit_bytes", default_query_limit_bytes)
.with_fact("query_share_bytes", query_share_bytes)
.with_fact("managed_capacity_bytes", managed_capacity_bytes));
}
let source_work_per_attempt = source_work_per_attempt.max(1);
let aux_source_work_capacity = max_active_attempts.checked_mul(source_work_per_attempt).ok_or_else(|| {
invalid(device_ordinal, "source_work_per_attempt", "auxiliary source-work capacity overflowed")
.with_fact("reason", "aux_source_work_overflow")
})?;
let required_slots = max_active_attempts.checked_add(aux_source_work_capacity).ok_or_else(|| {
invalid(device_ordinal, "stream_pool_size", "required stream capacity overflowed")
.with_fact("reason", "required_slot_overflow")
})?;
for (field, actual) in [("stream_pool_size", stream_pool_size), ("worker_job_capacity", worker_job_capacity)] {
if actual < required_slots {
return Err(invalid(device_ordinal, field, "must cover root and auxiliary source work")
.with_fact("reason", "below_required_slot_capacity")
.with_fact("required_slot_capacity", required_slots)
.with_fact("actual_capacity", actual));
}
}
Ok(Self {
device_ordinal,
device_capabilities,
configured_limit_bytes,
headroom_bytes,
backend_reserve_bytes,
usable_capacity_bytes,
managed_capacity_bytes,
default_query_limit_bytes,
max_active_attempts,
root_lane_count: max_active_attempts,
root_stream_count: max_active_attempts,
source_work_per_attempt,
aux_source_work_capacity,
stream_pool_size,
worker_job_capacity,
cache_cap_bytes,
})
}
}
fn invalid(device_ordinal: usize, field: &'static str, message: impl fmt::Display) -> QueryError {
QueryError::new(
QueryErrorCode::InvalidConfiguration,
OPERATION,
format!("device {device_ordinal} {field} {message}"),
)
.with_fact("device_ordinal", device_ordinal)
.with_fact("field", field)
}
#[cfg(test)]
mod tests {
use query_planning::admission::ExecutionCapability;
use super::*;
#[test]
fn resolves_heterogeneous_profiles_and_slot_math() {
let small = DeviceResourceProfile::resolve(DeviceResourceEnvelope {
device_ordinal: 0,
configured_limit_bytes: 1_000,
headroom_bytes: 100,
backend_reserve_bytes: 100,
default_query_limit_bytes: 100,
max_active_attempts: 2,
source_work_per_attempt: 3,
stream_pool_size: 8,
worker_job_capacity: 8,
device_capabilities: DeviceCapabilities::default(),
cache_cap_bytes: 0,
}).unwrap();
let large = DeviceResourceProfile::resolve(DeviceResourceEnvelope {
device_ordinal: 1,
configured_limit_bytes: 2_000,
headroom_bytes: 200,
backend_reserve_bytes: 100,
default_query_limit_bytes: 100,
max_active_attempts: 3,
source_work_per_attempt: 1,
stream_pool_size: 6,
worker_job_capacity: 6,
device_capabilities: DeviceCapabilities::default(),
cache_cap_bytes: 0,
})
.unwrap();
assert_eq!((small.usable_capacity_bytes, small.managed_capacity_bytes), (900, 800));
assert_eq!((small.root_lane_count, small.root_stream_count, small.aux_source_work_capacity), (2, 2, 6));
assert_eq!((large.usable_capacity_bytes, large.managed_capacity_bytes), (1_800, 1_700));
}
#[test]
fn rejects_invalid_capacity_and_query_limit_arithmetic() {
for (limit, headroom, reserve, reason) in [
(99, 100, 0, "configured_limit_below_headroom"),
(100, 0, 101, "backend_reserve_exceeds_usable_capacity"),
(100, 0, 100, "zero_managed_capacity"),
] {
let error = DeviceResourceProfile::resolve(DeviceResourceEnvelope {
device_ordinal: 7,
configured_limit_bytes: limit,
headroom_bytes: headroom,
backend_reserve_bytes: reserve,
default_query_limit_bytes: 0,
max_active_attempts: 1,
source_work_per_attempt: 1,
stream_pool_size: 2,
worker_job_capacity: 2,
device_capabilities: DeviceCapabilities::default(),
cache_cap_bytes: 0,
})
.unwrap_err();
assert_eq!(error.fact_value("device_ordinal"), Some("7"));
assert_eq!(error.fact_value("reason"), Some(reason));
}
let error = DeviceResourceProfile::resolve(DeviceResourceEnvelope {
device_ordinal: 7,
configured_limit_bytes: u64::MAX - 1,
headroom_bytes: 0,
backend_reserve_bytes: 0,
default_query_limit_bytes: u64::MAX,
max_active_attempts: 2,
source_work_per_attempt: 1,
stream_pool_size: 4,
worker_job_capacity: 4,
device_capabilities: DeviceCapabilities::default(),
cache_cap_bytes: 0,
})
.unwrap_err();
assert_eq!(error.fact_value("reason"), Some("default_query_limit_exceeds_managed_capacity"));
}
#[test]
fn validates_cache_cap_against_managed_capacity() {
let profile = DeviceResourceProfile::resolve(DeviceResourceEnvelope {
device_ordinal: 0,
configured_limit_bytes: 1_000,
headroom_bytes: 0,
backend_reserve_bytes: 0,
default_query_limit_bytes: 100,
max_active_attempts: 2,
source_work_per_attempt: 1,
stream_pool_size: 4,
worker_job_capacity: 4,
device_capabilities: DeviceCapabilities::default(),
cache_cap_bytes: 800,
})
.unwrap();
assert_eq!(profile.cache_cap_bytes, 800);
let error = DeviceResourceProfile::resolve(DeviceResourceEnvelope {
device_ordinal: 0,
configured_limit_bytes: 1_000,
headroom_bytes: 0,
backend_reserve_bytes: 0,
default_query_limit_bytes: 100,
max_active_attempts: 2,
source_work_per_attempt: 1,
stream_pool_size: 4,
worker_job_capacity: 4,
device_capabilities: DeviceCapabilities::default(),
cache_cap_bytes: 1_001,
})
.unwrap_err();
assert_eq!(error.fact_value("reason"), Some("cache_cap_exceeds_managed_capacity"));
}
#[test]
fn rejects_zero_default_query_limit_for_an_otherwise_valid_profile() {
let error = DeviceResourceProfile::resolve(DeviceResourceEnvelope {
device_ordinal: 7,
configured_limit_bytes: 100,
headroom_bytes: 0,
backend_reserve_bytes: 0,
default_query_limit_bytes: 0,
max_active_attempts: 1,
source_work_per_attempt: 1,
stream_pool_size: 2,
worker_job_capacity: 2,
device_capabilities: DeviceCapabilities::default(),
cache_cap_bytes: 0,
})
.unwrap_err();
assert_eq!(error.fact_value("field"), Some("default_query_limit_bytes"));
assert_eq!(error.fact_value("reason"), Some("zero"));
}
#[test]
fn rejects_query_limit_and_slot_capacity_below_requirements() {
let query_limit = DeviceResourceProfile::resolve(DeviceResourceEnvelope {
device_ordinal: 3,
configured_limit_bytes: 100,
headroom_bytes: 0,
backend_reserve_bytes: 0,
default_query_limit_bytes: 101,
max_active_attempts: 2,
source_work_per_attempt: 1,
stream_pool_size: 4,
worker_job_capacity: 4,
device_capabilities: DeviceCapabilities::default(),
cache_cap_bytes: 0,
})
.unwrap_err();
assert_eq!(query_limit.fact_value("reason"), Some("default_query_limit_exceeds_managed_capacity"));
let slots = DeviceResourceProfile::resolve(DeviceResourceEnvelope {
device_ordinal: 3,
configured_limit_bytes: 100,
headroom_bytes: 0,
backend_reserve_bytes: 0,
default_query_limit_bytes: 10,
max_active_attempts: 2,
source_work_per_attempt: 3,
stream_pool_size: 7,
worker_job_capacity: 7,
device_capabilities: DeviceCapabilities::default(),
cache_cap_bytes: 0,
})
.unwrap_err();
assert_eq!(slots.fact_value("reason"), Some("below_required_slot_capacity"));
}
#[test]
fn frozen_capability_metadata_requires_exact_matches() {
let required = ExecutionCapability::new("test_operator");
let other = ExecutionCapability::new("other_operator");
let restricted = DeviceCapabilities::only([DeviceCapability::new("test_operator")]);
assert!(restricted.supports(&required));
assert!(!restricted.supports(&other));
assert!(!DeviceCapabilities::only(std::iter::empty::<DeviceCapability>()).supports(&required));
}
#[test]
fn derives_default_query_limit_from_parallelism_target_and_cache_share() {
assert_eq!(DeviceResourceProfile::derive_default_query_limit_bytes(0, 24, 8, 2).unwrap(), 8);
assert_eq!(DeviceResourceProfile::derive_default_query_limit_bytes(0, 10, 0, 3).unwrap(), 3);
assert_eq!(DeviceResourceProfile::derive_default_query_limit_bytes(0, 1_000, 0, 1).unwrap(), 1_000);
let zero_share = DeviceResourceProfile::derive_default_query_limit_bytes(0, 10, 9, 3).unwrap_err();
assert_eq!(zero_share.fact_value("reason"), Some("derived_default_query_limit_is_zero"));
let no_share = DeviceResourceProfile::derive_default_query_limit_bytes(0, 8, 8, 2).unwrap_err();
assert_eq!(no_share.fact_value("reason"), Some("cache_cap_exceeds_managed_capacity"));
let no_parallelism = DeviceResourceProfile::derive_default_query_limit_bytes(0, 24, 8, 0).unwrap_err();
assert_eq!(no_parallelism.fact_value("reason"), Some("zero"));
}
#[test]
fn query_share_subtracts_the_cache_share() {
let envelope = |cache_cap_bytes, default_query_limit_bytes| DeviceResourceEnvelope {
device_ordinal: 0,
configured_limit_bytes: 1_000,
headroom_bytes: 0,
backend_reserve_bytes: 0,
default_query_limit_bytes,
max_active_attempts: 2,
source_work_per_attempt: 1,
stream_pool_size: 4,
worker_job_capacity: 4,
device_capabilities: DeviceCapabilities::default(),
cache_cap_bytes,
};
let profile = DeviceResourceProfile::resolve(envelope(300, 700)).unwrap();
assert_eq!(profile.query_share_bytes(), 700);
let beyond_share = DeviceResourceProfile::resolve(envelope(300, 701)).unwrap_err();
assert_eq!(beyond_share.fact_value("reason"), Some("default_query_limit_exceeds_managed_capacity"));
let no_share = DeviceResourceProfile::resolve(envelope(1_000, 1)).unwrap_err();
assert_eq!(no_share.fact_value("reason"), Some("cache_cap_exceeds_managed_capacity"));
}
}