diff --git a/nvitop/api/device.py b/nvitop/api/device.py index 6bdc546a..61d01ccf 100644 --- a/nvitop/api/device.py +++ b/nvitop/api/device.py @@ -189,7 +189,14 @@ def receive(self) -> int | NaType: return self.rx -_VALUE_OMITTED: str = object() # type: ignore[assignment] +# pylint: disable-next=missing-class-docstring,too-few-public-methods +class ValueOmitted: + def __repr__(self) -> str: + return '' + + +__VALUE_OMITTED: str = ValueOmitted() # type: ignore[assignment] +del ValueOmitted class Device: # pylint: disable=too-many-instance-attributes,too-many-public-methods @@ -499,7 +506,7 @@ def from_cuda_indices( @staticmethod def parse_cuda_visible_devices( - cuda_visible_devices: str | None = _VALUE_OMITTED, + cuda_visible_devices: str | None = __VALUE_OMITTED, ) -> list[int] | list[tuple[int, int]]: """Parse the given ``CUDA_VISIBLE_DEVICES`` value into a list of NVML device indices. @@ -525,7 +532,7 @@ def parse_cuda_visible_devices( return parse_cuda_visible_devices(cuda_visible_devices) @staticmethod - def normalize_cuda_visible_devices(cuda_visible_devices: str | None = _VALUE_OMITTED) -> str: + def normalize_cuda_visible_devices(cuda_visible_devices: str | None = __VALUE_OMITTED) -> str: """Parse the given ``CUDA_VISIBLE_DEVICES`` value and convert it into a comma-separated string of UUIDs. This is an alias of :func:`normalize_cuda_visible_devices`. @@ -2778,7 +2785,7 @@ def is_mig_device_uuid(uuid: str | None) -> bool: def parse_cuda_visible_devices( - cuda_visible_devices: str | None = _VALUE_OMITTED, + cuda_visible_devices: str | None = __VALUE_OMITTED, ) -> list[int] | list[tuple[int, int]]: """Parse the given ``CUDA_VISIBLE_DEVICES`` value into a list of NVML device indices. @@ -2831,13 +2838,13 @@ def parse_cuda_visible_devices( >>> parse_cuda_visible_devices('16') # invalid `CUDA_VISIBLE_DEVICES` (device ordinal out of range) [] """ # pylint: disable=line-too-long - if cuda_visible_devices is _VALUE_OMITTED: + if cuda_visible_devices is __VALUE_OMITTED: cuda_visible_devices = os.getenv('CUDA_VISIBLE_DEVICES', default=None) return _parse_cuda_visible_devices(cuda_visible_devices, format='index') -def normalize_cuda_visible_devices(cuda_visible_devices: str | None = _VALUE_OMITTED) -> str: +def normalize_cuda_visible_devices(cuda_visible_devices: str | None = __VALUE_OMITTED) -> str: """Parse the given ``CUDA_VISIBLE_DEVICES`` value and convert it into a comma-separated string of UUIDs. This function is aliased by :meth:`Device.normalize_cuda_visible_devices`. @@ -2888,7 +2895,7 @@ def normalize_cuda_visible_devices(cuda_visible_devices: str | None = _VALUE_OMI >>> normalize_cuda_visible_devices('16') # invalid `CUDA_VISIBLE_DEVICES` (device ordinal out of range) '' """ # pylint: disable=line-too-long - if cuda_visible_devices is _VALUE_OMITTED: + if cuda_visible_devices is __VALUE_OMITTED: cuda_visible_devices = os.getenv('CUDA_VISIBLE_DEVICES', default=None) return ','.join(_parse_cuda_visible_devices(cuda_visible_devices, format='uuid')) @@ -2904,14 +2911,17 @@ class _PhysicalDeviceAttrs(NamedTuple): support_mig_mode: bool -_PHYSICAL_DEVICE_ATTRS = None -_GLOBAL_PHYSICAL_DEVICE = None -_GLOBAL_PHYSICAL_DEVICE_LOCK = threading.RLock() +_PHYSICAL_DEVICE_ATTRS: OrderedDict[str, _PhysicalDeviceAttrs] | None = None +_GLOBAL_PHYSICAL_DEVICE: PhysicalDevice | None = None +_GLOBAL_PHYSICAL_DEVICE_LOCK: threading.RLock = threading.RLock() -def _get_all_physical_device_attrs() -> dict[str, _PhysicalDeviceAttrs]: +def _get_all_physical_device_attrs() -> OrderedDict[str, _PhysicalDeviceAttrs]: global _PHYSICAL_DEVICE_ATTRS # pylint: disable=global-statement + if _PHYSICAL_DEVICE_ATTRS is not None: + return _PHYSICAL_DEVICE_ATTRS + with _GLOBAL_PHYSICAL_DEVICE_LOCK: if _PHYSICAL_DEVICE_ATTRS is None: _PHYSICAL_DEVICE_ATTRS = OrderedDict( @@ -3099,7 +3109,7 @@ def strip_identifier(identifier: str) -> str: def _parse_cuda_visible_devices_to_uuids( - cuda_visible_devices: str | None = _VALUE_OMITTED, + cuda_visible_devices: str | None = __VALUE_OMITTED, verbose: bool = True, ) -> list[str]: """Parse the given ``CUDA_VISIBLE_DEVICES`` environment variable in a separate process and return a list of device UUIDs. @@ -3121,7 +3131,7 @@ def _parse_cuda_visible_devices_to_uuids( libcuda.CUDAError: If failed to parse the ``CUDA_VISIBLE_DEVICES`` environment variable. """ # pylint: disable=line-too-long - if cuda_visible_devices is _VALUE_OMITTED: + if cuda_visible_devices is __VALUE_OMITTED: cuda_visible_devices = os.getenv('CUDA_VISIBLE_DEVICES', default=None) # Do not inherit file descriptors and handles from the parent process @@ -3138,10 +3148,7 @@ def _parse_cuda_visible_devices_to_uuids( parser.start() parser.join() finally: - try: - parser.kill() # requires Python 3.7+ - except AttributeError: - pass + parser.kill() result = queue.get() if isinstance(result, Exception):