feat: Dynamic Notebook Image Selector (#259)

* Initial implementation: Dynamic Notebook Image Selector
* add jupyter-images, vscode-images, rstudio-images config with default
* change spawner file to Jinja template
* render spawner template with config
* seperate pushing the spawner from pushing the logos files
* feat:  integration test for notebook selector config
* feat: unit tests coverage
* refactor: use charm functionaility to get spawner_ui_config
* refactor: parametrize integration test
* refactor: use arrange-act-assert pattern for unit tests

---------

Co-authored-by: Andrew Scribner <[email protected]>

charms/jupyter-ui/config.yaml

@@ -19,3 +19,22 @@ options:
     type: boolean
     default: false
     description: Whether cookies should require HTTPS
+  jupyter-images:
+    type: string
+    default: |
+      - kubeflownotebookswg/jupyter-pytorch-full:v1.7.0
+      - kubeflownotebookswg/jupyter-pytorch-cuda-full:v1.7.0
+      - kubeflownotebookswg/jupyter-tensorflow-full:v1.7.0
+      - kubeflownotebookswg/jupyter-tensorflow-cuda-full:v1.7.0
+      - swr.cn-south-1.myhuaweicloud.com/mindspore/jupyter-mindspore:v1.6.1
+    description: list of image options for Jupyter Notebook
+  rstudio-images:
+    type: string
+    default: |
+      - kubeflownotebookswg/codeserver-python:v1.7.0
+    description: list of image options for RStudio
+  vscode-images:
+    type: string
+    default: |
+      - kubeflownotebookswg/rstudio-tidyverse:v1.7.0
+    description: list of image options for VSCode

charms/jupyter-ui/src/charm.py

@@ -7,12 +7,14 @@
 
 import logging
 from pathlib import Path
+from typing import List
 
 import yaml
 from charmed_kubeflow_chisme.exceptions import ErrorWithStatus
 from charmed_kubeflow_chisme.kubernetes import KubernetesResourceHandler
 from charmed_kubeflow_chisme.lightkube.batch import delete_many
 from charms.observability_libs.v1.kubernetes_service_patch import KubernetesServicePatch
+from jinja2 import Environment, FileSystemLoader
 from lightkube import ApiError
 from lightkube.generic_resource import load_in_cluster_generic_resources
 from lightkube.models.core_v1 import ServicePort
@@ -25,6 +27,10 @@
 K8S_RESOURCE_FILES = [
     "src/templates/auth_manifests.yaml.j2",
 ]
+JUPYTER_IMAGES_CONFIG = "jupyter-images"
+VSCODE_IMAGES_CONFIG = "vscode-images"
+RSTUDIO_IMAGES_CONFIG = "rstudio-images"
+JWA_CONFIG_FILE = "src/spawner_ui_config.yaml.j2"
 
 
 class CheckFailed(Exception):
@@ -167,15 +173,10 @@ def _update_layer(self) -> None:
             except ChangeError:
                 raise ErrorWithStatus("Failed to replan", BlockedStatus)
 
-    def _upload_files_to_container(self):
-        """Upload required files to container."""
-        with open("src/spawner_ui_config.yaml", "r") as ui_config:
-            file_content = ui_config.read()
-            self.container.push(
-                "/etc/config/spawner_ui_config.yaml",
-                file_content,
-                make_dirs=True,
-            )
+    def _upload_logos_files_to_container(self):
+        """Parses the logos-configmap.yaml file,
+        splits it into files as expected by the workload,
+        and pushes the files to the container"""
         for file_name, file_content in yaml.safe_load(
             Path("src/logos-configmap.yaml").read_text()
         )["data"].items():
@@ -195,6 +196,55 @@ def _deploy_k8s_resources(self) -> None:
             raise ErrorWithStatus("K8S resources creation failed", BlockedStatus)
         self.model.unit.status = MaintenanceStatus("K8S resources created")
 
+    def _get_from_config(self, config_key) -> List[str]:
+        """Returns the yaml value of the config stored in config_key."""
+        error_message = (
+            f"Cannot parse user-defined images from config "
+            f"`{config_key}` - ignoring this input."
+        )
+        try:
+            config = yaml.safe_load(self.model.config[config_key])
+        except yaml.YAMLError as err:
+            self.logger.warning(f"{error_message}  Got error: {err}")
+            return []
+        return config
+
+    def _render_jwa_file_with_images_config(
+        self, jupyter_images_config, vscode_images_config, rstudio_images_config
+    ):
+        """Renders the JWA configmap template with the user-set images in the juju config."""
+        environment = Environment(loader=FileSystemLoader("."))
+        template = environment.get_template(JWA_CONFIG_FILE)
+        content = template.render(
+            jupyter_images=jupyter_images_config,
+            vscode_images=vscode_images_config,
+            rstudio_images=rstudio_images_config,
+        )
+        return content
+
+    def _upload_jwa_file_to_container(self, file_content):
+        """Pushes the JWA spawner config file to the workload container."""
+        self.container.push(
+            "/etc/config/spawner_ui_config.yaml",
+            file_content,
+            make_dirs=True,
+        )
+
+    def _update_images_selector(self):
+        """Updates the images options that can be selected in the dropdown list."""
+        # get config
+        jupyter_images = self._get_from_config(JUPYTER_IMAGES_CONFIG)
+        vscode_images = self._get_from_config(VSCODE_IMAGES_CONFIG)
+        rstusio_images = self._get_from_config(RSTUDIO_IMAGES_CONFIG)
+        # render the jwa file
+        jwa_content = self._render_jwa_file_with_images_config(
+            jupyter_images_config=jupyter_images,
+            vscode_images_config=vscode_images,
+            rstudio_images_config=rstusio_images,
+        )
+        # push file
+        self._upload_jwa_file_to_container(jwa_content)
+
     def _on_install(self, _):
         """Perform installation only actions."""
         try:
@@ -210,7 +260,7 @@ def _on_pebble_ready(self, _):
             return
 
         # upload files to container
-        self._upload_files_to_container()
+        self._upload_logos_files_to_container()
 
         # proceed with other actions
         self.main(_)
@@ -271,6 +321,7 @@ def main(self, _) -> None:
             self._deploy_k8s_resources()
             if self._is_container_ready():
                 self._update_layer()
+                self._update_images_selector()
                 interfaces = self._get_interfaces()
                 self._configure_mesh(interfaces)
         except CheckFailed as err:

charms/jupyter-ui/src/spawner_ui_config.yaml → charms/jupyter-ui/src/spawner_ui_config.yaml.j2

@@ -18,24 +18,23 @@
 spawnerFormDefaults:
   image:
     # The container Image for the user's Jupyter Notebook
-    value: charmedkubeflow/jupyter-scipy:v1.7.0_20.04_1
+    value: {{ jupyter_images[0] }}
     # The list of available standard container Images
     options:
-    - charmedkubeflow/jupyter-scipy:v1.7.0_20.04_1
-    - charmedkubeflow/jupyter-pytorch-full:v1.7.0_20.04_1
-    - charmedkubeflow/jupyter-pytorch-cuda-full:v1.7.0_20.04_1
-    - charmedkubeflow/jupyter-tensorflow-full:v1.7.0_20.04_1
-    - charmedkubeflow/jupyter-tensorflow-cuda-full:v1.7.0_20.04_1
-    - swr.cn-south-1.myhuaweicloud.com/mindspore/jupyter-mindspore:v1.6.1
+    {% for image in jupyter_images -%}
+    - {{ image }}
+    {% endfor %}
   imageGroupOne:
     # The container Image for the user's Group One Server
     # The annotation `notebooks.kubeflow.org/http-rewrite-uri: /`
     # is applied to notebook in this group, configuring
     # the Istio rewrite for containers that host their web UI at `/`
-    value: kubeflownotebookswg/codeserver-python:v1.7.0
+    value: {{ vscode_images[0] }}
     # The list of available standard container Images
     options:
-    - kubeflownotebookswg/codeserver-python:v1.7.0
+    {% for image in vscode_images -%}
+    - {{ image }}
+    {% endfor %}
   imageGroupTwo:
     # The container Image for the user's Group Two Server
     # The annotation `notebooks.kubeflow.org/http-rewrite-uri: /`
@@ -44,10 +43,12 @@ spawnerFormDefaults:
     # The annotation `notebooks.kubeflow.org/http-headers-request-set`
     # is applied to notebook in this group, configuring Istio
     # to add the `X-RStudio-Root-Path` header to requests
-    value: kubeflownotebookswg/rstudio-tidyverse:v1.7.0
+    value: {{ rstudio_images[0] }}
     # The list of available standard container Images
     options:
-    - kubeflownotebookswg/rstudio-tidyverse:v1.7.0
+    {% for image in rstudio_images -%}
+    - {{ image }}
+    {% endfor %}
   # If true, hide registry and/or tag name in the image selection dropdown
   hideRegistry: true
   hideTag: false

charms/jupyter-ui/tests/integration/test_charm.py

@@ -4,6 +4,7 @@
 
 """Integration tests for Jupyter UI Operator/Charm."""
 
+import json
 import logging
 from pathlib import Path
 
@@ -17,6 +18,10 @@
 METADATA = yaml.safe_load(Path("./metadata.yaml").read_text())
 CONFIG = yaml.safe_load(Path("./config.yaml").read_text())
 APP_NAME = "jupyter-ui"
+JUPYTER_IMAGES_CONFIG = "jupyter-images"
+VSCODE_IMAGES_CONFIG = "vscode-images"
+RSTUDIO_IMAGES_CONFIG = "rstudio-images"
+PORT = CONFIG["options"]["port"]["default"]
 
 
 @pytest.mark.abort_on_fail
@@ -51,20 +56,51 @@ async def fetch_response(url):
     return result_status, str(result_text)
 
 
+async def get_unit_address(ops_test: OpsTest):
+    """Returns the unit address of jupyter-ui application."""
+    status = await ops_test.model.get_status()
+    jupyter_ui_units = status["applications"]["jupyter-ui"]["units"]
+    jupyter_ui_url = jupyter_ui_units["jupyter-ui/0"]["address"]
+    return jupyter_ui_url
+
+
 async def test_ui_is_accessible(ops_test: OpsTest):
     """Verify that UI is accessible."""
     # NOTE: This test is re-using deployment created in test_build_and_deploy()
     # NOTE: This test also tests Pebble checks since it uses the same URL.
-
-    status = await ops_test.model.get_status()
-    jupyter_ui_units = status["applications"]["jupyter-ui"]["units"]
-    jupyter_ui_url = jupyter_ui_units["jupyter-ui/0"]["address"]
+    jupyter_ui_url = await get_unit_address(ops_test)
 
     # obtain status and response text from Jupyter UI URL
-    port = CONFIG["options"]["port"]["default"]
-    result_status, result_text = await fetch_response(f"http://{jupyter_ui_url}:{port}")
+    result_status, result_text = await fetch_response(f"http://{jupyter_ui_url}:{PORT}")
 
     # verify that UI is accessible (NOTE: this also tests Pebble checks)
     assert result_status == 200
     assert len(result_text) > 0
     assert "Jupyter Management UI" in result_text
+
+
+@pytest.mark.parametrize(
+    "config_key,expected_images,yaml_key",
+    [
+        ("jupyter-images", ["jupyterimage1", "jupyterimage2"], "image"),
+        ("vscode-images", ["vscodeimage1", "vscodeimage2"], "imageGroupOne"),
+        ("rstudio-images", ["rstudioimage1", "rstudioimage2"], "imageGroupTwo"),
+    ],
+)
+async def test_notebook_image_selector(ops_test: OpsTest, config_key, expected_images, yaml_key):
+    """
+    Verify that setting the juju config for the 3 types of Notebook components
+    sets the notebook images selector list in the workload container,
+    with the same values in the configs.
+    """
+    await ops_test.model.applications[APP_NAME].set_config(
+        {config_key: yaml.dump(expected_images)}
+    )
+    await ops_test.model.wait_for_idle(
+        apps=[APP_NAME], status="active", raise_on_blocked=True, timeout=60 * 10, idle_period=30
+    )
+    jupyter_ui_url = await get_unit_address(ops_test)
+    response = await fetch_response(f"http://{jupyter_ui_url}:{PORT}/api/config")
+    response_json = json.loads(response[1])
+    actual_images = response_json["config"][yaml_key]["options"]
+    assert actual_images == expected_images