Enable configuration of New Notebook UI (#345)

* Enable configuration of New Notebook UI

This commit adds configuration for:
* the default number of GPUs and the GPUs available
* the default PodDefaults selected
* the Toleration configurations available
* the Affinity configurations available

These configurations are enabled through newly exposed charm configs.  These configs are lightly validated to ensure they're valid yaml, but not validated enough to ensure things like Tolerations or Affinities are proper Kubernetes yaml

charms/jupyter-ui/config.yaml

@@ -41,3 +41,74 @@ options:
     default: |
       - kubeflownotebookswg/rstudio-tidyverse:v1.8.0
     description: list of image options for RStudio
+  gpu-number-default:
+    type: int
+    default: 0
+    description: |
+      The number of GPUs that are selected by default in the New Notebook UI when creating a Notebook.
+  gpu-vendors:
+    type: string
+    default: '[{"limitsKey": "nvidia.com/gpu", "uiName": "NVIDIA"}, {"limitsKey": "amd.com/gpu", "uiName": "AMD"}]'
+    description: |
+        The GPU vendors that are selectable by users in the New Notebook UI when creating a Notebook.
+        Input is in JSON/YAML in the format defined by Kubeflow in:
+        https://github.com/kubeflow/kubeflow/blob/master/components/crud-web-apps/jupyter/manifests/base/configs/spawner_ui_config.yaml
+        Each item in the list should have keys:
+        - limitsKey: the key that corresponds to the GPU vendor resource in Kubernetes
+        - uiName: the name to be shown in the UI
+  gpu-vendors-default:
+    type: string
+    default: ""
+    description: |
+      The GPU vendor that is selected by default in the New Notebook UI when creating a Notebook.
+      This must be one of the limitsKey values from the gpu-vendors config.  Leave as an empty
+      string to select no GPU vendor by default
+  affinity-options:
+    type: string
+    default: "[]"
+    description: |
+        The Affinity configurations that are selectable by users in the New Notebook UI when creating a Notebook.
+        Input is in JSON/YAML in the format defined by Kubeflow in:
+        https://github.com/kubeflow/kubeflow/blob/master/components/crud-web-apps/jupyter/manifests/base/configs/spawner_ui_config.yaml
+        Each item in the list should have keys:
+        - configKey: an arbitrary key for the configuration
+        - displayName: the name to be shown in the UI
+        - affinity: the affinity configuration, as defined by Kubernetes: https://kubernetes.io/docs/concepts/scheduling-eviction/assign-pod-node/
+  affinity-options-default:
+    type: string
+    default: ""
+    description: |
+      The Affinity options that is selected by default in the New Notebook UI when creating a Notebook.
+      This must be one of the configKey values from the affinity-options config.  Leave as an empty
+      string to select no affinity by default
+  tolerations-options:
+    type: string
+    default: "[]"
+    description: |
+      The Toleration configurations that are selectable by users in the New Notebook UI when creating a Notebook.
+      Input is in JSON/YAML in the format defined by Kubeflow in:
+      https://github.com/kubeflow/kubeflow/blob/master/components/crud-web-apps/jupyter/manifests/base/configs/spawner_ui_config.yaml
+      Each item in the list should have keys:
+      - groupKey: an arbitrary key for the configuration
+      - displayName: the name to be shown in the UI
+      - tolerations: a list of Kubernetes tolerations, as defined in: https://kubernetes.io/docs/concepts/scheduling-eviction/taint-and-toleration/
+  tolerations-options-default:
+    type: string
+    default: ""
+    description: |
+      The Tolerations configuration that is selected by default in the New Notebook UI when creating a Notebook.
+      This must be one of the groupKey values from the tolerations-options config.  Leave as an empty
+      string to select no tolerations configuration by default
+  default-poddefaults:
+    type: string
+    # The default value allows users to access kfp from their Notebooks automatically
+    # Added from https://github.com/kubeflow/kubeflow/pull/6160 to fix
+    # https://github.com/canonical/bundle-kubeflow/issues/423.  This was not yet in
+    # upstream and if they go with something different we should consider syncing with
+    # upstream.
+    default: '["access-ml-pipeline"]'
+    description: |
+      The PodDefaults that are selected by default in the New Notebook UI when creating a new Notebook. 
+      Inputs is a JSON/YAML list of the names of the PodDefaults.
+      The New Notebook UI will always show all PodDefaults available to the user - this only defines
+      which PodDefaults are selected by default.

charms/jupyter-ui/requirements-integration.in

@@ -1,6 +1,8 @@
 aiohttp
+dpath
 # Pinning to <4.0 due to compatibility with the 3.1 controller version
 juju<4.0
 pytest
 pytest-operator
 pyyaml
+tenacity

charms/jupyter-ui/requirements-integration.txt

@@ -38,6 +38,8 @@ decorator==5.1.1
     # via
     #   ipdb
     #   ipython
+dpath==2.1.6
+    # via -r requirements-integration.in
 exceptiongroup==1.1.3
     # via pytest
 executing==1.2.0
@@ -170,6 +172,8 @@ six==1.16.0
     #   python-dateutil
 stack-data==0.6.2
     # via ipython
+tenacity==8.2.3
+    # via -r requirements-integration.in
 tomli==2.0.1
     # via
     #   ipdb

charms/jupyter-ui/src/charm.py

@@ -6,7 +6,7 @@
 
 import logging
 from pathlib import Path
-from typing import List
+from typing import Union
 
 import yaml
 from charmed_kubeflow_chisme.exceptions import ErrorWithStatus
@@ -26,14 +26,41 @@
 from ops.model import ActiveStatus, BlockedStatus, MaintenanceStatus, WaitingStatus
 from ops.pebble import ChangeError, Layer
 from serialized_data_interface import NoCompatibleVersions, NoVersionsListed, get_interfaces
+from yaml import YAMLError
+
+from config_validators import (
+    ConfigValidationError,
+    OptionsWithDefault,
+    parse_gpu_num,
+    validate_named_options_with_default,
+)
 
 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"
+GPU_NUMBER_CONFIG = "gpu-number-default"
+GPU_VENDORS_CONFIG = "gpu-vendors"
+GPU_VENDORS_CONFIG_DEFAULT = f"{GPU_VENDORS_CONFIG}-default"
+AFFINITY_OPTIONS_CONFIG = "affinity-options"
+AFFINITY_OPTIONS_CONFIG_DEFAULT = f"{AFFINITY_OPTIONS_CONFIG}-default"
+TOLERATIONS_OPTIONS_CONFIG = "tolerations-options"
+TOLERATIONS_OPTIONS_CONFIG_DEFAULT = f"{TOLERATIONS_OPTIONS_CONFIG}-default"
+DEFAULT_PODDEFAULTS_CONFIG = "default-poddefaults"
+JWA_CONFIG_FILE = "src/templates/spawner_ui_config.yaml.j2"
+
+IMAGE_CONFIGS = [
+    JUPYTER_IMAGES_CONFIG,
+    VSCODE_IMAGES_CONFIG,
+    RSTUDIO_IMAGES_CONFIG,
+]
+DEFAULT_WITH_OPTIONS_CONFIGS = [
+    GPU_VENDORS_CONFIG,
+    TOLERATIONS_OPTIONS_CONFIG,
+    AFFINITY_OPTIONS_CONFIG,
+]
 
 
 class CheckFailed(Exception):
@@ -213,29 +240,116 @@ 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]:
-        """Return 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."
-        )
+    def _get_from_config(self, key) -> Union[OptionsWithDefault, str]:
+        """Load, validate, render, and return the config value stored in self.model.config[key].
+
+        Different keys are parsed and validated differently.  Errors parsing a config result in
+        null values being returned and errors being logged - this should not raise an exception on
+        invalid input.
+        """
+        if key in IMAGE_CONFIGS:
+            return self._get_list_config(key)
+        elif key in DEFAULT_WITH_OPTIONS_CONFIGS:
+            return self._get_options_with_default_from_config(key)
+        elif key == DEFAULT_PODDEFAULTS_CONFIG:
+            # parsed the same as image configs
+            return self._get_list_config(key)
+        elif key == GPU_NUMBER_CONFIG:
+            return parse_gpu_num(self.model.config[key])
+        else:
+            return self.model.config[key]
+
+    def _get_list_config(self, key) -> OptionsWithDefault:
+        """Parse and return a config entry which should render to a list, like the image lists.
+
+        Returns a OptionsWithDefault with:
+            .options: the content of the config
+            .default: the first element of the list
+        """
+        error_message = f"Cannot parse list input from config '{key}` - ignoring this input."
         try:
-            config = yaml.safe_load(self.model.config[config_key])
+            options = yaml.safe_load(self.model.config[key])
+
+            # Empty yaml string, which resolves to None, should be treated as an empty list
+            if options is None:
+                options = []
+
+            # Check that we receive a list or tuple.  This filters out types that can be indexed but
+            # are not valid for this config (like strings or dicts).
+            if not isinstance(options, (tuple, list)):
+                self.logger.warning(
+                    f"{error_message}  Input must be a list or empty string. Got: '{options}'"
+                )
+                return OptionsWithDefault()
+
+            if len(options) > 0:
+                default = options[0]
+            else:
+                default = ""
+
+            return OptionsWithDefault(default=default, options=options)
         except yaml.YAMLError as err:
             self.logger.warning(f"{error_message}  Got error: {err}")
-            return []
-        return config
+            return OptionsWithDefault()
+
+    def _get_options_with_default_from_config(self, key) -> OptionsWithDefault:
+        """Return the input config for a config specified by a list of options and their default.
+
+        This is for options like the affinity, gpu, or tolerations options which consist of a list
+        of options dicts and a separate config specifying their default value.
+
+        This function handles any config parsing or validation errors, logging details and returning
+        and empty result in case of errors.
 
