feat(api-server): Enable configuration of launcher and driver images (#…

…455) Enable charm configuration of launcher and driver images used during pipeline steps. Those can be configured through config options `launcher-image` and `driver-image`. When unset, KFP uses the default values defined. This is based on upstream implementation introduced in kubeflow/pipelines#10269. Closes #452
canonical · May 22, 2024 · cbf33be · cbf33be
1 parent cbd8c58
commit cbf33be
Show file tree

Hide file tree

Showing 4 changed files with 105 additions and 50 deletions.
diff --git a/charms/kfp-api/config.yaml b/charms/kfp-api/config.yaml
@@ -49,3 +49,13 @@ options:
     type: string
     default: "mlpipeline"
     description: Default name of object storage bucket.
+  launcher-image:
+    type: string
+    # Source: https://github.com/kubeflow/pipelines/blob/2.0.5/backend/src/v2/compiler/argocompiler/container.go#L27
+    default: "gcr.io/ml-pipeline/kfp-launcher@sha256:80cf120abd125db84fa547640fd6386c4b2a26936e0c2b04a7d3634991a850a4"
+    description: Launcher image used during a pipeline's steps.
+  driver-image:
+    type: string
+    # Source: https://github.com/kubeflow/pipelines/blob/2.0.5/backend/src/v2/compiler/argocompiler/container.go#L29
+    default: "gcr.io/ml-pipeline/kfp-driver@sha256:8e60086b04d92b657898a310ca9757631d58547e76bbbb8bfc376d654bef1707"
+    description: Driver image used during a pipeline's steps.
diff --git a/charms/kfp-api/src/charm.py b/charms/kfp-api/src/charm.py
@@ -290,6 +290,8 @@ def _generate_environment(self) -> dict:
             "OBJECTSTORECONFIG_HOST": f"{object_storage['service']}.{object_storage['namespace']}",
             "OBJECTSTORECONFIG_PORT": str(object_storage["port"]),
             "OBJECTSTORECONFIG_REGION": "",
+            "V2_LAUNCHER_IMAGE": self.model.config["launcher-image"],
+            "V2_DRIVER_IMAGE": self.model.config["driver-image"],
         }
 
         return env_vars

diff --git a/charms/kfp-api/tests/unit/test_operator.py b/charms/kfp-api/tests/unit/test_operator.py
@@ -352,61 +352,18 @@ def test_install_with_all_inputs_and_pebble(
     ):
         """Test complete installation with all required relations and verify pebble layer."""
         harness.set_leader(True)
-        kfpapi_relation_name = "kfp-api"
         model_name = "kubeflow"
         service_port = "8888"
         harness.set_model_name(model_name)
         harness.update_config({"http-port": service_port})
 
         # Set up required relations
-
-        # mysql relation
-        mysql_data = {
-            "database": "database",
-            "host": "host",
-            "root_password": "root_password",
-            "port": "port",
-        }
-        mysql_rel_id = harness.add_relation("mysql", "mysql-provider")
-        harness.add_relation_unit(mysql_rel_id, "mysql-provider/0")
-        harness.update_relation_data(mysql_rel_id, "mysql-provider/0", mysql_data)
-
-        # object storage relation
-        objectstorage_data = {
-            "access-key": "access-key",
-            "namespace": "namespace",
-            "port": 1234,
-            "secret-key": "secret-key",
-            "secure": True,
-            "service": "service",
-        }
-        objectstorage_data_dict = {
-            "_supported_versions": "- v1",
-            "data": yaml.dump(objectstorage_data),
-        }
-        objectstorage_rel_id = harness.add_relation("object-storage", "storage-provider")
-        harness.add_relation_unit(objectstorage_rel_id, "storage-provider/0")
-        harness.update_relation_data(
-            objectstorage_rel_id, "storage-provider", objectstorage_data_dict
-        )
-
-        # kfp-viz relation
-        kfp_viz_data = {
-            "service-name": "viz-service",
-            "service-port": "1234",
-        }
-        kfp_viz_data_dict = {"_supported_versions": "- v1", "data": yaml.dump(kfp_viz_data)}
-        kfp_viz_id = harness.add_relation("kfp-viz", "kfp-viz")
-        harness.add_relation_unit(kfp_viz_id, "kfp-viz/0")
-        harness.update_relation_data(kfp_viz_id, "kfp-viz", kfp_viz_data_dict)
-
-        # example kfp-api provider relation
-        kfpapi_data = {
-            "_supported_versions": "- v1",
-        }
-        kfpapi_rel_id = harness.add_relation(kfpapi_relation_name, "kfp-api-subscriber")
-        harness.add_relation_unit(kfpapi_rel_id, "kfp-api-subscriber/0")
-        harness.update_relation_data(kfpapi_rel_id, "kfp-api-subscriber", kfpapi_data)
+        (
+            mysql_data,
+            objectstorage_data,
+            kfp_viz_data,
+            kfpapi_rel_id,
+        ) = self.setup_required_relations(harness)
 
         harness.begin_with_initial_hooks()
         harness.container_pebble_ready(KFP_API_CONTAINER_NAME)
@@ -483,12 +440,45 @@ def test_install_with_all_inputs_and_pebble(
             ),
             "OBJECTSTORECONFIG_PORT": str(objectstorage_data["port"]),
             "OBJECTSTORECONFIG_REGION": "",
