zenml-nightly 0.55.0.dev20240124__py3-none-any.whl → 0.72.0.dev20250116__py3-none-any.whl

zenml/integrations/kubernetes/__init__.py

@@ -24,6 +24,7 @@ from zenml.integrations.integration import Integration
 from zenml.stack import Flavor
 
 KUBERNETES_ORCHESTRATOR_FLAVOR = "kubernetes"
+KUBERNETES_STEP_OPERATOR_FLAVOR = "kubernetes"
 
 
 class KubernetesIntegration(Integration):
@@ -31,7 +32,9 @@ class KubernetesIntegration(Integration):
 
     NAME = KUBERNETES
     REQUIREMENTS = ["kubernetes>=21.7,<26"]
-
+    REQUIREMENTS_IGNORED_ON_UNINSTALL = [
+        "kfp", # it is used by many others
+    ]
     @classmethod
     def flavors(cls) -> List[Type[Flavor]]:
         """Declare the stack component flavors for the Kubernetes integration.
@@ -40,10 +43,10 @@ class KubernetesIntegration(Integration):
             List of new stack component flavors.
         """
         from zenml.integrations.kubernetes.flavors import (
-            KubernetesOrchestratorFlavor,
+            KubernetesOrchestratorFlavor, KubernetesStepOperatorFlavor
         )
 
-        return [KubernetesOrchestratorFlavor]
+        return [KubernetesOrchestratorFlavor, KubernetesStepOperatorFlavor]
 
 
 KubernetesIntegration.check_installation()

zenml/integrations/kubernetes/flavors/__init__.py

@@ -18,9 +18,17 @@ from zenml.integrations.kubernetes.flavors.kubernetes_orchestrator_flavor import
     KubernetesOrchestratorFlavor,
     KubernetesOrchestratorSettings,
 )
+from zenml.integrations.kubernetes.flavors.kubernetes_step_operator_flavor import (
+    KubernetesStepOperatorConfig,
+    KubernetesStepOperatorFlavor,
+    KubernetesStepOperatorSettings,
+)
 
 __all__ = [
     "KubernetesOrchestratorFlavor",
     "KubernetesOrchestratorConfig",
     "KubernetesOrchestratorSettings",
+    "KubernetesStepOperatorConfig",
+    "KubernetesStepOperatorFlavor",
+    "KubernetesStepOperatorSettings",
 ]

zenml/integrations/kubernetes/flavors/kubernetes_orchestrator_flavor.py

@@ -41,16 +41,24 @@ class KubernetesOrchestratorSettings(BaseSettings):
         service_account_name: Name of the service account to use for the
             orchestrator pod. If not provided, a new service account with "edit"
             permissions will be created.
-        pod_settings: Pod settings to apply.
+        step_pod_service_account_name: Name of the service account to use for the
+            step pods. If not provided, the default service account will be used.
+        privileged: If the container should be run in privileged mode.
+        pod_settings: Pod settings to apply to pods executing the steps.
+        orchestrator_pod_settings: Pod settings to apply to the pod which is
+            launching the actual steps.
     """
 
     synchronous: bool = True
     timeout: int = 0
     service_account_name: Optional[str] = None
+    step_pod_service_account_name: Optional[str] = None
+    privileged: bool = False
     pod_settings: Optional[KubernetesPodSettings] = None
+    orchestrator_pod_settings: Optional[KubernetesPodSettings] = None
 
 
-class KubernetesOrchestratorConfig(  # type: ignore[misc] # https://github.com/pydantic/pydantic/issues/4173
+class KubernetesOrchestratorConfig(
     BaseOrchestratorConfig, KubernetesOrchestratorSettings
 ):
     """Configuration for the Kubernetes orchestrator.
@@ -74,6 +82,9 @@ class KubernetesOrchestratorConfig(  # type: ignore[misc] # https://github.com/p
             containers).
         skip_local_validations: If `True`, the local validations will be
             skipped.
