Coverage for cuda / core / _host.py: 83.64%
55 statements
« prev ^ index » next coverage.py v7.14.0, created at 2026-05-22 01:37 +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 # Singleton cache keyed by (numa_id, is_numa_current).
37 _instances: ClassVar[dict[tuple[int | None, bool], Host]] = {}
38 _instances_lock: ClassVar[threading.Lock] = threading.Lock()
40 def __new__(cls, numa_id: int | None = None, *, is_numa_current: bool = False) -> Host:
41 if is_numa_current and numa_id is not None: 1ijklfhmnopgqtuvwrbexscyd
42 raise ValueError("numa_id and is_numa_current are mutually exclusive") 1s
43 if numa_id is not None and (isinstance(numa_id, bool) or not isinstance(numa_id, int) or numa_id < 0): 1ijklfhmnopgqtuvwrbexscyd
44 raise ValueError(f"numa_id must be a non-negative int, got {numa_id!r}") 1xy
45 return cls._get_or_create(numa_id, is_numa_current) 1ijklfhmnopgqtuvwrbescd
47 @classmethod
48 def _get_or_create(cls, numa_id: int | None, is_numa_current: bool) -> Host:
49 key = (numa_id, is_numa_current) 1ijklfhmnopgqtuvwrbescd
50 cache = cls._instances 1ijklfhmnopgqtuvwrbescd
51 inst = cache.get(key) 1ijklfhmnopgqtuvwrbescd
52 if inst is not None: 1ijklfhmnopgqtuvwrbescd
53 return inst 1ijklfhmnopgqtuvwrbs
54 with cls._instances_lock: 1becd
55 inst = cache.get(key) 1becd
56 if inst is None: 1becd
57 inst = object.__new__(cls) 1becd
58 inst._numa_id = numa_id 1becd
59 inst._is_numa_current = is_numa_current 1becd
60 cache[key] = inst 1becd
61 return inst 1becd
63 @property
64 def numa_id(self) -> int | None:
65 return self._numa_id 1iklfhmnopgqrcd
67 @property
68 def is_numa_current(self) -> bool:
69 return self._is_numa_current 1ijklhmnopgqrcd
71 @classmethod
72 def numa_current(cls) -> Host:
73 """Construct a ``Host`` referring to the calling thread's NUMA node."""
74 return cls(is_numa_current=True) 1jhbsc
76 def __eq__(self, other) -> bool:
77 if not isinstance(other, Host): 1fgb
78 return NotImplemented 1f
79 return self is other 1gb
81 def __hash__(self) -> int:
82 return hash((Host, self._numa_id, self._is_numa_current)) 1fb
84 def __reduce__(self):
85 if self._is_numa_current:
86 return (_reconstruct_numa_current, ())
87 return (Host, (self._numa_id,))
89 def __repr__(self) -> str:
90 if self.is_numa_current:
91 return "Host.numa_current()"
92 if self.numa_id is None:
93 return "Host()"
94 return f"Host(numa_id={self.numa_id})"
97def _reconstruct_numa_current() -> Host:
98 return Host.numa_current()