diff --git a/charms/jupyter-ui/config.yaml b/charms/jupyter-ui/config.yaml index ec886399..3ca367b4 100644 --- a/charms/jupyter-ui/config.yaml +++ b/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. diff --git a/charms/jupyter-ui/requirements-integration.in b/charms/jupyter-ui/requirements-integration.in index e8a32367..b4330a62 100644 --- a/charms/jupyter-ui/requirements-integration.in +++ b/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 diff --git a/charms/jupyter-ui/requirements-integration.txt b/charms/jupyter-ui/requirements-integration.txt index 66b48993..55cc8f1d 100644 --- a/charms/jupyter-ui/requirements-integration.txt +++ b/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 diff --git a/charms/jupyter-ui/src/charm.py b/charms/jupyter-ui/src/charm.py index 5fdbe303..aa326766 100755 --- a/charms/jupyter-ui/src/charm.py +++ b/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,6 +26,14 @@ 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", @@ -33,7 +41,26 @@ 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) diff --git a/charms/jupyter-ui/src/config_validators.py b/charms/jupyter-ui/src/config_validators.py new file mode 100644 index 00000000..82844da4 --- /dev/null +++ b/charms/jupyter-ui/src/config_validators.py @@ -0,0 +1,113 @@ +#!/usr/bin/env python3 +# Copyright 2023 Canonical Ltd. +# See LICENSE file for licensing details. + +"""Tools for validating configuration options.""" + +import dataclasses +from dataclasses import field +from typing import List, Union + +OPTIONS_LOOKUP = { + "gpu-vendors": { + "required_keys": ["limitsKey", "uiName"], + "option_index_key": "limitsKey", + }, + "affinity-options": { + "required_keys": ["configKey", "displayName", "affinity"], + "option_index_key": "configKey", + }, + "tolerations-options": { + "required_keys": ["groupKey", "displayName", "tolerations"], + "option_index_key": "groupKey", + }, +} + +VALID_GPU_NUMS = [0, 1, 2, 4, 8] + + +class ConfigValidationError(Exception): + """Raised when validate of a configuration option fails.""" + + pass + + +@dataclasses.dataclass +class OptionsWithDefault: + """A class to store configuration options with default values.""" + + default: str = "" + options: List[dict] = field(default_factory=list) + + +def validate_options_with_default( + default: Union[str, None], options: List[dict], option_index_key: str, required_keys: List[str] +) -> bool: + """Validate configuration specified by a list of options and their default. + + Validation function for options like the affinity, gpu, or tolerations options which accept + a list of options dicts, each with some required keys, and a default value that points at one + of those options by an index key. + + Raises ConfigValidationError if the configuration is invalid (missing a key, the default does + not exist in the list, etc), otherwise returns True. + + Args: + default: A key corresponding to the options entry that should be selected by default + options: A list of dictionaries, each containing the configuration options with some + required keys + option_index_key: The field in each `option` dict that is used as its index key + required_keys: A list of keys that each `option` dict must have + """ + for option in options: + if not isinstance(option, dict): + raise ConfigValidationError(f"Configuration option {option} is not a dictionary.") + for key in required_keys: + if key not in option: + raise ConfigValidationError( + f"Configuration option {option} missing required key: {key}" + ) + + if default and not any(default == option[option_index_key] for option in options): + raise ConfigValidationError( + f"Default selection {default} not found in the list of options." + ) + + return True + + +def validate_named_options_with_default( + default: Union[str, None], options: List[dict], name: str +) -> bool: + """Wrap validate_options_with_default to set up the validator by config name. + + This is a convenience function that automatically sets option_index_key and required_keys + for validate_options_with_default(). See validate_options_with_default() for more information. + + Args: + default: A key corresponding to the options entry that should be selected by default + options: A list of dictionaries, each containing the configuration options with some + required keys + name: the name of the configuration option to validate, for example "gpu-vendors" + """ + return validate_options_with_default( + default, + options, + OPTIONS_LOOKUP[name]["option_index_key"], + OPTIONS_LOOKUP[name]["required_keys"], + ) + + +def parse_gpu_num(num_gpu: str) -> str: + """Return the parsed value for the gpu-number-default configuration.""" + num_gpu = int(num_gpu) + if num_gpu == 0: + return "none" + try: + if num_gpu not in VALID_GPU_NUMS: + raise ConfigValidationError( + f"Invalid value for gpu-number-default: {num_gpu}. Must be one of {VALID_GPU_NUMS}." + ) + return str(num_gpu) + except ValueError: + raise ConfigValidationError("Invalid value for gpu-number-default.") diff --git a/charms/jupyter-ui/src/spawner_ui_config.yaml.j2 b/charms/jupyter-ui/src/templates/spawner_ui_config.yaml.j2 similarity index 85% rename from charms/jupyter-ui/src/spawner_ui_config.yaml.j2 rename to charms/jupyter-ui/src/templates/spawner_ui_config.yaml.j2 index 3609a7ec..88aba16a 100644 --- a/charms/jupyter-ui/src/spawner_ui_config.yaml.j2 +++ b/charms/jupyter-ui/src/templates/spawner_ui_config.yaml.j2 @@ -38,13 +38,19 @@ spawnerFormDefaults: ################################################################ image: # the default container image - value: {{ jupyter_images[0] }} + {%- if jupyter_images_default %} + value: {{ jupyter_images_default }} + {%- else %} + value: "" + {%- endif %} # the list of available container images in the dropdown + {%- if jupyter_images | length > 0 %} options: - {% for image in jupyter_images -%} - - {{ image }} - {% endfor %} + {{ jupyter_images | to_yaml | indent(4) }} + {%- else %} + options: [] + {%- endif %} ################################################################ # VSCode-like Container Images (Group 1) @@ -59,13 +65,19 @@ spawnerFormDefaults: ################################################################ imageGroupOne: # the default container image - value: {{ vscode_images[0] }} + {%- if vscode_images_default %} + value: {{ vscode_images_default }} + {%- else %} + value: "" + {%- endif %} # the list of available container images in the dropdown + {%- if vscode_images | length > 0 %} options: - {% for image in vscode_images -%} - - {{ image }} - {% endfor %} + {{ vscode_images | to_yaml | indent(4) }} + {%- else %} + options: [] + {%- endif %} ################################################################ # RStudio-like Container Images (Group 2) @@ -82,14 +94,19 @@ spawnerFormDefaults: ################################################################ imageGroupTwo: # the default container image - value: {{ rstudio_images[0] }} + {%- if rstudio_images_default %} + value: {{ rstudio_images_default }} + {%- else %} + value: "" + {%- endif %} # the list of available container images in the dropdown + {%- if rstudio_images | length > 0 %} options: - {% for image in rstudio_images -%} - - {{ image }} - {% endfor %} - + {{ rstudio_images | to_yaml | indent(4) }} + {%- else %} + options: [] + {%- endif %} ################################################################ # CPU Resources @@ -128,20 +145,25 @@ spawnerFormDefaults: value: # the `limitKey` of the default vendor # (to have no default, set as "") + {%- if gpu_vendors_default %} + vendor: {{ gpu_vendors_default }} + {%- else %} vendor: "" + {%- endif %} # the list of available vendors in the dropdown # `limitsKey` - what will be set as the actual limit # `uiName` - what will be displayed in the dropdown UI + {%- if gpu_vendors | length > 0 %} vendors: - - limitsKey: "nvidia.com/gpu" - uiName: "NVIDIA" - - limitsKey: "amd.com/gpu" - uiName: "AMD" + {{ gpu_vendors | to_yaml | indent(6) }} + {%- else %} + vendors: [] + {%- endif %} # the default value of the limit # (possible values: "none", "1", "2", "4", "8") - num: "none" + num: {{ gpu_number_default }} ################################################################ # Workspace Volumes @@ -202,10 +224,20 @@ spawnerFormDefaults: # the `configKey` of the default affinity config # (to have no default, set as "") # (if `readOnly`, the default `value` will be the only accessible option) + {%- if affinity_options_default %} + value: {{ affinity_options_default }} + {%- else %} value: "" + {%- endif %} # the list of available affinity configs in the dropdown + {%- if affinity_options | length > 0 %} + options: + {{ affinity_options | to_yaml | indent(4) }} + {%- else %} options: [] + {%- endif %} + #options: # - configKey: "dedicated_node_per_notebook" # displayName: "Dedicated Node Per Notebook" @@ -241,10 +273,20 @@ spawnerFormDefaults: # the `groupKey` of the default toleration group # (to have no default, set as "") # (if `readOnly`, the default `value` will be the only accessible option) + {%- if tolerations_options_default %} + value: {{ tolerations_options_default }} + {%- else %} value: "" + {%- endif %} # the list of available toleration groups in the dropdown + {%- if tolerations_options | length > 0 %} + options: + {{ tolerations_options | to_yaml | indent(4) }} + {%- else %} options: [] + {%- endif %} + #options: # - groupKey: "group_1" # displayName: "4 CPU 8Gb Mem at ~$X.XXX USD per day" @@ -295,13 +337,12 @@ spawnerFormDefaults: # the list of PodDefault names that are selected by default # (take care to ensure these PodDefaults exist in Profile Namespaces) + {%- if default_poddefaults | length > 0 %} value: - # 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. - # Auto-selects "Allow access to Kubeflow Pipelines" button in Notebook spawner UI - - access-ml-pipeline + {{ default_poddefaults | to_yaml | indent(4) }} + {%- else %} + value: [] + {%- endif %} ################################################################ # Environment diff --git a/charms/jupyter-ui/tests/integration/test_charm.py b/charms/jupyter-ui/tests/integration/test_charm.py index e53bedd0..c91cfe67 100644 --- a/charms/jupyter-ui/tests/integration/test_charm.py +++ b/charms/jupyter-ui/tests/integration/test_charm.py @@ -9,7 +9,9 @@ from pathlib import Path import aiohttp +import dpath import pytest +import tenacity import yaml from pytest_operator.plugin import OpsTest @@ -26,6 +28,38 @@ "kubeflow-userid": "", } +AFFINITY_OPTIONS = [ + { + "configKey": "test-affinity-config-1", + "displayName": "Test Affinity Config-1", + "affinity": dict( + nodeAffinity=dict( + requiredDuringSchedulingIgnoredDuringExecution=[ + dict( + matchExpressions=[ + dict(key="lifecycle", operator="In", values=["kubeflow-notebook-1"]) + ] + ) + ] + ) + ), + }, +] + +TOLERATIONS_OPTIONS = [ + { + "groupKey": "test-tolerations-group-1", + "displayName": "Test Tolerations Group 1", + "tolerations": [ + dict(effect="NoSchedule", key="dedicated", operator="Equal", value="big-machine") + ], + }, +] +DEFAULT_PODDEFAULTS = [ + "poddefault1", + "poddefault2", +] + @pytest.mark.abort_on_fail async def test_build_and_deploy(ops_test: OpsTest): @@ -84,28 +118,57 @@ async def test_ui_is_accessible(ops_test: OpsTest): @pytest.mark.parametrize( - "config_key,expected_images,yaml_key", + "config_key,config_value,yaml_path", [ - ("jupyter-images", ["jupyterimage1", "jupyterimage2"], "image"), - ("vscode-images", ["vscodeimage1", "vscodeimage2"], "imageGroupOne"), - ("rstudio-images", ["rstudioimage1", "rstudioimage2"], "imageGroupTwo"), + ("jupyter-images", ["jupyterimage1", "jupyterimagse2"], "config/image/options"), + ("vscode-images", ["vscodeimage1", "vscodeimagse2"], "config/imageGroupOne/options"), + ("rstudio-images", ["rstudioimage1", "rstudioismage2"], "config/imageGroupTwo/options"), + ("affinity-options", AFFINITY_OPTIONS, "config/affinityConfig/options"), + ("gpu-vendors", [{"limitsKey": "gpu1", "uiName": "GPsU 1"}], "config/gpus/value/vendors"), + ("tolerations-options", TOLERATIONS_OPTIONS, "config/tolerationGroup/options"), + ("default-poddefaults", DEFAULT_PODDEFAULTS, "config/configurations/value"), ], ) -async def test_notebook_image_selector(ops_test: OpsTest, config_key, expected_images, yaml_key): - """Test notebook image selector. - - 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. +async def test_notebook_configuration(ops_test: OpsTest, config_key, config_value, yaml_path): + """Test updating notebook configuration settings. + + Verify that setting the juju config for the default notebook configuration settings sets the + values in the Jupyter UI web form. + + Args: + config_key: The key in the charm config to set + config_value: The value to set the config key to + yaml_path: The path in the spawner_ui_config.yaml file that this config will be rendered to, + written as a "/" separated string (e.g. "config/image/options"). See dpath.get() + at https://github.com/dpath-maintainers/dpath-python?tab=readme-ov-file#searching + for more information on the path syntax. """ - 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 - ) + await ops_test.model.applications[APP_NAME].set_config({config_key: yaml.dump(config_value)}) + expected_images = config_value + + # To avoid waiting for a long idle_period between each of this series of tests, we do not use + # wait_for_idle. Instead we push the config and then try for 120 seconds to assert the config + # is updated. This ends up being faster than waiting for idle_period between each test. + jupyter_ui_url = await get_unit_address(ops_test) - response = await fetch_response(f"http://{jupyter_ui_url}:{PORT}/api/config", HEADERS) - response_json = json.loads(response[1]) - actual_images = response_json["config"][yaml_key]["options"] - assert actual_images == expected_images + logger.info("Polling the Jupyter UI API to check the configuration") + for attempt in RETRY_120_SECONDS: + logger.info("Testing whether the config has been updated") + with attempt: + try: + response = await fetch_response( + f"http://{jupyter_ui_url}:{PORT}/api/config", HEADERS + ) + response_json = json.loads(response[1]) + actual_config = dpath.get(response_json, yaml_path) + assert actual_config == expected_images + except AssertionError as e: + logger.info("Failed assertion that config is updated - will retry") + raise e + + +RETRY_120_SECONDS = tenacity.Retrying( + stop=tenacity.stop_after_delay(120), + wait=tenacity.wait_fixed(2), + reraise=True, +) diff --git a/charms/jupyter-ui/tests/unit/test_config_validators.py b/charms/jupyter-ui/tests/unit/test_config_validators.py new file mode 100644 index 00000000..581435c1 --- /dev/null +++ b/charms/jupyter-ui/tests/unit/test_config_validators.py @@ -0,0 +1,80 @@ +#!/usr/bin/env python3 +# Copyright 2023 Canonical Ltd. +# See LICENSE file for licensing details. + +"""Unit tests for the configuration validators.""" + +from contextlib import nullcontext as does_not_raise + +import pytest + +from config_validators import ( + ConfigValidationError, + validate_named_options_with_default, + validate_options_with_default, +) + +REQUIRED_KEYS = ["limitsKey", "uiName"] + + +@pytest.mark.parametrize( + "default, options, options_index_key, required_keys, context_raised", + [ + # Valid, parsable input + ( + "1", + [{"limitsKey": "1", "uiName": "a"}, {"limitsKey": "2", "uiName": "b"}], + "limitsKey", + REQUIRED_KEYS, + does_not_raise(), + ), + ( + "b", + [{"limitsKey": "1", "uiName": "a"}, {"limitsKey": "2", "uiName": "b"}], + "uiName", + REQUIRED_KEYS, + does_not_raise(), + ), + # One missing limitsKey + ( + None, + [{"limitsKey": "1", "uiName": "a"}, {"uiName": "b"}], + "limitsKey", + REQUIRED_KEYS, + pytest.raises(ConfigValidationError), + ), + # One missing uiName + ( + None, + [{"limitsKey": "1"}, {"limitsKey": "2", "uiName": "b"}], + "limitsKey", + REQUIRED_KEYS, + pytest.raises(ConfigValidationError), + ), + # Default not in vendors + ( + "not-in-list", + [{"limitsKey": "1", "uiName": "a"}, {"limitsKey": "2", "uiName": "b"}], + "limitsKey", + REQUIRED_KEYS, + pytest.raises(ConfigValidationError), + ), + ], +) +def test_validate_options_with_default( + default, options, options_index_key, required_keys, context_raised +): + """Test that validate_gpu_vendors_config raises an exception when a required key is missing.""" + # Test that the function raises a ConfigValidationError when a required key is missing. + with context_raised: + validate_options_with_default(default, options, options_index_key, required_keys) + + +def test_validate_named_options_with_default(): + """Test that validate_named_options_with_default passes with valid input. + + Tests using the gpu as an example case. + """ + validate_named_options_with_default( + "nvidia", [{"limitsKey": "nvidia", "uiName": "NVIDIA"}], name="gpu-vendors" + ) diff --git a/charms/jupyter-ui/tests/unit/test_operator.py b/charms/jupyter-ui/tests/unit/test_operator.py index 04e8ef74..f3c682d7 100644 --- a/charms/jupyter-ui/tests/unit/test_operator.py +++ b/charms/jupyter-ui/tests/unit/test_operator.py @@ -3,19 +3,101 @@ # """Unit tests for JupyterUI Charm.""" - +import copy import logging +from contextlib import nullcontext as does_not_raise from unittest.mock import MagicMock, patch import pytest import yaml +from lightkube.models.core_v1 import ( + Affinity, + NodeAffinity, + NodeSelector, + NodeSelectorRequirement, + NodeSelectorTerm, + Toleration, +) from ops.model import ActiveStatus, MaintenanceStatus, WaitingStatus from ops.testing import Harness from charm import JupyterUI +from config_validators import ConfigValidationError, OptionsWithDefault logger = logging.getLogger(__name__) +# Sample inputs for render_jwa_file tests +JUPYTER_IMAGES_CONFIG = ["jupyterimage1", "jupyterimage2"] +VSCODE_IMAGES_CONFIG = ["vscodeimage1", "vscodeimage2"] +RSTUDIO_IMAGES_CONFIG = ["rstudioimage1", "rstudioimage2"] +AFFINITY_OPTIONS_CONFIG = [ + { + "configKey": "test-affinity-config-1", + "displayName": "Test Affinity Config-1", + "affinity": Affinity( + nodeAffinity=NodeAffinity( + requiredDuringSchedulingIgnoredDuringExecution=NodeSelector( + [ + NodeSelectorTerm( + matchExpressions=[ + NodeSelectorRequirement( + key="lifecycle", operator="In", values=["kubeflow-notebook-1"] + ) + ] + ) + ] + ) + ) + ).to_dict(), + }, + { + "configKey": "test-affinity-config-2", + "displayName": "Test Affinity Config-2", + "affinity": Affinity( + nodeAffinity=NodeAffinity( + requiredDuringSchedulingIgnoredDuringExecution=NodeSelector( + [ + NodeSelectorTerm( + matchExpressions=[ + NodeSelectorRequirement( + key="lifecycle", operator="In", values=["kubeflow-notebook-2"] + ) + ] + ) + ] + ) + ) + ).to_dict(), + }, +] +GPU_VENDORS_CONFIG = [ + {"limitsKey": "nvidia", "uiName": "NVIDIA"}, +] +TOLERATIONS_OPTIONS_CONFIG = [ + { + "groupKey": "test-tolerations-group-1", + "displayName": "Test Tolerations Group 1", + "tolerations": [ + Toleration( + effect="NoSchedule", key="dedicated", operator="Equal", value="big-machine" + ).to_dict() + ], + }, + { + "groupKey": "test-tolerations-group-2", + "displayName": "Test Tolerations Group 2", + "tolerations": [ + Toleration( + effect="NoSchedule", key="dedicated", operator="Equal", value="big-machine" + ).to_dict() + ], + }, +] +DEFAULT_PODDEFAULTS_CONFIG = [ + "poddefault1", + "poddefault2", +] + @pytest.fixture(scope="function") def harness() -> Harness: @@ -52,6 +134,68 @@ def test_spawner_ui(self, k8s_resource_handler: MagicMock, harness: Harness): config_value = spawner_ui_config["spawnerFormDefaults"]["configurations"]["value"] assert config_value == ["access-ml-pipeline"] + @pytest.mark.parametrize( + "num_gpus, context_raised", + [ + (0, does_not_raise()), + (1, does_not_raise()), + (2, does_not_raise()), + (4, does_not_raise()), + (8, does_not_raise()), + ], + ) + @patch("charm.KubernetesServicePatch", lambda x, y, service_name: None) + @patch("charm.JupyterUI.k8s_resource_handler") + def test_spawner_ui_has_correct_num_gpu( + self, k8s_resource_handler: MagicMock, harness: Harness, num_gpus: int, context_raised + ): + """Test spawner UI. + + spawner_ui_config.yaml.j2 contains a number of changes that were done for Charmed + Kubeflow. This test is to validate those. If it fails, spawner_ui_config.yaml.j2 + should be reviewed and changes to this tests should be made, if required. + """ + harness.set_leader(True) + harness.update_config({"gpu-number-default": num_gpus}) + harness.begin_with_initial_hooks() + + spawner_ui_config = yaml.safe_load( + harness.charm.container.pull("/etc/config/spawner_ui_config.yaml") + ) + + # test for default configurations + # only single configuration value is currently set in the list of values + config_value = spawner_ui_config["spawnerFormDefaults"]["gpus"]["value"]["num"] + if num_gpus == 0: + assert config_value == "none" + else: + assert config_value == num_gpus + + @pytest.mark.parametrize( + "num_gpus, context_raised", + [ + # Invalid number + (3, pytest.raises(ConfigValidationError)), + # Nonsense input + ("adsda", pytest.raises(RuntimeError)), + ], + ) + @patch("charm.KubernetesServicePatch", lambda x, y, service_name: None) + @patch("charm.JupyterUI.k8s_resource_handler") + def test_spawner_ui_for_incorrect_gpu_number( + self, k8s_resource_handler: MagicMock, harness: Harness, num_gpus: int, context_raised + ): + """Test spawner UI. + + spawner_ui_config.yaml.j2 contains a number of changes that were done for Charmed + Kubeflow. This test is to validate those. If it fails, spawner_ui_config.yaml.j2 + should be reviewed and changes to this tests should be made, if required. + """ + harness.set_leader(True) + with context_raised: + harness.update_config({"gpu-number-default": num_gpus}) + harness.begin_with_initial_hooks() + @patch("charm.KubernetesServicePatch", lambda x, y, service_name: None) @patch("charm.JupyterUI.k8s_resource_handler") def test_not_leader(self, k8s_resource_handler: MagicMock, harness: Harness): @@ -147,59 +291,218 @@ def test_deploy_k8s_resources_success( assert isinstance(harness.charm.model.unit.status, MaintenanceStatus) @pytest.mark.parametrize( - "config_key,expected_images", + "config_key,expected_config_yaml", [ - ("jupyter-images", ["jupyterimage1", "jupyterimage2"]), - ("vscode-images", ["vscodeimage1", "vscodeimage2"]), - ("rstudio-images", ["rstudioimage1", "rstudioimage2"]), + ("jupyter-images", yaml.dump(["jupyterimage1", "jupyterimage2"])), + ("vscode-images", yaml.dump(["vscodeimage1", "vscodeimage2"])), + ("rstudio-images", yaml.dump(["rstudioimage1", "rstudioimage2"])), + ("jupyter-images", yaml.dump([])), + # Assert that we handle an empty string as if its an empty list + ("jupyter-images", ""), + # poddefaults inputs function like an image selector, so test them here too + ("default-poddefaults", yaml.dump(DEFAULT_PODDEFAULTS_CONFIG)), + ("default-poddefaults", ""), ], ) @patch("charm.KubernetesServicePatch", lambda x, y, service_name: None) @patch("charm.JupyterUI.k8s_resource_handler") - def test_notebook_selector_images_config( - self, k8s_resource_handler: MagicMock, harness: Harness, config_key, expected_images + def test_notebook_selector_config( + self, k8s_resource_handler: MagicMock, harness: Harness, config_key, expected_config_yaml ): - """Test that updating the images config works as expected. + """Test that updating the images config and poddefaults works as expected. The following should be tested: Jupyter images, VSCode images, and RStudio images. """ # Arrange - expected_images_yaml = yaml.dump(expected_images) + expected_config = yaml.safe_load(expected_config_yaml) + # Recast an empty input as an empty list to match the expected output + if expected_config is None: + expected_config = [] harness.set_leader(True) harness.begin() - harness.update_config({config_key: expected_images_yaml}) + harness.update_config({config_key: expected_config_yaml}) # Act - actual_images = harness.charm._get_from_config(config_key) + parsed_config = harness.charm._get_from_config(config_key) # Assert - assert actual_images == expected_images + assert parsed_config.options == expected_config + if expected_config: + assert parsed_config.default == expected_config[0] + else: + assert parsed_config.default == "" + @pytest.mark.parametrize( + "config_key,default_value,config_as_yaml", + [ + ("affinity-options", "test-affinity-config-1", yaml.dump(AFFINITY_OPTIONS_CONFIG)), + ("gpu-vendors", "nvidia", yaml.dump(GPU_VENDORS_CONFIG)), + ( + "tolerations-options", + "test-tolerations-group-1", + yaml.dump(TOLERATIONS_OPTIONS_CONFIG), + ), + ], + ) @patch("charm.KubernetesServicePatch", lambda x, y, service_name: None) @patch("charm.JupyterUI.k8s_resource_handler") - def test_render_jwa_file(self, k8s_resource_handler: MagicMock, harness: Harness): + def test_notebook_configurations( + self, + k8s_resource_handler: MagicMock, + harness: Harness, + config_key, + default_value, + config_as_yaml, + ): + """Test that updating the notebook configuration settings works as expected.""" + # Arrange + expected_config = yaml.safe_load(config_as_yaml) + # Recast an empty input as an empty list to match the expected output + if expected_config is None: + expected_config = [] + harness.set_leader(True) + harness.begin() + harness.update_config({config_key: config_as_yaml}) + harness.update_config({config_key + "-default": default_value}) + + # Act + parsed_config = harness.charm._get_from_config(config_key) + + # Assert + assert parsed_config.options == expected_config + assert parsed_config.default == default_value + + @pytest.mark.parametrize( + "render_jwa_file_with_images_config_args", + [ + # All options empty + ( + dict( + jupyter_images_config=OptionsWithDefault(), + vscode_images_config=OptionsWithDefault(), + rstudio_images_config=OptionsWithDefault(), + gpu_number_default=0, + gpu_vendors_config=OptionsWithDefault(), + affinity_options_config=OptionsWithDefault(), + tolerations_options_config=OptionsWithDefault(), + default_poddefaults_config=OptionsWithDefault(), + ) + ), + # All options with valid input + ( + dict( + jupyter_images_config=OptionsWithDefault( + default="jupyterimage1", options=["jupyterimage1", "jupyterimage2"] + ), + vscode_images_config=OptionsWithDefault( + default="vscodeimage1", options=["vscodeimage1", "vscodeimage2"] + ), + rstudio_images_config=OptionsWithDefault( + default="rstudioimage1", options=["rstudioimage1", "rstudioimage2"] + ), + gpu_number_default=1, + gpu_vendors_config=OptionsWithDefault( + default="nvidia", options=GPU_VENDORS_CONFIG + ), + affinity_options_config=OptionsWithDefault( + default="test-affinity-config-1", options=AFFINITY_OPTIONS_CONFIG + ), + tolerations_options_config=OptionsWithDefault( + default="test-tolerations-group-1", options=TOLERATIONS_OPTIONS_CONFIG + ), + default_poddefaults_config=OptionsWithDefault( + default="", options=DEFAULT_PODDEFAULTS_CONFIG + ), + ) + ), + ], + ) + @patch("charm.KubernetesServicePatch", lambda x, y, service_name: None) + @patch("charm.JupyterUI.k8s_resource_handler") + def test_render_jwa_file( + self, + k8s_resource_handler: MagicMock, + harness: Harness, + render_jwa_file_with_images_config_args, + ): """Tests the rendering of the jwa spawner file with the list of images.""" # Arrange - jupyter_images = ["jupyterimage1", "jupyterimage2"] - vscode_images = ["vscodeimage1", "vscodeimage2"] - rstudio_images = ["rstudioimage1", "rstudioimage2"] + render_args = render_jwa_file_with_images_config_args + + # Build the expected results + expected = copy.deepcopy(render_args) + harness.set_leader(True) harness.begin() # Act - actual_content_yaml = harness.charm._render_jwa_file_with_images_config( - jupyter_images, vscode_images, rstudio_images - ) + actual_content_yaml = harness.charm._render_jwa_spawner_inputs(**render_args) actual_content = yaml.safe_load(actual_content_yaml) - rendered_jupyter_images = actual_content["spawnerFormDefaults"]["image"]["options"] - rendered_vscode_images = actual_content["spawnerFormDefaults"]["imageGroupOne"]["options"] - rendered_rstudio_images = actual_content["spawnerFormDefaults"]["imageGroupTwo"]["options"] # Assert - assert rendered_jupyter_images == jupyter_images - assert rendered_vscode_images == vscode_images - assert rendered_rstudio_images == rstudio_images + assert ( + actual_content["spawnerFormDefaults"]["image"]["value"] + == expected["jupyter_images_config"].default + ) + assert ( + actual_content["spawnerFormDefaults"]["image"]["options"] + == expected["jupyter_images_config"].options + ) + + assert ( + actual_content["spawnerFormDefaults"]["imageGroupOne"]["value"] + == expected["vscode_images_config"].default + ) + assert ( + actual_content["spawnerFormDefaults"]["imageGroupOne"]["options"] + == expected["vscode_images_config"].options + ) + + assert ( + actual_content["spawnerFormDefaults"]["imageGroupTwo"]["value"] + == expected["rstudio_images_config"].default + ) + assert ( + actual_content["spawnerFormDefaults"]["imageGroupTwo"]["options"] + == expected["rstudio_images_config"].options + ) + + assert ( + actual_content["spawnerFormDefaults"]["gpus"]["value"]["vendor"] + == expected["gpu_vendors_config"].default + ) + assert ( + actual_content["spawnerFormDefaults"]["gpus"]["value"]["num"] + == expected["gpu_number_default"] + ) + assert ( + actual_content["spawnerFormDefaults"]["gpus"]["value"]["vendors"] + == expected["gpu_vendors_config"].options + ) + + assert ( + actual_content["spawnerFormDefaults"]["affinityConfig"]["value"] + == expected["affinity_options_config"].default + ) + assert ( + actual_content["spawnerFormDefaults"]["affinityConfig"]["options"] + == expected["affinity_options_config"].options + ) + + assert ( + actual_content["spawnerFormDefaults"]["tolerationGroup"]["value"] + == expected["tolerations_options_config"].default + ) + assert ( + actual_content["spawnerFormDefaults"]["tolerationGroup"]["options"] + == expected["tolerations_options_config"].options + ) + + assert ( + actual_content["spawnerFormDefaults"]["configurations"]["value"] + == expected["default_poddefaults_config"].options + ) @patch("charm.KubernetesServicePatch", lambda x, y, service_name: None) @patch("charm.JupyterUI.k8s_resource_handler") @@ -221,18 +524,23 @@ def test_upload_jwa_file(self, k8s_resource_handler: MagicMock, harness: Harness assert actual_config == test_config @pytest.mark.parametrize( - "config_key", - ["jupyter-images", "vscode-images", "rstudio-images"], + "config_key, yaml_string", + ( + ("jupyter-images", "{ not valid yaml"), + ("vscode-images", "{ not valid yaml"), + ("rstudio-images", "{ not valid yaml"), + ("jupyter-images", "A string"), + ("jupyter-images", "{}"), + ), ) @patch("charm.KubernetesServicePatch", lambda x, y, service_name: None) @patch("charm.JupyterUI.k8s_resource_handler") def test_failure_get_config( - self, k8s_resource_handler: MagicMock, harness: Harness, config_key + self, k8s_resource_handler: MagicMock, harness: Harness, config_key, yaml_string ): - """Tests that a warning is logged when a Notebook images config contains an invalid YAML.""" + """Tests that a warning is logged when a Notebook images config contains invalid input.""" # Arrange - invalid_yaml = "[ invalid yaml" - harness.update_config({config_key: invalid_yaml}) + harness.update_config({config_key: yaml_string}) harness.set_leader(True) harness.begin() harness.charm.logger = MagicMock()