Coverage for cuda/core/_host.py: 100.00%

58 statements  

« prev     ^ index     » next       coverage.py v7.16.0, created at 2026-09-10 02:27 +0000

1# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. 

2# SPDX-License-Identifier: Apache-2.0 

3 

4from __future__ import annotations 

5 

6import threading 

7from typing import ClassVar 

8 

9__all__ = ["Host"] 

10 

11 

12class Host: 

13 """Host (CPU) location for managed-memory operations. 

14 

15 Use one of the following forms: 

16 

17 * ``Host()`` — generic host (any NUMA node). 

18 * ``Host(numa_id=N)`` — specific NUMA node ``N``. 

19 * ``Host.numa_current()`` or ``Host(is_numa_current=True)`` — NUMA node 

20 of the calling thread. ``numa_id`` and ``is_numa_current`` are 

21 mutually exclusive. 

22 

23 ``Host`` is the symmetric counterpart of :class:`~cuda.core.Device` 

24 for managed-memory `prefetch`, `advise`, and `discard_prefetch` 

25 targets. Pass either a ``Device`` or a ``Host`` to those operations 

26 and to ``ManagedBuffer.preferred_location`` / ``accessed_by``. 

27 

28 ``Host`` is a singleton class, mirroring :class:`~cuda.core.Device`: 

29 constructor calls with the same arguments return the same instance, 

30 so ``Host() is Host()`` and ``Host(numa_id=1) is Host(numa_id=1)``. 

31 ``Host.numa_current()`` returns its own singleton, distinct from 

32 ``Host()`` because it represents a thread-relative location rather 

33 than a fixed one. 

34 """ 

35 

36 __slots__ = ("__weakref__", "_is_numa_current", "_numa_id") 

37 

38 _numa_id: int | None 

39 _is_numa_current: bool 

40 

41 # Singleton cache keyed by (numa_id, is_numa_current). 

42 _instances: ClassVar[dict[tuple[int | None, bool], Host]] = {} 

43 _instances_lock: ClassVar[threading.Lock] = threading.Lock() 

44 

45 def __new__(cls, numa_id: int | None = None, *, is_numa_current: bool = False) -> Host: 

46 if is_numa_current and numa_id is not None: 1rstijkuefgnmpvwxyozCDEFAdlGBqHhcb

47 raise ValueError("numa_id and is_numa_current are mutually exclusive") 1B

48 if numa_id is not None and (isinstance(numa_id, bool) or not isinstance(numa_id, int) or numa_id < 0): 1rstijkuefgnmpvwxyozCDEFAdlGBqHhcb

49 raise ValueError(f"numa_id must be a non-negative int, got {numa_id!r}") 1GH

50 return cls._get_or_create(numa_id, is_numa_current) 1rstijkuefgnmpvwxyozCDEFAdlBqhcb

51 

52 @classmethod 

53 def _get_or_create(cls, numa_id: int | None, is_numa_current: bool) -> Host: 

54 key = (numa_id, is_numa_current) 1rstijkuefgnmpvwxyozCDEFAdlBqhcb

55 cache = cls._instances 1rstijkuefgnmpvwxyozCDEFAdlBqhcb

56 inst = cache.get(key) 1rstijkuefgnmpvwxyozCDEFAdlBqhcb

57 if inst is not None: 1rstijkuefgnmpvwxyozCDEFAdlBqhcb

58 return inst 1rstuefgnmpvwxyozCDEFAdlBqhcb

59 with cls._instances_lock: 1ijkefgdlhcb

60 inst = cache.get(key) 1ijkefgdlhcb

61 if inst is None: 1ijkefgdlhcb

62 inst = object.__new__(cls) 1ijkefgdlhcb

63 inst._numa_id = numa_id 1ijkefgdlhcb

64 inst._is_numa_current = is_numa_current 1ijkefgdlhcb

65 cache[key] = inst 1ijkefgdlhcb

66 return inst 1ijkefgdlhcb

67 

68 @property 

69 def numa_id(self) -> int | None: 

70 """NUMA node ID, or ``None`` if not pinned to a specific NUMA node.""" 

71 return self._numa_id 1stjkufgnmpvwxyozAqhb

72 

73 @property 

74 def is_numa_current(self) -> bool: 

75 """``True`` if this ``Host`` represents the calling thread's NUMA node (constructed via :meth:`numa_current`).""" 

76 return self._is_numa_current 1rstijkuefgmpvwxyozAqhb

77 

78 @classmethod 

79 def numa_current(cls) -> Host: 

80 """Construct a ``Host`` referring to the calling thread's NUMA node.""" 

81 return cls(is_numa_current=True) 1riepdBqcb

82 

83 def __eq__(self, other: object) -> bool: 

84 if not isinstance(other, Host): 1nmod

85 return NotImplemented 1nm

86 return self is other 1od

87 

88 def __hash__(self) -> int: 

89 return hash((Host, self._numa_id, self._is_numa_current)) 1nmd

90 

91 def __reduce__(self) -> tuple[object, ...]: 

92 if self._is_numa_current: 1c

93 return (_reconstruct_numa_current, ()) 1c

94 return (Host, (self._numa_id,)) 1c

95 

96 def __repr__(self) -> str: 

97 if self.is_numa_current: 1b

98 return "Host.numa_current()" 1b

99 if self.numa_id is None: 1b

100 return "Host()" 1b

101 return f"Host(numa_id={self.numa_id})" 1b

102 

103 

104def _reconstruct_numa_current() -> Host: 

105 return Host.numa_current() 1c