#[cfg(test)]
use std::time::Duration;
use std::{fmt, time::Instant};
#[cfg(test)]
use query_planning::error::QueryErrorCode;
use query_planning::{
admission::AttemptCapabilityEvidence,
error::{QueryError, Result},
};
use serde::Serialize;
#[cfg(any(test, feature = "test-hooks"))]
use crate::attempt::AttemptOwner;
#[cfg(test)]
use crate::runtime::query_service::config::QueueBounds;
use crate::{
attempt::{AttemptAdmission, AttemptAdmissionBinding, AttemptLifecycleError},
observability::{AttemptId, RuntimeObservationHandle},
runtime::{
ExecutionCancellation,
cache::{CacheSourceEventSnapshot, DeviceCacheReport},
gpu::{DeviceLedgerReport, DeviceMemoryFailureReport},
gpu_execution_pool::{GpuExecutionPoolSnapshot, GpuExecutionQuerySnapshot},
},
};
#[cfg(feature = "qualification")]
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct DeviceQuarantineInjection {
device_ordinal: usize,
device_health_generation: u64,
newly_quarantined: bool,
}
#[cfg(feature = "qualification")]
impl DeviceQuarantineInjection {
pub(crate) fn new(device_ordinal: usize, device_health_generation: u64, newly_quarantined: bool) -> Self {
Self { device_ordinal, device_health_generation, newly_quarantined }
}
pub fn device_ordinal(self) -> usize {
self.device_ordinal
}
pub fn device_health_generation(self) -> u64 {
self.device_health_generation
}
pub fn newly_quarantined(self) -> bool {
self.newly_quarantined
}
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize)]
pub struct QueryServiceAdmissionSnapshot {
pub active_queries: usize,
pub queued_queries: usize,
pub queue_high_watermark: usize,
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct BackendDiagnostics {
pub admission: QueryServiceAdmissionSnapshot,
pub ledgers: Vec<DeviceLedgerReport>,
pub caches: Vec<DeviceCacheReport>,
pub cache_source_events: Vec<CacheSourceEventSnapshot>,
pub memory_failures: Vec<DeviceMemoryFailureReport>,
pub execution_pools: Vec<GpuExecutionPoolSnapshot>,
pub execution_queries: GpuExecutionQuerySnapshot,
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize)]
#[cfg(test)]
pub(crate) enum QueryAdmissionOutcome {
Capacity { queued_queries: Option<usize>, queue_capacity: Option<usize>, max_queue_wait: Option<Duration> },
BudgetUnavailable { requested_cap_bytes: Option<u64>, largest_eligible_cap_bytes: Option<u64> },
CapabilityRejected,
InvalidInput,
Cancelled,
ExecutionFailure,
}
#[cfg(test)]
impl QueryAdmissionOutcome {
#[cfg(test)]
pub(crate) fn from_error(error: &QueryError) -> Self {
Self::from_error_with_optional_queue_bounds(error, None)
}
#[cfg(test)]
pub(crate) fn from_error_with_queue_bounds(error: &QueryError, queue_bounds: &QueueBounds) -> Self {
Self::from_error_with_optional_queue_bounds(error, Some(queue_bounds))
}
#[cfg(test)]
pub(crate) fn from_execution_error(error: &QueryError) -> Self {
if error.code() == QueryErrorCode::Cancelled { Self::Cancelled } else { Self::ExecutionFailure }
}
#[cfg(test)]
fn from_error_with_optional_queue_bounds(error: &QueryError, queue_bounds: Option<&QueueBounds>) -> Self {
match error.code() {
QueryErrorCode::QueryBudgetUnavailable => Self::BudgetUnavailable {
requested_cap_bytes: fact_u64(error, "requested_cap_bytes"),
largest_eligible_cap_bytes: fact_u64(error, "largest_eligible_cap_bytes"),
},
QueryErrorCode::ServiceOverloaded | QueryErrorCode::MemoryReservationFailed => Self::Capacity {
queued_queries: fact_usize(error, "queued_queries"),
queue_capacity: fact_usize(error, "max_waiters")
.or_else(|| queue_bounds.map(|bounds| bounds.max_waiters)),
max_queue_wait: queue_bounds.map(|bounds| bounds.max_wait),
},
QueryErrorCode::UnsupportedPipelineContract
| QueryErrorCode::UnsupportedOperator
| QueryErrorCode::UnsupportedExpression
| QueryErrorCode::UnsupportedType
| QueryErrorCode::CudfUnsupported => Self::CapabilityRejected,
QueryErrorCode::InvalidPlan
| QueryErrorCode::InvalidAttemptCorrelation
| QueryErrorCode::InvalidConfiguration
| QueryErrorCode::GraphInputInvalid
| QueryErrorCode::CudfInvalidInput => Self::InvalidInput,
QueryErrorCode::Cancelled => Self::Cancelled,
QueryErrorCode::CudfTemporaryFailure
| QueryErrorCode::CudfPersistentFailure
| QueryErrorCode::AlgorithmDidNotConverge
| QueryErrorCode::CudfDependencyFailure
| QueryErrorCode::SourceAccessDenied
| QueryErrorCode::SourceUnauthenticated
| QueryErrorCode::SourceObjectNotFound
| QueryErrorCode::SourceTemporaryFailure
| QueryErrorCode::CudfInvariantViolation
| QueryErrorCode::QueryMemoryLimitExceeded
| QueryErrorCode::DeviceOutOfMemory
| QueryErrorCode::DeviceAllocationFailed
| QueryErrorCode::DeviceFault
| QueryErrorCode::CleanupFailed
| QueryErrorCode::InternalInvariantViolation
| QueryErrorCode::SizeOverflow
| QueryErrorCode::HostMaterializationLimitExceeded
| QueryErrorCode::Internal => Self::ExecutionFailure,
}
}
#[cfg(test)]
pub(crate) fn queued_queries(&self) -> Option<usize> {
match self {
Self::Capacity { queued_queries, .. } => *queued_queries,
_ => None,
}
}
#[cfg(test)]
pub(crate) fn queue_capacity(&self) -> Option<usize> {
match self {
Self::Capacity { queue_capacity, .. } => *queue_capacity,
_ => None,
}
}
#[cfg(test)]
pub(crate) fn max_queue_wait(&self) -> Option<Duration> {
match self {
Self::Capacity { max_queue_wait, .. } => *max_queue_wait,
_ => None,
}
}
}
#[cfg(test)]
fn fact_usize(error: &QueryError, key: &str) -> Option<usize> {
error.fact_value(key)?.parse().ok()
}
#[cfg(test)]
fn fact_u64(error: &QueryError, key: &str) -> Option<u64> {
error.fact_value(key)?.parse().ok()
}
#[cfg(test)]
mod outcome_tests {
use super::*;
fn error(code: QueryErrorCode) -> QueryError {
QueryError::new(code, "query_engine.test", "test error")
}
#[test]
fn capacity_projection_preserves_queue_facts_and_configured_wait_bound() {
let error = error(QueryErrorCode::ServiceOverloaded)
.with_fact("reason", "admission_queue_full")
.with_fact("queued_queries", 7)
.with_fact("max_waiters", 8);
let bounds = QueueBounds { max_waiters: 8, max_wait: Duration::from_secs(11) };
let outcome = QueryAdmissionOutcome::from_error_with_queue_bounds(&error, &bounds);
assert_eq!(outcome.queued_queries(), Some(7));
assert_eq!(outcome.queue_capacity(), Some(8));
assert_eq!(outcome.max_queue_wait(), Some(Duration::from_secs(11)));
assert!(matches!(outcome, QueryAdmissionOutcome::Capacity { .. }));
}
#[test]
fn admission_projection_is_exhaustive_and_does_not_parse_reason_strings() {
let cases = [
(QueryErrorCode::MemoryReservationFailed, "capacity"),
(QueryErrorCode::QueryBudgetUnavailable, "budget_unavailable"),
(QueryErrorCode::UnsupportedPipelineContract, "capability"),
(QueryErrorCode::InvalidPlan, "invalid"),
(QueryErrorCode::InvalidConfiguration, "invalid"),
(QueryErrorCode::Cancelled, "cancelled"),
(QueryErrorCode::CudfInvariantViolation, "execution"),
];
for (code, expected) in cases {
let outcome = QueryAdmissionOutcome::from_error(&error(code).with_fact("reason", "unrelated_reason_name"));
let actual = match outcome {
QueryAdmissionOutcome::Capacity { .. } => "capacity",
QueryAdmissionOutcome::BudgetUnavailable { .. } => "budget_unavailable",
QueryAdmissionOutcome::CapabilityRejected => "capability",
QueryAdmissionOutcome::InvalidInput => "invalid",
QueryAdmissionOutcome::Cancelled => "cancelled",
QueryAdmissionOutcome::ExecutionFailure => "execution",
};
assert_eq!(actual, expected, "classification for {code:?}");
}
let bounds = QueueBounds { max_waiters: 8, max_wait: Duration::from_secs(11) };
let budget = QueryAdmissionOutcome::from_error_with_queue_bounds(
&error(QueryErrorCode::QueryBudgetUnavailable).with_fact("queued_queries", 7).with_fact("max_waiters", 8),
&bounds,
);
assert_eq!(budget.queued_queries(), None);
assert_eq!(budget.queue_capacity(), None);
assert_eq!(budget.max_queue_wait(), None);
}
#[test]
fn capacity_projection_leaves_unknown_queue_metrics_unknown() {
let outcome = QueryAdmissionOutcome::from_error(&error(QueryErrorCode::MemoryReservationFailed));
assert!(matches!(
outcome,
QueryAdmissionOutcome::Capacity { queued_queries: None, queue_capacity: None, max_queue_wait: None }
));
let bounds = QueueBounds { max_waiters: 4, max_wait: Duration::from_secs(3) };
let outcome = QueryAdmissionOutcome::from_error_with_queue_bounds(
&error(QueryErrorCode::MemoryReservationFailed),
&bounds,
);
assert_eq!(outcome.queued_queries(), None);
assert_eq!(outcome.queue_capacity(), Some(4));
assert_eq!(outcome.max_queue_wait(), Some(Duration::from_secs(3)));
}
#[test]
fn execution_projection_does_not_reclassify_post_admission_resource_errors() {
for code in [QueryErrorCode::MemoryReservationFailed, QueryErrorCode::DeviceAllocationFailed] {
assert_eq!(
QueryAdmissionOutcome::from_execution_error(&error(code)),
QueryAdmissionOutcome::ExecutionFailure,
"post-admission {code:?} is an execution failure"
);
}
assert_eq!(
QueryAdmissionOutcome::from_execution_error(&error(QueryErrorCode::Cancelled)),
QueryAdmissionOutcome::Cancelled
);
}
}
#[derive(Debug)]
#[non_exhaustive]
pub struct QueryAdmissionRequest {
requested_query_cap_bytes: Option<u64>,
pub(in crate::runtime::query_service) attempt_requirement: Option<AttemptCapabilityEvidence>,
planned_capability_evidence: Option<AttemptCapabilityEvidence>,
pub(crate) queue_deadline: Option<Instant>,
pub(crate) execution_cancellation: ExecutionCancellation,
attempt: AttemptAdmissionBinding,
runtime_observation_activity: Option<Result<QueryActivityLease>>,
#[cfg(any(test, feature = "test-hooks"))]
pub(in crate::runtime::query_service) test_attempt: Option<AttemptOwner>,
}
pub(in crate::runtime::query_service) struct AdmissionObservation {
pub(in crate::runtime::query_service) handle: RuntimeObservationHandle,
pub(in crate::runtime::query_service) activity: Option<QueryActivityLease>,
pub(in crate::runtime::query_service) activity_error: Option<QueryError>,
pub(in crate::runtime::query_service) attempt_id: AttemptId,
}
pub struct QueryActivityLease {
_lease: Box<dyn Send + Sync>,
}
impl QueryActivityLease {
pub fn new(lease: impl Send + Sync + 'static) -> Self {
Self { _lease: Box::new(lease) }
}
}
impl fmt::Debug for QueryActivityLease {
fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
formatter.debug_struct("QueryActivityLease").finish_non_exhaustive()
}
}
impl QueryAdmissionRequest {
pub fn new(attempt: AttemptAdmissionBinding) -> Self {
Self {
requested_query_cap_bytes: None,
attempt_requirement: None,
planned_capability_evidence: None,
queue_deadline: None,
execution_cancellation: ExecutionCancellation::default(),
attempt,
runtime_observation_activity: None,
#[cfg(any(test, feature = "test-hooks"))]
test_attempt: None,
}
}
#[cfg(any(test, feature = "test-hooks"))]
#[doc(hidden)]
pub fn test_attempt() -> Self {
let owner = AttemptOwner::native_for_test();
let mut request = Self::new(owner.admission_binding());
request.test_attempt = Some(owner);
request
}
#[cfg(any(test, feature = "test-hooks"))]
pub(in crate::runtime::query_service) fn split_test_attempt(mut self) -> (Self, Option<AttemptOwner>) {
let owner = self.test_attempt.take();
(self, owner)
}
pub fn with_query_cap_bytes(mut self, query_cap_bytes: u64) -> Self {
self.requested_query_cap_bytes = Some(query_cap_bytes);
self
}
pub(in crate::runtime::query_service) const fn requested_query_cap_bytes(&self) -> Option<u64> {
self.requested_query_cap_bytes
}
pub fn with_query_plan_capability_evidence(mut self, evidence: AttemptCapabilityEvidence) -> Self {
self.planned_capability_evidence = Some(evidence);
self
}
#[cfg(test)]
pub(crate) fn with_attempt_requirement(mut self, attempt_requirement: AttemptCapabilityEvidence) -> Self {
self.attempt_requirement = Some(attempt_requirement);
self
}
pub(in crate::runtime::query_service) fn resolved_attempt_requirement(&self) -> Option<AttemptCapabilityEvidence> {
self.planned_capability_evidence.clone().or_else(|| self.attempt_requirement.clone())
}
pub(in crate::runtime::query_service) fn submitted_plan_capability_evidence(
&self,
) -> Option<AttemptCapabilityEvidence> {
self.planned_capability_evidence.clone()
}
pub fn with_queue_deadline(mut self, queue_deadline: Instant) -> Self {
self.queue_deadline = Some(queue_deadline);
self
}
pub fn with_execution_cancellation(mut self, execution_cancellation: ExecutionCancellation) -> Self {
self.execution_cancellation = execution_cancellation;
self
}
#[cfg(test)]
pub(crate) fn with_runtime_observation_activity(mut self, activity: QueryActivityLease) -> Self {
self.runtime_observation_activity = Some(Ok(activity));
self
}
#[cfg(test)]
pub(crate) fn with_runtime_observation_activity_result(mut self, activity: Result<QueryActivityLease>) -> Self {
self.runtime_observation_activity = Some(activity);
self
}
pub(in crate::runtime::query_service) fn into_admission_observation(self) -> AdmissionObservation {
let (activity, activity_error) = match self.runtime_observation_activity {
Some(Ok(activity)) => (Some(activity), None),
Some(Err(error)) => (None, Some(error)),
None => (None, None),
};
AdmissionObservation {
handle: self.attempt.observation_handle(),
activity,
activity_error,
attempt_id: self.attempt.attempt_id(),
}
}
pub(in crate::runtime::query_service) fn claim_attempt_admission(
&self,
) -> std::result::Result<AttemptAdmission, AttemptLifecycleError> {
self.attempt.claim()
}
}
#[cfg(test)]
mod tests {
use std::sync::Arc;
use super::*;
use crate::{
attempt::{AttemptBackendCounters, BeginAttemptRequest},
observability::{AttemptDisposition, RuntimeObservationConfig},
runtime::BackendLifecycle,
};
#[test]
fn attempt_binding_attaches_observation_and_release_together() {
let lifecycle = BackendLifecycle::new();
let attempt = AttemptOwner::begin(
BeginAttemptRequest {
observation: RuntimeObservationConfig::bounded(8, 8 * 1024),
..BeginAttemptRequest::default()
},
lifecycle.try_begin_observation().expect("attempt observation permit"),
Arc::new(AttemptBackendCounters::default()),
None,
);
let attempt_id = attempt.owner.attempt_id();
attempt.owner.handle().set_disposition(AttemptDisposition::Native).expect("native disposition");
let request = QueryAdmissionRequest::new(attempt.owner.admission_binding());
assert!(request.claim_attempt_admission().is_ok());
assert!(matches!(request.claim_attempt_admission(), Err(AttemptLifecycleError::AdmissionAlreadyClaimed)));
let observation = request.into_admission_observation();
assert_eq!(observation.attempt_id, attempt_id);
}
}