-    def _render_jwa_file_with_images_config(
-        self, jupyter_images_config, vscode_images_config, rstudio_images_config
+        Returns a OptionsWithDefault with:
+            .options: the content of this config
+            .default: the option selected by f'{key}-default'
+        """
+        default_key = f"{key}-default"
+        try:
+            default = self.model.config[default_key]
+            options = self.model.config[key]
+            options = yaml.safe_load(options)
+            # Convert anything empty to an empty list
+            if not options:
+                options = []
+            validate_named_options_with_default(default, options, name=key)
+            return OptionsWithDefault(default=default, options=options)
+        except (YAMLError, ConfigValidationError) as e:
+            self.logger.warning(f"Failed to parse {key} config:\n{e}")
+            return OptionsWithDefault()
+
+    @staticmethod
+    def _render_jwa_spawner_inputs(
+        jupyter_images_config: OptionsWithDefault,
+        vscode_images_config: OptionsWithDefault,
+        rstudio_images_config: OptionsWithDefault,
+        gpu_number_default: str,
+        gpu_vendors_config: OptionsWithDefault,
+        affinity_options_config: OptionsWithDefault,
+        tolerations_options_config: OptionsWithDefault,
+        default_poddefaults_config: OptionsWithDefault,
     ):
         """Render the JWA configmap template with the user-set images in the juju config."""
         environment = Environment(loader=FileSystemLoader("."))
+        # Add a filter to render yaml with proper formatting
+        environment.filters["to_yaml"] = _to_yaml
         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,
+            jupyter_images=jupyter_images_config.options,
+            jupyter_images_default=jupyter_images_config.default,
+            vscode_images=vscode_images_config.options,
+            vscode_images_default=vscode_images_config.default,
+            rstudio_images=rstudio_images_config.options,
+            rstudio_images_default=rstudio_images_config.default,
+            gpu_number_default=gpu_number_default,
+            gpu_vendors=gpu_vendors_config.options,
+            gpu_vendors_default=gpu_vendors_config.default,
+            affinity_options=affinity_options_config.options,
+            affinity_options_default=affinity_options_config.default,
+            tolerations_options=tolerations_options_config.options,
+            tolerations_options_default=tolerations_options_config.default,
+            default_poddefaults=default_poddefaults_config.options,
         )
         return content
 
@@ -247,17 +361,27 @@ def _upload_jwa_file_to_container(self, file_content):
             make_dirs=True,
         )
 
-    def _update_images_selector(self):
+    def _update_spawner_ui_config(self):
         """Update 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)
