Merge pull request #41 from lsst-sqre/tickets/DM-29532

[DM-29532] Code cleanup

dev-values.yaml

@@ -91,22 +91,21 @@ jupyterhub:
 
 config:
   base_url: "https://minikube.lsst.codes"
-  options_form:
-    images_url: "http://cachemachine.cachemachine.svc.cluster.local/cachemachine/jupyter/available"
-    images: []
-    sizes:
-      - name: Tiny
-        cpu: .5
-        ram: 1536M
-      - name: Small
-        cpu: 1
-        ram: 3072M
-      - name: Medium
-        cpu: 2
-        ram: 6144M
-      - name: Large
-        cpu: 4
-        ram: 12288M
+  images_url: "http://cachemachine.cachemachine.svc.cluster.local/cachemachine/jupyter/available"
+  pinned_images: []
+  sizes:
+    - name: Tiny
+      cpu: 0.5
+      ram: 1536M
+    - name: Small
+      cpu: 1
+      ram: 3072M
+    - name: Medium
+      cpu: 2
+      ram: 6144M
+    - name: Large
+      cpu: 4
+      ram: 12288M
   user_resources:
     - apiVersion: v1
       kind: Namespace
@@ -127,7 +126,7 @@ config:
         SODA_ROUTE: /api/image/soda
         WORKFLOW_ROUTE: /wf
         NO_SUDO: 'TRUE'
-        AUTO_REPO_URLS: https://github.com/lsst-sqre/notebook-demo
+        AUTO_REPO_URLS: "https://github.com/lsst-sqre/notebook-demo"
         EXTERNAL_GROUPS: "{{ external_groups }}"
         EXTERNAL_USER: "{{ user }}"
         EXTERNAL_UID: "{{ uid }}"

src/nublado2/auth.py

@@ -9,8 +9,8 @@
 from jupyterhub.utils import url_path_join
 from tornado import web
 
+from nublado2.http import get_session
 from nublado2.nublado_config import NubladoConfig
-from nublado2.options import session
 
 if TYPE_CHECKING:
     from typing import Any, Dict, List, Optional, Tuple, Type, Union
@@ -151,10 +151,11 @@ async def _build_auth_info(headers: HTTPHeaders) -> Dict[str, Any]:
             raise web.HTTPError(401, "No request token")
 
         # Retrieve the token metadata.
-        base_url = NubladoConfig().get().get("base_url")
+        base_url = NubladoConfig().base_url
         if not base_url:
             raise web.HTTPError(500, "base_url not set in configuration")
         api_url = url_path_join(base_url, "/auth/analyze")
+        session = await get_session()
         resp = await session.post(api_url, data={"token": token})
         if resp.status != 200:
             raise web.HTTPError(500, "Cannot reach token analysis API")

src/nublado2/hooks.py

@@ -20,8 +20,8 @@ async def pre_spawn(self, spawner: Spawner) -> None:
         )
 
         spawner.image = options.image_info.reference
-        spawner.mem_limit = options.ram
-        spawner.cpu_limit = options.cpu
+        spawner.mem_limit = options.size.ram
+        spawner.cpu_limit = options.size.cpu
 
         auth_state = await spawner.user.get_auth_state()
 

src/nublado2/http.py

@@ -0,0 +1,17 @@
+from aiohttp import ClientSession
+
+_session = None
+
+
+async def get_session() -> ClientSession:
+    """This is the way to retrieve a ClientSession to make HTTP requests.
+
+    ClientSession needs to be created inside an async function, so by
+    calling this, you ensure it exists, or create it if it doesn't.
+
+    Since there are some connection pools, we don't want to be creating
+    these all the time.  Better to just reuse one."""
+    global _session
+    if not _session:
+        _session = ClientSession()
+    return _session

src/nublado2/hub_config.py

@@ -16,7 +16,7 @@ def configure(self, c: JupyterHub) -> None:
         self.log.info("Configuring JupyterHub Nublado2 style")
         self.log.debug(f"JupyterHub configuration starting as: {c}")
 
-        nc = NubladoConfig().get()
+        nc = NubladoConfig()
         self.log.debug(f"Nublado Config is:\n{nc}")
 
         c.JupyterHub.hub_connect_url = self._get_hub_connect_url()

