zenml-nightly 0.75.0.dev20250318__py3-none-any.whl → 0.80.0.dev20250322__py3-none-any.whl

zenml/integrations/discord/__init__.py

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

zenml/integrations/evidently/__init__.py

@@ -87,5 +87,3 @@ class EvidentlyIntegration(Integration):
 
         return [EvidentlyDataValidatorFlavor]
 
-
-EvidentlyIntegration.check_installation()

zenml/integrations/facets/__init__.py

@@ -45,5 +45,3 @@ class FacetsIntegration(Integration):
         return cls.REQUIREMENTS + \
             PandasIntegration.get_requirements(target_os=target_os)
 
-
-FacetsIntegration.check_installation()

zenml/integrations/feast/__init__.py

@@ -60,5 +60,3 @@ class FeastIntegration(Integration):
         return cls.REQUIREMENTS + \
             PandasIntegration.get_requirements(target_os=target_os)
 
-
-FeastIntegration.check_installation()

zenml/integrations/gcp/__init__.py

@@ -84,6 +84,3 @@ class GcpIntegration(Integration):
             VertexOrchestratorFlavor,
             VertexStepOperatorFlavor,
         ]
-
-
-GcpIntegration.check_installation()

zenml/integrations/github/__init__.py

@@ -39,5 +39,3 @@ class GitHubIntegration(Integration):
 
         return [GithubWebhookEventSourceFlavor]
 
-
-GitHubIntegration.check_installation()

zenml/integrations/github/code_repositories/github_code_repository.py

@@ -237,12 +237,17 @@ class GitHubCodeRepository(BaseCodeRepository):
         Returns:
             Whether the remote url is correct.
         """
-        https_url = f"https://{self.config.host}/{self.config.owner}/{self.config.repository}.git"
+        host = self.config.host or "github.com"
+        host = host.rstrip("/")
+
+        https_url = (
+            f"https://{host}/{self.config.owner}/{self.config.repository}.git"
+        )
         if url == https_url:
             return True
 
         ssh_regex = re.compile(
-            f".*@{self.config.host}:{self.config.owner}/{self.config.repository}.git"
+            f".*@{host}:{self.config.owner}/{self.config.repository}.git"
         )
         if ssh_regex.fullmatch(url):
             return True

zenml/integrations/gitlab/__init__.py

@@ -25,4 +25,3 @@ class GitLabIntegration(Integration):
     REQUIREMENTS: List[str] = ["python-gitlab"]
 
 
-GitLabIntegration.check_installation()

zenml/integrations/gitlab/code_repositories/gitlab_code_repository.py

@@ -176,10 +176,13 @@ class GitLabCodeRepository(BaseCodeRepository):
         Returns:
             Whether the remote url is correct.
         """
+        host = self.config.host or "gitlab.com"
+        host = host.rstrip("/")
+
         parsed_url = urlparse(url)
         if (
             parsed_url.scheme == "https"
-            and parsed_url.hostname == self.config.host
+            and parsed_url.hostname == host
             and parsed_url.path
             == f"/{self.config.group}/{self.config.project}.git"
         ):
@@ -188,7 +191,7 @@ class GitLabCodeRepository(BaseCodeRepository):
         ssh_regex = re.compile(
             r"^(?P<scheme_with_delimiter>ssh://)?"
             r"(?P<userinfo>git)"
-            f"@{self.config.host}:"
+            f"@{host}:"
             r"(?P<port>\d+)?"
             r"(?(scheme_with_delimiter)/|/?)"
             f"{self.config.group}/{self.config.project}.git$",

zenml/integrations/great_expectations/__init__.py

@@ -66,6 +66,3 @@ class GreatExpectationsIntegration(Integration):
 
         return cls.REQUIREMENTS + \
             PandasIntegration.get_requirements(target_os=target_os)
-
-
-GreatExpectationsIntegration.check_installation()

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/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):
@@ -221,45 +223,52 @@ class KubernetesOrchestrator(ContainerizedOrchestrator):
             msg = f"'{self.name}' Kubernetes orchestrator error: "
 
             if not self.connector:
