From de03499975b614cc5d278f14a3261a868f359cc4 Mon Sep 17 00:00:00 2001 From: Aadesh Rasal Date: Sun, 6 Apr 2025 20:28:24 +0100 Subject: [PATCH 1/7] Initial Commit From 815403534b93ab04d02db8bbd84447eec7555105 Mon Sep 17 00:00:00 2001 From: Aadesh Rasal Date: Sun, 13 Apr 2025 22:24:31 +0100 Subject: [PATCH 2/7] Boilerplate code added for kubeconnector in kaniko image builder stack_component --- .../kaniko_service_connectors.py | 54 +++++++++++++++++++ 1 file changed, 54 insertions(+) create mode 100644 src/zenml/integrations/kaniko/service_connectors/kaniko_service_connectors.py diff --git a/src/zenml/integrations/kaniko/service_connectors/kaniko_service_connectors.py b/src/zenml/integrations/kaniko/service_connectors/kaniko_service_connectors.py new file mode 100644 index 00000000000..c344019600d --- /dev/null +++ b/src/zenml/integrations/kaniko/service_connectors/kaniko_service_connectors.py @@ -0,0 +1,54 @@ +import base64 +import datetime +import json +import os +import re +from typing import Any, Dict, List, Optional, Tuple, cast + +from pydantic import Field + +from zenml.constants import ( + DOCKER_REGISTRY_RESOURCE_TYPE, + KUBERNETES_CLUSTER_RESOURCE_TYPE, +) + +from zenml.logger import get_logger +from zenml.exceptions import AuthorizationException + +from zenml.models import ( + AuthenticationMethodModel, + ResourceTypeModel, + ServiceConnectorTypeModel, +) +from zenml.service_connectors.docker_service_connector import ( + DockerAuthenticationMethods, + DockerConfiguration, + DockerServiceConnector, +) +from zenml.service_connectors.service_connector import ( + AuthenticationConfig, + ServiceConnector, +) + +class KanikoCredeintials(AuthenticationConfig): + """ + Kaniko credentials for Docker registry authentication. + """ + + # The Docker registry URL + registry_url: str = Field( + description="The URL of the Docker registry.", + title="Docker Registry URL", + ) + + # The username for Docker registry authentication + username: str = Field( + description="The username for Docker registry authentication.", + title="Username", + ) + + # The password for Docker registry authentication + password: str = Field( + description="The password for Docker registry authentication.", + title="Password", + ) \ No newline at end of file From b6bdfc3dd995591113ea55f08b5c8d5b78309a07 Mon Sep 17 00:00:00 2001 From: Aadesh Rasal Date: Sun, 13 Apr 2025 22:39:34 +0100 Subject: [PATCH 3/7] Updated the kubeservice connector in kaniko image builder stack --- .../kaniko_service_connectors.py | 47 ++++++++++++------- 1 file changed, 29 insertions(+), 18 deletions(-) diff --git a/src/zenml/integrations/kaniko/service_connectors/kaniko_service_connectors.py b/src/zenml/integrations/kaniko/service_connectors/kaniko_service_connectors.py index c344019600d..5b65a5d4408 100644 --- a/src/zenml/integrations/kaniko/service_connectors/kaniko_service_connectors.py +++ b/src/zenml/integrations/kaniko/service_connectors/kaniko_service_connectors.py @@ -4,7 +4,6 @@ import os import re from typing import Any, Dict, List, Optional, Tuple, cast - from pydantic import Field from zenml.constants import ( @@ -30,25 +29,37 @@ ServiceConnector, ) -class KanikoCredeintials(AuthenticationConfig): - """ - Kaniko credentials for Docker registry authentication. - """ +class KanikoConnectorConfig(AuthenticationConfig): + """Kubernetes connection configuration for Kaniko.""" - # The Docker registry URL - registry_url: str = Field( - description="The URL of the Docker registry.", - title="Docker Registry URL", + api_token: str = Field( + description="Kubernetes API token for authentication.", + title="API token", + secret=True, + default=None ) - # The username for Docker registry authentication - username: str = Field( - description="The username for Docker registry authentication.", - title="Username", + service_account_name: str = Field( + description="Kubernetes service account name for authentication.", + title="Service Account Name", + default=None, + ) + kubeconfig: Optional[str] = Field( + description="Content of the kubecofig file,", + title="Kubeconfig", + secret=True, + default=None, ) - # The password for Docker registry authentication - password: str = Field( - description="The password for Docker registry authentication.", - title="Password", - ) \ No newline at end of file +class KubernetesKanikoServiceConnector(ServiceConnector): + """Kubernetes Service Connector for Kaniko.""" + + config: KanikoConnectorConfig + + @classmethod + def _get_connector_type(cls): + return ServiceConnectorTypeModel( + name="kubernetes-kaniko", + type="kubernetes", + description="Kubernetes Service Connector for Kaniko.", + ) \ No newline at end of file From a2bfb2e8e9edb6f5e677fc83a841a283d159537b Mon Sep 17 00:00:00 2001 From: Aadesh Rasal Date: Mon, 28 Apr 2025 00:11:47 +0100 Subject: [PATCH 4/7] Adding kuberntes connector support for kaniko image builder Flavour + image builder --- .../flavors/kaniko_image_builder_flavor.py | 11 +++- .../image_builders/kaniko_image_builder.py | 59 ++++++++++++++----- 2 files changed, 54 insertions(+), 16 deletions(-) diff --git a/src/zenml/integrations/kaniko/flavors/kaniko_image_builder_flavor.py b/src/zenml/integrations/kaniko/flavors/kaniko_image_builder_flavor.py index b7733bd29a4..597160f902d 100644 --- a/src/zenml/integrations/kaniko/flavors/kaniko_image_builder_flavor.py +++ b/src/zenml/integrations/kaniko/flavors/kaniko_image_builder_flavor.py @@ -19,7 +19,8 @@ from zenml.image_builders import BaseImageBuilderConfig, BaseImageBuilderFlavor from zenml.integrations.kaniko import KANIKO_IMAGE_BUILDER_FLAVOR - +from zenml.models import ServiceConnectorRequirements +from zenml.constants import KUBERNETES_CLUSTER_RESOURCE_TYPE if TYPE_CHECKING: from zenml.integrations.kaniko.image_builders import KanikoImageBuilder @@ -156,3 +157,11 @@ def implementation_class(self) -> Type["KanikoImageBuilder"]: from zenml.integrations.kaniko.image_builders import KanikoImageBuilder return KanikoImageBuilder + @property + def service_connector_requirements(self) -> Optional[ServiceConnectorRequirements]: + """Service connector requirements. + + Returns: + The service connector requirements. + """ + return ServiceConnectorRequirements(resource_type=KUBERNETES_CLUSTER_RESOURCE_TYPE) diff --git a/src/zenml/integrations/kaniko/image_builders/kaniko_image_builder.py b/src/zenml/integrations/kaniko/image_builders/kaniko_image_builder.py index 22074472521..ea4a6bcf793 100644 --- a/src/zenml/integrations/kaniko/image_builders/kaniko_image_builder.py +++ b/src/zenml/integrations/kaniko/image_builders/kaniko_image_builder.py @@ -26,6 +26,9 @@ from zenml.logger import get_logger from zenml.stack import StackValidator from zenml.utils.archivable import ArchiveType +from zenml.service_connectors.service_connector import ServiceConnector +import kubernetes.client +from kubernetes.client.rest import ApiException if TYPE_CHECKING: from zenml.container_registries import BaseContainerRegistry @@ -202,9 +205,7 @@ def _generate_spec_overrides( optional_spec_args: Dict[str, Any] = {} if self.config.service_account_name: - optional_spec_args["serviceAccountName"] = ( - self.config.service_account_name - ) + optional_spec_args["serviceAccountName"] = self.config.service_account_name return { "apiVersion": "v1", @@ -226,6 +227,42 @@ def _generate_spec_overrides( }, } + def _run_kaniko_build_with_kubeconnect( + self, + pod_name: str, + spec_overrides: Dict[str, Any], + build_context: "BuildContext", + ) -> None: + """Runs the kaniko build in kubernetes with provided connectors. + Args: + pod_name: Name of the Pod that should be created to run the build. + spec_overrides: Pod spec override values. + build_context: The build context. + Raises: + RuntimeError: If the process running the Kaniko build failed. + """ + connector = self.get_connector() + if connector: + logger.info("Using kubernetes connector to run the kaniko build.") + api_client = connector.connect() + + if not isinstance(api_client, kubernetes.client.api_client.ApiClient): + raise RuntimeError(f"Expected ApiClient, got {type(api_client)}") + core_api = kubernetes.client.CoreV1Api(api_client) + + # Second step define the kaniko pod spec + container = kubernetes.client.V1Container( + name=pod_name, + image=self.config.excutor_image, + args=[ + f"--destination={self.config.target_image}", + "--container=tar://stdin", + "--dockerfile=DockerFile", + "--verbosity=info", + ], + vloume_mounts=self.config.volume_mounts, + env=self.config.env, + ) def _run_kaniko_build( self, pod_name: str, @@ -267,9 +304,7 @@ def _run_kaniko_build( stdin=subprocess.PIPE, ) as p: if not self.config.store_context_in_artifact_store: - self._write_build_context( - process=p, build_context=build_context - ) + self._write_build_context(process=p, build_context=build_context) try: return_code = p.wait() @@ -284,9 +319,7 @@ def _run_kaniko_build( ) @staticmethod - def _write_build_context( - process: BytePopen, build_context: "BuildContext" - ) -> None: + def _write_build_context(process: BytePopen, build_context: "BuildContext") -> None: """Writes the build context to the process stdin. Args: @@ -376,14 +409,10 @@ def _check_prerequisites() -> None: RuntimeError: If any of the prerequisites are not installed. """ if not shutil.which("kubectl"): - raise RuntimeError( - "`kubectl` is required to run the Kaniko image builder." - ) + raise RuntimeError("`kubectl` is required to run the Kaniko image builder.") @staticmethod - def _verify_image_name( - image_name_with_tag: str, image_name_with_sha: str - ) -> None: + def _verify_image_name(image_name_with_tag: str, image_name_with_sha: str) -> None: """Verifies the name/sha of the pushed image. Args: From c821d8e79badf6aee155b694e3f9ea49095d35c3 Mon Sep 17 00:00:00 2001 From: Aadesh Rasal Date: Mon, 5 May 2025 22:02:06 +0100 Subject: [PATCH 5/7] Complete implementation of kubernetes connection support to kaniko image builder --- .../image_builders/kaniko_image_builder.py | 50 ++++++++++++++++++- 1 file changed, 48 insertions(+), 2 deletions(-) diff --git a/src/zenml/integrations/kaniko/image_builders/kaniko_image_builder.py b/src/zenml/integrations/kaniko/image_builders/kaniko_image_builder.py index ea4a6bcf793..5ab21cdf22f 100644 --- a/src/zenml/integrations/kaniko/image_builders/kaniko_image_builder.py +++ b/src/zenml/integrations/kaniko/image_builders/kaniko_image_builder.py @@ -28,6 +28,7 @@ from zenml.utils.archivable import ArchiveType from zenml.service_connectors.service_connector import ServiceConnector import kubernetes.client +import kubernetes from kubernetes.client.rest import ApiException if TYPE_CHECKING: @@ -227,7 +228,7 @@ def _generate_spec_overrides( }, } - def _run_kaniko_build_with_kubeconnect( + def _run_kaniko_build( self, pod_name: str, spec_overrides: Dict[str, Any], @@ -263,7 +264,52 @@ def _run_kaniko_build_with_kubeconnect( vloume_mounts=self.config.volume_mounts, env=self.config.env, ) - def _run_kaniko_build( + pod_spec = kubernetes.client.V1PodSpec( + containers=[container], + restart_policy="Never", + service_account_name=self.config.service_account_name + ) + pod = kubernetes.client.V1Pod(metadata=kubernetes.client.V1ObjectMeta(name = pod_name), + spec = pod_spec) + + # now lets create the pod + try: + core_api.create_namespace( + namespace=self.config.kubernetes_namespace, + body=pod + ) + logger.info(f"Kaniko pod {pod_name} created") + except ApiException as e: + raise RuntimeError(f"Failed to create the pod {pod_name}: {e}") + # Stream the context if needed + if not self.config.store_context_in_artifact_store: + logger.info("streaming build context into kaniko pod.") + kubernetes.k8s_utils.stream_file_to_pod( + core_api, + namespcae=self.config.kubernetes_namespace, + pod_name=pod_name, + container_name="kaniko", + source_path=build_context, + destination_parh="/workspace" + ) + # wait for the pod completion + kubernetes.k8s_utils.wait_pod( + core_api, + pod_name=pod_name, + namespace=self.config.kubernetes_namespace, + exit_condition_lambda = kubernetes.k8s_utils.pod_is_done, + timeout_sec = self.config.pod_running_timeout + ) + + #cleanup the pod + logger.info(f"Deleting the pod {pod_name}") + core_api.delete_namespaced_pod(name=pod_name, + namespace=self.config.kubernetes_namespace + ) + else: + logger.info("Connector not found continuing build with kubectl") + self._run_kaniko_build_kubectl(pod_name, spec_overrides, build_context) + def _run_kaniko_build_kubectl( self, pod_name: str, spec_overrides: Dict[str, Any], From 57f4dcc880445695a49559b25c3e9185de0bac2d Mon Sep 17 00:00:00 2001 From: Aadesh Rasal Date: Sat, 20 Sep 2025 17:47:00 +0100 Subject: [PATCH 6/7] Alignment fix for the function --- .../integrations/kaniko/image_builders/kaniko_image_builder.py | 1 + 1 file changed, 1 insertion(+) diff --git a/src/zenml/integrations/kaniko/image_builders/kaniko_image_builder.py b/src/zenml/integrations/kaniko/image_builders/kaniko_image_builder.py index 5ab21cdf22f..dcbd47b1260 100644 --- a/src/zenml/integrations/kaniko/image_builders/kaniko_image_builder.py +++ b/src/zenml/integrations/kaniko/image_builders/kaniko_image_builder.py @@ -309,6 +309,7 @@ def _run_kaniko_build( else: logger.info("Connector not found continuing build with kubectl") self._run_kaniko_build_kubectl(pod_name, spec_overrides, build_context) + def _run_kaniko_build_kubectl( self, pod_name: str, From edfe689cd8426b6f1a5907e1fea9bbfb4574e5c9 Mon Sep 17 00:00:00 2001 From: Aadesh Rasal Date: Sat, 18 Oct 2025 17:50:59 +0100 Subject: [PATCH 7/7] Tests for new kaniko image builder --- ...iko_image_builder_connector_integration.py | 244 ++++++++++++++++++ 1 file changed, 244 insertions(+) create mode 100644 tests/integration/integrations/kaniko/image_builders/test_kaniko_image_builder_connector_integration.py diff --git a/tests/integration/integrations/kaniko/image_builders/test_kaniko_image_builder_connector_integration.py b/tests/integration/integrations/kaniko/image_builders/test_kaniko_image_builder_connector_integration.py new file mode 100644 index 00000000000..857ece7ba2d --- /dev/null +++ b/tests/integration/integrations/kaniko/image_builders/test_kaniko_image_builder_connector_integration.py @@ -0,0 +1,244 @@ +# Copyright (c) ZenML GmbH 2025. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at: +# +# https://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express +# or implied. See the License for the specific language governing +# permissions and limitations under the License. +"""Integration tests for Kaniko image builder Kubernetes connector enhancement.""" + +from contextlib import ExitStack as does_not_raise +from datetime import datetime +from typing import Optional +from unittest.mock import Mock, patch +from uuid import uuid4 + +from zenml.enums import StackComponentType +from zenml.integrations.kaniko.flavors.kaniko_image_builder_flavor import ( + KanikoImageBuilderConfig, +) +from zenml.integrations.kaniko.image_builders import KanikoImageBuilder + + +def _get_kaniko_image_builder( + config: Optional[KanikoImageBuilderConfig] = None, +) -> KanikoImageBuilder: + """Helper function to get a Kaniko image builder.""" + if config is None: + config = KanikoImageBuilderConfig( + kubernetes_context="test-context", + kubernetes_namespace="test-namespace", + ) + + return KanikoImageBuilder( + name="test-kaniko", + id=uuid4(), + config=config, + flavor="kaniko", + type=StackComponentType.IMAGE_BUILDER, + user=uuid4(), + created=datetime.now(), + updated=datetime.now(), + ) + +def _patch_k8s_clients(mocker): + """Helper function to patch k8s clients.""" + mock_connector = Mock() + mock_api_client = Mock() + mock_core_api = Mock() + + mock_connector.connect.return_value = mock_api_client + mocker.patch("kubernetes.client.CoreV1Api", return_value=mock_core_api) + mocker.patch("kubernetes.client.V1Container") + mocker.patch("kubernetes.client.V1PodSpec") + mocker.patch("kubernetes.client.V1Pod") + mocker.patch("kubernetes.client.V1ObjectMeta") + + # Mock k8s_utils + mocker.patch("kubernetes.k8s_utils.stream_file_to_pod") + mocker.patch("kubernetes.k8s_utils.wait_pod") + mocker.patch("kubernetes.k8s_utils.pod_is_done") + + return mock_connector, mock_api_client, mock_core_api + + +@patch("subprocess.Popen") +def test_connector_fallback_to_kubectl_integration(mock_popen, mocker): + """Test complete integration flow with connector fallback to kubectl.""" + image_builder = _get_kaniko_image_builder() + pod_name = "test-pod" + spec_overrides = { + "apiVersion": "v1", + "spec": { + "containers": [{ + "name": pod_name, + "image": "gcr.io/kaniko-project/executor:v1.9.1", + "args": ["--dockerfile=Dockerfile"], + }] + } + } + build_context = Mock() + + # Mock process for kubectl + mock_process = Mock() + mock_process.stdin = Mock() + mock_process.wait.return_value = 0 # Success + mock_popen.return_value.__enter__.return_value = mock_process + + mock_write_context = mocker.patch.object(image_builder, '_write_build_context') + + with patch.object(image_builder, 'get_connector', return_value=None): + image_builder._run_kaniko_build(pod_name, spec_overrides, build_context) + + # Verify kubectl command was constructed correctly + mock_popen.assert_called_once() + args, kwargs = mock_popen.call_args + command = args[0] + + expected_command_parts = [ + "kubectl", + "--context", "test-context", + "--namespace", "test-namespace", + "run", pod_name, + "--stdin", "true", + "--restart", "Never", + ] + + for part in expected_command_parts: + assert part in command + + +def test_connector_vs_kubectl_consistency(mocker): + """Test that connector and kubectl paths produce consistent results.""" + config = KanikoImageBuilderConfig( + kubernetes_context="test-context", + kubernetes_namespace="test-namespace", + executor_image="custom-kaniko:latest", + service_account_name="test-sa", + env=[{"name": "TEST_ENV", "value": "test"}], + volume_mounts=[{"name": "test-vol", "mountPath": "/test"}], + ) + image_builder = _get_kaniko_image_builder(config) + + # Test connector path + mock_connector, mock_api_client, mock_core_api = _patch_k8s_clients(mocker) + + with patch.object(image_builder, 'get_connector', return_value=mock_connector): + try: + image_builder._run_kaniko_build("test-pod", {}, Mock()) + connector_success = True + except Exception: + connector_success = False + + # Test kubectl fallback path + mock_kubectl_build = mocker.patch.object(image_builder, '_run_kaniko_build_kubectl') + + with patch.object(image_builder, 'get_connector', return_value=None): + image_builder._run_kaniko_build("test-pod", {}, Mock()) + kubectl_called = mock_kubectl_build.called + + # Both paths should be available and functional + assert connector_success or kubectl_called, "At least one execution path should work" + + +def test_complete_build_flow_with_connector(mocker): + """Test complete build flow using Kubernetes connector.""" + # Setup for a realistic scenario + config = KanikoImageBuilderConfig( + kubernetes_context="prod-cluster", + kubernetes_namespace="zenml-builds", + executor_image="gcr.io/kaniko-project/executor:v1.9.1", + service_account_name="kaniko-builder", + pod_running_timeout=600, + store_context_in_artifact_store=False, + env=[ + {"name": "DOCKER_CONFIG", "value": "/kaniko/.docker"}, + {"name": "AWS_REGION", "value": "us-west-2"} + ], + volume_mounts=[ + {"name": "docker-config", "mountPath": "/kaniko/.docker"}, + {"name": "aws-credentials", "mountPath": "/root/.aws"} + ], + volumes=[ + {"name": "docker-config", "secret": {"secretName": "docker-config"}}, + {"name": "aws-credentials", "secret": {"secretName": "aws-credentials"}} + ], + executor_args=["--cache=true", "--compressed-caching=false"] + ) + + image_builder = _get_kaniko_image_builder(config) + pod_name = "kaniko-build-12345678" + spec_overrides = image_builder._generate_spec_overrides( + pod_name=pod_name, + image_name="my-app:latest", + context="tar://stdin" + ) + build_context = Mock() + + # Mock all Kubernetes components + mock_connector, mock_api_client, mock_core_api = _patch_k8s_clients(mocker) + mock_logger = mocker.patch("zenml.integrations.kaniko.image_builders.kaniko_image_builder.logger") + + with patch.object(image_builder, 'get_connector', return_value=mock_connector): + # Execute complete flow + image_builder._run_kaniko_build(pod_name, spec_overrides, build_context) + + # Verify all steps executed + assert mock_connector.connect.called + assert mock_core_api.create_namespace.called + assert mock_core_api.delete_namespaced_pod.called + + # Verify logging + mock_logger.info.assert_any_call("Using kubernetes connector to run the kaniko build.") + mock_logger.info.assert_any_call(f"Kaniko pod {pod_name} created") + mock_logger.info.assert_any_call(f"Deleting the pod {pod_name}") + + +def test_config_compatibility_with_connector_enhancement(): + """Test that existing configs remain compatible with connector enhancement.""" + # Test with minimal config (backward compatibility) + minimal_config = KanikoImageBuilderConfig( + kubernetes_context="test-context" + ) + image_builder = _get_kaniko_image_builder(minimal_config) + + # Should work without errors + assert image_builder.config.kubernetes_context == "test-context" + assert image_builder.config.kubernetes_namespace == "zenml-kaniko" # Default + assert not image_builder.config.store_context_in_artifact_store # Default + + # Test with full config + full_config = KanikoImageBuilderConfig( + kubernetes_context="prod-context", + kubernetes_namespace="custom-namespace", + executor_image="custom-kaniko:latest", + service_account_name="custom-sa", + pod_running_timeout=900, + store_context_in_artifact_store=True, + env=[{"name": "CUSTOM_ENV", "value": "custom"}], + env_from=[{"secretRef": {"name": "custom-secret"}}], + volume_mounts=[{"name": "custom-vol", "mountPath": "/custom"}], + volumes=[{"name": "custom-vol", "emptyDir": {}}], + executor_args=["--custom-flag=value"] + ) + + image_builder_full = _get_kaniko_image_builder(full_config) + + # All config values should be preserved + assert image_builder_full.config.kubernetes_context == "prod-context" + assert image_builder_full.config.kubernetes_namespace == "custom-namespace" + assert image_builder_full.config.executor_image == "custom-kaniko:latest" + assert image_builder_full.config.service_account_name == "custom-sa" + assert image_builder_full.config.pod_running_timeout == 900 + assert image_builder_full.config.store_context_in_artifact_store == True + assert len(image_builder_full.config.env) == 1 + assert len(image_builder_full.config.env_from) == 1 + assert len(image_builder_full.config.volume_mounts) == 1 + assert len(image_builder_full.config.volumes) == 1 + assert len(image_builder_full.config.executor_args) == 1