zenml-nightly 0.75.0.dev20250317__py3-none-any.whl → 0.80.0.dev20250321__py3-none-any.whl

zenml/integrations/huggingface/__init__.py

@@ -75,4 +75,3 @@ class HuggingfaceIntegration(Integration):
         return [HuggingFaceModelDeployerFlavor]
 
 
-HuggingfaceIntegration.check_installation()

zenml/integrations/huggingface/services/huggingface_deployment.py

@@ -28,11 +28,13 @@ from huggingface_hub.errors import HfHubHTTPError
 from pydantic import Field
 
 from zenml.client import Client
+from zenml.enums import ServiceState
 from zenml.integrations.huggingface.flavors.huggingface_model_deployer_flavor import (
     HuggingFaceBaseConfig,
 )
 from zenml.logger import get_logger
-from zenml.services import ServiceState, ServiceStatus, ServiceType
+from zenml.models.v2.misc.service import ServiceType
+from zenml.services import ServiceStatus
 from zenml.services.service import BaseDeploymentService, ServiceConfig
 
 logger = get_logger(__name__)

zenml/integrations/hyperai/__init__.py

@@ -49,5 +49,3 @@ class HyperAIIntegration(Integration):
 
         return [HyperAIOrchestratorFlavor]
 
-
-HyperAIIntegration.check_installation()

zenml/integrations/kaniko/__init__.py

@@ -38,5 +38,3 @@ class KanikoIntegration(Integration):
 
         return [KanikoImageBuilderFlavor]
 
-
-KanikoIntegration.check_installation()

zenml/integrations/kubeflow/__init__.py

@@ -49,4 +49,3 @@ class KubeflowIntegration(Integration):
         return [KubeflowOrchestratorFlavor]
 
 
-KubeflowIntegration.check_installation()

zenml/integrations/kubeflow/orchestrators/kubeflow_orchestrator.py

@@ -578,6 +578,12 @@ class KubeflowOrchestrator(ContainerizedOrchestrator):
                             "Volume mounts are set but not supported in "
                             "Kubeflow with Kubeflow Pipelines 2.x. Ignoring..."
                         )