src/nublado2/imageinfo.py

@@ -1,4 +1,4 @@
-from dataclasses import dataclass, field
+from dataclasses import dataclass
 from typing import Dict
 
 FIELD_DELIMITER = "|"
@@ -31,14 +31,18 @@ class ImageInfo:
     Example: "sha256:419c4b7e14603711b25fa9e0569460a753c4b2449fe275bb5f89743b01794a30"  # noqa: E501
     """
 
-    packed_string: str = field(init=False, default="")
-    """packed_string is the form in which the image info is packed into the
-    JupyterHub options form and in which is is returned as the form
-    selection.  It is specification, display_name, and digest
-    concatenated with the pipe character.
+    @property
+    def packed_string(self) -> str:
+        """packed_string is the form in which the image info is packed into the
+        JupyterHub options form and in which is is returned as the form
+        selection.  It is specification, display_name, and digest
+        concatenated with the pipe character.
 
-    Example: "registry.hub.docker.com/lsstsqre/sciplat-lab:w_2021_13|Weekly 13|sha256:419c4b7e14603711b25fa9e0569460a753c4b2449fe275bb5f89743b01794a30"  # noqa: E501
-    """
+        Example: "registry.hub.docker.com/lsstsqre/sciplat-lab:w_2021_13|Weekly 13|sha256:419c4b7e14603711b25fa9e0569460a753c4b2449fe275bb5f89743b01794a30"  # noqa: E501
+        """
+        return FIELD_DELIMITER.join(
+            [self.reference, self.display_name, self.digest]
+        )
 
     @classmethod
     def from_cachemachine_entry(cls, entry: CachemachineEntry) -> "ImageInfo":
@@ -67,11 +71,3 @@ def from_packed_string(cls, packed_string: str) -> "ImageInfo":
         return cls(
             reference=fields[0], display_name=fields[1], digest=fields[2]
         )
-
-    def __post_init__(self) -> None:
-        object.__setattr__(self, "packed_string", self._pack())
-
-    def _pack(self) -> str:
-        return FIELD_DELIMITER.join(
-            [self.reference, self.display_name, self.digest]
-        )

src/nublado2/labsize.py

@@ -0,0 +1,20 @@
+from dataclasses import dataclass
+
+
+@dataclass(frozen=True)
+class LabSize:
+    """The cpu and ram settings for a lab container."""
+
+    cpu: float
+    """Number of virtual CPUs to allocate for this lab.
+
+    This can be a partial number, such as 2.5 or .5 vCPUs."""
+
+    name: str
+    """The name referring to this pairing of cpu and ram."""
+
+    ram: str
+    """Amount of memory to allocate for this lab.
+
+    This is a string with special characters for units, such as
+    2048M, or 2G."""

src/nublado2/nublado_config.py

@@ -2,26 +2,64 @@
 
 __all__ = ["NubladoConfig"]
 
-from typing import Any, Dict, Tuple
+from typing import Any, Dict, List, Tuple
 
 from ruamel import yaml
 from ruamel.yaml import RoundTripLoader
-from traitlets.config import LoggingConfigurable
 
+from nublado2.imageinfo import ImageInfo
+from nublado2.labsize import LabSize
 
-class NubladoConfig(LoggingConfigurable):
-    def get(self) -> Dict[str, Any]:
+
+class NubladoConfig:
+    def __init__(self) -> None:
+        """Load the nublado_config.yaml file from disk.
+
+        This file normally comes from mounting a configmap with the
+        nublado_config.yaml mounted into the hub container."""
         with open("/etc/jupyterhub/nublado_config.yaml") as f:
-            nc = yaml.load(f.read(), Loader=RoundTripLoader)
+            self._config = yaml.load(f.read(), Loader=RoundTripLoader)
+
+        self._sizes = {
+            s.name: s
+            for s in [
+                LabSize(float(s["cpu"]), s["name"], s["ram"])
+                for s in self._config["options_form"]["sizes"]
+            ]
+        }
+
+    @property
+    def base_url(self) -> str:
+        """Base URL for the environment, like https://data.lsst.cloud"""
+        return self._config["base_url"]
+
+    @property
+    def images_url(self) -> str:
+        """URL to fetch list of images to show in options form.
+
+        Generally, this is a link to the cachemachine service."""
+        return self._config["images_url"]
 
