From 92e133f8de1e98b088ac050fef2a3ac5049b41f6 Mon Sep 17 00:00:00 2001 From: Michael Schuster Date: Wed, 2 Jul 2025 11:19:12 +0200 Subject: [PATCH 1/3] Always build pipeline image for Kubernetes orchestrator --- .../orchestrators/kubernetes_orchestrator.py | 17 ++++++++++++++++- .../orchestrators/containerized_orchestrator.py | 16 ++++++++++++++++ 2 files changed, 32 insertions(+), 1 deletion(-) diff --git a/src/zenml/integrations/kubernetes/orchestrators/kubernetes_orchestrator.py b/src/zenml/integrations/kubernetes/orchestrators/kubernetes_orchestrator.py index c3e008fbe1..bf0d1988d5 100644 --- a/src/zenml/integrations/kubernetes/orchestrators/kubernetes_orchestrator.py +++ b/src/zenml/integrations/kubernetes/orchestrators/kubernetes_orchestrator.py @@ -70,7 +70,10 @@ from zenml.stack import StackValidator if TYPE_CHECKING: - from zenml.models import PipelineDeploymentResponse, PipelineRunResponse + from zenml.models import ( + PipelineDeploymentResponse, + PipelineRunResponse, + ) from zenml.stack import Stack logger = get_logger(__name__) @@ -84,6 +87,18 @@ class KubernetesOrchestrator(ContainerizedOrchestrator): _k8s_client: Optional[k8s_client.ApiClient] = None + @property + def always_build_pipeline_image(self) -> bool: + """Whether to always build the pipeline image. + + Returns: + Whether to always build the pipeline image. + """ + # Users might want to use a different image for the orchestrator pod, + # which is why make sure to always build the pipeline image even if + # all steps have a custom build. + return True + def get_kube_client( self, incluster: Optional[bool] = None ) -> k8s_client.ApiClient: diff --git a/src/zenml/orchestrators/containerized_orchestrator.py b/src/zenml/orchestrators/containerized_orchestrator.py index d9e2d1fabd..c18f715caf 100644 --- a/src/zenml/orchestrators/containerized_orchestrator.py +++ b/src/zenml/orchestrators/containerized_orchestrator.py @@ -53,6 +53,15 @@ def get_image( component_key=ORCHESTRATOR_DOCKER_IMAGE_KEY, step=step_name ) + @property + def always_build_pipeline_image(self) -> bool: + """Whether to always build the pipeline image. + + Returns: + Whether to always build the pipeline image. + """ + return False + def get_docker_builds( self, deployment: "PipelineDeploymentBase" ) -> List["BuildConfiguration"]: @@ -87,4 +96,11 @@ def get_docker_builds( builds.append(pipeline_build) included_pipeline_build = True + if not included_pipeline_build and self.always_build_pipeline_image: + pipeline_build = BuildConfiguration( + key=ORCHESTRATOR_DOCKER_IMAGE_KEY, + settings=pipeline_settings, + ) + builds.append(pipeline_build) + return builds From 41319d783d45e3d6db56546d920320e3b79718f7 Mon Sep 17 00:00:00 2001 From: Michael Schuster Date: Wed, 2 Jul 2025 11:54:32 +0200 Subject: [PATCH 2/3] Make configurable --- .../flavors/kubernetes_orchestrator_flavor.py | 3 +++ .../orchestrators/kubernetes_orchestrator.py | 14 ++++++++------ .../orchestrators/containerized_orchestrator.py | 13 ++++++++----- 3 files changed, 19 insertions(+), 11 deletions(-) diff --git a/src/zenml/integrations/kubernetes/flavors/kubernetes_orchestrator_flavor.py b/src/zenml/integrations/kubernetes/flavors/kubernetes_orchestrator_flavor.py index 7f2bd57f15..a35fe36b50 100644 --- a/src/zenml/integrations/kubernetes/flavors/kubernetes_orchestrator_flavor.py +++ b/src/zenml/integrations/kubernetes/flavors/kubernetes_orchestrator_flavor.py @@ -69,6 +69,8 @@ class KubernetesOrchestratorSettings(BaseSettings): scheduling a pipeline. prevent_orchestrator_pod_caching: If `True`, the orchestrator pod will not try to compute cached steps before starting the step pods. + always_build_pipeline_image: If `True`, the orchestrator will always + build the pipeline image, even if all steps have a custom build. """ synchronous: bool = True @@ -88,6 +90,7 @@ class KubernetesOrchestratorSettings(BaseSettings): failed_jobs_history_limit: Optional[NonNegativeInt] = None ttl_seconds_after_finished: Optional[NonNegativeInt] = None prevent_orchestrator_pod_caching: bool = False + always_build_pipeline_image: bool = False class KubernetesOrchestratorConfig( diff --git a/src/zenml/integrations/kubernetes/orchestrators/kubernetes_orchestrator.py b/src/zenml/integrations/kubernetes/orchestrators/kubernetes_orchestrator.py index bf0d1988d5..498c45ff06 100644 --- a/src/zenml/integrations/kubernetes/orchestrators/kubernetes_orchestrator.py +++ b/src/zenml/integrations/kubernetes/orchestrators/kubernetes_orchestrator.py @@ -71,6 +71,7 @@ if TYPE_CHECKING: from zenml.models import ( + PipelineDeploymentBase, PipelineDeploymentResponse, PipelineRunResponse, ) @@ -87,17 +88,18 @@ class KubernetesOrchestrator(ContainerizedOrchestrator): _k8s_client: Optional[k8s_client.ApiClient] = None - @property - def always_build_pipeline_image(self) -> bool: + def should_build_pipeline_image( + self, deployment: "PipelineDeploymentBase" + ) -> bool: """Whether to always build the pipeline image. Returns: Whether to always build the pipeline image. """ - # Users might want to use a different image for the orchestrator pod, - # which is why make sure to always build the pipeline image even if - # all steps have a custom build. - return True + settings = cast( + KubernetesOrchestratorSettings, self.get_settings(deployment) + ) + return settings.always_build_pipeline_image def get_kube_client( self, incluster: Optional[bool] = None diff --git a/src/zenml/orchestrators/containerized_orchestrator.py b/src/zenml/orchestrators/containerized_orchestrator.py index c18f715caf..58db7949c3 100644 --- a/src/zenml/orchestrators/containerized_orchestrator.py +++ b/src/zenml/orchestrators/containerized_orchestrator.py @@ -53,12 +53,13 @@ def get_image( component_key=ORCHESTRATOR_DOCKER_IMAGE_KEY, step=step_name ) - @property - def always_build_pipeline_image(self) -> bool: - """Whether to always build the pipeline image. + def should_build_pipeline_image( + self, deployment: "PipelineDeploymentBase" + ) -> bool: + """Whether to build the pipeline image. Returns: - Whether to always build the pipeline image. + Whether to build the pipeline image. """ return False @@ -96,7 +97,9 @@ def get_docker_builds( builds.append(pipeline_build) included_pipeline_build = True - if not included_pipeline_build and self.always_build_pipeline_image: + if not included_pipeline_build and self.should_build_pipeline_image( + deployment + ): pipeline_build = BuildConfiguration( key=ORCHESTRATOR_DOCKER_IMAGE_KEY, settings=pipeline_settings, From 61d8e1096ca4839bf8b0d456c22d2fa3bc575408 Mon Sep 17 00:00:00 2001 From: Michael Schuster Date: Thu, 3 Jul 2025 16:23:51 +0200 Subject: [PATCH 3/3] Docstrings --- .../kubernetes/orchestrators/kubernetes_orchestrator.py | 3 +++ src/zenml/orchestrators/containerized_orchestrator.py | 3 +++ 2 files changed, 6 insertions(+) diff --git a/src/zenml/integrations/kubernetes/orchestrators/kubernetes_orchestrator.py b/src/zenml/integrations/kubernetes/orchestrators/kubernetes_orchestrator.py index 498c45ff06..9d72d61a09 100644 --- a/src/zenml/integrations/kubernetes/orchestrators/kubernetes_orchestrator.py +++ b/src/zenml/integrations/kubernetes/orchestrators/kubernetes_orchestrator.py @@ -93,6 +93,9 @@ def should_build_pipeline_image( ) -> bool: """Whether to always build the pipeline image. + Args: + deployment: The pipeline deployment. + Returns: Whether to always build the pipeline image. """ diff --git a/src/zenml/orchestrators/containerized_orchestrator.py b/src/zenml/orchestrators/containerized_orchestrator.py index 58db7949c3..7659131e0d 100644 --- a/src/zenml/orchestrators/containerized_orchestrator.py +++ b/src/zenml/orchestrators/containerized_orchestrator.py @@ -58,6 +58,9 @@ def should_build_pipeline_image( ) -> bool: """Whether to build the pipeline image. + Args: + deployment: The pipeline deployment. + Returns: Whether to build the pipeline image. """