+        parallel_step_startup_waiting_period: How long to wait in between
+            starting parallel steps. This can be used to distribute server
+            load when running pipelines with a huge amount of parallel steps.
     """
 
     incluster: bool = False
@@ -81,6 +92,7 @@ class KubernetesOrchestratorConfig(  # type: ignore[misc] # https://github.com/p
     kubernetes_namespace: str = "zenml"
     local: bool = False
     skip_local_validations: bool = False
+    parallel_step_startup_waiting_period: Optional[float] = None
 
     @property
     def is_remote(self) -> bool:
@@ -104,6 +116,35 @@ class KubernetesOrchestratorConfig(  # type: ignore[misc] # https://github.com/p
         """
         return self.local
 
+    @property
+    def is_synchronous(self) -> bool:
+        """Whether the orchestrator runs synchronous or not.
+
+        Returns:
+            Whether the orchestrator runs synchronous or not.
+        """
+        return self.synchronous
+
+    @property
+    def is_schedulable(self) -> bool:
+        """Whether the orchestrator is schedulable or not.
+
+        Returns:
+            Whether the orchestrator is schedulable or not.
+        """
+        return True
+
+    @property
+    def supports_client_side_caching(self) -> bool:
+        """Whether the orchestrator supports client side caching.
+
+        Returns:
+            Whether the orchestrator supports client side caching.
+        """
+        # The Kubernetes orchestrator starts step pods from a pipeline pod.
+        # This is currently not supported when using client-side caching.
+        return False
+
 
 class KubernetesOrchestratorFlavor(BaseOrchestratorFlavor):
     """Kubernetes orchestrator flavor."""

zenml/integrations/kubernetes/flavors/kubernetes_step_operator_flavor.py