+                    if pod_settings.env or pod_settings.env_from:
+                        logger.warning(
+                            "Environment variables are set but not supported "
+                            "in Kubeflow with Kubeflow Pipelines 2.x. "
+                            "Ignoring..."
+                        )
 
                     # apply pod settings
                     if (

zenml/integrations/kubernetes/__init__.py

@@ -48,5 +48,3 @@ class KubernetesIntegration(Integration):
 
         return [KubernetesOrchestratorFlavor, KubernetesStepOperatorFlavor]
 
-
-KubernetesIntegration.check_installation()

zenml/integrations/kubernetes/flavors/kubernetes_orchestrator_flavor.py

@@ -85,6 +85,9 @@ class KubernetesOrchestratorConfig(
         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.
+        pass_zenml_token_as_secret: If `True`, the ZenML token will be passed
+            as a Kubernetes secret to the pods. For this to work, the Kubernetes
+            client must have permissions to create secrets in the namespace.
     """
 
     incluster: bool = False
@@ -93,6 +96,7 @@ class KubernetesOrchestratorConfig(
     local: bool = False
     skip_local_validations: bool = False
     parallel_step_startup_waiting_period: Optional[float] = None
+    pass_zenml_token_as_secret: bool = False
 
     @property
     def is_remote(self) -> bool:

zenml/integrations/kubernetes/orchestrators/kube_utils.py

@@ -34,7 +34,7 @@ Adjusted from https://github.com/tensorflow/tfx/blob/master/tfx/utils/kube_utils
 import enum
 import re
 import time
-from typing import Any, Callable, Optional, TypeVar, cast
+from typing import Any, Callable, Dict, Optional, TypeVar, cast
 
 from kubernetes import client as k8s_client
 from kubernetes import config as k8s_config
@@ -43,6 +43,7 @@ from kubernetes.client.rest import ApiException
 from zenml.integrations.kubernetes.orchestrators.manifest_utils import (
     build_namespace_manifest,
     build_role_binding_manifest_for_service_account,
+    build_secret_manifest,
     build_service_account_manifest,
 )
 from zenml.logger import get_logger
@@ -371,3 +372,90 @@ def create_namespace(core_api: k8s_client.CoreV1Api, namespace: str) -> None:
     """
     manifest = build_namespace_manifest(namespace)
     _if_not_exists(core_api.create_namespace)(body=manifest)
+
+
+def create_secret(
+    core_api: k8s_client.CoreV1Api,
+    namespace: str,
+    secret_name: str,
+    data: Dict[str, Optional[str]],
+) -> None:
+    """Create a Kubernetes secret.
+
+    Args:
+        core_api: Client of Core V1 API of Kubernetes API.
+        namespace: The namespace in which to create the secret.
+        secret_name: The name of the secret to create.
+        data: The secret data.
+    """
+    core_api.create_namespaced_secret(
+        namespace=namespace,
+        body=build_secret_manifest(name=secret_name, data=data),
+    )
+
+
+def update_secret(
+    core_api: k8s_client.CoreV1Api,
+    namespace: str,
+    secret_name: str,
+    data: Dict[str, Optional[str]],
+) -> None:
+    """Update a Kubernetes secret.
+
+    Args:
+        core_api: Client of Core V1 API of Kubernetes API.
+        namespace: The namespace in which to update the secret.
+        secret_name: The name of the secret to update.
+        data: The secret data. If the value is None, the key will be removed
+            from the secret.
+    """
+    core_api.patch_namespaced_secret(
+        namespace=namespace,
+        name=secret_name,
+        body=build_secret_manifest(name=secret_name, data=data),
+    )
+
+
+def create_or_update_secret(
+    core_api: k8s_client.CoreV1Api,
+    namespace: str,
+    secret_name: str,
+    data: Dict[str, Optional[str]],
+) -> None:
+    """Create a Kubernetes secret if it doesn't exist, or update it if it does.
+
+    Args:
+        core_api: Client of Core V1 API of Kubernetes API.
+        namespace: The namespace in which to create or update the secret.
+        secret_name: The name of the secret to create or update.
+        data: The secret data. If the value is None, the key will be removed
+            from the secret.
+
+    Raises:
+        ApiException: If the secret creation failed for any reason other than
+            the secret already existing.
+    """
+    try:
+        create_secret(core_api, namespace, secret_name, data)
+    except ApiException as e:
+        if e.status != 409:
+            raise
+        update_secret(core_api, namespace, secret_name, data)
+
+
+def delete_secret(
+    core_api: k8s_client.CoreV1Api,
+    namespace: str,
+    secret_name: str,
+) -> None:
+    """Delete a Kubernetes secret.
+
+    Args:
+        core_api: Client of Core V1 API of Kubernetes API.
+        namespace: The namespace in which to delete the secret.
+        secret_name: The name of the secret to delete.
+    """
+    core_api.delete_namespaced_secret(
+        name=secret_name,
+        namespace=namespace,
+    )

zenml/integrations/kubernetes/orchestrators/kubernetes_orchestrator.py

@@ -41,6 +41,7 @@ from typing import (
     Type,
     cast,
 )
+from uuid import UUID
 
 from kubernetes import client as k8s_client
 from kubernetes import config as k8s_config
@@ -72,6 +73,7 @@ if TYPE_CHECKING:
 logger = get_logger(__name__)
 
 ENV_ZENML_KUBERNETES_RUN_ID = "ZENML_KUBERNETES_RUN_ID"
+KUBERNETES_SECRET_TOKEN_KEY_NAME = "zenml_api_token"
 
 
 class KubernetesOrchestrator(ContainerizedOrchestrator):
@@ -368,6 +370,17 @@ class KubernetesOrchestrator(ContainerizedOrchestrator):
 
         return pod_settings
 
+    def get_token_secret_name(self, deployment_id: UUID) -> str:
+        """Returns the name of the secret that contains the ZenML token.
+
+        Args:
+            deployment_id: The ID of the deployment.
+
+        Returns:
+            The name of the secret that contains the ZenML token.
+        """
+        return f"zenml-token-{deployment_id}"
+
     def prepare_or_run_pipeline(
         self,
         deployment: "PipelineDeploymentResponse",
@@ -439,6 +452,38 @@ class KubernetesOrchestrator(ContainerizedOrchestrator):
         # Authorize pod to run Kubernetes commands inside the cluster.
         service_account_name = self._get_service_account_name(settings)
 
+        # 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,
+        )
+
+        if self.config.pass_zenml_token_as_secret:
+            secret_name = self.get_token_secret_name(deployment.id)
+            token = environment.pop("ZENML_STORE_API_TOKEN")
+            kube_utils.create_or_update_secret(
+                core_api=self._k8s_core_api,
+                namespace=self.config.kubernetes_namespace,
+                secret_name=secret_name,
+                data={KUBERNETES_SECRET_TOKEN_KEY_NAME: token},
+            )
+            orchestrator_pod_settings.env.append(
+                {
+                    "name": "ZENML_STORE_API_TOKEN",
+                    "valueFrom": {
+                        "secretKeyRef": {
+                            "name": secret_name,
+                            "key": KUBERNETES_SECRET_TOKEN_KEY_NAME,
+                        }
+                    },
+                }
+            )
+
         # Schedule as CRON job if CRON schedule is given.
         if deployment.schedule:
             if not deployment.schedule.cron_expression:
@@ -458,7 +503,7 @@ class KubernetesOrchestrator(ContainerizedOrchestrator):
                 args=args,
                 service_account_name=service_account_name,
                 privileged=False,
-                pod_settings=settings.orchestrator_pod_settings,
+                pod_settings=orchestrator_pod_settings,
                 env=environment,
                 mount_local_stores=self.config.is_local,
             )
@@ -472,57 +517,46 @@ class KubernetesOrchestrator(ContainerizedOrchestrator):
                 f'`"{cron_expression}"`.'
             )
             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,
