DM-37758: Replace Cachemachine with JupyterLab Controller

requirements/dev.txt

@@ -1,5 +1,5 @@
 #
-# This file is autogenerated by pip-compile with Python 3.10
+# This file is autogenerated by pip-compile with Python 3.11
 # by the following command:
 #
 #    pip-compile --generate-hashes --output-file=requirements/dev.txt requirements/dev.in
@@ -204,10 +204,6 @@ distlib==0.3.6 \
     --hash=sha256:14bad2d9b04d3a36127ac97f30b12a19268f211063d8f8ee4f47108896e11b46 \
     --hash=sha256:f35c4b692542ca110de7ef0bea44d73981caeb34ca0b9b6b2e6d7790dda8f80e
     # via virtualenv
-exceptiongroup==1.1.0 \
-    --hash=sha256:327cbda3da756e2de031a3107b81ab7b3770a602c4d16ca618298c526f4bec1e \
-    --hash=sha256:bcb67d800a4497e1b404c2dd44fca47d3b7a5e5433dbab67f96c1a685cdfdf23
-    # via pytest
 filelock==3.9.0 \
     --hash=sha256:7b319f24340b51f55a2bf7a12ac0755a9b03e718311dac567a0f4f7fabd2f5de \
     --hash=sha256:f58d535af89bb9ad5cd4df046f741f8553a418c01a7856bf0d173bbc9f6bd16d
@@ -310,9 +306,9 @@ httpx==0.23.3 \
     #   -c requirements/main.txt
     #   -r requirements/dev.in
     #   pytest-httpx
-identify==2.5.16 \
-    --hash=sha256:832832a58ecc1b8f33d5e8cb4f7d3db2f5c7fbe922dfee5f958b48fed691501a \
-    --hash=sha256:c47acedfe6495b1c603ed7e93469b26e839cab38db4155113f36f718f8b3dc47
+identify==2.5.17 \
+    --hash=sha256:7d526dd1283555aafcc91539acc061d8f6f59adb0a7bba462735b0a318bff7ed \
+    --hash=sha256:93cc61a861052de9d4c541a7acb7e3dcc9c11b398a2144f6e52ae5285f5f4f06
     # via pre-commit
 idna==3.4 \
     --hash=sha256:814f528e8dead7d329833b91c5faa87d60bf71824cd12a7530b5526063d02cb4 \
@@ -460,9 +456,9 @@ pluggy==1.0.0 \
     --hash=sha256:4224373bacce55f955a878bf9cfa763c1e360858e330072059e10bad68531159 \
     --hash=sha256:74134bbf457f031a36d68416e1509f34bd5ccc019f0bcc952c7b909d06b37bd3
     # via pytest
-pre-commit==3.0.2 \
-    --hash=sha256:aa97fa71e7ab48225538e1e91a6b26e483029e6de64824f04760c32557bc91d7 \
-    --hash=sha256:f448d5224c70e196a6c6f87961d2333dfdc49988ebbf660477f9efe991c03597
+pre-commit==3.0.3 \
+    --hash=sha256:4187e74fda38f0f700256fb2f757774385503b04292047d0899fc913207f314b \
+    --hash=sha256:83e2e8cc5cbb3691cff9474494816918d865120768aa36c9eda6185126667d21
     # via -r requirements/dev.in
 pytest==7.2.1 \
     --hash=sha256:c7c6ca206e93355074ae32f7403e8ea12163b1163c976fee7d4d84027c162be5 \
@@ -543,16 +539,9 @@ sniffio==1.3.0 \
     #   asgi-lifespan
     #   httpcore
     #   httpx
-tomli==2.0.1 \
-    --hash=sha256:939de3e7a6161af0c887ef91b7d41a53e7c5a1ca976325f429cb46ea9bc30ecc \
-    --hash=sha256:de526c12914f0c550d15924c62d72abc48d6fe7364aa87328337a31007fe8a4f
-    # via
-    #   coverage
-    #   mypy
-    #   pytest
-types-pyyaml==6.0.12.3 \
-    --hash=sha256:17ce17b3ead8f06e416a3b1d5b8ddc6cb82a422bb200254dd8b469434b045ffc \
-    --hash=sha256:879700e9f215afb20ab5f849590418ab500989f83a57e635689e1d50ccc63f0c
+types-pyyaml==6.0.12.4 \
+    --hash=sha256:ade6e328a5a3df816c47c912c2e1e946ae2bace90744aa73111ee6834b03a314 \
+    --hash=sha256:de3bacfc4e0772d9b1baf007c37354f3c34c8952e90307d5155b6de0fc183a67
     # via -r requirements/dev.in
 types-requests==2.28.11.8 \
     --hash=sha256:61960554baca0008ae7e2db2bd3b322ca9a144d3e80ce270f5fb640817e40994 \