+        jupyter_images_config = self._get_from_config(JUPYTER_IMAGES_CONFIG)
+        vscode_images_config = self._get_from_config(VSCODE_IMAGES_CONFIG)
+        rstusio_images_config = self._get_from_config(RSTUDIO_IMAGES_CONFIG)
+        gpu_number_default = self._get_from_config(GPU_NUMBER_CONFIG)
+        gpu_vendors_config = self._get_from_config(GPU_VENDORS_CONFIG)
+        affinity_options_config = self._get_from_config(AFFINITY_OPTIONS_CONFIG)
+        tolerations_options_config = self._get_from_config(TOLERATIONS_OPTIONS_CONFIG)
+        default_poddefaults = self._get_from_config(DEFAULT_PODDEFAULTS_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,
+        jwa_content = self._render_jwa_spawner_inputs(
+            jupyter_images_config=jupyter_images_config,
+            vscode_images_config=vscode_images_config,
+            rstudio_images_config=rstusio_images_config,
+            gpu_number_default=gpu_number_default,
+            gpu_vendors_config=gpu_vendors_config,
+            affinity_options_config=affinity_options_config,
+            tolerations_options_config=tolerations_options_config,
+            default_poddefaults_config=default_poddefaults,
         )
         # push file
         self._upload_jwa_file_to_container(jwa_content)
@@ -338,7 +462,7 @@ def main(self, _) -> None:
             self._deploy_k8s_resources()
             if self._is_container_ready():
                 self._update_layer()
-                self._update_images_selector()
+                self._update_spawner_ui_config()
                 interfaces = self._get_interfaces()
                 self._configure_mesh(interfaces)
         except CheckFailed as err:
@@ -348,8 +472,13 @@ def main(self, _) -> None:
         self.model.unit.status = ActiveStatus()
 
 
-#
-# Start main
-#
+def _to_yaml(data: str) -> str:
+    """Jinja filter to convert data to formatted yaml.
+
+    This is used in the jinja template to format the yaml in the template.
+    """
+    return yaml.safe_dump(data)
+
+
 if __name__ == "__main__":
     main(JupyterUI)