Merge pull request #8 from analysiscenter/run_notebook_improvements

`run_notebook` improvements

Author: SergeyTsimfer

nbtools/__init__.py

@@ -4,4 +4,4 @@
 from .run_notebook import run_notebook
 from .pylint_notebook import pylint_notebook
 
-__version__ = '0.9.9'
+__version__ = '0.9.10'

nbtools/core.py

@@ -162,6 +162,8 @@ def get_available_gpus(n=1, min_free_memory=0.9, max_processes=2, verbose=False,
             print(f'Device {i} | Free memory: {fraction_free:4.2f} | '
                   f'Number of running processes: {num_processes:>2} | Free: {consider_available}')
 
+    nvidia_smi.nvmlShutdown()
+
     if isinstance(n, str) and n.startswith('max'):
         n = len(available_devices)
 
@@ -187,7 +189,7 @@ def get_gpu_free_memory(index):
     nvidia_smi.nvmlDeviceGetCount()
     handle = nvidia_smi.nvmlDeviceGetHandleByIndex(index)
     info = nvidia_smi.nvmlDeviceGetMemoryInfo(handle)
-
+    nvidia_smi.nvmlShutdown()
     return info.free / info.total
 
 def set_gpus(n=1, min_free_memory=0.9, max_processes=2, verbose=False, raise_error=False):
@@ -225,3 +227,31 @@ def set_gpus(n=1, min_free_memory=0.9, max_processes=2, verbose=False, raise_err
     if verbose:
         print(f'{newline}`CUDA_VISIBLE_DEVICES` set to "{str_devices}"')
     return devices
+
+def free_gpus(devices=None):
+    """ Terminate all processes on gpu devices.
+
+    Parameters
+    ----------
+    devices : iterable of ints
+        Device indices to terminate processes.
+        If None, than free all available gpus.
+    """
+    import nvidia_smi
+    import psutil
+
+    nvidia_smi.nvmlInit()
+
+    if devices is None:
+        if 'CUDA_VISIBLE_DEVICES' in os.environ.keys():
+            devices = [int(d) for d in os.environ["CUDA_VISIBLE_DEVICES"].split(',')]
+        else:
+            devices = range(0, nvidia_smi.nvmlDeviceGetCount())
+
+    for device_index in devices:
+        handle = nvidia_smi.nvmlDeviceGetHandleByIndex(device_index)
+
+        for proc in nvidia_smi.nvmlDeviceGetComputeRunningProcesses(handle):
+            psutil.Process(proc.pid).terminate()
+
+    nvidia_smi.nvmlShutdown()

nbtools/nbstat/cli.py

@@ -118,6 +118,12 @@ def output_looped(inspector, name, formatter, view_args,
                         elif inkey == 'm':
                             if formatter[Resource.DEVICE_UTIL]:
                                 formatter[Resource.DEVICE_UTIL_MA] = not formatter[Resource.DEVICE_UTIL_MA]
+                        elif inkey.code == 265:
+                            formatter[Resource.PID] = not formatter[Resource.PID]
+                        elif inkey.code == 266:
+                            formatter[Resource.CPU] = not formatter[Resource.CPU]
+                        elif inkey.code == 267:
+                            formatter[Resource.TYPE] = not formatter[Resource.TYPE]
                         elif inkey == 'q':
                             raise KeyboardInterrupt
                         else:

nbtools/nbstat/resource.py

@@ -1,5 +1,9 @@
 """ Resource -- a class to describe a property of an entry. Refer to class documentation for more. """
-from enum import Enum, auto
+from enum import Enum
+
+# `enum.Auto` is bugged in Python 3.11
+int_generator = iter(range(0, 10000))
+auto = lambda: next(int_generator)
 
 
 class Resource(Enum):

nbtools/nbstat/resource_entry.py

@@ -50,12 +50,20 @@ def to_format_data(self, resource, terminal, **kwargs):
                         Resource.DEVICE_PROCESS_PID,]:
             pass
 
-        elif resource in [Resource.NAME, Resource.TYPE]:
+        elif resource == Resource.NAME:
             if data is not None:
-                if 'zombie' in data or 'container' in data:
-                    style = terminal.red
                 if '/' in data:
                     data = '~' + data.split('/')[-1]
+                if len(data) >= 60:
+                    data = data.replace('.ipynb', '').replace('.py', '')
+                    data = data[:30] + '[...]' + data[-30]
+
+        elif resource == Resource.TYPE:
+            if data is not None:
+                if 'zombie' in data or 'containerd' in data:
+                    style = terminal.red
+                if 'run_notebook' in data:
+                    style = terminal.green
 
         elif resource == Resource.CREATE_TIME:
             if data is not None:

nbtools/nbstat/resource_inspector.py

@@ -13,12 +13,14 @@
 from .resource import Resource
 from .resource_table import ResourceTable
 from .utils import format_memory, pid_to_name, pid_to_ngid, true_len, true_rjust, true_center, FiniteList
+from ..run_notebook import get_run_notebook_name
 
 
 
 KERNEL_ID_SEARCHER   = re.compile('kernel-(.*).json').search
 VSCODE_KEY_SEARCHER  = re.compile('key=b"(.*)"').search
 SCRIPT_NAME_SEARCHER = re.compile('python.* (.*).py').search
+RUN_NOTEBOOK_PATH_SEARCHER = re.compile('/tmp/.*.json.*--HistoryManager.hist_file=:memory:.*').search
 
 
 class ResourceInspector:
@@ -232,6 +234,7 @@ def get_process_table(self, formatter=None):
                     kernel_id = KERNEL_ID_SEARCHER(cmdline)
                     vscode_key = VSCODE_KEY_SEARCHER(cmdline)
                     script_name = SCRIPT_NAME_SEARCHER(cmdline)
+                    run_notebook_path = RUN_NOTEBOOK_PATH_SEARCHER(cmdline)
 
                     if kernel_id:
                         # The name will be changed by data from `notebook_table`.
@@ -245,6 +248,11 @@ def get_process_table(self, formatter=None):
                         type_ = 'vscode'
                         name = vscode_key.group(1).split('-')[0] + '.ipynb'
                         path = kernel_id = vscode_key.group(1)
+                    elif run_notebook_path:
+                        type_ = 'run_notebook'
+                        name = get_run_notebook_name(process.ppid())
+                        path = cwd
+                        kernel_id = None
                     elif script_name:
                         type_ = 'script'
                         name = script_name.group(1) + '.py'
@@ -258,7 +266,10 @@ def get_process_table(self, formatter=None):
 
                     # PYTHON_PPID = PPID if parent is Python process else -1
                     ppid = process.ppid()
-                    if ppid in python_pids:
+                    if type_ == 'run_notebook':
+                        # Spawned by `run_notebook` function of the library
+                        python_ppid = ppid
+                    elif ppid in python_pids:
                         # Spawned by one of other Python processes
                         type_ = 'subprocess'
                         python_ppid = ppid
@@ -495,6 +506,12 @@ def get_view(self, name='nbstat', formatter=None, index_condition=None, force_st
         else:
             raise ValueError('Wrong name of view to get!')
 
+        # Filter some processes
+        if 'nb' in name and table:
+            bad_names = ['lsp_server']
+            function = lambda index_value, _: not any(name in index_value for name in bad_names)
+            table.filter_on_index(function, inplace=True)
+
         # Filter index of the table by a regular expression
         if table and index_condition is not None:
             function = lambda index_value, _: bool(re.search(index_condition, str(index_value)))