Merge pull request #63 from lsst-sqre/tickets/DM-30978

Shut down kernel after notebook execution
lsst-sqre · Jul 13, 2021 · be88186 · be88186
2 parents f0cc39e + 14a4bcb
commit be88186
Show file tree

Hide file tree

Showing 3 changed files with 38 additions and 17 deletions.
diff --git a/src/mobu/jupyterclient.py b/src/mobu/jupyterclient.py
@@ -172,6 +172,16 @@ async def create_kernel(self, kernel_name: str = "LSST") -> str:
             response = await r.json()
             return response["id"]
 
+    async def delete_kernel(self, kernel_id: str) -> None:
+        kernel_url = (
+            self.jupyter_url
+            + f"user/{self.user.username}/api/kernels/{kernel_id}"
+        )
+        async with self.session.delete(kernel_url, raise_for_status=True) as r:
+            if r.status != 204:
+                self.log.warning(f"Delete kernel {kernel_id}: {r}")
+            return
+
     async def run_python(self, kernel_id: str, code: str) -> str:
         kernel_url = (
             self.jupyter_url

diff --git a/src/mobu/jupyterpythonloop.py b/src/mobu/jupyterpythonloop.py
@@ -13,36 +13,44 @@
 from mobu.jupyterclient import JupyterClient
 from mobu.jupyterloginloop import JupyterLoginLoop
 
+MAX_EXECUTIONS = 20
+SLEEP_TIME = 1
+
 
 @dataclass
 class JupyterPythonLoop(JupyterLoginLoop):
     async def run(self) -> None:
         logger = self.monkey.log
         logger.info("Starting up...")
 
-        client = JupyterClient(self.monkey.user, logger, self.options)
-        self._client = client
+        self._client = JupyterClient(self.monkey.user, logger, self.options)
         self.start_event("hub_login")
         await self._client.hub_login()
         self.stop_current_event()
         self.start_event("ensure_lab")
-        await client.ensure_lab()
-        self.stop_current_event()
-        self.start_event("create_kernel")
-        kernel = await client.create_kernel()
+        await self._client.ensure_lab()
         self.stop_current_event()
-
         while True:
-            self.start_event("execute_code")
-            code_str = "print(2+2)"
-            reply = await client.run_python(kernel, code_str)
-            sw = self.get_current_event()
-            if sw is not None:
-                sw.annotation = {"code": code_str, "result": reply}
+            logger.info("create_kernel")
+            self.start_event("create_kernel")
+            kernel = await self._client.create_kernel()
             self.stop_current_event()
-            logger.info(f"{code_str} -> {reply}")
-            self.start_event("lab_wait")
-            await asyncio.sleep(60)
+
+            for count in range(MAX_EXECUTIONS):
+                self.start_event("execute_code")
+                code_str = "print(2+2)"
+                reply = await self._client.run_python(kernel, code_str)
+                sw = self.get_current_event()
+                if sw is not None:
+                    sw.annotation = {"code": code_str, "result": reply}
+                self.stop_current_event()
+                logger.info(f"{code_str} -> {reply}")
+                self.start_event("lab_wait")
+                await asyncio.sleep(SLEEP_TIME)
+                self.stop_current_event()
+            logger.info("delete_kernel")
+            self.start_event("delete_kernel")
+            await self._client.delete_kernel(kernel)
             self.stop_current_event()
 
     def dump(self) -> dict:

diff --git a/src/mobu/notebookrunner.py b/src/mobu/notebookrunner.py
@@ -99,7 +99,10 @@ async def run(self) -> None:
                                 }
                             self.stop_current_event()
                             logger.info(f"Result:\n{reply}\n")
-
+                    logger.info(f"Deleting kernel {kernel}")
+                    self.start_event("delete_kernel")
+                    await self._client.delete_kernel(kernel)
+                    self.stop_current_event()
                     logger.info(
                         f"Success running notebook: {self.notebook.name}"
                     )