diff --git a/nvitop/api/collector.py b/nvitop/api/collector.py index 5887cae..4496a4d 100644 --- a/nvitop/api/collector.py +++ b/nvitop/api/collector.py @@ -250,7 +250,7 @@ def collect_in_background( raise ValueError(f'Invalid argument interval={interval!r}') interval = min(interval, collector.interval) - def target(): + def target() -> None: if on_start is not None: on_start(collector) try: @@ -402,10 +402,7 @@ class ResourceMetricCollector: # pylint: disable=too-many-instance-attributes if devices is None: devices = Device.all() - if root_pids is None: - root_pids = {os.getpid()} - else: - root_pids = set(root_pids) + root_pids = {os.getpid()} if root_pids is None else set(root_pids) self.interval = interval diff --git a/nvitop/api/device.py b/nvitop/api/device.py index 72c8f6e..1dfff37 100644 --- a/nvitop/api/device.py +++ b/nvitop/api/device.py @@ -2019,7 +2019,7 @@ class MigDevice(Device): # pylint: disable=too-many-instance-attributes return snapshot - SNAPSHOT_KEYS = Device.SNAPSHOT_KEYS + ['gpu_instance_id', 'compute_instance_id'] + SNAPSHOT_KEYS = [*Device.SNAPSHOT_KEYS, 'gpu_instance_id', 'compute_instance_id'] class CudaDevice(Device): @@ -2387,15 +2387,14 @@ def normalize_cuda_visible_devices(cuda_visible_devices: Optional[str] = _VALUE_ # Helper functions ################################################################################# -_PhysicalDeviceAttrs = NamedTuple( - 'PhysicalDeviceAttrs', - [ - ('index', int), - ('name', str), - ('uuid', str), - ('support_mig_mode', bool), - ], -) + +class _PhysicalDeviceAttrs(NamedTuple): + index: int + name: str + uuid: str + support_mig_mode: bool + + _PHYSICAL_DEVICE_ATTRS = None _GLOBAL_PHYSICAL_DEVICE = None _GLOBAL_PHYSICAL_DEVICE_LOCK = threading.RLock() @@ -2551,7 +2550,7 @@ def _parse_cuda_visible_devices( # pylint: disable=too-many-branches,too-many-s def _parse_cuda_visible_devices_to_uuids( cuda_visible_devices: Optional[str] = _VALUE_OMITTED, - verbose=True, + verbose: bool = True, ) -> List[str]: """Parse the given ``CUDA_VISIBLE_DEVICES`` environment variable in a separate process and return a list of device UUIDs. diff --git a/nvitop/api/host.py b/nvitop/api/host.py index 0c79f2e..778dc4f 100644 --- a/nvitop/api/host.py +++ b/nvitop/api/host.py @@ -40,7 +40,7 @@ __all__[__all__.index('Error')] = 'PsutilError' PsutilError = Error # make alias # noqa: F405 -del Error # noqa: F405,F821 # pylint: disable=undefined-variable +del Error # noqa: F821 # pylint: disable=undefined-variable cpu_percent = _ttl_cache(ttl=0.25)(_psutil.cpu_percent) diff --git a/nvitop/api/libcudart.py b/nvitop/api/libcudart.py index 65e6bbc..2f89c84 100644 --- a/nvitop/api/libcudart.py +++ b/nvitop/api/libcudart.py @@ -483,6 +483,12 @@ def __LoadCudaLibrary() -> None: # pylint: disable=too-many-branches _glob.iglob(_os.path.join(candidate_dir, f'cudart{bits}*.dll')) ) + # Normalize paths and remove duplicates + candidate_paths = list( + dict.fromkeys( + _os.path.normpath(_os.path.normcase(p)) for p in candidate_paths + ) + ) for lib_filename in candidate_paths: try: __cudaLib = _ctypes.CDLL(lib_filename) diff --git a/nvitop/api/libnvml.py b/nvitop/api/libnvml.py index 9d9ac80..f76952d 100644 --- a/nvitop/api/libnvml.py +++ b/nvitop/api/libnvml.py @@ -391,10 +391,7 @@ def nvmlQuery( retval = func(*args, **kwargs) except NVMLError_FunctionNotFound as e2: if not ignore_function_not_found: - if func.__name__ == '': - identifier = _inspect.getsource(func) - else: - identifier = repr(func) + identifier = _inspect.getsource(func) if func.__name__ == '' else repr(func) with __lock: if ( identifier not in UNKNOWN_FUNCTIONS diff --git a/nvitop/api/process.py b/nvitop/api/process.py index 1805b39..5d2a9cc 100644 --- a/nvitop/api/process.py +++ b/nvitop/api/process.py @@ -100,7 +100,9 @@ _RAISE = object() _USE_FALLBACK_WHEN_RAISE = threading.local() # see also `GpuProcess.failsafe` -def auto_garbage_clean(fallback=_RAISE): +def auto_garbage_clean( + fallback: Any = _RAISE, +) -> Callable[[Callable[..., Any]], Callable[..., Any]]: """Remove the object references in the instance cache if the method call fails (the process is gone). The fallback value will be used with `:meth:`GpuProcess.failsafe`` context manager, otherwise @@ -979,7 +981,7 @@ class GpuProcess: # pylint: disable=too-many-instance-attributes,too-many-publi @classmethod def take_snapshots( # batched version of `as_snapshot` - cls, gpu_processes: Iterable['GpuProcess'], *, failsafe=False + cls, gpu_processes: Iterable['GpuProcess'], *, failsafe: bool = False ) -> List[Snapshot]: """Take snapshots for a list of :class:`GpuProcess` instances. diff --git a/nvitop/callbacks/tensorboard.py b/nvitop/callbacks/tensorboard.py index 9e90d28..04c7b91 100644 --- a/nvitop/callbacks/tensorboard.py +++ b/nvitop/callbacks/tensorboard.py @@ -17,8 +17,28 @@ # pylint: disable=missing-module-docstring +from typing import TYPE_CHECKING, Dict, Optional, Union -def add_scalar_dict(writer, main_tag, tag_scalar_dict, global_step=None, walltime=None): + +if TYPE_CHECKING: + import numpy as np + + try: + from tensorboard.summary import Writer as SummaryWriter + except ImportError: + try: + from tensorboardX import SummaryWriter + except ImportError: + pass + + +def add_scalar_dict( + writer: 'SummaryWriter', + main_tag: str, + tag_scalar_dict: Dict[str, Union[int, float, 'np.floating']], + global_step: Optional[Union[int, 'np.integer']] = None, + walltime: Optional[float] = None, +) -> None: """Add a batch of scalars to the writer. Batched version of ``writer.add_scalar``. diff --git a/nvitop/cli.py b/nvitop/cli.py index 5f301cb..8334399 100644 --- a/nvitop/cli.py +++ b/nvitop/cli.py @@ -19,13 +19,14 @@ NVITOP_MONITOR_MODE = set( ) -def parse_arguments(): # pylint: disable=too-many-branches,too-many-statements +# pylint: disable=too-many-branches,too-many-statements +def parse_arguments() -> argparse.Namespace: """Parse command-line arguments for ``nvtiop``.""" coloring_rules = '{} < th1 %% <= {} < th2 %% <= {}'.format( colored('light', 'green'), colored('moderate', 'yellow'), colored('heavy', 'red') ) - def posint(argstring): + def posint(argstring: str) -> int: num = int(argstring) if num <= 0: raise ValueError @@ -251,7 +252,8 @@ def parse_arguments(): # pylint: disable=too-many-branches,too-many-statements return args -def main(): # pylint: disable=too-many-branches,too-many-statements,too-many-locals +# pylint: disable-next=too-many-branches,too-many-statements,too-many-locals +def main() -> None: """Main function for ``nvitop`` CLI.""" args = parse_arguments() @@ -272,10 +274,7 @@ def main(): # pylint: disable=too-many-branches,too-many-statements,too-many-lo if hasattr(args, 'monitor') and args.monitor is None: mode = NVITOP_MONITOR_MODE.intersection({'auto', 'full', 'compact'}) - if len(mode) != 1: - mode = 'auto' - else: - mode = mode.pop() + mode = 'auto' if len(mode) != 1 else mode.pop() args.monitor = mode if not setlocale_utf8(): diff --git a/nvitop/select.py b/nvitop/select.py index 76418d8..6f082d6 100644 --- a/nvitop/select.py +++ b/nvitop/select.py @@ -60,7 +60,7 @@ import math import os import sys import warnings -from typing import Iterable, List, Optional, Tuple, Union +from typing import Any, Iterable, List, Optional, Tuple, Union from nvitop.api import Device, GpuProcess, colored, human2bytes, libnvml from nvitop.version import __version__ @@ -76,8 +76,7 @@ except ModuleNotFoundError: TTY = sys.stdout.isatty() -# pylint: disable-next=too-many-branches,too-many-statements,too-many-locals -def select_devices( +def select_devices( # pylint: disable=too-many-branches,too-many-statements,too-many-locals,unused-argument devices: Iterable[Device] = None, *, format: str = 'index', # pylint: disable=redefined-builtin @@ -91,7 +90,7 @@ def select_devices( tolerance: int = 0, # in percentage free_accounts: List[str] = None, sort: bool = True, - **kwargs, # fmt: skip # pylint: disable=unused-argument + **kwargs: Any, ) -> Union[List[int], List[Tuple[int, int]], List[str]]: """Select a subset of devices satisfying the specified criteria. @@ -163,7 +162,7 @@ def select_devices( if isinstance(min_total_memory, str): min_total_memory = human2bytes(min_total_memory) - available_devices = [] # type: Iterable[DeviceSnapshot] + available_devices = [] for device in devices: available_devices.extend(dev.as_snapshot() for dev in device.to_leaf_devices()) for device in available_devices: @@ -261,20 +260,16 @@ def select_devices( if format == 'device': return [device.real for device in available_devices] - if format == 'uuid': - identifiers = [device.uuid for device in available_devices] # type: List[str] - else: - identifiers = [ - device.index for device in available_devices - ] # type: List[int, Tuple[int, int]] - return identifiers + return [device.uuid for device in available_devices] + return [device.index for device in available_devices] -def parse_arguments(): # pylint: disable=too-many-branches,too-many-statements +# pylint: disable-next=too-many-branches,too-many-statements +def parse_arguments() -> argparse.Namespace: """Parse command-line arguments for ``nvisel``.""" - def non_negint(argstring): + def non_negint(argstring: str) -> int: num = int(argstring) if num < 0: raise ValueError @@ -489,7 +484,7 @@ def parse_arguments(): # pylint: disable=too-many-branches,too-many-statements return args -def main(): +def main() -> None: """Main function for ``nvisel`` CLI.""" args = parse_arguments()