From cbf33bec5895eaca9b8a9471e50af507ff9ad9a2 Mon Sep 17 00:00:00 2001 From: Orfeas Kourkakis Date: Wed, 22 May 2024 15:37:00 +0300 Subject: [PATCH] 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 --- charms/kfp-api/config.yaml | 10 ++ charms/kfp-api/src/charm.py | 2 + charms/kfp-api/tests/unit/test_operator.py | 141 ++++++++++++++------- tools/get-images.sh | 2 +- 4 files changed, 105 insertions(+), 50 deletions(-) diff --git a/charms/kfp-api/config.yaml b/charms/kfp-api/config.yaml index bce24b57..1bb47e61 100644 --- 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 index 9d2ec7f6..b43ba876 100755 --- 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 index 4f11df06..4e48a15a 100644 --- 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 index c9c5812c..c3159c6e 100755 --- 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[@]}"