basic E2E for autogenerating JupyterNotebooks and associating them with runs on Portal

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

Author: Lance-Drane

.gitignore

@@ -1,4 +1,5 @@
 # temporary files
+*.tmp
 *.tmp.*
 
 

examples-proposed/004-time-loop/README.md

@@ -19,5 +19,5 @@ pip install -e .
 To run the code, run:
 
 ```bash
-ips.py --config=sim.conf --platform=platform.conf
+./run.sh
 ```

examples-proposed/004-time-loop/mymodule/components.py

@@ -29,6 +29,15 @@ def step(self, timestamp=0.0):
             # TODO - perhaps monitor timestep does not need to be called every step, but only every 20 steps?
             self.services.call(monitor, 'step', t)
 
+        # create notebook here
+        NOTEBOOK_NAME = 'full_state.ipynb'
+        jupyter_files = self.services.get_staged_jupyterhub_files()
+        self.services.create_jupyterhub_notebook(jupyter_files, NOTEBOOK_NAME)
+        # NOTE: depending on the names of the files, you may have to use a custom mapping function to get the tag
+        # You MUST store the tag somewhere in the file name
+        tags = jupyter_files
+        self.services.portal_register_jupyter_notebook(NOTEBOOK_NAME, tags)
+
         self.services.call(worker, 'finalize', 0)
 
 
@@ -77,8 +86,11 @@ def step(self, timestamp=0.0, **keywords):
         self.services.stage_state()
 
         state_file = self.services.get_config_param('STATE_FILES')
-        with open(state_file, 'r') as f:
-            data = json.load(f)
+        with open(state_file, 'rb') as f:
+            data = f.read()
 
+        # example of updating Jupyter state
+        _jupyterhub_state_file = self.services.jupyterhub_make_state(state_file, timestamp)
+        # if you wanted to create a notebook per timestep, call send_portal_data with _jupyterhub_state_file as the argument.
         print('SEND PORTAL DATA', timestamp, data, file=stderr)
-        # self.services.send_portal_data(timestamp, data)
+        self.services.send_portal_data(timestamp, data)

examples-proposed/004-time-loop/run.sh

@@ -1,2 +1,2 @@
 rm -rf sim
-PYTHONPATH=$PWD ips.py --config=sim.conf --platform=platform.conf --log=ips.log #--debug --verbose
+PYTHONPATH=$PWD PSCRATCH=${PSCRATCH:-/tmp} ips.py --config=sim.conf --platform=platform.conf --log=ips.log #--debug --verbose

examples-proposed/004-time-loop/sim.conf

@@ -12,7 +12,18 @@ TAG = tag
 USER_W3_DIR = $PWD/www
 USER_W3_BASEURL =
 
-PORTAL_URL = http://lb.ipsportal.development.svc.spin.nersc.org
+PORTAL_URL = https://lb.ipsportal.development.svc.spin.nersc.org
+
+# OPTIONAL
+# The BASE DIRECTORY of your machine's JupyterHub web server directory. This is used strictly for moving files around on the machine itself.
+# This MUST be an absolute directory
+# if executing portal on a web server also running JupyterHub, you can set this value to manually add your file to this location.
+# NOTE: "/ipsframework/runs/${RUNID}/${TIMESTEP}.ipynb" will automatically be prepended to the file, guaranteeing that files will NEVER be overwritten.
+JUPYTERHUB_DIR = ${PSCRATCH}
+# OPTIONAL
+# if PORTAL_URL is defined + JUPYTERHUB_DIR is defined + this value is defined, the IPSFramework can send a direct link to the JupyterHub file.
+# In the framework, we don't make any assumptions about how JupyterHub's base URL or root server path is configured; you'll generally want to provide the full URL up to the JupyterHub base directory.
+JUPYTERHUB_URL = https://jupyter.nersc.gov/user/${USER}/perlmutter-login-node-base/lab/tree${PSCRATCH}
 
 STATE_FILES = state.json
 STATE_WORK_DIR = $SIM_ROOT/state
@@ -79,5 +90,5 @@ STATE_WORK_DIR = $SIM_ROOT/state
 [TIME_LOOP]
     MODE = REGULAR
     START = 1
-    FINISH = 100
-    NSTEP = 100
+    FINISH = 30
+    NSTEP = 30

ipsframework/jupyter.py

@@ -0,0 +1,52 @@
+"""
+This module is designed to help generate JupyterNotebooks to be used with IPS Portal analysis.
+Some parts of the script will need direction from users on the Framework side to generate.
+
+Note that this module is currently biased towards working with NERSC (jupyter.nersc.gov), so will attempt to import specific libraries.
+
+To see available libraries on NERSC, run:
+  !pip list
+
+...in a shell on Jupyter NERSC.
+"""
+
+from os.path import sep
+from typing import List
+
+import nbformat as nbf
+
+
+def _get_multi_state_file_notebook_code(state_file_paths: List[str]) -> str:
+    """TODO this is currently just an example."""
+    return f"""FILES = [{','.join([f"'data{sep}{file}'" for file in state_file_paths])}]
+mapping = {{}}
+for file in FILES:
+    with open(file, 'rb') as f:
+        mapping[file] = f.read()
+print(mapping)
+"""
+
+
+def _get_nb_v4(code: str):
+    """Returns an nbf.v4 object"""
+    nb = nbf.v4.new_notebook()
+    text = '# AUTOGENERATED from IPS Framework'
+
+    nb['cells'] = [nbf.v4.new_markdown_cell(text), nbf.v4.new_code_cell(code)]
+    return nb
+
+
+def create_multi_state_file_notebook(state_file_paths: List[str], notebook_path: str):
+    """
+    TODO
+
+    Writes notebook which will try to load multiple state files
+
+    for now, will just store all data in a dictionary of filepath str to raw bytes
+    """
+
+    code = _get_multi_state_file_notebook_code(state_file_paths)
+    nb = _get_nb_v4(code)
+
+    with open(notebook_path, 'w') as f:
+        nbf.write(nb, f)