@@ -0,0 +1,166 @@
+#  Copyright (c) ZenML GmbH 2024. 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.
+"""Kubernetes step operator flavor."""
+
+from typing import TYPE_CHECKING, Optional, Type
+
+from zenml.config.base_settings import BaseSettings
+from zenml.constants import KUBERNETES_CLUSTER_RESOURCE_TYPE
+from zenml.integrations.kubernetes import KUBERNETES_STEP_OPERATOR_FLAVOR
+from zenml.integrations.kubernetes.pod_settings import KubernetesPodSettings
+from zenml.models import ServiceConnectorRequirements
+from zenml.step_operators import BaseStepOperatorConfig, BaseStepOperatorFlavor
+
+if TYPE_CHECKING:
+    from zenml.integrations.kubernetes.step_operators import (
+        KubernetesStepOperator,
+    )
+
+
+class KubernetesStepOperatorSettings(BaseSettings):
+    """Settings for the Kubernetes step operator.
+
+    Attributes:
+        pod_settings: Pod settings to apply to pods executing the steps.
+        service_account_name: Name of the service account to use for the pod.
+        privileged: If the container should be run in privileged mode.
+    """
+
+    pod_settings: Optional[KubernetesPodSettings] = None
+    service_account_name: Optional[str] = None
+    privileged: bool = False
+
+
+class KubernetesStepOperatorConfig(
+    BaseStepOperatorConfig, KubernetesStepOperatorSettings
+):
+    """Configuration for the Kubernetes step operator.
+
+    Attributes:
+        kubernetes_namespace: Name of the Kubernetes namespace to be used.
+        incluster: If `True`, the step operator will run the pipeline inside the
+            same cluster in which the orchestrator is running. For this to work,
+            the pod running the orchestrator needs permissions to create new
+            pods. If set, the `kubernetes_context` config option is ignored. If
+            the stack component is linked to a Kubernetes service connector,
+            this field is ignored.
+        kubernetes_context: Name of a Kubernetes context to run pipelines in.
+            If the stack component is linked to a Kubernetes service connector,
+            this field is ignored. Otherwise, it is mandatory.
+    """
+
+    kubernetes_namespace: str = "zenml"
+    incluster: bool = False
+    kubernetes_context: Optional[str] = None
+
+    @property
+    def is_remote(self) -> bool:
+        """Checks if this stack component is running remotely.
+
+        This designation is used to determine if the stack component can be
+        used with a local ZenML database or if it requires a remote ZenML
+        server.
+
+        Returns:
+            True if this config is for a remote component, False otherwise.
+        """
+        return True
+
+    @property
+    def is_local(self) -> bool:
+        """Checks if this stack component is running locally.
+
+        Returns:
+            True if this config is for a local component, False otherwise.
+        """
+        return False
+
+
+class KubernetesStepOperatorFlavor(BaseStepOperatorFlavor):
+    """Kubernetes step operator flavor."""
+
+    @property
+    def name(self) -> str:
+        """Name of the flavor.
+
+        Returns:
+            The name of the flavor.
+        """
+        return KUBERNETES_STEP_OPERATOR_FLAVOR
+
+    @property
+    def service_connector_requirements(
+        self,
+    ) -> Optional[ServiceConnectorRequirements]:
+        """Service connector resource requirements for service connectors.
+
+        Specifies resource requirements that are used to filter the available
+        service connector types that are compatible with this flavor.
+
+        Returns:
+            Requirements for compatible service connectors, if a service
+            connector is required for this flavor.
+        """
+        return ServiceConnectorRequirements(
+            resource_type=KUBERNETES_CLUSTER_RESOURCE_TYPE,
+        )
+
+    @property
+    def docs_url(self) -> Optional[str]:
+        """A url to point at docs explaining this flavor.
+
+        Returns:
+            A flavor docs url.
+        """
+        return self.generate_default_docs_url()
+
+    @property
+    def sdk_docs_url(self) -> Optional[str]:
+        """A url to point at SDK docs explaining this flavor.
+
+        Returns:
+            A flavor SDK docs url.
+        """
+        return self.generate_default_sdk_docs_url()
+
+    @property
+    def logo_url(self) -> str:
+        """A url to represent the flavor in the dashboard.
+
+        Returns:
+            The flavor logo.
+        """
+        return "https://public-flavor-logos.s3.eu-central-1.amazonaws.com/step_operator/kubernetes.png"
+
+    @property
+    def config_class(self) -> Type[KubernetesStepOperatorConfig]:
+        """Returns `KubernetesStepOperatorConfig` config class.
+
+        Returns:
+                The config class.
+        """
+        return KubernetesStepOperatorConfig
+
+    @property
+    def implementation_class(self) -> Type["KubernetesStepOperator"]:
+        """Implementation class for this flavor.
+
+        Returns:
+            The implementation class.
+        """
+        from zenml.integrations.kubernetes.step_operators import (
+            KubernetesStepOperator,
+        )
+
+        return KubernetesStepOperator

zenml/integrations/kubernetes/orchestrators/kube_utils.py

@@ -42,8 +42,8 @@ from kubernetes import config as k8s_config
 from kubernetes.client.rest import ApiException
 
 from zenml.integrations.kubernetes.orchestrators.manifest_utils import (
-    build_cluster_role_binding_manifest_for_service_account,
     build_namespace_manifest,
+    build_role_binding_manifest_for_service_account,
     build_service_account_manifest,
 )
 from zenml.logger import get_logger
