Merge pull request #57 from lsst-sqre/tickets/DM-30007

[DM-30007] Rework lab environment creation

src/nublado2/nublado_config.py

@@ -49,6 +49,11 @@ def images_url(self) -> str:
         Generally, this is a link to the cachemachine service."""
         return self._config["images_url"]
 
+    @property
+    def lab_environment(self) -> Dict[str, str]:
+        """Environment variable settings for the lab (possibly templates)."""
+        return dict(self._config.get("lab_environment", {}))
+
     @property
     def pinned_images(self) -> List[ImageInfo]:
         """List of images to keep pinned in the options form."""
@@ -73,11 +78,6 @@ def user_resources_template(self) -> str:
         """Retrieve a copy of the lab resources templates."""
         return self._config.get("user_resources_template")
 
-    @property
-    def custom_resources_template(self) -> str:
-        """Retrieve a copy of the lab custom resource templates."""
-        return self._config.get("custom_resources_template")
-
     @property
     def volumes(self) -> List[Dict[str, Any]]:
         return list(self._config["volumes"])

src/nublado2/resourcemgr.py

@@ -1,24 +1,33 @@
+from __future__ import annotations
+
+from io import StringIO
+from typing import TYPE_CHECKING
+
 import aiohttp
 from jinja2 import Template
-from jupyterhub.spawner import Spawner
 from kubernetes import client, config
 from kubernetes.utils import create_from_dict
-from ruamel import yaml
-from ruamel.yaml import RoundTripLoader
+from ruamel.yaml import YAML
 from traitlets.config import LoggingConfigurable
 
 from nublado2.crdparser import CRDParser
 from nublado2.nublado_config import NubladoConfig
 from nublado2.provisioner import Provisioner
-from nublado2.selectedoptions import SelectedOptions
+
+if TYPE_CHECKING:
+    from typing import Any, Dict
+
+    from jupyterhub.spawner import Spawner
+
+    from nublado2.selectedoptions import SelectedOptions
 
 
 class ResourceManager(LoggingConfigurable):
     # These k8s clients don't copy well with locks, connection,
     # pools, locks, etc.  Copying seems to happen under the hood of the
     # LoggingConfigurable base class, so just have them be class variables.
     # Should be safe to share these, and better to have fewer of them.
-    k8s_api = client.api_client.ApiClient()
+    k8s_api = client.ApiClient()
     custom_api = client.CustomObjectsApi()
     k8s_client = client.CoreV1Api()
 
@@ -30,85 +39,87 @@ def __init__(self) -> None:
             headers={"Authorization": f"Bearer {token}"}
         )
         self.provisioner = Provisioner(self.http_client)
+        self.yaml = YAML()
+        self.yaml.indent(mapping=2, sequence=4, offset=2)
 
     async def create_user_resources(
         self, spawner: Spawner, options: SelectedOptions
     ) -> None:
         """Create the user resources for this spawning session."""
         await self.provisioner.provision_homedir(spawner)
-
         try:
-            auth_state = await spawner.user.get_auth_state()
-            self.log.debug(f"Auth state={auth_state}")
-            groups = auth_state["groups"]
-
-            # Build a comma separated list of group:gid
-            # ex: group1:1000,group2:1001,group3:1002
-            external_groups = ",".join(
-                [f'{g["name"]}:{g["id"]}' for g in groups]
-            )
-
-            # Retrieve image tag and corresponding hash (if any)
-            # These come back from the options form as one-item lists
-
-            template_values = {
-                "user_namespace": spawner.namespace,
-                "user": spawner.user.name,
-                "uid": auth_state["uid"],
-                "token": auth_state["token"],
-                "groups": groups,
-                "external_groups": external_groups,
-                "base_url": self.nublado_config.base_url,
-                "dask_yaml": await self._build_dask_template(spawner),
-                "options": options,
-                "labels": spawner.common_labels,
-                "annotations": spawner.extra_annotations,
-                "nublado_base_url": spawner.hub.base_url,
-                "butler_secret_path": self.nublado_config.butler_secret_path,
-            }
-
-            self.log.debug(f"Template values={template_values}")
-            self.log.debug("Template:")
-            self.log.debug(self.nublado_config.user_resources_template)
-            t = Template(self.nublado_config.user_resources_template)
-            templated_user_resources = t.render(template_values)
-            self.log.debug("Generated user resources:")
-            self.log.debug(templated_user_resources)
-
-            user_resources = yaml.load(
-                templated_user_resources, Loader=RoundTripLoader
-            )
-
-            for r in user_resources:
-                self.log.debug(f"Creating: {r}")
-                create_from_dict(self.k8s_api, r)
+            await self._create_kubernetes_resources(spawner, options)
         except Exception:
             self.log.exception("Exception creating user resource!")
             raise
-        try:
-            # CRDs cannot be created with create_from_dict:
+
+    def _create_lab_environment_configmap(
+        self, spawner: Spawner, template_values: Dict[str, Any]
+    ) -> None:
+        """Create the ConfigMap that holds environment settings for the lab."""
+        environment = {}
+        for variable, template in self.nublado_config.lab_environment.items():
+            value = Template(template).render(template_values)
+            environment[variable] = value
+
+        self.log.debug(f"Creating environment ConfigMap with {environment}")
+        body = client.V1ConfigMap(
+            api_version="v1",
+            kind="ConfigMap",
+            metadata=client.V1ObjectMeta(
+                name="lab-environment",
+                namespace=spawner.namespace,
+                annotations=spawner.extra_annotations,
+                labels=spawner.common_labels,
+            ),
+            data=environment,
+        )
+        self.k8s_client.create_namespaced_config_map(spawner.namespace, body)
+
+    async def _create_kubernetes_resources(
+        self, spawner: Spawner, options: SelectedOptions
+    ) -> None:
+        template_values = await self._build_template_values(spawner, options)
+
+        # Construct the lab environment ConfigMap.  This is constructed from
+        # configuration settings and doesn't use a resource template like
+        # other resources.
+        self._create_lab_environment_configmap(spawner, template_values)
+
+        # Generate the list of additional user resources from the template.
+        self.log.debug("Template:")
+        self.log.debug(self.nublado_config.user_resources_template)
+        t = Template(self.nublado_config.user_resources_template)
+        templated_user_resources = t.render(template_values)
+        self.log.debug("Generated user resources:")
+        self.log.debug(templated_user_resources)
+        resources = self.yaml.load(templated_user_resources)
+
+        # Add in the standard labels and annotations common to every resource
+        # and create the resources.
+        for resource in resources:
+            if "metadata" not in resource:
+                resource["metadata"] = {}
+            resource["metadata"]["annotations"] = spawner.extra_annotations
+            resource["metadata"]["labels"] = spawner.common_labels
+
+            # Custom resources cannot be created by create_from_dict:
             # https://github.com/kubernetes-client/python/issues/740
-            ct = Template(self.nublado_config.custom_resources_template)
-            templated_custom_resources = ct.render(template_values)
-            self.log.debug("Generated custom resources:")
-            self.log.debug(templated_custom_resources)
-            custom_resources = yaml.load(
-                templated_custom_resources, Loader=RoundTripLoader
-            )
-            for cr in custom_resources:
-                self.log.debug(f"Creating: {cr}")
-                crd_parser = CRDParser.from_crd_body(cr)
-                self.log.debug(f"CRD_Parser: {crd_parser}")
+            #
+            # Detect those from the apiVersion field and handle them
+            # specially.
+            api_version = resource["apiVersion"]
+            if "." in api_version and ".k8s.io/" not in api_version:
+                crd_parser = CRDParser.from_crd_body(resource)
                 self.custom_api.create_namespaced_custom_object(
-                    body=cr,
+                    body=resource,
                     group=crd_parser.group,
                     version=crd_parser.version,
                     namespace=spawner.namespace,
                     plural=crd_parser.plural,
                 )
-        except Exception:
-            self.log.exception("Exception creating custom resource!")
-            raise
+            else:
+                create_from_dict(self.k8s_api, resource)
 
     async def _build_dask_template(self, spawner: Spawner) -> str:
         """Build a template for dask workers from the jupyter pod manifest."""
@@ -131,15 +142,41 @@ async def _build_dask_template(self, spawner: Spawner) -> str:
         # This will take the python model names and transform
         # them to the names kubernetes expects, which to_dict
         # alone doesn't.
-        dask_yaml = yaml.dump(
-            self.k8s_api.sanitize_for_serialization(dask_template)
+        dask_yaml_stream = StringIO()
+        self.yaml.dump(
+            self.k8s_api.sanitize_for_serialization(dask_template),
+            dask_yaml_stream,
         )
+        return dask_yaml_stream.getvalue()
 
-        if not dask_yaml:
-            # This is mostly to help with the typing.
-            raise Exception("Dask template ended up empty.")
-        else:
-            return dask_yaml
+    async def _build_template_values(
+        self, spawner: Spawner, options: SelectedOptions
+    ) -> Dict[str, Any]:
+        """Construct the template variables for Jinja templating."""
+        auth_state = await spawner.user.get_auth_state()
+        self.log.debug(f"Auth state={auth_state}")
+        groups = auth_state["groups"]
+
+        # Build a comma separated list of group:gid
+        # ex: group1:1000,group2:1001,group3:1002
+        external_groups = ",".join([f'{g["name"]}:{g["id"]}' for g in groups])
+
+        # Define the template variables.
+        template_values = {
+            "user_namespace": spawner.namespace,
+            "user": spawner.user.name,
+            "uid": auth_state["uid"],
+            "token": auth_state["token"],
+            "groups": groups,
+            "external_groups": external_groups,
+            "base_url": self.nublado_config.base_url,
+            "dask_yaml": await self._build_dask_template(spawner),
+            "options": options,
+            "nublado_base_url": spawner.hub.base_url,
+            "butler_secret_path": self.nublado_config.butler_secret_path,
+        }
+        self.log.debug(f"Template values={template_values}")
+        return template_values
 
     def delete_user_resources(self, namespace: str) -> None:
         """Clean up a jupyterlab by deleting the whole namespace.

tests/provisioner_test.py

@@ -24,7 +24,7 @@ async def config_mock() -> AsyncGenerator:
     with patch("nublado2.resourcemgr.NubladoConfig") as mock:
         mock.return_value = MagicMock()
         mock.return_value.base_url = "https://data.example.com/"
-        mock.return_valid.gafaelfawr_token = "admin-token"
+        mock.return_value.gafaelfawr_token = "admin-token"
         with patch("nublado2.provisioner.NubladoConfig") as mock:
             mock.return_value = MagicMock()
             mock.return_value.base_url = "https://data.example.com/"