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

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``. A 

27 ``Host`` may also be returned by ``ManagedBuffer.last_prefetch_location``. 

28 

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

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

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

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

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

34 than a fixed one. 

35 """ 

36 

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

38 

39 _numa_id: int | None 

40 _is_numa_current: bool 

41 

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

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

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

45 

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

47 if is_numa_current and numa_id is not None: 36 ctx1mnopqrstuvgCflwhixyzjkDEFGAcHIBdJeba

48 raise ValueError("numa_id and is_numa_current are mutually exclusive") 1 ctx1B

49 if numa_id is not None and (isinstance(numa_id, bool) or not isinstance(numa_id, int) or numa_id < 0): 36 ctx1mnopqrstuvgCflwhixyzjkDEFGAcHIBdJeba

50 raise ValueError(f"numa_id must be a non-negative int, got {numa_id!r}") 2 ctx1IJ

51 return cls._get_or_create(numa_id, is_numa_current) 34 ctx1mnopqrstuvgCflwhixyzjkDEFGAcHBdeba

52 

53 @classmethod 

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

55 key = (numa_id, is_numa_current) 34 ctx1mnopqrstuvgCflwhixyzjkDEFGAcHBdeba

56 cache = cls._instances 34 ctx1mnopqrstuvgCflwhixyzjkDEFGAcHBdeba

57 inst = cache.get(key) 34 ctx1mnopqrstuvgCflwhixyzjkDEFGAcHBdeba

58 if inst is not None: 34 ctx1mnopqrstuvgCflwhixyzjkDEFGAcHBdeba

59 return inst 33 ctx1mnopqrstuvgCflwhixyzjkDEFGAcHBdba

60 with cls._instances_lock: 5 ctx1cdeba

61 inst = cache.get(key) 5 ctx1cdeba

62 if inst is None: 5 ctx1cdeba

63 inst = object.__new__(cls) 5 ctx1cdeba

64 inst._numa_id = numa_id 5 ctx1cdeba

65 inst._is_numa_current = is_numa_current 5 ctx1cdeba

66 cache[key] = inst 5 ctx1cdeba

67 return inst 5 ctx1cdeba

68 

69 @property 

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

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

72 return self._numa_id 22 ctx1noqrsuvgflwhixyzjkAdea

73 

74 @property 

75 def is_numa_current(self) -> bool: 

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

77 return self._is_numa_current 24 ctx1mnopqrstuvflwhixyzjkAdea

78 

79 @classmethod 

80 def numa_current(cls) -> Host: 

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

82 return cls(is_numa_current=True) 9 ctx1mptlcBdba

83 

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

85 if not isinstance(other, Host): 7 ctx1gfhijkc

86 return NotImplemented 2 ctx1gf

87 return self is other 5 ctx1hijkc

88 

89 def __hash__(self) -> int: 

90 return hash((Host, self._numa_id, self._is_numa_current)) 4 ctx1gCfc

91 

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

93 if self._is_numa_current: 1 ctx1b

94 return (_reconstruct_numa_current, ()) 1 ctx1b

95 return (Host, (self._numa_id,)) 1 ctx1b

96 

97 def __repr__(self) -> str: 

98 if self.is_numa_current: 1 ctx1a

99 return "Host.numa_current()" 1 ctx1a

100 if self.numa_id is None: 1 ctx1a

101 return "Host()" 1 ctx1a

102 return f"Host(numa_id={self.numa_id})" 1 ctx1a

103 

104 

105def _reconstruct_numa_current() -> Host: 

106 return Host.numa_current() 1 ctx1b