+            "V2_LAUNCHER_IMAGE": "gcr.io/ml-pipeline/kfp-launcher@sha256:80cf120abd125db84fa547640fd6386c4b2a26936e0c2b04a7d3634991a850a4",  # noqa: E501
+            "V2_DRIVER_IMAGE": "gcr.io/ml-pipeline/kfp-driver@sha256:8e60086b04d92b657898a310ca9757631d58547e76bbbb8bfc376d654bef1707",  # noqa: E501
         }
         test_env = pebble_plan_info["services"][KFP_API_SERVICE_NAME]["environment"]
 
         assert test_env == expected_env
         assert model_name == test_env["POD_NAMESPACE"]
 
+    @patch("charm.KubernetesServicePatch", lambda x, y: None)
+    @patch("charm.KfpApiOperator.k8s_resource_handler")
+    def test_launcher_driver_images_config(
+        self,
+        k8s_resource_handler: MagicMock,
+        harness: Harness,
+    ):
+        """Test complete installation with all required relations and verify pebble layer."""
+        harness.set_leader(True)
+        model_name = "kubeflow"
+        service_port = "8888"
+        harness.set_model_name(model_name)
+        harness.update_config({"http-port": service_port})
+        harness.update_config({"launcher-image": "fake-launcher-image"})
+        harness.update_config({"driver-image": "fake-driver-image"})
+
+        # Set up required relations
+        self.setup_required_relations(harness)
+
+        harness.begin_with_initial_hooks()
+        harness.container_pebble_ready(KFP_API_CONTAINER_NAME)
+
+        # test Pebble
+        pebble_plan = harness.get_container_pebble_plan(KFP_API_CONTAINER_NAME)
+        pebble_plan_info = pebble_plan.to_dict()
+        test_env = pebble_plan_info["services"][KFP_API_SERVICE_NAME]["environment"]
+
+        assert test_env["V2_LAUNCHER_IMAGE"] == "fake-launcher-image"
+        assert test_env["V2_DRIVER_IMAGE"] == "fake-driver-image"
+        assert model_name == test_env["POD_NAMESPACE"]
+
     @patch("charm.KubernetesServicePatch", lambda x, y: None)
     @patch("charm.KfpApiOperator._apply_k8s_resources")
     @patch("charm.KfpApiOperator._check_status")
@@ -728,3 +718,56 @@ def return_krh_with_mocked_lightkube(*args, **kwargs):
         )
         assert len(minio_service.spec.ports) == 1
         assert minio_service.spec.ports[0].targetPort == objectstorage_data["port"]
+
+    def setup_required_relations(self, harness: Harness):
+        kfpapi_relation_name = "kfp-api"
+
+        # mysql relation
+        mysql_data = {
+            "database": "database",
+            "host": "host",
+            "root_password": "root_password",
+            "port": "port",
+        }
+        mysql_rel_id = harness.add_relation("mysql", "mysql-provider")
+        harness.add_relation_unit(mysql_rel_id, "mysql-provider/0")
+        harness.update_relation_data(mysql_rel_id, "mysql-provider/0", mysql_data)
+
+        # object storage relation
+        objectstorage_data = {
+            "access-key": "access-key",
+            "namespace": "namespace",
+            "port": 1234,
+            "secret-key": "secret-key",
+            "secure": True,
+            "service": "service",
+        }
+        objectstorage_data_dict = {
+            "_supported_versions": "- v1",
+            "data": yaml.dump(objectstorage_data),
+        }
+        objectstorage_rel_id = harness.add_relation("object-storage", "storage-provider")
+        harness.add_relation_unit(objectstorage_rel_id, "storage-provider/0")
+        harness.update_relation_data(
+            objectstorage_rel_id, "storage-provider", objectstorage_data_dict
+        )
+
+        # kfp-viz relation
+        kfp_viz_data = {
+            "service-name": "viz-service",
+            "service-port": "1234",
+        }
+        kfp_viz_data_dict = {"_supported_versions": "- v1", "data": yaml.dump(kfp_viz_data)}
+        kfp_viz_id = harness.add_relation("kfp-viz", "kfp-viz")
+        harness.add_relation_unit(kfp_viz_id, "kfp-viz/0")
+        harness.update_relation_data(kfp_viz_id, "kfp-viz", kfp_viz_data_dict)
+
+        # example kfp-api provider relation
+        kfpapi_data = {
+            "_supported_versions": "- v1",
+        }
+        kfpapi_rel_id = harness.add_relation(kfpapi_relation_name, "kfp-api-subscriber")
+        harness.add_relation_unit(kfpapi_rel_id, "kfp-api-subscriber/0")
+        harness.update_relation_data(kfpapi_rel_id, "kfp-api-subscriber", kfpapi_data)
+
+        return mysql_data, objectstorage_data, kfp_viz_data, kfpapi_rel_id
diff --git a/tools/get-images.sh b/tools/get-images.sh
@@ -5,5 +5,5 @@
 # dynamic list
 IMAGE_LIST=()
 IMAGE_LIST+=($(find -type f -name metadata.yaml -exec yq '.resources | to_entries | .[] | .value | ."upstream-source"' {} \;))
-IMAGE_LIST+=($(find -type f -name config.yaml -exec yq '.options.cache-image |  select(.) | .default' {} \;))
+IMAGE_LIST+=($(find -type f -name config.yaml -exec yq '.options.*-image |  select(.) | .default' {} \;))
 printf "%s\n" "${IMAGE_LIST[@]}"