-            pod_name=pod_name,
-            pipeline_name=pipeline_name,
-            image_name=image,
-            command=command,
-            args=args,
-            privileged=False,
-            pod_settings=orchestrator_pod_settings,
-            service_account_name=service_account_name,
-            env=environment,
-            mount_local_stores=self.config.is_local,
-        )
-
-        self._k8s_core_api.create_namespaced_pod(
-            namespace=self.config.kubernetes_namespace,
-            body=pod_manifest,
-        )
-
-        # Wait for the orchestrator pod to finish and stream logs.
-        if settings.synchronous:
-            logger.info("Waiting for Kubernetes orchestrator pod...")
-            kube_utils.wait_pod(
-                kube_client_fn=self.get_kube_client,
+        else:
+            # Create and run the orchestrator pod.
+            pod_manifest = build_pod_manifest(
+                run_name=orchestrator_run_name,
                 pod_name=pod_name,
-                namespace=self.config.kubernetes_namespace,
-                exit_condition_lambda=kube_utils.pod_is_done,
-                timeout_sec=settings.timeout,
-                stream_logs=True,
+                pipeline_name=pipeline_name,
+                image_name=image,
+                command=command,
+                args=args,
+                privileged=False,
+                pod_settings=orchestrator_pod_settings,
+                service_account_name=service_account_name,
+                env=environment,
+                mount_local_stores=self.config.is_local,
             )
-        else:
-            logger.info(
-                f"Orchestration started asynchronously in pod "
-                f"`{self.config.kubernetes_namespace}:{pod_name}`. "
-                f"Run the following command to inspect the logs: "
-                f"`kubectl logs {pod_name} -n {self.config.kubernetes_namespace}`."
+
+            self._k8s_core_api.create_namespaced_pod(
+                namespace=self.config.kubernetes_namespace,
+                body=pod_manifest,
             )
 
+            # Wait for the orchestrator pod to finish and stream logs.
+            if settings.synchronous:
+                logger.info("Waiting for Kubernetes orchestrator pod...")
+                kube_utils.wait_pod(
+                    kube_client_fn=self.get_kube_client,
+                    pod_name=pod_name,
+                    namespace=self.config.kubernetes_namespace,
+                    exit_condition_lambda=kube_utils.pod_is_done,
+                    timeout_sec=settings.timeout,
+                    stream_logs=True,
+                )
+            else:
+                logger.info(
+                    f"Orchestration started asynchronously in pod "
+                    f"`{self.config.kubernetes_namespace}:{pod_name}`. "
+                    f"Run the following command to inspect the logs: "
+                    f"`kubectl logs {pod_name} -n {self.config.kubernetes_namespace}`."
+                )
+
     def _get_service_account_name(
         self, settings: KubernetesOrchestratorSettings
     ) -> str:

zenml/integrations/kubernetes/orchestrators/kubernetes_orchestrator_entrypoint.py

@@ -28,6 +28,7 @@ from zenml.integrations.kubernetes.flavors.kubernetes_orchestrator_flavor import
 from zenml.integrations.kubernetes.orchestrators import kube_utils
 from zenml.integrations.kubernetes.orchestrators.kubernetes_orchestrator import (
     ENV_ZENML_KUBERNETES_RUN_ID,
+    KUBERNETES_SECRET_TOKEN_KEY_NAME,
     KubernetesOrchestrator,
 )
 from zenml.integrations.kubernetes.orchestrators.manifest_utils import (
@@ -82,6 +83,9 @@ def main() -> None:
     kube_client = orchestrator.get_kube_client(incluster=True)
     core_api = k8s_client.CoreV1Api(kube_client)
 
+    env = get_config_environment_vars()
+    env[ENV_ZENML_KUBERNETES_RUN_ID] = orchestrator_run_id
+
     def run_step_on_kubernetes(step_name: str) -> None:
         """Run a pipeline step in a separate Kubernetes pod.
 
@@ -115,9 +119,6 @@ def main() -> None:
             orchestrator_settings
         )
 
-        env = get_config_environment_vars()
-        env[ENV_ZENML_KUBERNETES_RUN_ID] = orchestrator_run_id
-
         # We set some default minimum memory resource requests for the step pod
         # here if the user has not specified any, because the step pod takes up
         # some memory resources itself and, if not specified, the pod will be
@@ -128,6 +129,23 @@ def main() -> None:
             pod_settings=settings.pod_settings,
         )
 
+        if orchestrator.config.pass_zenml_token_as_secret:
+            env.pop("ZENML_STORE_API_TOKEN", None)
+            secret_name = orchestrator.get_token_secret_name(
+                deployment_config.id
+            )
+            pod_settings.env.append(
+                {
+                    "name": "ZENML_STORE_API_TOKEN",
+                    "valueFrom": {
+                        "secretKeyRef": {
+                            "name": secret_name,
+                            "key": KUBERNETES_SECRET_TOKEN_KEY_NAME,
+                        }
+                    },
+                }
+            )
+
         # Define Kubernetes pod manifest.
         pod_manifest = build_pod_manifest(
             pod_name=pod_name,
@@ -166,13 +184,29 @@ def main() -> None:
     parallel_node_startup_waiting_period = (
         orchestrator.config.parallel_step_startup_waiting_period or 0.0
     )
-    ThreadedDagRunner(
-        dag=pipeline_dag,
-        run_fn=run_step_on_kubernetes,
-        parallel_node_startup_waiting_period=parallel_node_startup_waiting_period,
-    ).run()
-
-    logger.info("Orchestration pod completed.")
+    try:
+        ThreadedDagRunner(
+            dag=pipeline_dag,
+            run_fn=run_step_on_kubernetes,
+            parallel_node_startup_waiting_period=parallel_node_startup_waiting_period,
+        ).run()
+        logger.info("Orchestration pod completed.")
+    finally:
+        if (
+            orchestrator.config.pass_zenml_token_as_secret
+            and deployment_config.schedule is None
+        ):
+            secret_name = orchestrator.get_token_secret_name(
+                deployment_config.id
+            )
+            try:
+                kube_utils.delete_secret(
+                    core_api=core_api,
+                    namespace=args.kubernetes_namespace,
+                    secret_name=secret_name,
+                )
+            except k8s_client.rest.ApiException as e:
+                logger.error(f"Error cleaning up secret {secret_name}: {e}")
 
 
 if __name__ == "__main__":

zenml/integrations/kubernetes/orchestrators/manifest_utils.py

@@ -13,9 +13,10 @@
 #  permissions and limitations under the License.
 """Utility functions for building manifests for k8s pods."""
 
+import base64
 import os
 import sys
-from typing import Any, Dict, List, Optional
+from typing import Any, Dict, List, Mapping, Optional
 
 from kubernetes import client as k8s_client
 
@@ -222,6 +223,18 @@ def add_pod_settings(
             else:
                 container.volume_mounts = settings.volume_mounts
 
+        if settings.env:
+            if container.env:
+                container.env.extend(settings.env)
+            else:
+                container.env = settings.env
+
+        if settings.env_from:
+            if container.env_from:
+                container.env_from.extend(settings.env_from)
+            else:
+                container.env_from = settings.env_from
+
     if settings.volumes:
         if pod_spec.volumes:
             pod_spec.volumes.extend(settings.volumes)
@@ -378,3 +391,34 @@ def build_namespace_manifest(namespace: str) -> Dict[str, Any]:
             "name": namespace,
         },
     }
+
+
+def build_secret_manifest(
+    name: str,
+    data: Mapping[str, Optional[str]],
+    secret_type: str = "Opaque",
+) -> Dict[str, Any]:
+    """Builds a Kubernetes secret manifest.
+
+    Args:
+        name: Name of the secret.
+        data: The secret data.
+        secret_type: The secret type.
+
+    Returns:
+        The secret manifest.
+    """
+    encoded_data = {
+        key: base64.b64encode(value.encode()).decode() if value else None
+        for key, value in data.items()
+    }
+
+    return {
+        "apiVersion": "v1",
+        "kind": "Secret",
+        "metadata": {
+            "name": name,
+        },
+        "type": secret_type,
+        "data": encoded_data,
+    }

zenml/integrations/kubernetes/pod_settings.py

@@ -35,6 +35,8 @@ class KubernetesPodSettings(BaseSettings):
         host_ipc: Whether to enable host IPC for the pod.
         image_pull_secrets: Image pull secrets to use for the pod.
         labels: Labels to apply to the pod.
+        env: Environment variables to apply to the container.
+        env_from: Environment variables to apply to the container.
     """
 
     node_selectors: Dict[str, str] = {}
@@ -47,6 +49,8 @@ class KubernetesPodSettings(BaseSettings):
     host_ipc: bool = False
     image_pull_secrets: List[str] = []
     labels: Dict[str, str] = {}
+    env: List[Dict[str, Any]] = []
+    env_from: List[Dict[str, Any]] = []
 
     @field_validator("volumes", mode="before")
     @classmethod
@@ -155,3 +159,51 @@ class KubernetesPodSettings(BaseSettings):
             return serialization_utils.serialize_kubernetes_model(value)
         else:
             return value
+
+    @field_validator("env", mode="before")
+    @classmethod
+    def _convert_env(cls, value: Any) -> Any:
+        """Converts Kubernetes EnvVar to a dict.
+
+        Args:
+            value: The env value.
+
+        Returns:
+            The converted value.
+        """
+        from kubernetes.client.models import V1EnvVar
+
+        result = []
+        for element in value:
+            if isinstance(element, V1EnvVar):
+                result.append(
+                    serialization_utils.serialize_kubernetes_model(element)
+                )
+            else:
+                result.append(element)
+
+        return result
+
+    @field_validator("env_from", mode="before")
+    @classmethod
+    def _convert_env_from(cls, value: Any) -> Any:
+        """Converts Kubernetes EnvFromSource to a dict.
+
+        Args:
+            value: The env from value.
+
+        Returns:
+            The converted value.
+        """
+        from kubernetes.client.models import V1EnvFromSource
+
+        result = []
+        for element in value:
+            if isinstance(element, V1EnvFromSource):
+                result.append(
+                    serialization_utils.serialize_kubernetes_model(element)
+                )
+            else:
+                result.append(element)
+
+        return result

zenml/integrations/label_studio/__init__.py

@@ -42,5 +42,3 @@ class LabelStudioIntegration(Integration):
 
         return [LabelStudioAnnotatorFlavor]
 
-
-LabelStudioIntegration.check_installation()

zenml/integrations/langchain/__init__.py

@@ -38,5 +38,3 @@ class LangchainIntegration(Integration):
         """Activates the integration."""
         from zenml.integrations.langchain import materializers  # noqa
 
-
-LangchainIntegration.check_installation()

zenml/integrations/lightgbm/__init__.py

@@ -30,4 +30,3 @@ class LightGBMIntegration(Integration):
         from zenml.integrations.lightgbm import materializers  # noqa
 
 
-LightGBMIntegration.check_installation()

zenml/integrations/lightning/__init__.py

@@ -44,5 +44,3 @@ class LightningIntegration(Integration):
         return [
             LightningOrchestratorFlavor,
         ]
-
-LightningIntegration.check_installation()

zenml/integrations/mlflow/__init__.py

@@ -58,7 +58,7 @@ class MlflowIntegration(Integration):
         from zenml.integrations.pandas import PandasIntegration
 
         reqs = [
-            "mlflow>=2.1.1,<3",
+            "mlflow>=2.1.1,<2.21.0",
             # TODO: remove this requirement once rapidjson is fixed
             "python-rapidjson<1.15",
             # When you do:
@@ -111,5 +111,3 @@ class MlflowIntegration(Integration):
             MLFlowModelRegistryFlavor,
         ]
 
-
-MlflowIntegration.check_installation()

zenml/integrations/mlflow/services/mlflow_deployment.py

@@ -30,6 +30,7 @@ from zenml.integrations.mlflow.experiment_trackers.mlflow_experiment_tracker imp
     MLFlowExperimentTracker,
 )
 from zenml.logger import get_logger
+from zenml.models.v2.misc.service import ServiceType
 from zenml.services import (
     HTTPEndpointHealthMonitor,
     HTTPEndpointHealthMonitorConfig,
@@ -38,7 +39,6 @@ from zenml.services import (
     LocalDaemonServiceEndpoint,
     LocalDaemonServiceEndpointConfig,
     ServiceEndpointProtocol,
-    ServiceType,
 )
 from zenml.services.service import BaseDeploymentService
 

zenml/integrations/modal/__init__.py

@@ -43,4 +43,3 @@ class ModalIntegration(Integration):
         return [ModalStepOperatorFlavor]
 
 
-ModalIntegration.check_installation()

zenml/integrations/neptune/__init__.py

@@ -46,5 +46,3 @@ class NeptuneIntegration(Integration):
             NeptuneExperimentTrackerFlavor,
         ]
 
-
-NeptuneIntegration.check_installation()  # this checks if the requirements are installed

zenml/integrations/neural_prophet/__init__.py

@@ -33,5 +33,3 @@ class NeuralProphetIntegration(Integration):
         """Activates the integration."""
         from zenml.integrations.neural_prophet import materializers  # noqa
 
-
-NeuralProphetIntegration.check_installation()

zenml/integrations/numpy/__init__.py

@@ -28,5 +28,3 @@ class NumpyIntegration(Integration):
         """Activates the integration."""
         from zenml.integrations.numpy import materializers  # noqa
 
-
-NumpyIntegration.check_installation()