Merge pull request #12 from lsst-sqre/tickets/DM-26625

[DM-26625] Robo-Simon

.github/workflows/ci.yaml

@@ -12,7 +12,7 @@ jobs:
       - name: Set up Python
         uses: actions/setup-python@v1
         with:
-          python-version: 3.7
+          python-version: 3.8
 
       - name: Install tox
         run: pip install tox

Dockerfile

@@ -14,7 +14,7 @@
 #   - Runs a non-root user.
 #   - Sets up the entrypoint and port.
 
-FROM python:3.7-slim-buster AS base-image
+FROM python:3.8-slim-buster AS base-image
 
 # Update system packages
 COPY scripts/install-base-packages.sh .

pyproject.toml

@@ -95,5 +95,5 @@ exclude = '''
 include_trailing_comma = true
 multi_line_output = 3
 known_first_party = ["mobu", "tests"]
-known_third_party = ["aiohttp", "aiojobs", "click", "jinja2", "jwt", "pyvo", "requests", "safir", "setuptools", "structlog"]
+known_third_party = ["aiohttp", "aiojobs", "click", "git", "jinja2", "jwt", "pyvo", "requests", "safir", "setuptools", "structlog"]
 skip = ["docs/conf.py"]

requirements/main.in

@@ -12,6 +12,7 @@ aiojobs
 cchardet
 click
 cryptography
+gitpython
 importlib_metadata
 jinja2
 pyjwt

scripts/lsptestuser05.json

@@ -1,6 +1,6 @@
 {
   "username": "lsptestuser05",
   "uidnumber": 60185,
-  "business": "JupyterLoginLoop",
+  "business": "NotebookRunner",
   "restart": true
 }

src/mobu/jupyterclient.py

@@ -9,6 +9,7 @@
 
 import asyncio
 import random
+import re
 import string
 from dataclasses import dataclass
 from http.cookies import BaseCookie
@@ -21,6 +22,12 @@
 from mobu.user import User
 
 
+class NotebookException(Exception):
+    """Passing an error back from a remote notebook session."""
+
+    pass
+
+
 @dataclass
 class JupyterClient:
     log: BoundLoggerLazyProxy
@@ -48,6 +55,12 @@ def __init__(self, user: User, log: BoundLoggerLazyProxy):
             BaseCookie({"_xsrf": self.xsrftoken})
         )
 
+    __ansi_reg_exp = re.compile(r"(\x9B|\x1B\[)[0-?]*[ -\/]*[@-~]")
+
+    @classmethod
+    def _ansi_escape(cls, line: str) -> str:
+        return cls.__ansi_reg_exp.sub("", line)
+
     async def hub_login(self) -> None:
         async with self.session.get(self.jupyter_url + "hub/login") as r:
             if r.status != 200:
@@ -188,14 +201,24 @@ async def run_python(self, kernel_id: str, code: str) -> str:
 
             while True:
                 r = await ws.receive_json()
+                self.log.debug(f"Recieved kernel message: {r}")
                 msg_type = r["msg_type"]
                 if msg_type == "error":
-                    raise Exception(f"Error running python {r}")
+                    error_message = "".join(r["content"]["traceback"])
+                    raise NotebookException(self._ansi_escape(error_message))
                 elif (
                     msg_type == "stream"
                     and msg_id == r["parent_header"]["msg_id"]
                 ):
                     return r["content"]["text"]
+                elif msg_type == "execute_reply":
+                    status = r["content"]["status"]
+                    if status == "ok":
+                        return ""
+                    else:
+                        raise NotebookException(
+                            f"Error content status is {status}"
+                        )
 
     def dump(self) -> dict:
         return {

src/mobu/monkeybusinessfactory.py

@@ -10,6 +10,7 @@
 from mobu.jupyterloginloop import JupyterLoginLoop
 from mobu.jupyterpythonloop import JupyterPythonLoop
 from mobu.monkey import Monkey
+from mobu.notebookrunner import NotebookRunner
 from mobu.querymonkey import QueryMonkey
 from mobu.user import User
 
@@ -30,6 +31,7 @@ def create(body: Dict) -> Monkey:
             Business,
             JupyterLoginLoop,
             JupyterPythonLoop,
+            NotebookRunner,
             QueryMonkey,
         ]
 

src/mobu/notebookrunner.py

@@ -0,0 +1,120 @@
+"""NotebookRunner logic for mobu.
+
+This business pattern will clone a git repo full
+of notebooks, randomly pick the notebooks, and run
+them on the remote jupyter lab."""
+
+__all__ = [
+    "NotebookRunner",
+]
+
+import json
+import os
+from dataclasses import dataclass, field
+from pathlib import Path
+from tempfile import TemporaryDirectory
+from typing import Iterator
+
+import git
+
+from mobu.business import Business
+from mobu.jupyterclient import JupyterClient, NotebookException
+
+REPO_URL = "https://github.com/lsst-sqre/notebook-demo.git"
+REPO_BRANCH = "prod"
+
+
+@dataclass
+class NotebookRunner(Business):
+    success_count: int = 0
+    failure_count: int = 0
+    _client: JupyterClient = field(init=False)
+    _failed_notebooks: list = field(init=False, default_factory=list)
+    _repo_dir: TemporaryDirectory = field(
+        init=False, default_factory=TemporaryDirectory
+    )
+    _repo: git.Repo = field(init=False, default=None)
+    _notebook_iterator: Iterator = field(init=False)
+    notebook: os.DirEntry = field(init=False)
+    code: str = field(init=False)
+
+    async def run(self) -> None:
+        try:
+            logger = self.monkey.log
+
+            self._client = JupyterClient(self.monkey.user, logger)
+
+            if not self._repo:
+                self._repo = git.Repo.clone_from(
+                    REPO_URL, self._repo_dir.name, branch=REPO_BRANCH
+                )
+
+            self._notebook_iterator = os.scandir(self._repo_dir.name)
+
+            logger.info("Repository cloned and ready")
+
+            await self._client.hub_login()
+
+            while True:
+                self._next_notebook()
+
+                if self.success_count % 100 == 0:
+                    await self._client.delete_lab()
+
+                await self._client.ensure_lab()
+
+                if self.notebook.path.endswith(".ipynb"):
+                    logger.info(f"Starting notebook: {self.notebook.name}")
+                    notebook_text = Path(self.notebook.path).read_text()
+                    cells = json.loads(notebook_text)["cells"]
+
+                    kernel = await self._client.create_kernel(
+                        kernel_name="LSST"
+                    )
+
+                    for cell in cells:
+                        if cell["cell_type"] == "code":
+                            self.code = "".join(cell["source"])
+                            logger.info("Executing:\n%s\n", self.code)
+                            reply = await self._client.run_python(
+                                kernel, self.code
+                            )
+
+                            if reply:
+                                logger.info(f"Response:\n{reply}\n")
+
+                    logger.info(
+                        f"Success running notebook: {self.notebook.name}"
+                    )
+
+                self.success_count += 1
+
+        except NotebookException as e:
+            logger.error(f"Error running notebook: {self.notebook.name}")
+            self._failed_notebooks.append(self.notebook.name)
+            self.failure_count += 1
+            raise NotebookException(
+                f"Running {self.notebook.name}: '"
+                f"```{self.code}``` generated: ```{e}```"
+            )
+
+    def dump(self) -> dict:
+        return {
+            "name": "NotebookRunner",
+            "current_notebook": self.notebook.name,
+            "running_code": self.code,
+            "failed_notebooks": self._failed_notebooks,
+            "failure_count": self.failure_count,
+            "success_count": self.success_count,
+            "jupyter_client": self._client.dump(),
+        }
+
+    def _next_notebook(self) -> None:
+        try:
+            self.notebook = next(self._notebook_iterator)
+        except StopIteration:
+            self.monkey.log.info(
+                "Done with this cycle of notebooks, recreating lab."
+            )
+            self._notebook_iterator = os.scandir(self._repo_dir.name)
+            self._next_notebook()