simplify Jupyter API

Signed-off-by: Lance-Drane <[email protected]>

Author: Lance-Drane

bokeh-plots.ipynb

@@ -8,12 +8,12 @@
    "outputs": [],
    "source": [
     "# Notebook template, the IPS Framework will add a cell before this one\n",
-    "# defining IPS_STATE_FILES as a list of state file paths.\n",
+    "# defining ANALYSIS_FILES as a list of state file paths.\n",
     "\n",
     "# In this example, this notebook is generated during the time loop.\n",
     "\n",
     "mapping = {}\n",
-    "for file in IPS_STATE_FILES:\n",
+    "for file in ANALYSIS_FILES:\n",
     "    with open(file, 'rb') as f:\n",
     "        mapping[file] = f.read()\n",
     "print(mapping)\n"

examples-proposed/004-time-loop/sim/input_dir/basic.ipynb

@@ -8,7 +8,7 @@
    "outputs": [],
    "source": [
     "# Notebook template, the IPS Framework will add a cell before this one\n",
-    "# defining IPS_STATE_FILES as a list of state file paths.\n",
+    "# defining ANALYSIS_FILES as a list of state file paths.\n",
     "\n",
     "# In this example, this notebook is only generated at the end of the run.\n",
     "\n",
@@ -21,10 +21,10 @@
     "\n",
     "DATA = []\n",
     "# create DATA list, will depend on user input type (i.e. 'hdf5', 'json')\n",
-    "for file in IPS_STATE_FILES:\n",
+    "for file in ANALYSIS_FILES:\n",
     "    with open(file, 'rb') as f:\n",
     "        DATA.append(json.load(f))\n",
-    "x = [float(f.rpartition('/')[2]) for f in IPS_STATE_FILES]\n",
+    "x = [float(f.rpartition('/')[2]) for f in ANALYSIS_FILES]\n",
     "\n",
     "COLORS = ['red', 'green', 'blue']\n",
     "\n",

examples-proposed/004-time-loop/sim/input_dir/bokeh-plots.ipynb

@@ -22,6 +22,8 @@
 It is written to a notebook cell on initializing it, and is searched for when adding a data file to it.
 """
 
+DIRECTORY_VARIABLE_NAME = 'DATA_DIR'
+
 
 def replace_last(source_string: str, old: str, new: str) -> str:
     """Attempt to replace the last occurence of 'old' with 'new' in 'source_string', searching from the right.
@@ -32,17 +34,17 @@ def replace_last(source_string: str, old: str, new: str) -> str:
     return f'{head}{new}{tail}'
 
 
-def _initial_jupyter_file_notebook_cell(dest: str, variable: str) -> str:
+def _initial_jupyter_file_notebook_cell(dest: str, files_variable_name: str) -> str:
     return f"""{HOOK}
 
 import os
 
 # NOTE: directory should be sim_name plus the run id from the Portal
-IPS_DATA_DIR = '{str(Path(dest).parent / 'data') + sep}'
+{DIRECTORY_VARIABLE_NAME} = '{str(Path(dest).parent / 'data') + sep}'
 # Uncomment below line to implicitly use any state files saved in the data directory, note that the IPS framework explicitly lists out each file used
-#{variable} = os.listdir('data')
+#{files_variable_name} = os.listdir('data')
 # files created during the run
-{variable} = [
+{files_variable_name} = [
 ]
 """
 
@@ -61,10 +63,15 @@ def initialize_jupyter_notebook(dest: str, src: str, variable_name: str, index:
     # to avoid conversion, use as_version=nbf.NO_CONVERT
     nb: nbf.NotebookNode = nbf.read(src, as_version=4)
 
-    header = '# Next cell generated by IPS Framework'
     nb['cells'] = (
-        nb['cells'][:index]
-        + [nbf.v4.new_markdown_cell(header), nbf.v4.new_code_cell(_initial_jupyter_file_notebook_cell(dest, variable_name))]
+        # warning notification for users inspecting the file, unused programatically
+        [nbf.v4.new_markdown_cell('# WARNING: Do not manually modify this file until the IPS simulation is complete.')]
+        + nb['cells'][:index]
+        + [
+            # explicitly mark the IPS cell for users inspecting the file, unused programatically
+            nbf.v4.new_markdown_cell('## Next cell generated by IPS Framework'),
+            nbf.v4.new_code_cell(_initial_jupyter_file_notebook_cell(dest, variable_name)),
+        ]
         + nb['cells'][index:]
     )
 
@@ -88,21 +95,21 @@ def add_data_file_to_notebook(dest: str, data_file: str, index: Optional[int] =
         raise Exception('Cannot find IPS notebook node')
     ips_cell: str = nb['cells'][index]['source']
 
-    if ips_cell.find(f"f'{{IPS_DATA_DIR}}{data_file}',\n]") != -1:
+    if ips_cell.find(f"f'{{{DIRECTORY_VARIABLE_NAME}}}{data_file}',\n]") != -1:
         # The data file is already referenced in the notebook, so there's nothing else to do
         return
 
     # data file does not exist, so we need to add it
     # search from right of string for the ']' character, should work assuming user does not modify the cell past the variable definition
-    result = replace_last(ips_cell, ']', f"f'{{IPS_DATA_DIR}}{data_file}',\n]")
+    result = replace_last(ips_cell, ']', f"f'{{{DIRECTORY_VARIABLE_NAME}}}{data_file}',\n]")
     nb['cells'][index]['source'] = result
 
     with open(dest, 'w') as f:
         nbf.write(nb, f)
 
 
 def remove_data_file_from_notebook(dest: str, data_file: str, index: Optional[int] = None):
-    """Remove data file from the notebook list.
+    """Remove a specific data file from the notebook list.
 
     Params:
       - dest: path to notebook which will be modified
@@ -116,7 +123,7 @@ def remove_data_file_from_notebook(dest: str, data_file: str, index: Optional[in
         raise Exception('Cannot find IPS notebook node')
     ips_cell: str = nb['cells'][index]['source']
 
-    head, sep, tail = ips_cell.rpartition(f"f'{{IPS_DATA_DIR}}{data_file}',\n")
+    head, sep, tail = ips_cell.rpartition(f"f'{{{DIRECTORY_VARIABLE_NAME}}}{data_file}',\n")
     if sep == '':
         # existing match not found, so there's nothing left to remove
         return
@@ -133,7 +140,7 @@ def remove_last_data_file_from_notebook(dest: str, index: Optional[int] = None)
     Note that this function assumes the notebook maintains a specific format.
 
     Returns:
-      - None if there were no data entries in the notebook,
+      - None if there were no data entries in the notebook, the name of the file removed (without the directory) as a string if there was
     """
     nb: nbf.NotebookNode = nbf.read(dest, as_version=4)
     if index is None:
@@ -142,7 +149,7 @@ def remove_last_data_file_from_notebook(dest: str, index: Optional[int] = None)
         raise Exception('Cannot find IPS notebook node')
     ips_cell: str = nb['cells'][index]['source']
 
-    search_hook = "f'{IPS_DATA_DIR}"
+    search_hook = f"f'{{{DIRECTORY_VARIABLE_NAME}}}"
 
     start_index = ips_cell.rfind(search_hook)
     if start_index == -1:

ipsframework/jupyter.py

@@ -29,7 +29,7 @@
 from . import ipsutil, messages
 from .cca_es_spec import initialize_event_service
 from .ips_es_spec import eventManager
-from .jupyter import add_data_file_to_notebook, initialize_jupyter_notebook, remove_data_file_from_notebook, remove_last_data_file_from_notebook
+from .jupyter import add_data_file_to_notebook, initialize_jupyter_notebook, remove_last_data_file_from_notebook
 from .taskManager import TaskInit
 
 RunningTask = namedtuple('RunningTask', ['process', 'start_time', 'timeout', 'nproc', 'cores_allocated', 'command', 'binary', 'args'])
@@ -1903,7 +1903,7 @@ def initialize_jupyter_notebook(
         self,
         dest_notebook_name: str,
         source_notebook_path: str,
-        variable_name: str = 'IPS_STATE_FILES',
+        variable_name: str = 'ANALYSIS_FILES',
         cell_to_modify: int = 0,
     ) -> None:
         """Loads a notebook from source_notebook_path, adds a cell to load the data, and then saves it to source_notebook_path. Will also try to register the notebook with the IPS Portal, if available.
@@ -1913,7 +1913,7 @@ def initialize_jupyter_notebook(
         Params:
           - dest_notebook_name: name of the JupyterNotebook you want to write (do not include file paths).
           - source_notebook_path: location you want to load the source notebook from
-          - variable_name: name of the variable you want to load files from (default: "IPS_STATE_FILES")
+          - variable_name: name of the variable you want to load files from (default: "ANALYSIS_FILES")
           - cell_to_modify: which cell in the JupyterNotebook you want to add the data call to (0-indexed).
                (This will not overwrite any cells, just appends.)
                By default, the data listing will happen in the FIRST cell.
@@ -1942,7 +1942,7 @@ def initialize_jupyter_notebook(
         self.publish('_IPS_MONITOR', 'PORTAL_REGISTER_NOTEBOOK', event_data)
         self._send_monitor_event('IPS_PORTAL_REGISTER_NOTEBOOK', f'URL = {url}')
 
-    def add_data_file_to_notebook(self, state_file_path: str, timestamp: float, notebook_name: str, index: Optional[int] = None):
+    def add_data_file_to_notebook(self, state_file_path: str, timestamp: float, notebook_name: str, replace: bool = False, index: Optional[int] = None):
         """Add data file to JupyterHub directory, and reference it in the notebook.
 
         This function assumes that a notebook has already been created with intialize_jupyter_notebook. Using this function does not call the IPS Portal.
@@ -1951,83 +1951,30 @@ def add_data_file_to_notebook(self, state_file_path: str, timestamp: float, note
         - state_file_path: location of the current state file we want to copy to the Jupyter directory
         - timestamp: label to assign to the data (currently must be a floating point value)
         - notebook_name: name of notebook which will be modified. Note that this path is relative to the JupyterHub directory.
+        - replace: If True, replace the last data file added with the new data file. If False, simply append the new data file.
         - index: optional index of the IPS notebook cell. If not provided, the IPS Framework will attempt to automatically find the cell it created,
             which should work for every usecase where you don't anticipate modifying the notebook until after the run is complete.
         """
-
         if not self._jupyterhub_dir:
             if not self._init_jupyter():
                 # TODO generic exception
                 raise Exception('Unable to initialize base JupyterHub dir')
 
-        file_parts = state_file_path.split('.')
-        if len(file_parts) > 2:  # name of the file could just be a floating point value with no extension
-            extension = f'.{file_parts[-1]}'
-        else:
-            extension = ''
-
-        state_file_name = f'{timestamp}{extension}'
-        jupyter_data_dir = os.path.join(self._jupyterhub_dir, 'data', state_file_name)
+        data_file_name = f'{timestamp}_{os.path.basename(state_file_path)}'
+        jupyter_data_dir = os.path.join(self._jupyterhub_dir, 'data', data_file_name)
         # this may raise an OSError, it is the responsibility of the caller to handle it.
         shutil.copyfile(state_file_path, jupyter_data_dir)
 
-        # TODO - maybe add flag which allows us to replace old state files
-        add_data_file_to_notebook(f'{self._jupyterhub_dir}{notebook_name}', state_file_name, index)
-
-    def remove_data_file_from_notebook(self, state_file_path: str, timestamp: float, notebook_name: str, index: Optional[int] = None):
-        """Remove data file from JupyterHub data directory and from being referenced in the notebook.
-
-        This function assumes that a notebook has already been created with intialize_jupyter_notebook. Using this function does not call the IPS Portal.
-
-        Params:
-        - state_file_path: location of the current state file we want to copy to the Jupyter directory
-        - timestamp: label to assign to the data (currently must be a floating point value)
-        - notebook_name: name of notebook which will be modified. Note that this path is relative to the JupyterHub directory.
-        - index: optional index of the IPS notebook cell. If not provided, the IPS Framework will attempt to automatically find the cell it created,
-            which should work for every usecase where you don't anticipate modifying the notebook until after the run is complete.
-        """
-
-        if not self._jupyterhub_dir:
-            if not self._init_jupyter():
-                # TODO generic exception
-                raise Exception('Unable to initialize base JupyterHub dir')
+        if replace:
+            # first try to remove the reference from the Jupyter Notebook
+            filename_to_remove = remove_last_data_file_from_notebook(f'{self._jupyterhub_dir}{notebook_name}', index)
+            if filename_to_remove is not None:
+                # now remove the state file from the filesyste,
+                file_to_remove = os.path.join(self._jupyterhub_dir, 'data', filename_to_remove)
+                shutil.rmtree(file_to_remove, ignore_errors=True)
 
-        file_parts = state_file_path.split('.')
-        if len(file_parts) > 2:  # name of the file could just be a floating point value with no extension
-            extension = f'.{file_parts[-1]}'
-        else:
-            extension = ''
-
-        state_file_name = f'{timestamp}{extension}'
-        jupyter_data_dir = os.path.join(self._jupyterhub_dir, 'data', state_file_name)
-
-        # if this errors out, we can safely ignore them
-        shutil.rmtree(jupyter_data_dir, ignore_errors=True)
-
-        # TODO - maybe add flag which allows us to replace old state files
-        remove_data_file_from_notebook(f'{self._jupyterhub_dir}{notebook_name}', state_file_name, index)
-
-    def remove_last_data_file_from_notebook(self, notebook_name: str, index: Optional[int] = None):
-        """Remove the last added data file from a notebook and from the filesystem.
-
-        This function assumes that a notebook has already been created with intialize_jupyter_notebook. Using this function does not call the IPS Portal.
-
-        Params:
-        - notebook_name: name of notebook which will be modified. Note that this path is relative to the JupyterHub directory.
-        - index: optional index of the IPS notebook cell. If not provided, the IPS Framework will attempt to automatically find the cell it created,
-            which should work for every usecase where you don't anticipate modifying the notebook until after the run is complete.
-        """
-
-        if not self._jupyterhub_dir:
-            if not self._init_jupyter():
-                # TODO generic exception
-                raise Exception('Unable to initialize base JupyterHub dir')
-
-        last_state_file = remove_last_data_file_from_notebook(notebook_name, index)
-        if last_state_file is None:
-            return
-        data_file = os.path.join(self._jupyterhub_dir, 'data', last_state_file)
-        shutil.rmtree(data_file, ignore_errors=True)
+        # add newest data file to notebook
+        add_data_file_to_notebook(f'{self._jupyterhub_dir}{notebook_name}', data_file_name, index)
 
     def publish(self, topicName: str, eventName: str, eventBody: Any):
         """