query_runtime/runtime/gpu/
memory_failure.rs

1//! Signal-only device memory allocation-failure counters.
2//!
3//! Signal-only means the counters inform reports and never admission: the
4//! ledger decides capacity from its own integer accounting, and a failed
5//! allocation is already a typed error at its source. The counters are
6//! written by the callback the device memory manager invokes inside RMM's
7//! failure handling, so that callback allocates nothing and takes no lock.
8//!
9//! The manager and its callback are process-lifetime, while backends are
10//! constructed per generation. The registry below therefore keeps one
11//! tracker per ordinal before installation (`pending`) and rebinds it to the
12//! manager's installation generation afterwards (`installed`). A later backend
13//! recovers the tracker the live callback already points at, because the
14//! callback would never update a second counter.
15
16use std::{
17    collections::HashMap,
18    sync::{
19        Arc, Mutex, OnceLock,
20        atomic::{AtomicBool, AtomicU64, Ordering},
21    },
22};
23
24use cudf::{DeviceMemoryFailureSource, DeviceMemoryPoolUsage};
25use query_planning::error::{QueryError, QueryErrorCode, Result};
26use serde::Serialize;
27
28const TRACKER_LOOKUP_OPERATION: &str = "query_engine.gpu.memory_failure_tracker.lookup";
29
30#[derive(Debug, Default)]
31struct DeviceMemoryFailureTrackerRegistryState {
32    pending: HashMap<usize, DeviceMemoryFailureTracker>,
33    installed: HashMap<(usize, u64), DeviceMemoryFailureTracker>,
34}
35
36type DeviceMemoryFailureTrackerRegistry = Mutex<DeviceMemoryFailureTrackerRegistryState>;
37
38fn device_memory_failure_trackers() -> &'static DeviceMemoryFailureTrackerRegistry {
39    static TRACKERS: OnceLock<DeviceMemoryFailureTrackerRegistry> = OnceLock::new();
40    TRACKERS.get_or_init(|| Mutex::new(DeviceMemoryFailureTrackerRegistryState::default()))
41}
42
43/// Per-device allocation-failure counters observed by the device memory
44/// manager failure callback.
45#[derive(Debug, Clone)]
46pub(crate) struct DeviceMemoryFailureTracker {
47    inner: Arc<DeviceMemoryFailureTrackerInner>,
48}
49
50#[derive(Debug)]
51struct DeviceMemoryFailureTrackerInner {
52    device_ordinal: usize,
53    /// Zero until the tracker is bound to a manager installation; bound once
54    /// and never rebound to a different generation.
55    installation_generation: AtomicU64,
56    /// Whether the live callback updates this tracker. Mutually exclusive
57    /// with `callback_install_unavailable`; both false before installation.
58    callback_installed: AtomicBool,
59    callback_install_unavailable: AtomicBool,
60    allocation_failure_count: AtomicU64,
61    total_failed_allocation_bytes: AtomicU64,
62    largest_failed_allocation_bytes: AtomicU64,
63    pool_usage_observation_count: AtomicU64,
64    last_failed_allocation_bytes: AtomicU64,
65    last_pool_reserved_current_bytes: AtomicU64,
66    last_pool_used_current_bytes: AtomicU64,
67}
68
69/// Point-in-time device memory allocation-failure report.
70#[derive(Debug, Clone, PartialEq, Eq, Serialize)]
71pub struct DeviceMemoryFailureReport {
72    /// CUDA ordinal monitored by the installed allocation-failure callback.
73    pub device_ordinal: usize,
74    /// Whether the native memory manager updates this report's tracker.
75    pub callback_installed: bool,
76    /// Whether callback installation was unavailable; excludes `callback_installed`.
77    pub callback_install_unavailable: bool,
78    /// Allocation failures observed over the memory-manager installation lifetime.
79    pub allocation_failure_count: u64,
80    /// Sum of failed allocation request sizes in bytes.
81    pub total_failed_allocation_bytes: u64,
82    /// Largest failed allocation request size in bytes; zero before any failure.
83    pub largest_failed_allocation_bytes: u64,
84    /// Always zero because this memory-manager callback declines allocator retries.
85    pub trim_retry_count: u64,
86    /// Failure callbacks carrying a CUDA async-pool usage sample.
87    pub pool_usage_observation_count: u64,
88    /// Most recent failed request size in bytes; zero before any failure.
89    pub last_failed_allocation_bytes: u64,
90    /// Pool-reserved bytes from the last callback carrying a pool sample.
91    /// Zero before the first sample; consult `pool_usage_observation_count`.
92    pub last_pool_reserved_current_bytes: u64,
93    /// Pool-used bytes from the last callback carrying a pool sample.
94    /// Zero before the first sample; consult `pool_usage_observation_count`.
95    pub last_pool_used_current_bytes: u64,
96}
97
98impl DeviceMemoryFailureTracker {
99    /// Creates a zeroed tracker for one CUDA device ordinal.
100    pub(crate) fn new(device_ordinal: usize) -> Self {
101        Self {
102            inner: Arc::new(DeviceMemoryFailureTrackerInner {
103                device_ordinal,
104                installation_generation: AtomicU64::new(0),
105                callback_installed: AtomicBool::new(false),
106                callback_install_unavailable: AtomicBool::new(false),
107                allocation_failure_count: AtomicU64::new(0),
108                total_failed_allocation_bytes: AtomicU64::new(0),
109                largest_failed_allocation_bytes: AtomicU64::new(0),
110                pool_usage_observation_count: AtomicU64::new(0),
111                last_failed_allocation_bytes: AtomicU64::new(0),
112                last_pool_reserved_current_bytes: AtomicU64::new(0),
113                last_pool_used_current_bytes: AtomicU64::new(0),
114            }),
115        }
116    }
117
118    /// Recovers the tracker referenced by an already-installed callback, or
119    /// creates an unbound candidate only while the process manager is still
120    /// uninstalled. This prevents a fresh counter from being substituted for
121    /// the callback's original `(device, installation_generation)` owner.
122    pub(crate) fn for_installation_state(
123        device_ordinal: usize,
124        state: cudf::DeviceMemoryManagerInstallationState,
125    ) -> Result<Self> {
126        Self::for_installation_state_in_registry(device_ordinal, state, device_memory_failure_trackers())
127    }
128
129    fn for_installation_state_in_registry(
130        device_ordinal: usize,
131        state: cudf::DeviceMemoryManagerInstallationState,
132        trackers: &DeviceMemoryFailureTrackerRegistry,
133    ) -> Result<Self> {
134        let mut trackers = trackers.lock().map_err(|_error| tracker_registry_poisoned(device_ordinal))?;
135        match state.status {
136            cudf::DeviceMemoryManagerInstallationStatus::Uninstalled => {
137                if trackers.installed.keys().any(|(installed_device, _)| *installed_device == device_ordinal) {
138                    return Err(QueryError::new(
139                        QueryErrorCode::Internal,
140                        TRACKER_LOOKUP_OPERATION,
141                        "tracker remains bound while the device manager reports uninstalled",
142                    )
143                    .with_fact("device_ordinal", device_ordinal)
144                    .with_fact("reason", "tracker_generation_without_installation"));
145                }
146                let tracker = trackers.pending.entry(device_ordinal).or_insert_with(|| Self::new(device_ordinal));
147                if tracker.installation_generation().is_some() {
148                    return Err(QueryError::new(
149                        QueryErrorCode::Internal,
150                        TRACKER_LOOKUP_OPERATION,
151                        "pending tracker is bound while the device manager reports uninstalled",
152                    )
153                    .with_fact("device_ordinal", device_ordinal)
154                    .with_fact("reason", "pending_tracker_generation_without_installation"));
155                }
156                Ok(tracker.clone())
157            },
158            cudf::DeviceMemoryManagerInstallationStatus::Owned
159            | cudf::DeviceMemoryManagerInstallationStatus::Incomplete => {
160                let generation = state.installation_generation.ok_or_else(|| {
161                    QueryError::new(
162                        QueryErrorCode::Internal,
163                        TRACKER_LOOKUP_OPERATION,
164                        "owned device-manager callback has no installation generation",
165                    )
166                    .with_fact("device_ordinal", device_ordinal)
167                    .with_fact("reason", "installed_callback_generation_missing")
168                })?;
169                if let Some(tracker) = trackers.installed.get(&(device_ordinal, generation)).cloned() {
170                    tracker.require_installation_generation(generation)?;
171                    return Ok(tracker);
172                }
173                if let Some((_, existing_generation)) =
174                    trackers.installed.keys().find(|(installed_device, _)| *installed_device == device_ordinal).copied()
175                {
176                    return Err(QueryError::new(
177                        QueryErrorCode::CudfPersistentFailure,
178                        TRACKER_LOOKUP_OPERATION,
179                        "installed device-manager callback generation does not match the recoverable tracker",
180                    )
181                    .with_fact("device_ordinal", device_ordinal)
182                    .with_fact("expected_installation_generation", existing_generation)
183                    .with_fact("observed_installation_generation", generation)
184                    .with_fact("reason", "callback_tracker_generation_mismatch"));
185                }
186                let tracker = trackers.pending.remove(&device_ordinal).ok_or_else(|| {
187                    QueryError::new(
188                        QueryErrorCode::Internal,
189                        TRACKER_LOOKUP_OPERATION,
190                        "installed device-manager callback has no recoverable original tracker",
191                    )
192                    .with_fact("device_ordinal", device_ordinal)
193                    .with_fact("installation_generation", generation)
194                    .with_fact("reason", "installed_callback_tracker_missing")
195                })?;
196                if state.status == cudf::DeviceMemoryManagerInstallationStatus::Incomplete
197                    && tracker.installation_generation().is_none()
198                {
199                    // The native/Rust cuDF installation record commits callback
200                    // identity before attempting the current-resource switch.
201                    // A retryable switch failure therefore binds the original
202                    // pre-install candidate here before a later owner retries.
203                    tracker.bind_installation_generation(generation)?;
204                } else {
205                    tracker.require_installation_generation(generation)?;
206                }
207                trackers.installed.insert((device_ordinal, generation), tracker.clone());
208                Ok(tracker)
209            },
210            _ => Err(QueryError::new(
211                QueryErrorCode::CudfPersistentFailure,
212                TRACKER_LOOKUP_OPERATION,
213                "device-manager callback installation is permanently unusable",
214            )
215            .with_fact("device_ordinal", device_ordinal)
216            .with_fact("installation_status", format!("{:?}", state.status))
217            .with_fact("installation_retry", format!("{:?}", state.retry))
218            .with_fact("installation_cause", format!("{:?}", state.cause))
219            .with_fact("reason", "installed_callback_unrecoverable")),
220        }
221    }
222
223    #[cfg(test)]
224    fn for_device_in_registry(device_ordinal: usize, trackers: &DeviceMemoryFailureTrackerRegistry) -> Result<Self> {
225        let mut trackers = trackers.lock().map_err(|_error| tracker_registry_poisoned(device_ordinal))?;
226        if let Some(tracker) = trackers
227            .installed
228            .iter()
229            .find_map(|((installed_device, _), tracker)| (*installed_device == device_ordinal).then(|| tracker.clone()))
230        {
231            return Ok(tracker);
232        }
233        Ok(trackers.pending.entry(device_ordinal).or_insert_with(|| Self::new(device_ordinal)).clone())
234    }
235
236    /// Moves this tracker from the pending to the installed registry under
237    /// `generation`, after a successful manager installation. Idempotent for
238    /// the same tracker; a different tracker already bound to the generation,
239    /// or a pending candidate that is not this tracker, is refused so the
240    /// callback and the reported counters cannot diverge.
241    pub(crate) fn commit_installation_generation(&self, generation: u64) -> Result<()> {
242        self.bind_installation_generation(generation)?;
243        let device_ordinal = self.device_ordinal();
244        let mut trackers =
245            device_memory_failure_trackers().lock().map_err(|_error| tracker_registry_poisoned(device_ordinal))?;
246        if let Some(existing) = trackers.installed.get(&(device_ordinal, generation)) {
247            if Arc::ptr_eq(&existing.inner, &self.inner) {
248                return Ok(());
249            }
250            return Err(QueryError::new(
251                QueryErrorCode::CudfPersistentFailure,
252                TRACKER_LOOKUP_OPERATION,
253                "installation generation is already bound to a different callback tracker",
254            )
255            .with_fact("device_ordinal", device_ordinal)
256            .with_fact("installation_generation", generation)
257            .with_fact("reason", "callback_tracker_identity_mismatch"));
258        }
259        let pending = trackers.pending.remove(&device_ordinal).ok_or_else(|| {
260            QueryError::new(
261                QueryErrorCode::CudfPersistentFailure,
262                TRACKER_LOOKUP_OPERATION,
263                "installed callback tracker was not the registered pre-install candidate",
264            )
265            .with_fact("device_ordinal", device_ordinal)
266            .with_fact("installation_generation", generation)
267            .with_fact("reason", "pending_callback_tracker_missing")
268        })?;
269        if !Arc::ptr_eq(&pending.inner, &self.inner) {
270            trackers.pending.insert(device_ordinal, pending);
271            return Err(QueryError::new(
272                QueryErrorCode::CudfPersistentFailure,
273                TRACKER_LOOKUP_OPERATION,
274                "pending callback tracker does not match the installed callback tracker",
275            )
276            .with_fact("device_ordinal", device_ordinal)
277            .with_fact("installation_generation", generation)
278            .with_fact("reason", "callback_tracker_identity_mismatch"));
279        }
280        trackers.installed.insert((device_ordinal, generation), self.clone());
281        Ok(())
282    }
283
284    pub(crate) fn bind_installation_generation(&self, generation: u64) -> Result<()> {
285        if generation == 0 {
286            return Err(QueryError::new(
287                QueryErrorCode::Internal,
288                TRACKER_LOOKUP_OPERATION,
289                "callback installation generation must be nonzero",
290            )
291            .with_fact("device_ordinal", self.device_ordinal())
292            .with_fact("reason", "zero_installation_generation"));
293        }
294        match self.inner.installation_generation.compare_exchange(0, generation, Ordering::AcqRel, Ordering::Acquire) {
295            Ok(_) => Ok(()),
296            Err(existing) if existing == generation => Ok(()),
297            Err(existing) => Err(QueryError::new(
298                QueryErrorCode::CudfPersistentFailure,
299                TRACKER_LOOKUP_OPERATION,
300                "callback tracker is bound to a different installation generation",
301            )
302            .with_fact("device_ordinal", self.device_ordinal())
303            .with_fact("expected_installation_generation", existing)
304            .with_fact("observed_installation_generation", generation)
305            .with_fact("reason", "callback_tracker_generation_mismatch")),
306        }
307    }
308
309    fn require_installation_generation(&self, generation: u64) -> Result<()> {
310        match self.installation_generation() {
311            Some(existing) if existing == generation => Ok(()),
312            Some(existing) => Err(QueryError::new(
313                QueryErrorCode::CudfPersistentFailure,
314                TRACKER_LOOKUP_OPERATION,
315                "installed callback generation does not match its original tracker",
316            )
317            .with_fact("device_ordinal", self.device_ordinal())
318            .with_fact("expected_installation_generation", existing)
319            .with_fact("observed_installation_generation", generation)
320            .with_fact("reason", "callback_tracker_generation_mismatch")),
321            None => Err(QueryError::new(
322                QueryErrorCode::CudfPersistentFailure,
323                TRACKER_LOOKUP_OPERATION,
324                "installed callback's tracker is not bound to an installation generation",
325            )
326            .with_fact("device_ordinal", self.device_ordinal())
327            .with_fact("observed_installation_generation", generation)
328            .with_fact("reason", "callback_tracker_generation_unbound")),
329        }
330    }
331
332    fn installation_generation(&self) -> Option<u64> {
333        match self.inner.installation_generation.load(Ordering::Acquire) {
334            0 => None,
335            generation => Some(generation),
336        }
337    }
338
339    /// CUDA device ordinal this tracker describes.
340    pub fn device_ordinal(&self) -> usize {
341        self.inner.device_ordinal
342    }
343
344    /// Returns a re-entrancy-safe callback for cudf/RMM.
345    ///
346    /// The callback records counters without deciding allocation retry or
347    /// reclamation. It must remain allocation-free and lock-free because RMM
348    /// invokes it inside allocation failure handling.
349    pub(crate) fn failure_callback(&self) -> Arc<dyn Fn(cudf::DeviceMemoryFailureInfo) + Send + Sync + 'static> {
350        let tracker = self.clone();
351        Arc::new(move |info| {
352            if info.source != DeviceMemoryFailureSource::Physical {
353                return;
354            }
355            tracker.record_allocation_failure(info.requested_bytes, info.pool_usage);
356        })
357    }
358
359    pub(crate) fn mark_callback_installed(&self) {
360        self.inner.callback_installed.store(true, Ordering::Release);
361        self.inner.callback_install_unavailable.store(false, Ordering::Release);
362    }
363
364    pub(crate) fn mark_callback_install_unavailable(&self) {
365        self.inner.callback_installed.store(false, Ordering::Release);
366        self.inner.callback_install_unavailable.store(true, Ordering::Release);
367    }
368
369    fn record_allocation_failure(&self, bytes: u64, pool_usage: Option<DeviceMemoryPoolUsage>) {
370        self.inner.allocation_failure_count.fetch_add(1, Ordering::Relaxed);
371        self.inner.total_failed_allocation_bytes.fetch_add(bytes, Ordering::Relaxed);
372        self.inner.last_failed_allocation_bytes.store(bytes, Ordering::Relaxed);
373        update_max(&self.inner.largest_failed_allocation_bytes, bytes);
374        if let Some(pool_usage) = pool_usage {
375            self.inner.pool_usage_observation_count.fetch_add(1, Ordering::Relaxed);
376            self.inner.last_pool_reserved_current_bytes.store(pool_usage.reserved_current_bytes, Ordering::Relaxed);
377            self.inner.last_pool_used_current_bytes.store(pool_usage.used_current_bytes, Ordering::Relaxed);
378        }
379    }
380
381    pub(crate) fn report(&self) -> DeviceMemoryFailureReport {
382        DeviceMemoryFailureReport {
383            device_ordinal: self.inner.device_ordinal,
384            callback_installed: self.inner.callback_installed.load(Ordering::Acquire),
385            callback_install_unavailable: self.inner.callback_install_unavailable.load(Ordering::Acquire),
386            allocation_failure_count: self.inner.allocation_failure_count.load(Ordering::Acquire),
387            total_failed_allocation_bytes: self.inner.total_failed_allocation_bytes.load(Ordering::Acquire),
388            largest_failed_allocation_bytes: self.inner.largest_failed_allocation_bytes.load(Ordering::Acquire),
389            trim_retry_count: 0,
390            pool_usage_observation_count: self.inner.pool_usage_observation_count.load(Ordering::Acquire),
391            last_failed_allocation_bytes: self.inner.last_failed_allocation_bytes.load(Ordering::Acquire),
392            last_pool_reserved_current_bytes: self.inner.last_pool_reserved_current_bytes.load(Ordering::Acquire),
393            last_pool_used_current_bytes: self.inner.last_pool_used_current_bytes.load(Ordering::Acquire),
394        }
395    }
396}
397
398fn tracker_registry_poisoned(device_ordinal: usize) -> QueryError {
399    QueryError::new(
400        QueryErrorCode::Internal,
401        TRACKER_LOOKUP_OPERATION,
402        "device memory failure-tracker registry mutex was poisoned",
403    )
404    .with_fact("device_ordinal", device_ordinal)
405    .with_fact("reason", "tracker_registry_poisoned")
406}
407
408fn update_max(target: &AtomicU64, candidate: u64) {
409    let mut current = target.load(Ordering::Relaxed);
410    while candidate > current {
411        match target.compare_exchange_weak(current, candidate, Ordering::Relaxed, Ordering::Relaxed) {
412            Ok(_) => return,
413            Err(next) => current = next,
414        }
415    }
416}
417
418#[cfg(test)]
419mod tests {
420    use std::panic::{AssertUnwindSafe, catch_unwind};
421
422    use super::*;
423
424    #[test]
425    fn failure_callback_records_counters_and_declines_retry() {
426        let tracker = DeviceMemoryFailureTracker::new(3);
427
428        tracker.record_allocation_failure(10, None);
429        tracker.record_allocation_failure(
430            40,
431            Some(DeviceMemoryPoolUsage {
432                reserved_current_bytes: 400,
433                reserved_high_bytes: 800,
434                used_current_bytes: 100,
435                used_high_bytes: 700,
436            }),
437        );
438        tracker.record_allocation_failure(20, None);
439
440        let report = tracker.report();
441        assert_eq!(report.device_ordinal, 3);
442        assert!(!report.callback_installed);
443        assert_eq!(report.allocation_failure_count, 3);
444        assert_eq!(report.total_failed_allocation_bytes, 70);
445        assert_eq!(report.largest_failed_allocation_bytes, 40);
446        assert_eq!(report.trim_retry_count, 0);
447        assert_eq!(report.pool_usage_observation_count, 1);
448        assert_eq!(report.last_failed_allocation_bytes, 20);
449        assert_eq!(report.last_pool_reserved_current_bytes, 400);
450        assert_eq!(report.last_pool_used_current_bytes, 100);
451    }
452
453    #[test]
454    fn failure_callback_declines_reservation_limit_without_recording_device_failure() {
455        let tracker = DeviceMemoryFailureTracker::new(3);
456        let callback = tracker.failure_callback();
457
458        callback(cudf::DeviceMemoryFailureInfo {
459            device: cudf::DeviceId::DEFAULT,
460            source: DeviceMemoryFailureSource::ReservationLimit,
461            allocation_domain: None,
462            requested_bytes: 10,
463            callback_attempt: 0,
464            trim_retry_count: 0,
465            pool_usage: Some(DeviceMemoryPoolUsage {
466                reserved_current_bytes: 1_000,
467                reserved_high_bytes: 1_000,
468                used_current_bytes: 0,
469                used_high_bytes: 0,
470            }),
471        });
472
473        let report = tracker.report();
474        assert_eq!(report.allocation_failure_count, 0);
475        assert_eq!(report.trim_retry_count, 0);
476        assert_eq!(report.pool_usage_observation_count, 0);
477    }
478
479    #[test]
480    fn failure_callback_observes_physical_failure_with_idle_pool() {
481        let tracker = DeviceMemoryFailureTracker::new(3);
482        let callback = tracker.failure_callback();
483
484        callback(cudf::DeviceMemoryFailureInfo {
485            device: cudf::DeviceId::DEFAULT,
486            source: DeviceMemoryFailureSource::Physical,
487            allocation_domain: None,
488            requested_bytes: 10,
489            callback_attempt: 0,
490            trim_retry_count: 0,
491            pool_usage: Some(DeviceMemoryPoolUsage {
492                reserved_current_bytes: 1_000,
493                reserved_high_bytes: 1_000,
494                used_current_bytes: 0,
495                used_high_bytes: 0,
496            }),
497        });
498
499        let report = tracker.report();
500        assert_eq!(report.allocation_failure_count, 1);
501        assert_eq!(report.trim_retry_count, 0);
502        assert_eq!(report.pool_usage_observation_count, 1);
503        assert_eq!(report.last_failed_allocation_bytes, 10);
504    }
505
506    #[test]
507    fn compatible_generation_lookup_recovers_original_callback_tracker() {
508        let trackers = Mutex::new(DeviceMemoryFailureTrackerRegistryState::default());
509        let first_generation = DeviceMemoryFailureTracker::for_device_in_registry(3, &trackers).unwrap();
510        let original_identity = Arc::downgrade(&first_generation.inner);
511        let original_callback = first_generation.failure_callback();
512        drop(first_generation);
513
514        let second_generation = DeviceMemoryFailureTracker::for_device_in_registry(3, &trackers).unwrap();
515        let original_inner = original_identity.upgrade().expect("registry must retain the original tracker");
516        assert!(Arc::ptr_eq(&original_inner, &second_generation.inner));
517
518        original_callback(cudf::DeviceMemoryFailureInfo {
519            device: cudf::DeviceId::DEFAULT,
520            source: DeviceMemoryFailureSource::Physical,
521            allocation_domain: None,
522            requested_bytes: 64,
523            callback_attempt: 0,
524            trim_retry_count: 0,
525            pool_usage: None,
526        });
527
528        let report = second_generation.report();
529        assert_eq!(report.allocation_failure_count, 1);
530        assert_eq!(report.total_failed_allocation_bytes, 64);
531        assert_eq!(report.last_failed_allocation_bytes, 64);
532    }
533
534    #[test]
535    fn retryable_incomplete_install_binds_the_original_preinstall_tracker() {
536        let trackers = Mutex::new(DeviceMemoryFailureTrackerRegistryState::default());
537        let original = DeviceMemoryFailureTracker::for_device_in_registry(3, &trackers).unwrap();
538        assert_eq!(original.installation_generation(), None);
539        let state = cudf::DeviceMemoryManagerInstallationState {
540            device_ordinal: 3,
541            status: cudf::DeviceMemoryManagerInstallationStatus::Incomplete,
542            retry: cudf::DeviceMemoryManagerInstallationRetry::Retryable,
543            cause: cudf::DeviceMemoryManagerInstallationCause::CurrentResourceInstallFailed,
544            effective_options: Some(cudf::DeviceMemoryManagerOptions::default()),
545            identity: Some(cudf::DeviceMemoryManagerInstallationIdentity::new([1; 16], 1, [2; 16])),
546            installation_generation: Some(41),
547            failure_callback_installed: true,
548            current_resource_installed: false,
549            failure_callback_panicked: false,
550        };
551
552        let recovered = DeviceMemoryFailureTracker::for_installation_state_in_registry(3, state, &trackers).unwrap();
553
554        assert!(Arc::ptr_eq(&original.inner, &recovered.inner));
555        assert_eq!(recovered.installation_generation(), Some(41));
556        let registry = trackers.lock().unwrap();
557        assert!(!registry.pending.contains_key(&3));
558        assert!(registry.installed.get(&(3, 41)).is_some_and(|tracker| Arc::ptr_eq(&tracker.inner, &original.inner)));
559    }
560
561    #[test]
562    fn tracker_registry_poison_fails_closed() {
563        let trackers = Mutex::new(DeviceMemoryFailureTrackerRegistryState::default());
564        let poison = catch_unwind(AssertUnwindSafe(|| {
565            let _guard = trackers.lock().unwrap();
566            panic!("poison test tracker registry");
567        }));
568        assert!(poison.is_err());
569
570        let error = DeviceMemoryFailureTracker::for_device_in_registry(7, &trackers).unwrap_err();
571
572        assert_eq!(error.code(), QueryErrorCode::Internal);
573        assert_eq!(error.operation(), TRACKER_LOOKUP_OPERATION);
574        assert_eq!(error.fact_value("device_ordinal"), Some("7"));
575        assert_eq!(error.fact_value("reason"), Some("tracker_registry_poisoned"));
576    }
577}
query_runtime/runtime/gpu/memory_failure.rs

query_runtime/runtime/gpu/
memory_failure.rs