ipsframework/portalBridge.py

@@ -98,6 +98,38 @@ def send_post_data(conn: Connection, stop: EventType, url: str):
             break
 
 
+def send_put_jupyter_url(conn: Connection, stop: EventType, url: str):
+    fail_count = 0
+
+    http = urllib3.PoolManager(retries=urllib3.util.Retry(3, backoff_factor=0.25))
+
+    while True:
+        if conn.poll(0.1):
+            next_val: dict[str, Any] = conn.recv()
+            # TODO - consider using multipart/form-data instead
+            try:
+                resp = http.request(
+                    'PUT',
+                    url,
+                    body=json.dumps({'url': next_val['url'], 'tags': next_val['tags'], 'portal_runid': next_val['portal_runid']}).encode(),
+                    headers={
+                        'Content-Type': 'application/json',
+                    },
+                )
+            except urllib3.exceptions.MaxRetryError as e:
+                fail_count += 1
+                conn.send((999, str(e)))
+            else:
+                conn.send((resp.status, resp.data.decode()))
+                fail_count = 0
+
+            if fail_count >= 3:
+                conn.send((-1, 'Too many consecutive failed connections'))
+                break
+        elif stop.is_set():
+            break
+
+
 class PortalBridge(Component):
     """
     Framework component to communicate with the SWIM web portal.
@@ -142,6 +174,10 @@ def __init__(self, services, config):
         self.data_childProcess = None
         self.data_childProcessStop = None
         self.data_parent_conn = None
+        self.dataurl_first_event = True
+        self.dataurl_childProcess = None
+        self.dataurl_childProcessStop = None
+        self.dataurl_parent_conn = None
         self.mpo = None
         self.mpo_name_counter = defaultdict(lambda: 0)
         self.counter = 0
@@ -232,6 +268,10 @@ def process_event(self, topicName, theEvent):
             self.send_data(sim_data, portal_data)
             return
 
+        if portal_data['eventtype'] == 'PORTAL_REGISTER_NOTEBOOK':
+            self.send_notebook_url(sim_data, portal_data)
+            return
+
         if portal_data['eventtype'] == 'IPS_SET_MONITOR_URL':
             sim_data.monitor_url = portal_data['vizurl']
         elif sim_data.monitor_url:
@@ -345,7 +385,7 @@ def send_data(self, sim_data, event_data):
             if self.data_first_event:  # First time, launch sendPost.py daemon
                 self.data_parent_conn, child_conn = Pipe()
                 self.data_childProcessStop = Event()
-                self.data_childProcess = Process(target=send_post_data, args=(child_conn, self.childProcessStop, self.portal_url + '/api/data'))
+                self.data_childProcess = Process(target=send_post_data, args=(child_conn, self.data_childProcessStop, self.portal_url + '/api/data'))
                 self.data_childProcess.start()
                 self.data_first_event = False
 
@@ -357,9 +397,9 @@ def send_data(self, sim_data, event_data):
             self.check_data_send_post_responses()
 
     def check_data_send_post_responses(self):
-        while self.parent_conn.poll():
+        while self.data_parent_conn.poll():
             try:
-                code, msg = self.parent_conn.recv()
+                code, msg = self.data_parent_conn.recv()
             except (EOFError, OSError):
                 break
 
@@ -380,6 +420,49 @@ def check_data_send_post_responses(self):
             else:
                 self.services.error('Portal Error: %d %s', code, msg)
 
+    def send_notebook_url(self, sim_data, event_data):
+        """
+        Send notebook contents
+        """
+        if self.portal_url:
+            if self.dataurl_first_event:  # First time, launch sendPost.py daemon
+                self.dataurl_parent_conn, child_conn = Pipe()
+                self.dataurl_childProcessStop = Event()
+                self.dataurl_childProcess = Process(
+                    target=send_put_jupyter_url, args=(child_conn, self.dataurl_childProcessStop, self.portal_url + '/api/data/add_url')
+                )
+                self.dataurl_childProcess.start()
+                self.dataurl_first_event = False
+
+            try:
+                self.dataurl_parent_conn.send(event_data)
+            except OSError:
+                pass
+
+            while self.dataurl_parent_conn.poll():
+                try:
+                    code, msg = self.dataurl_parent_conn.recv()
+                except (EOFError, OSError):
+                    break
+
+                print('PUT RESPONSE', code, msg)
+                try:
+                    data = json.loads(msg)
+                    if 'runid' in data:
+                        self.services.info('Run Portal URL = %s/%s', self.portal_url, data.get('runid'))
+
+                    msg = json.dumps(data)
+                except (TypeError, json.decoder.JSONDecodeError):
+                    pass
+                if code == -1:
+                    # disable portal, stop trying to send more data
+                    self.portal_url = None
+                    self.services.error('Disabling portal because: %s', msg)
+                elif code < 400:
+                    self.services.debug('Portal Response: %d %s', code, msg)
+                else:
+                    self.services.error('Portal Error: %d %s', code, msg)
+
     def send_mpo_data(self, event_data, sim_data):  # pragma: no cover
         def md5(fname):
             "Courtesy of stackoverflow 3431825"