Skip to content

Commit

Permalink
style: miscellaneous style housekeeping
Browse files Browse the repository at this point in the history
  • Loading branch information
XuehaiPan committed Nov 5, 2023
1 parent 1cff66b commit 7eb0d20
Showing 1 changed file with 24 additions and 17 deletions.
41 changes: 24 additions & 17 deletions nvitop/api/device.py
Original file line number Diff line number Diff line change
Expand Up @@ -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>'


__VALUE_OMITTED: str = ValueOmitted() # type: ignore[assignment]
del ValueOmitted


class Device: # pylint: disable=too-many-instance-attributes,too-many-public-methods
Expand Down Expand Up @@ -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.
Expand All @@ -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`.
Expand Down Expand Up @@ -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.
Expand Down Expand Up @@ -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`.
Expand Down Expand Up @@ -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'))
Expand All @@ -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(
Expand Down Expand Up @@ -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.
Expand All @@ -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
Expand All @@ -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):
Expand Down

0 comments on commit 7eb0d20

Please sign in to comment.