-        self.log.debug(f"Loaded Nublado Config:\n{nc}")
-        return nc
+    @property
+    def pinned_images(self) -> List[ImageInfo]:
+        """List of images to keep pinned in the options form."""
+        return [
+            ImageInfo.from_cachemachine_entry(i)
+            for i in self._config["pinned_images"]
+        ]
 
-    def lookup_size(self, name: str) -> Tuple[float, str]:
-        sizes = self.get()["options_form"]["sizes"]
+    @property
+    def signing_key(self) -> str:
+        """Retrieve the gafaelfawr signing key to mint tokens."""
+        with open("/etc/keys/signing_key.pem", "r") as f:
+            return f.read()
 
-        for s in sizes:
-            if s["name"] == name:
-                return (float(s["cpu"]), s["ram"])
+    @property
+    def sizes(self) -> Dict[str, LabSize]:
+        """Retrieve a copy of the sizes a lab can spawn as."""
+        return dict(self._sizes)
 
-        raise ValueError(f"Size {name} not found")
+    @property
+    def user_resources(self) -> Tuple[Any, ...]:
+        """Retrieve a copy of the lab resources templates."""
+        return tuple(self._config.get("user_resources", []))

src/nublado2/options.py

@@ -1,10 +1,10 @@
-from typing import Dict, List, Optional
+from typing import List, Optional, Tuple
 
-from aiohttp import ClientSession
 from jinja2 import Template
 from jupyterhub.spawner import Spawner
 from traitlets.config import LoggingConfigurable
 
+from nublado2.http import get_session
 from nublado2.imageinfo import ImageInfo
 from nublado2.nublado_config import NubladoConfig
 
@@ -73,47 +73,42 @@
 """
 )
 
-# Don't have this be a member of NubladoOptions, we should
-# share this connection pool.  Also the LoggingConfigurable
-# will try to pickle it to json, and it can't pickle a session.
-session = ClientSession()
-
 
 class NubladoOptions(LoggingConfigurable):
     async def show_options_form(self, spawner: Spawner) -> str:
-        options_config = NubladoConfig().get()["options_form"]
-        sizes = options_config["sizes"]
-
-        images_url = options_config.get("images_url")
+        nc = NubladoConfig()
 
-        cachemachine_response = await self._get_images_from_url(images_url)
+        (cached_images, all_images) = await self._get_images_from_url(
+            nc.images_url
+        )
+        cached_images.extend(nc.pinned_images)
 
-        all_imageinfos = [
-            ImageInfo.from_cachemachine_entry(img)
-            for img in cachemachine_response["all"]
-        ]
-        # Start with the cachemachine response, then extend it with
-        #  contents of options_config
-        cached_images = cachemachine_response["images"]
-        cached_images.extend(options_config["images"])
-        cached_imageinfos = [
-            ImageInfo.from_cachemachine_entry(img) for img in cached_images
-        ]
         return options_template.render(
             dropdown_sentinel=DROPDOWN_SENTINEL_VALUE,
-            cached_images=cached_imageinfos,
-            all_images=all_imageinfos,
-            sizes=sizes,
+            cached_images=cached_images,
+            all_images=all_images,
+            sizes=nc.sizes.values(),
         )
 
     async def _get_images_from_url(
         self, url: Optional[str]
-    ) -> Dict[str, List[Dict[str, str]]]:
+    ) -> Tuple[List[ImageInfo], List[ImageInfo]]:
         if not url:
-            return {"all": [], "images": []}
+            return ([], [])
 
+        session = await get_session()
         r = await session.get(url)
         if r.status != 200:
             raise Exception(f"Error {r.status} from {url}")
 
-        return await r.json()
+        body = await r.json()
+
+        cached_images = [
+            ImageInfo.from_cachemachine_entry(img) for img in body["images"]
+        ]
+
+        all_images = [
+            ImageInfo.from_cachemachine_entry(img) for img in body["all"]
+        ]
+
+        return (cached_images, all_images)