Coverage for cuda/core/_host.py: 100.00%
57 statements
« prev ^ index » next coverage.py v7.14.1, created at 2026-06-13 01:38 +0000
1# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
2# SPDX-License-Identifier: Apache-2.0
4from __future__ import annotations
6import threading
7from typing import ClassVar
10class Host:
11 """Host (CPU) location for managed-memory operations.
13 Use one of the following forms:
15 * ``Host()`` — generic host (any NUMA node).
16 * ``Host(numa_id=N)`` — specific NUMA node ``N``.
17 * ``Host.numa_current()`` or ``Host(is_numa_current=True)`` — NUMA node
18 of the calling thread. ``numa_id`` and ``is_numa_current`` are
19 mutually exclusive.
21 ``Host`` is the symmetric counterpart of :class:`~cuda.core.Device`
22 for managed-memory `prefetch`, `advise`, and `discard_prefetch`
23 targets. Pass either a ``Device`` or a ``Host`` to those operations
24 and to ``ManagedBuffer.preferred_location`` / ``accessed_by``.
26 ``Host`` is a singleton class, mirroring :class:`~cuda.core.Device`:
27 constructor calls with the same arguments return the same instance,
28 so ``Host() is Host()`` and ``Host(numa_id=1) is Host(numa_id=1)``.
29 ``Host.numa_current()`` returns its own singleton, distinct from
30 ``Host()`` because it represents a thread-relative location rather
31 than a fixed one.
32 """
34 __slots__ = ("__weakref__", "_is_numa_current", "_numa_id")
36 _numa_id: int | None
37 _is_numa_current: bool
39 # Singleton cache keyed by (numa_id, is_numa_current).
40 _instances: ClassVar[dict[tuple[int | None, bool], Host]] = {}
41 _instances_lock: ClassVar[threading.Lock] = threading.Lock()
43 def __new__(cls, numa_id: int | None = None, *, is_numa_current: bool = False) -> Host:
44 if is_numa_current and numa_id is not None: 1lfgmienopqjhuvwxrdyzskAtca
45 raise ValueError("numa_id and is_numa_current are mutually exclusive") 1s
46 if numa_id is not None and (isinstance(numa_id, bool) or not isinstance(numa_id, int) or numa_id < 0): 1lfgmienopqjhuvwxrdyzskAtca
47 raise ValueError(f"numa_id must be a non-negative int, got {numa_id!r}") 1zA
48 return cls._get_or_create(numa_id, is_numa_current) 1lfgmienopqjhuvwxrdysktca
50 @classmethod
51 def _get_or_create(cls, numa_id: int | None, is_numa_current: bool) -> Host:
52 key = (numa_id, is_numa_current) 1lfgmienopqjhuvwxrdysktca
53 cache = cls._instances 1lfgmienopqjhuvwxrdysktca
54 inst = cache.get(key) 1lfgmienopqjhuvwxrdysktca
55 if inst is not None: 1lfgmienopqjhuvwxrdysktca
56 return inst 1lmienopqjuvwxrdysktca
57 with cls._instances_lock: 1fgehdca
58 inst = cache.get(key) 1fgehdca
59 if inst is None: 1fgehdca
60 inst = object.__new__(cls) 1fgehdca
61 inst._numa_id = numa_id 1fgehdca
62 inst._is_numa_current = is_numa_current 1fgehdca
63 cache[key] = inst 1fgehdca
64 return inst 1fgehdca
66 @property
67 def numa_id(self) -> int | None:
68 """NUMA node ID, or ``None`` if not pinned to a specific NUMA node."""
69 return self._numa_id 1lgmienopqjhrkta
71 @property
72 def is_numa_current(self) -> bool:
73 """``True`` if this ``Host`` represents the calling thread's NUMA node (constructed via :meth:`numa_current`)."""
74 return self._is_numa_current 1lfgmenopqjhrkta
76 @classmethod
77 def numa_current(cls) -> Host:
78 """Construct a ``Host`` referring to the calling thread's NUMA node."""
79 return cls(is_numa_current=True) 1fedskca
81 def __eq__(self, other: object) -> bool:
82 if not isinstance(other, Host): 1ijd
83 return NotImplemented 1i
84 return self is other 1jd
86 def __hash__(self) -> int:
87 return hash((Host, self._numa_id, self._is_numa_current)) 1id
89 def __reduce__(self) -> tuple[object, ...]:
90 if self._is_numa_current: 1c
91 return (_reconstruct_numa_current, ()) 1c
92 return (Host, (self._numa_id,)) 1c
94 def __repr__(self) -> str:
95 if self.is_numa_current: 1a
96 return "Host.numa_current()" 1a
97 if self.numa_id is None: 1a
98 return "Host()" 1a
99 return f"Host(numa_id={self.numa_id})" 1a
102def _reconstruct_numa_current() -> Host:
103 return Host.numa_current() 1c