Support functionality to override default images for kfp-profile-cont…

…roller (#416)

* Support functionality to override default images for kfp-profile-controller

charms/kfp-profile-controller/config.yaml

@@ -0,0 +1,12 @@
+# Copyright 2024 Canonical Ltd.
+# See LICENSE file for licensing details.
+
+options:
+  custom_images:
+    type: string
+    default: | 
+      visualization_server : ''
+      frontend_image : ''
+    description: >
+      YAML or JSON formatted input defining images to use in Katib
+      For usage details, see https://github.com/canonical/kfp-operators.

charms/kfp-profile-controller/src/charm.py

@@ -11,8 +11,10 @@
 import logging
 from base64 import b64encode
 from pathlib import Path
+from typing import Dict
 
 import lightkube
+import yaml
 from charmed_kubeflow_chisme.components import (
     ContainerFileTemplate,
     SdiRelationDataReceiverComponent,
@@ -33,6 +35,10 @@
     KfpProfileControllerPebbleService,
 )
 
+DEFAULT_IMAGES_FILE = "src/default-custom-images.json"
+with open(DEFAULT_IMAGES_FILE, "r") as json_file:
+    DEFAULT_IMAGES = json.load(json_file)
+
 logger = logging.getLogger(__name__)
 
 DecoratorController = create_global_resource(
@@ -45,7 +51,7 @@
     "src/templates/secrets.yaml.j2",
 ]
 KFP_DEFAULT_PIPELINE_ROOT = ""
-KFP_IMAGES_VERSION = "2.0.3"
+KFP_IMAGES_VERSION = "2.0.3"  # Remember to change this version also in default-custom-images.json
 # This service name must be the Service from the mlmd-operator
 # FIXME: leaving it hardcoded now, but we should share this
 # host and port through relation data
@@ -56,6 +62,35 @@
 SYNC_CODE_DESTINATION_PATH = "/hooks/sync.py"
 
 
+def parse_images_config(config: str) -> Dict:
+    """
+    Parse a YAML config-defined images list.
+
+    This function takes a YAML-formatted string 'config' containing a list of images
+    and returns a dictionary representing the images.
+
+    Args:
+        config (str): YAML-formatted string representing a list of images.
+
+    Returns:
+        Dict: A list of images.
+    """
+    error_message = (
+        f"Cannot parse a config-defined images list from config '{config}' - this"
+        "config input will be ignored."
+    )
+    if not config:
+        return []
+    try:
+        images = yaml.safe_load(config)
+    except yaml.YAMLError as err:
+        logger.warning(
+            f"{error_message}  Got error: {err}, while parsing the custom_image config."
+        )
+        raise err
+    return images
+
+
 class KfpProfileControllerOperator(CharmBase):
     """Charm for the Kubeflow Pipelines Profile controller.
 
@@ -64,6 +99,10 @@ class KfpProfileControllerOperator(CharmBase):
 
     def __init__(self, *args):
         super().__init__(*args)
+        self.images = self.get_images(
+            DEFAULT_IMAGES,
+            parse_images_config(self.model.config["custom_images"]),
+        )
 
         # expose controller's port
         http_port = ServicePort(CONTROLLER_PORT, name="http")
@@ -156,6 +195,10 @@ def __init__(self, *args):
                     CONTROLLER_PORT=CONTROLLER_PORT,
                     METADATA_GRPC_SERVICE_HOST=METADATA_GRPC_SERVICE_HOST,
                     METADATA_GRPC_SERVICE_PORT=METADATA_GRPC_SERVICE_PORT,
+                    VISUALIZATION_SERVER_IMAGE=self.images["visualization_server__image"],
+                    VISUALIZATION_SERVER_TAG=self.images["visualization_server__version"],
+                    FRONTEND_IMAGE=self.images["frontend__image"],
+                    FRONTEND_TAG=self.images["frontend__version"],
                 ),
             ),
             depends_on=[
@@ -167,6 +210,45 @@ def __init__(self, *args):
 
         self.charm_reconciler.install_default_event_handlers()
 
+    def get_images(
+        self, default_images: Dict[str, str], custom_images: Dict[str, str]
+    ) -> Dict[str, str]:
+        """
+        Combine default images with custom images.
+
+        This function takes two dictionaries, 'default_images' and 'custom_images',
+        representing the default set of images and the custom set of images respectively.
+        It combines the custom images into the default image list, overriding any matching
+        image names from the default list with the custom ones.
+
+        Args:
+            default_images (Dict[str, str]): A dictionary containing the default image names
+                as keys and their corresponding default image URIs as values.
+            custom_images (Dict[str, str]): A dictionary containing the custom image names
+                as keys and their corresponding custom image URIs as values.
+
+        Returns:
+            Dict[str, str]: A dictionary representing the combined images, where image names
+            from the custom_images override any matching image names from the default_images.
+        """
+        images = default_images
+        for image_name, custom_image in custom_images.items():
+            if custom_image:
+                if image_name in images:
+                    images[image_name] = custom_image
+                else:
+                    self.log.warning(f"image_name {image_name} not in image list, ignoring.")
+
+        # This are special cases comfigmap where they need to be split into image and version
+        for image_name in [
+            "visualization_server",
+            "frontend",
+        ]:
+            images[f"{image_name}__image"], images[f"{image_name}__version"] = images[
+                image_name
+            ].rsplit(":", 1)
+        return images
+
 
 if __name__ == "__main__":
     main(KfpProfileControllerOperator)

charms/kfp-profile-controller/src/components/pebble_components.py

@@ -23,6 +23,10 @@ class KfpProfileControllerInputs:
     CONTROLLER_PORT: int
     METADATA_GRPC_SERVICE_HOST: str
     METADATA_GRPC_SERVICE_PORT: str
+    VISUALIZATION_SERVER_IMAGE: str
+    VISUALIZATION_SERVER_TAG: str
+    FRONTEND_IMAGE: str
+    FRONTEND_TAG: str
 
 
 class KfpProfileControllerPebbleService(PebbleServiceComponent):
@@ -61,6 +65,10 @@ def get_layer(self) -> Layer:
                             "CONTROLLER_PORT": inputs.CONTROLLER_PORT,
                             "METADATA_GRPC_SERVICE_HOST": inputs.METADATA_GRPC_SERVICE_HOST,
                             "METADATA_GRPC_SERVICE_PORT": inputs.METADATA_GRPC_SERVICE_PORT,
+                            "VISUALIZATION_SERVER_IMAGE": inputs.VISUALIZATION_SERVER_IMAGE,
+                            "VISUALIZATION_SERVER_TAG": inputs.VISUALIZATION_SERVER_TAG,
+                            "FRONTEND_IMAGE": inputs.FRONTEND_IMAGE,
+                            "FRONTEND_TAG": inputs.FRONTEND_TAG,
                         },
                     }
                 }

charms/kfp-profile-controller/src/default-custom-images.json

@@ -0,0 +1,4 @@
+{
+    "visualization_server": "gcr.io/ml-pipeline/visualization-server:2.0.3",
+    "frontend": "gcr.io/ml-pipeline/frontend:2.0.3"
+}

charms/kfp-profile-controller/tests/integration/test_charm.py

@@ -1,6 +1,7 @@
 # Copyright 2023 Canonical Ltd.
 # See LICENSE file for licensing details.
 
+import json
 import logging
 from base64 import b64decode
 from pathlib import Path
@@ -22,6 +23,8 @@
 
 MINIO_APP_NAME = "minio"
 MINIO_CONFIG = {"access-key": "minio", "secret-key": "minio-secret-key"}
+CUSTOM_FRONTEND_IMAGE = "gcr.io/ml-pipeline/frontend:latest"
+CUSTOM_VISUALISATION_IMAGE = "gcr.io/ml-pipeline/visualization-server:latest"
 
 PodDefault = create_namespaced_resource(
     group="kubeflow.org", version="v1alpha1", kind="PodDefault", plural="poddefaults"
@@ -163,6 +166,34 @@ def validate_profile_resources(
     assert expected_label_value == namespace.metadata.labels[expected_label]
 
 
+@retry(
+    wait=wait_exponential(multiplier=1, min=1, max=10),
+    stop=stop_after_delay(30),
+    reraise=True,
+)
+def validate_profile_deployments_with_custom_images(
+    lightkube_client: lightkube.Client,
+    profile_name: str,
+    frontend_image: str,
+    visualisation_image: str,
+):
+    """Tests if profile's deployment have correct images"""
+    # Get deployments
+    pipeline_ui_deployment = lightkube_client.get(
+        Deployment, name="ml-pipeline-ui-artifact", namespace=profile_name
+    )
+    visualization_server_deployment = lightkube_client.get(
+        Deployment, name="ml-pipeline-visualizationserver", namespace=profile_name
+    )
+
+    # Assert images
+    assert pipeline_ui_deployment.spec.template.spec.containers[0].image == frontend_image
+    assert (
+        visualization_server_deployment.spec.template.spec.containers[0].image
+        == visualisation_image
+    )
+
+
 async def test_model_resources(ops_test: OpsTest):
     """Tests if the resources associated with secret's namespace were created.
 
@@ -216,3 +247,24 @@ async def test_sync_webhook(lightkube_client: lightkube.Client, profile: str):
     ]
     for resource, name in desired_resources:
         lightkube_client.get(resource, name=name, namespace=profile)
+
+
+async def test_change_custom_images(
+    ops_test: OpsTest, lightkube_client: lightkube.Client, profile: str
+):
+    """Tests that updating images deployed to user Namespaces works as expected."""
+    custom_images = {
+        "visualization_server": CUSTOM_VISUALISATION_IMAGE,
+        "frontend": CUSTOM_FRONTEND_IMAGE,
+    }
+    await ops_test.model.applications[CHARM_NAME].set_config(
+        {"custom_images": json.dumps(custom_images)}
+    )
+
+    await ops_test.model.wait_for_idle(
+        apps=[CHARM_NAME], status="active", raise_on_blocked=True, timeout=300
+    )
+
+    validate_profile_deployments_with_custom_images(
+        lightkube_client, profile, CUSTOM_FRONTEND_IMAGE, CUSTOM_VISUALISATION_IMAGE
+    )

charms/kfp-profile-controller/tests/unit/test_operator.py

@@ -39,6 +39,10 @@
     "MINIO_NAMESPACE": MOCK_OBJECT_STORAGE_DATA["namespace"],
     "MINIO_PORT": MOCK_OBJECT_STORAGE_DATA["port"],
     "MINIO_SECRET_KEY": MOCK_OBJECT_STORAGE_DATA["secret-key"],
+    "FRONTEND_IMAGE": "gcr.io/ml-pipeline/frontend",
+    "FRONTEND_TAG": KFP_IMAGES_VERSION,
+    "VISUALIZATION_SERVER_IMAGE": "gcr.io/ml-pipeline/visualization-server",
+    "VISUALIZATION_SERVER_TAG": KFP_IMAGES_VERSION,
 }