Refactor nbexecute to do file-ops in sync context

Author: u8sand

appyter/render/nbexecute.py

@@ -13,7 +13,7 @@
 from appyter.cli import cli
 from appyter.ext.emitter import json_emitter_factory
 from appyter.ext.fsspec.core import url_to_chroot_fs
-from appyter.ext.asyncio.helpers import ensure_async_contextmanager, ensure_sync
+from appyter.ext.asyncio.helpers import ensure_sync, ensure_async
 from appyter.ext.nbclient import NotebookClientIOPubHook
 from appyter.parse.nb import nb_from_ipynb_io, nb_to_ipynb_io, nb_to_json
 from appyter.ext.click import click_option_setenv, click_argument_setenv
@@ -33,9 +33,10 @@ async def iopub_hook(cell, cell_index):
     await emit({ 'type': 'cell', 'data': [cell, cell_index] })
   return iopub_hook
 
-async def nbexecute_async(ipynb='', emit=json_emitter_factory(sys.stdout), cwd='', subscribe=None, fuse=False):
+def nbexecute_sync(ipynb='', emit=json_emitter_factory(sys.stdout), cwd='', subscribe=None, fuse=False):
   logger.info('starting')
   assert callable(emit), 'Emit must be callable'
+  emit_sync = ensure_sync(emit)
   with fsspec.open(str(URI(cwd).join(ipynb)), 'r') as fr:
     nb = nb_from_ipynb_io(fr)
   #