requirements/main.txt

@@ -1,5 +1,5 @@
 #
-# This file is autogenerated by pip-compile with Python 3.10
+# This file is autogenerated by pip-compile with Python 3.11
 # by the following command:
 #
 #    pip-compile --generate-hashes --output-file=requirements/main.txt requirements/main.in

src/mobu/business/jupyterloginloop.py

@@ -22,7 +22,7 @@
 )
 from ..jupyterclient import JupyterClient
 from ..models.business import BusinessConfig, BusinessData
-from ..models.jupyter import JupyterImage
+from ..models.jupyter import ControllerImage
 from ..models.user import AuthenticatedUser
 from .base import Business
 
@@ -69,7 +69,7 @@ def __init__(
         user: AuthenticatedUser,
     ) -> None:
         super().__init__(logger, business_config, user)
-        self.image: Optional[JupyterImage] = None
+        self.image: Optional[ControllerImage] = None
         self._client = JupyterClient(user, logger, business_config.jupyter)
 
     async def close(self) -> None:

src/mobu/config.py

@@ -37,17 +37,6 @@ class Configuration:
     Set with the ``ENVIRONMENT_URL`` environment variable.
     """
 
-    cachemachine_image_policy: str = os.getenv(
-        "CACHEMACHINE_IMAGE_POLICY", "available"
-    )
-    """Whether to use the images available on all nodes, or the images
-    desired by cachemachine.  In instances where image streaming is enabled,
-    and therefore pulls are fast, ``desired`` is preferred.  The default is
-    ``available``.
-
-    Set with the ``CACHEMACHINE_IMAGE_POLICY`` environment variable.
-    """
-
     gafaelfawr_token: Optional[str] = os.getenv("GAFAELFAWR_TOKEN")
     """The Gafaelfawr admin token to use to create user tokens.
 

src/mobu/controller.py

@@ -0,0 +1,83 @@
+"""Client for the JupyterLab Controller service."""
+
+from __future__ import annotations
+
+from aiohttp import ClientSession
+
+from .config import config
+from .exceptions import ControllerError
+from .models.jupyter import ControllerImage, ControllerImages
+
+
+class ControllerClient:
+    """Query the JupyterLab Controller service for image information.
+
+    The JupyterLab Controller is canonical for the available images and
+    details such as which image is recommended and what the latest weeklies
+    are.  This client queries it and returns the image that matches some
+    selection criteria.
+
+    This will be modified into a form suitable for POSTing to the Controller
+    to create the requested lab.
+    """
+
+    def __init__(
+        self, session: ClientSession, token: str, username: str
+    ) -> None:
+        self._session = session
+        self._token = token
+        self._username = username
+        self._url = config.environment_url + "/nublado/spawner/v1/images"
+
+    async def get_latest_weekly(self) -> ControllerImage:
+        """Image for the latest weekly version.
+
+        Returns
+        -------
+        image : `mobu.models.jupyter.ControllerImage`
+            The corresponding image.
+
+        Raises
+        ------
+        mobu.exceptions.ControllerError
+            Some error occurred talking to JupyterLab Controller or it does
+            not have a latest weekly image.
+        """
+        images = await self._get_images()
+        if images.latest_weekly is None:
+            raise ControllerError(
+                self._username, "No latest weekly image found"
+            )
+        return images.latest_weekly
+
+    async def get_recommended(self) -> ControllerImage:
+        """Path suitable for image pulling for the latest recommended version.
+
+        Returns
+        -------
+        path : `mobu.models.jupyter.ControllerImage`
+            The corresponding image.
+
+        Raises
+        ------
+        mobu.exceptions.ControllerError
+            Some error occurred talking to JupyterLab Controller or it does
+            not have a recommended image.
+        """
+        images = await self._get_images()
+        if images.recommended is None:
+            raise ControllerError(self._username, "No recommended image found")
+        return images.recommended
+
+    async def _get_images(self) -> ControllerImages:
+        headers = {"Authorization": f"bearer {self._token}"}
+        async with self._session.get(self._url, headers=headers) as r:
+            if r.status != 200:
+                msg = f"Cannot get image status: {r.status} {r.reason}"
+                raise ControllerError(self._username, msg)
+            try:
+                data = await r.json()
+                return ControllerImages.parse_obj(data)
+            except Exception as e:
+                msg = f"Invalid response: {type(e).__name__}: {str(e)}"
+                raise ControllerError(self._username, msg)

src/mobu/exceptions.py