@@ -94,18 +94,62 @@ def load_kube_config(
         k8s_config.load_kube_config(context=context)
 
 
-def sanitize_pod_name(pod_name: str) -> str:
+def calculate_max_pod_name_length_for_namespace(namespace: str) -> int:
+    """Calculate the max pod length for a certain namespace.
+
+    Args:
+        namespace: The namespace in which the pod will be created.
+
+    Returns:
+        The maximum pod name length.
+    """
+    # Kubernetes allows Pod names to have 253 characters. However, when
+    # creating a pod they try to create a log file which is called
+    # <NAMESPACE>_<POD_NAME>_<UUID>, which adds additional characters and
+    # runs into filesystem limitations for filename lengths (255). We therefore
+    # subtract the length of a UUID (36), the two underscores and the
+    # namespace length from the max filename length.
+    return 255 - 38 - len(namespace)
+
+
+def sanitize_pod_name(pod_name: str, namespace: str) -> str:
     """Sanitize pod names so they conform to Kubernetes pod naming convention.
 
     Args:
         pod_name: Arbitrary input pod name.
+        namespace: Namespace in which the Pod will be created.
 
     Returns:
         Sanitized pod name.
     """
+    # https://kubernetes.io/docs/concepts/overview/working-with-objects/names/#dns-subdomain-names
     pod_name = re.sub(r"[^a-z0-9-]", "-", pod_name.lower())
     pod_name = re.sub(r"^[-]+", "", pod_name)
-    return re.sub(r"[-]+", "-", pod_name)
+    pod_name = re.sub(r"[-]+$", "", pod_name)
+    pod_name = re.sub(r"[-]+", "-", pod_name)
+
+    allowed_length = calculate_max_pod_name_length_for_namespace(
+        namespace=namespace
+    )
+    return pod_name[:allowed_length]
+
+
+def sanitize_label(label: str) -> str:
+    """Sanitize a label for a Kubernetes resource.
+
+    Args:
+        label: The label to sanitize.
+
+    Returns:
+        The sanitized label.
+    """
+    # https://kubernetes.io/docs/concepts/overview/working-with-objects/names/#rfc-1035-label-names
+    label = re.sub(r"[^a-z0-9-]", "-", label.lower())
+    label = re.sub(r"^[-]+", "", label)
+    label = re.sub(r"[-]+$", "", label)
+    label = re.sub(r"[-]+", "-", label)
+
+    return label[:63]
 
 
 def pod_is_not_pending(pod: k8s_client.V1Pod) -> bool:
@@ -169,7 +213,7 @@ def get_pod(
 
 
 def wait_pod(
-    core_api: k8s_client.CoreV1Api,
+    kube_client_fn: Callable[[], k8s_client.ApiClient],
     pod_name: str,
     namespace: str,
     exit_condition_lambda: Callable[[k8s_client.V1Pod], bool],
@@ -180,7 +224,11 @@ def wait_pod(
     """Wait for a pod to meet an exit condition.
 
     Args:
-        core_api: Client of `CoreV1Api` of Kubernetes API.
+        kube_client_fn: the kube client fn is a function that is called
+            periodically and is used to get a `CoreV1Api` client for
+            the Kubernetes API. It should cache the client to avoid
+            unnecessary overhead but should also instantiate a new client if
+            the previous one is using credentials that are about to expire.
         pod_name: The name of the pod.
         namespace: The namespace of the pod.
         exit_condition_lambda: A lambda
@@ -210,6 +258,9 @@ def wait_pod(
     logged_lines = 0
 
     while True:
+        kube_client = kube_client_fn()
+        core_api = k8s_client.CoreV1Api(kube_client)
+
         resp = get_pod(core_api, pod_name, namespace)
 
         # Stream logs to `zenml.logger.info()`.
@@ -218,8 +269,11 @@ def wait_pod(
             response = core_api.read_namespaced_pod_log(
                 name=pod_name,
                 namespace=namespace,
+                _preload_content=False,
             )
-            logs = response.splitlines()
+            raw_data = response.data
+            decoded_log = raw_data.decode("utf-8", errors="replace")
+            logs = decoded_log.splitlines()
             if len(logs) > logged_lines:
                 for line in logs[logged_lines:]:
                     logger.info(line)
@@ -278,7 +332,7 @@ def create_edit_service_account(
     rbac_api: k8s_client.RbacAuthorizationV1Api,
     service_account_name: str,
     namespace: str,
-    cluster_role_binding_name: str = "zenml-edit",
+    role_binding_name: str = "zenml-edit",
 ) -> None:
     """Create a new Kubernetes service account with "edit" rights.
 
@@ -287,16 +341,17 @@ def create_edit_service_account(
         rbac_api: Client of Rbac Authorization V1 API of Kubernetes API.
         service_account_name: Name of the service account.
         namespace: Kubernetes namespace. Defaults to "default".
-        cluster_role_binding_name: Name of the cluster role binding.
-            Defaults to "zenml-edit".
+        role_binding_name: Name of the role binding. Defaults to "zenml-edit".
     """
-    crb_manifest = build_cluster_role_binding_manifest_for_service_account(
-        name=cluster_role_binding_name,
+    rb_manifest = build_role_binding_manifest_for_service_account(
+        name=role_binding_name,
         role_name="edit",
         service_account_name=service_account_name,
         namespace=namespace,
     )
-    _if_not_exists(rbac_api.create_cluster_role_binding)(body=crb_manifest)
+    _if_not_exists(rbac_api.create_namespaced_role_binding)(
+        namespace=namespace, body=rb_manifest
+    )
 
     sa_manifest = build_service_account_manifest(
         name=service_account_name, namespace=namespace

zenml/integrations/kubernetes/orchestrators/kubernetes_orchestrator.py

@@ -31,14 +31,22 @@
 """Kubernetes-native orchestrator."""
 
 import os
-from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple, Type, cast
+from typing import (
+    TYPE_CHECKING,
+    Any,
+    Dict,
+    List,
+    Optional,
+    Tuple,
+    Type,
+    cast,
+)
 
 from kubernetes import client as k8s_client
 from kubernetes import config as k8s_config
 
 from zenml.config.base_settings import BaseSettings
 from zenml.enums import StackComponentType
-from zenml.environment import Environment
 from zenml.integrations.kubernetes.flavors.kubernetes_orchestrator_flavor import (
     KubernetesOrchestratorConfig,
     KubernetesOrchestratorSettings,
@@ -51,6 +59,7 @@ from zenml.integrations.kubernetes.orchestrators.manifest_utils import (
     build_cron_job_manifest,
     build_pod_manifest,
 )
+from zenml.integrations.kubernetes.pod_settings import KubernetesPodSettings
 from zenml.logger import get_logger
 from zenml.orchestrators import ContainerizedOrchestrator
 from zenml.orchestrators.utils import get_orchestrator_run_name
@@ -88,6 +97,14 @@ class KubernetesOrchestrator(ContainerizedOrchestrator):
         if incluster is None:
             incluster = self.config.incluster
 
+        if incluster:
+            kube_utils.load_kube_config(
+                incluster=incluster,
+                context=self.config.kubernetes_context,
+            )
+            self._k8s_client = k8s_client.ApiClient()
+            return self._k8s_client
+
         # Refresh the client also if the connector has expired
         if self._k8s_client and not self.connector_has_expired():
             return self._k8s_client
@@ -200,11 +217,10 @@ class KubernetesOrchestrator(ContainerizedOrchestrator):
             # this, but just in case
             assert container_registry is not None
 
-            connector = self.get_connector()
             kubernetes_context = self.config.kubernetes_context
             msg = f"'{self.name}' Kubernetes orchestrator error: "
 
-            if not connector:
+            if not self.connector:
                 if not kubernetes_context:
                     return False, (
                         f"{msg}you must either link this stack component to a "
@@ -313,6 +329,45 @@ class KubernetesOrchestrator(ContainerizedOrchestrator):
             custom_validation_function=_validate_local_requirements,
         )
 
+    @classmethod
+    def apply_default_resource_requests(
+        cls,
+        memory: str,
+        cpu: Optional[str] = None,
+        pod_settings: Optional[KubernetesPodSettings] = None,
+    ) -> KubernetesPodSettings:
+        """Applies default resource requests to a pod settings object.
+
+        Args:
+            memory: The memory resource request.
+            cpu: The CPU resource request.
+            pod_settings: The pod settings to update. A new one will be created
+                if not provided.
+
+        Returns:
+            The new or updated pod settings.
+        """
+        resources = {
+            "requests": {"memory": memory},
+        }
+        if cpu:
+            resources["requests"]["cpu"] = cpu
+        if not pod_settings:
+            pod_settings = KubernetesPodSettings(resources=resources)
+        elif not pod_settings.resources:
+            # We can't update the pod settings in place (because it's a frozen
+            # pydantic model), so we have to create a new one.
+            pod_settings = KubernetesPodSettings(
+                **pod_settings.model_dump(exclude_unset=True),
+                resources=resources,
+            )
+        else:
+            set_requests = pod_settings.resources.get("requests", {})
+            resources["requests"].update(set_requests)
+            pod_settings.resources["requests"] = resources["requests"]
+
+        return pod_settings
+
     def prepare_or_run_pipeline(
         self,
         deployment: "PipelineDeploymentResponse",
@@ -328,19 +383,8 @@ class KubernetesOrchestrator(ContainerizedOrchestrator):
                 environment.
 
         Raises:
-            RuntimeError: If trying to run from a Jupyter notebook.
+            RuntimeError: If the Kubernetes orchestrator is not configured.
         """
-        # First check whether the code is running in a notebook.
-        if Environment.in_notebook():
-            raise RuntimeError(
-                "The Kubernetes orchestrator cannot run pipelines in a notebook "
-                "environment. The reason is that it is non-trivial to create "
-                "a Docker image of a notebook. Please consider refactoring "
-                "your notebook cells into separate scripts in a Python module "
-                "and run the code outside of a notebook when using this "
-                "orchestrator."
-            )
-
         for step_name, step in deployment.step_configurations.items():
             if self.requires_resources_in_orchestration_environment(step):
                 logger.warning(
@@ -351,8 +395,19 @@ class KubernetesOrchestrator(ContainerizedOrchestrator):
                 )
 
         pipeline_name = deployment.pipeline_configuration.name
-        orchestrator_run_name = get_orchestrator_run_name(pipeline_name)
-        pod_name = kube_utils.sanitize_pod_name(orchestrator_run_name)
+
+        # We already make sure the orchestrator run name has the correct length
+        # to make sure we don't cut off the randomized suffix later when
+        # sanitizing the pod name. This avoids any pod naming collisions.
+        max_length = kube_utils.calculate_max_pod_name_length_for_namespace(
+            namespace=self.config.kubernetes_namespace
+        )
+        orchestrator_run_name = get_orchestrator_run_name(
+            pipeline_name, max_length=max_length
+        )
+        pod_name = kube_utils.sanitize_pod_name(
+            orchestrator_run_name, namespace=self.config.kubernetes_namespace
+        )
 
         assert stack.container_registry
 
@@ -402,7 +457,8 @@ class KubernetesOrchestrator(ContainerizedOrchestrator):
                 command=command,
                 args=args,
                 service_account_name=service_account_name,
-                settings=settings,
+                privileged=False,
+                pod_settings=settings.orchestrator_pod_settings,
                 env=environment,
                 mount_local_stores=self.config.is_local,
             )
@@ -417,6 +473,17 @@ class KubernetesOrchestrator(ContainerizedOrchestrator):
             )
             return
 
+        # We set some default minimum resource requests for the orchestrator pod
+        # here if the user has not specified any, because the orchestrator pod
+        # takes up some memory resources itself and, if not specified, the pod
+        # will be scheduled on any node regardless of available memory and risk
+        # negatively impacting or even crashing the node due to memory pressure.
+        orchestrator_pod_settings = self.apply_default_resource_requests(
+            memory="400Mi",
+            cpu="100m",
+            pod_settings=settings.orchestrator_pod_settings,
+        )
+
         # Create and run the orchestrator pod.
         pod_manifest = build_pod_manifest(
             run_name=orchestrator_run_name,
@@ -425,8 +492,9 @@ class KubernetesOrchestrator(ContainerizedOrchestrator):
             image_name=image,
             command=command,
             args=args,
+            privileged=False,
+            pod_settings=orchestrator_pod_settings,
             service_account_name=service_account_name,
-            settings=settings,
             env=environment,
             mount_local_stores=self.config.is_local,
         )
@@ -440,7 +508,7 @@ class KubernetesOrchestrator(ContainerizedOrchestrator):
         if settings.synchronous:
             logger.info("Waiting for Kubernetes orchestrator pod...")
             kube_utils.wait_pod(
-                core_api=self._k8s_core_api,
+                kube_client_fn=self.get_kube_client,
                 pod_name=pod_name,
                 namespace=self.config.kubernetes_namespace,
                 exit_condition_lambda=kube_utils.pod_is_done,