@@ -62,78 +63,30 @@ async def nbexecute_async(ipynb='', emit=json_emitter_factory(sys.stdout), cwd='
     del nb.metadata['execution_info']
   #
   if 'completed' in nb.metadata['appyter']['nbexecute']:
-    await emit({ 'type': 'error', 'data': f"Execution already completed at {nb.metadata['appyter']['nbexecute']['completed']}" })
+    emit_sync({ 'type': 'error', 'data': f"Execution already completed at {nb.metadata['appyter']['nbexecute']['completed']}" })
     return
   elif 'started' in nb.metadata['appyter']['nbexecute']:
-    await emit({ 'type': 'error', 'data': f"Execution already started at {nb.metadata['appyter']['nbexecute']['started']}" })
+    emit_sync({ 'type': 'error', 'data': f"Execution already started at {nb.metadata['appyter']['nbexecute']['started']}" })
     return
   #
-  await emit({ 'type': 'status', 'data': 'Starting' })
+  emit_sync({ 'type': 'status', 'data': 'Starting' })
   state = dict(progress=0, status='Starting')
   #
   try:
     files = nb.metadata['appyter']['nbconstruct'].get('files')
     logger.debug(f"{cwd=} {files=}")
-    async with ensure_async_contextmanager(url_to_chroot_fs(cwd, pathmap=files)) as fs:
-      async with ensure_async_contextmanager(fs.mount(fuse=fuse)) as mnt:
+    with url_to_chroot_fs(cwd, pathmap=files) as fs:
+      with fs.mount(fuse=fuse) as mnt:
         # setup execution_info with start time
         nb.metadata['appyter']['nbexecute']['started'] = datetime.datetime.now().replace(tzinfo=datetime.timezone.utc).isoformat()
         with (mnt/ipynb).open('w') as fw:
           nb_to_ipynb_io(nb, fw)
-        #
-        logger.info('initializing')
-        if callable(subscribe):
-          await subscribe(lambda: dict(nb=nb_to_json(nb), **state))
-        #
-        try:
-          iopub_hook = iopub_hook_factory(nb, emit)
-          client = NotebookClientIOPubHook(
-            nb,
-            allow_errors=True,
-            timeout=None,
-            kernel_name='python3',
-            resources={ 'metadata': {'path': str(mnt) } },
-            iopub_hook=iopub_hook,
-          )
-          await emit({ 'type': 'nb', 'data': nb_to_json(nb) })
-          try:
-            async with client.async_setup_kernel(
-              env=dict(
-                { k: v for k, v in os.environ.items() if not k.startswith('APPYTER_') },
-                PYTHONPATH=':'.join(sys.path),
-              ),
-            ):
-              logger.info('executing')
-              state.update(status='Executing...', progress=0)
-              await emit({ 'type': 'status', 'data': state['status'] })
-              await emit({ 'type': 'progress', 'data': state['progress'] })
-              client.set_widgets_metadata()
-              n_cells = len(nb.cells)
-              exec_count = 1
-              for index, cell in enumerate(nb.cells):
-                logger.debug(f"executing cell {index}")
-                cell = await client.async_execute_cell(
-                  cell, index,
-                  execution_count=exec_count,
-                )
-                if cell_is_code(cell):
-                  if cell_has_error(cell):
-                    raise Exception('Cell execution error on cell %d' % (exec_count))
-                  exec_count += 1
-                if index < n_cells-1:
-                  state['progress'] = index + 1
-                  await emit({ 'type': 'progress', 'data': state['progress'] })
-                else:
-                  state['status'] = 'Success'
-                  await emit({ 'type': 'status', 'data': state['status'] })
-          finally:
-            client.set_widgets_metadata()
-        except asyncio.CancelledError:
-          logger.info('cancelled')
-          raise
-        except Exception as e:
+        exc, ret = ensure_sync(nbexecute_async_inner)(nb, mnt, emit=emit, subscribe=subscribe)
+        nb = ret['nb']
+        state = ret['state']
+        if exc:
           logger.info(f"execution error: {traceback.format_exc()}")
-          await emit({ 'type': 'error', 'data': str(e) })
+          emit_sync({ 'type': 'error', 'data': str(exc) })
         # Save execution completion time
         logger.info('saving')
         nb.metadata['appyter']['nbexecute']['completed'] = datetime.datetime.now().replace(tzinfo=datetime.timezone.utc).isoformat()
@@ -153,19 +106,70 @@ async def nbexecute_async(ipynb='', emit=json_emitter_factory(sys.stdout), cwd='
           nb_to_ipynb_io(nb, fw)
         #
         logger.info('finalized')
-  except asyncio.CancelledError:
-    logger.info(f"outer cancelled")
-    raise
-  except Exception as e:
+  except Exception:
     logger.error(traceback.format_exc())
     if state['status'] != 'Success':
-      await emit({ 'type': 'error', 'data': f"Error occured while executing" })
+      emit_sync({ 'type': 'error', 'data': f"Error occured while executing" })
     else:
-      await emit({ 'type': 'error', 'data': f"Error occured while finalizing" })
+      emit_sync({ 'type': 'error', 'data': f"Error occured while finalizing" })
     raise
   finally:
     logger.info('complete')
-  #
+
+async def nbexecute_async_inner(nb, mnt, emit=None, subscribe=None):
+  state = dict(progress=0, status='Starting')
+  try:
+    logger.info('initializing')
+    if callable(subscribe):
+      await subscribe(lambda: dict(nb=nb_to_json(nb), **state))
+    iopub_hook = iopub_hook_factory(nb, emit)
+    client = NotebookClientIOPubHook(
+      nb,
+      allow_errors=True,
+      timeout=None,
+      kernel_name='python3',
+      resources={ 'metadata': {'path': str(mnt) } },
+      iopub_hook=iopub_hook,
+    )
+    await emit({ 'type': 'nb', 'data': nb_to_json(nb) })
+    try:
+      async with client.async_setup_kernel(
+        env=dict(
+          { k: v for k, v in os.environ.items() if not k.startswith('APPYTER_') },
+          PYTHONPATH=':'.join(sys.path),
+        ),
+      ):
+        logger.info('executing')
+        state.update(status='Executing...', progress=0)
+        await emit({ 'type': 'status', 'data': state['status'] })
+        await emit({ 'type': 'progress', 'data': state['progress'] })
+        client.set_widgets_metadata()
+        n_cells = len(nb.cells)
+        exec_count = 1
+        for index, cell in enumerate(nb.cells):
+          logger.debug(f"executing cell {index}")
+          cell = await client.async_execute_cell(
+            cell, index,
+            execution_count=exec_count,
+          )
+          if cell_is_code(cell):
+            if cell_has_error(cell):
+              raise Exception('Cell execution error on cell %d' % (exec_count))
+            exec_count += 1
+          if index < n_cells-1:
+            state['progress'] = index + 1
+            await emit({ 'type': 'progress', 'data': state['progress'] })
+          else:
+            state['status'] = 'Success'
+            await emit({ 'type': 'status', 'data': state['status'] })
+    finally:
+      client.set_widgets_metadata()
+  except Exception as exc:
+    return exc, dict(nb=nb, state=state)
+  else:
+    return None, dict(nb=nb, state=state)
+
+nbexecute_async = ensure_async(nbexecute_sync)
 
 @cli.command(help='Execute a jupyter notebook on the command line asynchronously')
 @click.option('-s', type=str, metavar='URI', default='file:///dev/stdout', help='Status stream')