@@ -10,7 +10,7 @@
 from .constants import DATE_FORMAT
 
 __all__ = [
-    "CachemachineError",
+    "ControllerError",
     "CodeExecutionError",
     "FlockNotFoundException",
     "JupyterError",
@@ -111,11 +111,11 @@ def common_fields(self) -> List[Dict[str, Union[str, bool]]]:
         return fields
 
 
-class CachemachineError(SlackError):
-    """Failed to obtain a valid image list from cachemachine."""
+class ControllerError(SlackError):
+    """Failed to obtain a valid image list from JupyterLab Controller."""
 
     def __init__(self, user: str, msg: str) -> None:
-        super().__init__(user, f"Cachemachine error: {msg}")
+        super().__init__(user, f"JupyterLab Controller error: {msg}")
 
 
 class CodeExecutionError(SlackError):

src/mobu/jupyterclient.py

@@ -21,6 +21,7 @@
     Awaitable,
     Callable,
     Dict,
+    List,
     Optional,
     TypeVar,
     cast,
@@ -38,15 +39,15 @@
 from aiohttp.client import _RequestContextManager, _WSRequestContextManager
 from structlog import BoundLogger
 
-from .cachemachine import CachemachineClient
 from .config import config
+from .controller import ControllerClient
 from .exceptions import (
     CodeExecutionError,
     JupyterError,
     JupyterResponseError,
     JupyterWebSocketError,
 )
-from .models.jupyter import JupyterConfig, JupyterImage, JupyterImageClass
+from .models.jupyter import ControllerImage, JupyterConfig, JupyterImageClass
 from .models.user import AuthenticatedUser
 
 __all__ = ["JupyterClient", "JupyterLabSession"]
@@ -242,10 +243,10 @@ def __init__(
         session.cookie_jar.update_cookies(BaseCookie({"_xsrf": xsrftoken}))
         self.session = JupyterClientSession(session, user.token)
 
-        # We also send the XSRF token to cachemachine because of how we're
-        # sharing the session, but that shouldn't matter.
+        # We also send the XSRF token to the JupyterLabController because
+        # of how we're sharing the session, but that shouldn't matter.
         assert config.gafaelfawr_token
-        self.cachemachine = CachemachineClient(
+        self.controller = ControllerClient(
             session, config.gafaelfawr_token, self.user.username
         )
 
@@ -295,17 +296,21 @@ async def is_lab_stopped(self, final: bool = False) -> bool:
         return result
 
     @_convert_exception
-    async def spawn_lab(self) -> JupyterImage:
+    async def spawn_lab(self) -> ControllerImage:
         spawn_url = self.jupyter_url + "hub/spawn"
 
         # Determine what image to spawn.
         if self.config.image_class == JupyterImageClass.RECOMMENDED:
-            image = await self.cachemachine.get_recommended()
+            image = await self.controller.get_recommended()
         elif self.config.image_class == JupyterImageClass.LATEST_WEEKLY:
-            image = await self.cachemachine.get_latest_weekly()
+            image = await self.controller.get_latest_weekly()
         else:
             assert self.config.image_reference
-            image = JupyterImage.from_reference(self.config.image_reference)
+            ref = self.config.image_reference
+            image = ControllerImage(
+                path=ref, name=ref, digest="unknown", tags={}
+            )
+
         msg = f"Spawning lab image {image.name} for {self.user.username}"
         self.log.info(msg)
 
@@ -501,10 +506,12 @@ def _remove_ansi_escapes(string: str) -> str:
         """
         return _ANSI_REGEX.sub("", string)
 
-    def _build_jupyter_spawn_form(self, image: JupyterImage) -> Dict[str, str]:
+    def _build_jupyter_spawn_form(
+        self, image: ControllerImage
+    ) -> Dict[str, List[str]]:
         """Construct the form to submit to the JupyterHub login page."""
         return {
-            "image_list": str(image),
-            "image_dropdown": "use_image_from_dropdown",
-            "size": self.config.image_size,
+            "image_list": [image.path],
+            "image_dropdown": [image.path],  # Not used
+            "size": [self.config.image_size],
         }

src/mobu/models/business.py

@@ -5,7 +5,7 @@
 from pydantic import BaseModel, Field
 
 from ..constants import NOTEBOOK_REPO_BRANCH, NOTEBOOK_REPO_URL
-from .jupyter import JupyterConfig, JupyterImage
+from .jupyter import ControllerImage, JupyterConfig
 from .timings import StopwatchData
 
 __all__ = ["BusinessConfig", "BusinessData"]
@@ -189,7 +189,7 @@ class BusinessData(BaseModel):
 
     timings: List[StopwatchData] = Field(..., title="Timings of events")
 
-    image: Optional[JupyterImage] = Field(
+    image: Optional[ControllerImage] = Field(
         None,
         title="JupyterLab image information",
         description="Will only be present when there is an active Jupyter lab",