-                if not kubernetes_context:
+                if kubernetes_context:
+                    contexts, active_context = self.get_kubernetes_contexts()
+
+                    if kubernetes_context not in contexts:
+                        return False, (
+                            f"{msg}could not find a Kubernetes context named "
+                            f"'{kubernetes_context}' in the local "
+                            "Kubernetes configuration. Please make sure that "
+                            "the Kubernetes cluster is running and that the "
+                            "kubeconfig file is configured correctly. To list "
+                            "all configured contexts, run:\n\n"
+                            "  `kubectl config get-contexts`\n"
+                        )
+                    if kubernetes_context != active_context:
+                        logger.warning(
+                            f"{msg}the Kubernetes context "
+                            f"'{kubernetes_context}' configured for the "
+                            f"Kubernetes orchestrator is not the same as the "
+                            f"active context in the local Kubernetes "
+                            f"configuration. If this is not deliberate, you "
+                            f"should update the orchestrator's "
+                            f"`kubernetes_context` field by running:\n\n"
+                            f"  `zenml orchestrator update {self.name} "
+                            f"--kubernetes_context={active_context}`\n"
+                            f"To list all configured contexts, run:\n\n"
+                            f"  `kubectl config get-contexts`\n"
+                            f"To set the active context to be the same as the "
+                            f"one configured in the Kubernetes orchestrator "
+                            f"and silence this warning, run:\n\n"
+                            f"  `kubectl config use-context "
+                            f"{kubernetes_context}`\n"
+                        )
+                elif self.config.incluster:
+                    # No service connector or kubernetes_context is needed when
+                    # the orchestrator is being used from within a Kubernetes
+                    # cluster.
+                    pass
+                else:
                     return False, (
-                        f"{msg}you must either link this stack component to a "
+                        f"{msg}you must either link this orchestrator to a "
                         "Kubernetes service connector (see the 'zenml "
-                        "orchestrator connect' CLI command) or explicitly set "
+                        "orchestrator connect' CLI command), explicitly set "
                         "the `kubernetes_context` attribute to the name of the "
                         "Kubernetes config context pointing to the cluster "
-                        "where you would like to run pipelines."
-                    )
-
-                contexts, active_context = self.get_kubernetes_contexts()
-
-                if kubernetes_context not in contexts:
-                    return False, (
-                        f"{msg}could not find a Kubernetes context named "
-                        f"'{kubernetes_context}' in the local "
-                        "Kubernetes configuration. Please make sure that the "
-                        "Kubernetes cluster is running and that the kubeconfig "
-                        "file is configured correctly. To list all configured "
-                        "contexts, run:\n\n"
-                        "  `kubectl config get-contexts`\n"
-                    )
-                if kubernetes_context != active_context:
-                    logger.warning(
-                        f"{msg}the Kubernetes context '{kubernetes_context}' "  # nosec
-                        f"configured for the Kubernetes orchestrator is not "
-                        f"the same as the active context in the local "
-                        f"Kubernetes configuration. If this is not deliberate,"
-                        f" you should update the orchestrator's "
-                        f"`kubernetes_context` field by running:\n\n"
-                        f"  `zenml orchestrator update {self.name} "
-                        f"--kubernetes_context={active_context}`\n"
-                        f"To list all configured contexts, run:\n\n"
-                        f"  `kubectl config get-contexts`\n"
-                        f"To set the active context to be the same as the one "
-                        f"configured in the Kubernetes orchestrator and "
-                        f"silence this warning, run:\n\n"
-                        f"  `kubectl config use-context "
-                        f"{kubernetes_context}`\n"
+                        "where you would like to run pipelines, or set the "
+                        "`incluster` attribute to `True`."
                     )
 
             silence_local_validations_msg = (
@@ -368,6 +377,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 +459,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 +510,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 +524,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__":