Merge pull request #5 from analysiscenter/mask_extra_code_run_notebook

Add `mask_extra_code` option `into run_notebook`

Author: SergeyTsimfer

nbtools/__init__.py

@@ -1,6 +1,7 @@
 """ Init file. """
-from .pylint_notebook import pylint_notebook
+#pylint: disable=wildcard-import
+from .core import *
 from .run_notebook import run_notebook
-from .core import StringWithDisabledRepr, notebook_to_script
+from .pylint_notebook import pylint_notebook
 
 __version__ = '0.9.7'

nbtools/core.py

@@ -1,15 +1,16 @@
-""" !!. """
+""" Core utility functions to work with Jupyter Notebooks. """
 #pylint: disable=import-outside-toplevel
 import os
 import re
 import json
+import warnings
 
 
 
 class StringWithDisabledRepr(str):
     """ String with disabled repr. Used to avoid cluttering repr from function outputs. """
     def __repr__(self):
-        """ !!. """
+        """ Shorten the repr of a string. """
         return f'<StringWithDisabledRepr at {hex(id(self))}. Use `str`/`print` explicitly!>'
 
 
@@ -70,9 +71,8 @@ def get_notebook_name():
     return os.path.splitext(get_notebook_path())[0].split('/')[-1]
 
 
-
 def notebook_to_script(path_script, path_notebook=None, ignore_markdown=True, return_info=False):
-    """ !!. """
+    """ Convert a notebook to a script. """
     import nbformat #pylint: disable=import-outside-toplevel
     path_notebook = path_notebook or get_notebook_path()
     if path_notebook is None:
@@ -113,3 +113,115 @@ def notebook_to_script(path_script, path_notebook=None, ignore_markdown=True, re
         return {'code': StringWithDisabledRepr(code),
                 'cell_line_numbers': cell_line_numbers}
     return None
+
+
+
+def get_available_gpus(n=1, min_free_memory=0.9, max_processes=2, verbose=False, raise_error=False):
+    """ Select `n` gpus from available and free devices.
+
+    Parameters
+    ----------
+    n : int, str
+        If `max`, then use maximum number of available devices.
+        If int, then number of devices to select.
+    min_free_memory : float
+        Minimum percentage of free memory on a device to consider it free.
+    max_processes : int
+        Maximum amount of computed processes on a device to consider it free.
+    verbose : bool
+        Whether to show individual device information.
+    raise_error : bool
+        Whether to raise an exception if not enough devices are available.
+
+    Returns
+    -------
+    List with indices of available GPUs
+    """
+    try:
+        import nvidia_smi
+    except ImportError as exception:
+        raise ImportError('Install Python interface for nvidia_smi') from exception
+
+    nvidia_smi.nvmlInit()
+    n_devices = nvidia_smi.nvmlDeviceGetCount()
+
+    available_devices, memory_usage = [], []
+    for i in range(n_devices):
+        handle = nvidia_smi.nvmlDeviceGetHandleByIndex(i)
+        info = nvidia_smi.nvmlDeviceGetMemoryInfo(handle)
+
+        fraction_free = info.free / info.total
+        num_processes = len(nvidia_smi.nvmlDeviceGetComputeRunningProcesses(handle))
+
+        consider_available = (fraction_free > min_free_memory) & (num_processes <= max_processes)
+        if consider_available:
+            available_devices.append(i)
+            memory_usage.append(fraction_free)
+
+        if verbose:
+            print(f'Device {i} | Free memory: {fraction_free:4.2f} | '
+                  f'Number of running processes: {num_processes:>2} | Free: {consider_available}')
+
+    if isinstance(n, str) and n.startswith('max'):
+        n = len(available_devices)
+
+    if len(available_devices) < n:
+        msg = f'Not enough free devices: requested {n}, found {len(available_devices)}'
+        if raise_error:
+            raise ValueError(msg)
+        warnings.warn(msg, RuntimeWarning)
+
+    # Argsort of `memory_usage` in a descending order
+    indices = sorted(range(len(available_devices)), key=memory_usage.__getitem__, reverse=True)
+    available_devices = [available_devices[i] for i in indices]
+    return sorted(available_devices[:n])
+
+def get_gpu_free_memory(index):
+    """ Get free memory of a device. """
+    try:
+        import nvidia_smi
+    except ImportError as exception:
+        raise ImportError('Install Python interface for nvidia_smi') from exception
+
+    nvidia_smi.nvmlInit()
+    nvidia_smi.nvmlDeviceGetCount()
+    handle = nvidia_smi.nvmlDeviceGetHandleByIndex(index)
+    info = nvidia_smi.nvmlDeviceGetMemoryInfo(handle)
+
+    return info.free / info.total
+
+def set_gpus(n=1, min_free_memory=0.9, max_processes=2, verbose=False, raise_error=False):
+    """ Set the `CUDA_VISIBLE_DEVICES` variable to `n` available devices.
+
+    Parameters
+    ----------
+    n : int, str
+        If `max`, then use maximum number of available devices.
+        If int, then number of devices to select.
+    min_free_memory : float
+        Minimum percentage of free memory on a device to consider it free.
+    max_processes : int
+        Maximum amount of computed processes on a device to consider it free.
+    verbose : bool or int
+        Whether to show individual device information.
+        If 0 or False, then no information is displayed.
+        If 1 or True, then display the value assigned to `CUDA_VISIBLE_DEVICES` variable.
+        If 2, then display memory and process information for each device.
+    raise_error : bool
+        Whether to raise an exception if not enough devices are available.
+    """
+    #pylint: disable=consider-iterating-dictionary
+    if 'CUDA_VISIBLE_DEVICES' in os.environ.keys():
+        str_devices = os.environ["CUDA_VISIBLE_DEVICES"]
+        warnings.warn(f'`CUDA_VISIBLE_DEVICES` is already set to "{str_devices}"!')
+        return [int(d) for d in str_devices.split(',')]
+
+    devices = get_available_gpus(n=n, min_free_memory=min_free_memory, max_processes=max_processes,
+                                 verbose=(verbose==2), raise_error=raise_error)
+    str_devices = ','.join(str(i) for i in devices)
+    os.environ['CUDA_VISIBLE_DEVICES'] = str_devices
+
+    newline = "\n" if verbose==2 else ""
+    if verbose:
+        print(f'{newline}`CUDA_VISIBLE_DEVICES` set to "{str_devices}"')
+    return devices

nbtools/run_notebook.py

@@ -28,6 +28,9 @@
 
         locals().update(inputs)
 """
+YAML_IMPORT = """import yaml"""
+INPUTS_DISPLAY = """print(yaml.dump(inputs))"""
+
 INPUTS_CODE_CELL = dedent(INPUTS_CODE_CELL)
 
 # Save notebook outputs
@@ -43,12 +46,15 @@
     with shelve.open(out_path_db) as notebook_db:
         notebook_db['outputs'] = output
 """
+OUTPUTS_DISPLAY = """print(yaml.dump(output))"""
+
 OUTPUTS_CODE_CELL = dedent(OUTPUTS_CODE_CELL)
 
 
 def run_notebook(path, inputs=None, outputs=None, inputs_pos=1, working_dir = './', execute_kwargs=None,
                  out_path_db=None, out_path_ipynb=None, out_path_html=None, remove_db='always', add_timestamp=True,
-                 hide_code_cells=False, display_links=True, raise_exception=False, return_notebook=False):
+                 hide_code_cells=False, mask_extra_code=True, display_links=True,
+                 raise_exception=False, return_notebook=False):
     """ Execute a Jupyter Notebook programmatically.
     Heavily inspired by https://github.com/tritemio/nbrun.
 
@@ -110,6 +116,9 @@ def run_notebook(path, inputs=None, outputs=None, inputs_pos=1, working_dir = '.
         Whether to add a cell with execution information at the beginning of the saved notebook.
     hide_code_cells : bool, optional
         Whether to hide the code cells in the saved notebook.
+    mask_extra_code : bool, optional
+        Whether to mask database reading and dumping code.
+        For more, see :func:`~.mask_inputs_reading` and :func`~.mask_outputs_dumping` docstrings.
     display_links : bool, optional
         Whether to display links to the executed notebook and html at execution.
     raise_exception : bool, optional
@@ -182,6 +191,9 @@ def run_notebook(path, inputs=None, outputs=None, inputs_pos=1, working_dir = '.
             notebook_db.update(inputs)
 
         code = CELL_INSERT_COMMENT + DB_CONNECT_CODE_CELL.format(repr(out_path_db)) + INPUTS_CODE_CELL
+        if mask_extra_code:
+            code += YAML_IMPORT + '\n' + INPUTS_DISPLAY
+
         notebook['cells'].insert(inputs_pos, nbformat.v4.new_code_cell(code))
 
     if outputs is not None:
@@ -191,6 +203,12 @@ def run_notebook(path, inputs=None, outputs=None, inputs_pos=1, working_dir = '.
         code = CELL_INSERT_COMMENT + \
                (DB_CONNECT_CODE_CELL.format(repr(out_path_db)) if not inputs else "") + \
                OUTPUTS_CODE_CELL.format(outputs)
+
+        if mask_extra_code:
+            if inputs is None:
+                code = YAML_IMPORT + '\n' + code
+            code += OUTPUTS_DISPLAY
+
         output_cell = nbformat.v4.new_code_cell(code)
         notebook['cells'].append(output_cell)
 
@@ -240,13 +258,23 @@ def run_notebook(path, inputs=None, outputs=None, inputs_pos=1, working_dir = '.
         if outputs is not None:
             exec_res['outputs'] = outputs_values
 
+        # Notebook postprocessing: add timestamp, mask db reading/dumping code
         if add_timestamp:
             timestamp = (f"**Executed:** {time.ctime(start_time)}<br>"
                          f"**Duration:** {time.strftime('%H:%M:%S', time.gmtime(time.time() - start_time))}<br>"
                          f"**Autogenerated from:** [{path}]\n\n---")
             timestamp_cell = nbformat.v4.new_markdown_cell(timestamp)
             notebook['cells'].insert(0, timestamp_cell)
 
+        if mask_extra_code:
+            if inputs is not None:
+                pos = inputs_pos + 1 if add_timestamp else inputs_pos
+                mask_inputs_reading(notebook=notebook, pos=pos)
+
+            if outputs is not None:
+                pos = len(notebook['cells']) - 1
+                mask_outputs_dumping(notebook=notebook, pos=pos)
+
         # Save the executed notebook/HTML to disk
         if out_path_ipynb is not None:
             save_notebook(notebook=notebook, out_path_ipynb=out_path_ipynb, display_link=display_links)
@@ -257,6 +285,63 @@ def run_notebook(path, inputs=None, outputs=None, inputs_pos=1, working_dir = '.
             exec_res['notebook'] = notebook
         return exec_res
 
+# Mask functions for database operations cells
+def mask_inputs_reading(notebook, pos):
+    """ Replace database reading by variables initialization.
+
+    Result is a code cell with the following view:
+    .. code-block:: python
+
+        varible_name_1 = varible_value_1
+        varible_name_2 = varible_value_2
+        ...
+    """
+    import nbformat
+
+    execution_count = notebook['cells'][pos]['execution_count']
+
+    code_mask = str(notebook['cells'][pos]['outputs'][0]['text']).split('\n')
+    code_mask = [variable.replace(':', ' =', 1) for variable in code_mask]
+    code_mask = '\n'.join(code_mask)[:-2]
+
+    cell_mask = nbformat.v4.new_code_cell(source=code_mask, execution_count=execution_count)
+    notebook['cells'][pos] = cell_mask
+
+def mask_outputs_dumping(notebook, pos):
+    """Replace database dumping by printing outputs.
+
+    Result is a code cell with the following view and corresponding output:
+    .. code-block:: python
+
+        print(varible_name_1)
+        print(varible_name_2)
+        ...
+    """
+    import nbformat
+
+    execution_count = notebook['cells'][pos]['execution_count']
+    outputs_variables = str(notebook['cells'][pos]['outputs'][0]['text']).split('\n')
+
+    code_mask = ''
+    text_mask = ''
+
+    for variable in outputs_variables:
+        separator_pos = int(variable.find(':'))
+
+        if separator_pos != -1:
+            variable_name = variable[:separator_pos]
+            variable_value = variable[separator_pos+2:] + '\n'
+
+            code_mask += f'print({variable_name})\n'
+            text_mask += variable_value
+
+    code_mask = code_mask[:-1]
+    outputs_mask =  [nbformat.v4.new_output(text=text_mask, name='stdout', output_type='stream')]
+
+    cell_mask = nbformat.v4.new_code_cell(source=code_mask, execution_count=execution_count, outputs=outputs_mask)
+    notebook['cells'][pos] = cell_mask
+
+
 # Save notebook functions
 def save_notebook(notebook, out_path_ipynb, display_link):
     """ Save an instance of :class:`nbformat.notebooknode.NotebookNode` as ipynb file."""