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

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
 
@@ -390,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/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()

zenml/integrations/openai/__init__.py

@@ -24,4 +24,3 @@ class OpenAIIntegration(Integration):
     REQUIREMENTS = ["openai>=1.0.0"]
 
 
-OpenAIIntegration.check_installation()

zenml/integrations/pandas/__init__.py

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

zenml/integrations/pigeon/__init__.py

@@ -41,4 +41,3 @@ class PigeonIntegration(Integration):
         return [PigeonAnnotatorFlavor]
 
 
-PigeonIntegration.check_installation()

zenml/integrations/pillow/__init__.py

@@ -27,6 +27,3 @@ class PillowIntegration(Integration):
     def activate(cls) -> None:
         """Activates the integration."""
         from zenml.integrations.pillow import materializers  # noqa
-
-
-PillowIntegration.check_installation()

zenml/integrations/polars/__init__.py

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

zenml/integrations/prodigy/__init__.py

@@ -45,4 +45,3 @@ class ProdigyIntegration(Integration):
         return [ProdigyAnnotatorFlavor]
 
 
-ProdigyIntegration.check_installation()

zenml/integrations/pycaret/__init__.py

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

zenml/integrations/pytorch/__init__.py

@@ -31,4 +31,3 @@ class PytorchIntegration(Integration):
         from zenml.integrations.pytorch import materializers  # noqa
 
 
-PytorchIntegration.check_installation()

zenml/integrations/pytorch_lightning/__init__.py

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

zenml/integrations/registry.py

@@ -13,6 +13,8 @@
 #  permissions and limitations under the License.
 """Implementation of a registry to track ZenML integrations."""
 
+import importlib
+import os
 from typing import TYPE_CHECKING, Any, Dict, List, Optional, Type
 
 from zenml.exceptions import IntegrationError
@@ -30,6 +32,7 @@ class IntegrationRegistry(object):
     def __init__(self) -> None:
         """Initializing the integration registry."""
         self._integrations: Dict[str, Type["Integration"]] = {}
+        self._initialized = False
 
     @property
     def integrations(self) -> Dict[str, Type["Integration"]]:
@@ -38,6 +41,7 @@ class IntegrationRegistry(object):
         Returns:
             A dict of integration key to type of `Integration`.
         """
+        self._initialize()
         return self._integrations
 
     @integrations.setter
@@ -68,8 +72,33 @@ class IntegrationRegistry(object):
         """
         self._integrations[key] = type_
 
+    def _initialize(self) -> None:
+        """Method to register all integrations."""
+        if self._initialized:
+            return
+        self._initialized = True
+
+        # Load all submodules in the integrations module
+        integrations_dir = os.path.dirname(__file__)
+        for file in os.listdir(integrations_dir):
+            full_path = os.path.join(integrations_dir, file)
+            # Skip anything that isn't a directory
+            if not os.path.isdir(full_path):
+                continue
+            # Skip anything that doesn't have a __init__.py file
+            if not os.path.exists(os.path.join(full_path, "__init__.py")):
+                continue
+            # Import the module
+            module_path = f"zenml.integrations.{file}"
+            try:
+                importlib.import_module(module_path)
+            except ImportError:
+                logger.exception(f"Failed to import module `{module_path}`.")
+                continue
+
     def activate_integrations(self) -> None:
         """Method to activate the integrations with are registered in the registry."""
+        self._initialize()
         for name, integration in self._integrations.items():
             if integration.check_installation():
                 logger.debug(f"Activating integration `{name}`...")
@@ -85,6 +114,7 @@ class IntegrationRegistry(object):
         Returns:
             A list of all possible integrations.
         """
+        self._initialize()
         return [name for name in self._integrations]
 
     def select_integration_requirements(
@@ -104,6 +134,7 @@ class IntegrationRegistry(object):
         Raises:
             KeyError: If the integration is not found.
         """
+        self._initialize()
         if integration_name:
             if integration_name in self.list_integration_names:
                 return self._integrations[integration_name].get_requirements(
@@ -141,6 +172,7 @@ class IntegrationRegistry(object):
         Raises:
             KeyError: If the integration is not found.
         """
+        self._initialize()
         if integration_name:
             if integration_name in self.list_integration_names:
                 return self._integrations[
@@ -173,6 +205,7 @@ class IntegrationRegistry(object):
         Raises:
             KeyError: If the integration is not found.
         """
+        self._initialize()
         if integration_name in self.list_integration_names:
             return self._integrations[integration_name].check_installation()
         elif not integration_name:
@@ -194,6 +227,7 @@ class IntegrationRegistry(object):
         Returns:
             List of installed integrations.
         """
+        self._initialize()
         return [
             name
             for name, integration in integration_registry.integrations.items()

zenml/integrations/s3/__init__.py

@@ -58,5 +58,3 @@ class S3Integration(Integration):
 
         return [S3ArtifactStoreFlavor]
 
-
-S3Integration.check_installation()

zenml/integrations/scipy/__init__.py

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

zenml/integrations/seldon/__init__.py

@@ -67,5 +67,3 @@ class SeldonIntegration(Integration):
         return cls.REQUIREMENTS + \
             NumpyIntegration.get_requirements(target_os=target_os)
 
-
-SeldonIntegration.check_installation()

zenml/integrations/seldon/services/seldon_deployment.py

@@ -22,6 +22,7 @@ import requests
 from pydantic import Field, ValidationError
 
 from zenml import __version__
+from zenml.enums import ServiceState
 from zenml.integrations.seldon.seldon_client import (
     SeldonClient,
     SeldonDeployment,
@@ -30,9 +31,9 @@ from zenml.integrations.seldon.seldon_client import (
     SeldonResourceRequirements,
 )
 from zenml.logger import get_logger
+from zenml.models.v2.misc.service import ServiceType
 from zenml.services.service import BaseDeploymentService, ServiceConfig
-from zenml.services.service_status import ServiceState, ServiceStatus
-from zenml.services.service_type import ServiceType
+from zenml.services.service_status import ServiceStatus
 
 logger = get_logger(__name__)
 

zenml/integrations/sklearn/__init__.py

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

zenml/integrations/skypilot/orchestrators/skypilot_base_vm_orchestrator.py

@@ -318,30 +318,36 @@ class SkypilotBaseOrchestrator(ContainerizedOrchestrator):
                     disk_tier=settings.disk_tier,
                 )
             )
-            # Set the cluster name
+            # Do not detach run if logs are being streamed
+            # Otherwise, the logs will not be streamed after the task is submitted
+            # Could also be a parameter in the settings to control this behavior
+            detach_run = not settings.stream_logs
+
+            launch_new_cluster = True
             if settings.cluster_name:
-                sky.exec(
-                    task,
-                    settings.cluster_name,
-                    down=down,
-                    stream_logs=settings.stream_logs,
-                    backend=None,
-                    detach_run=True,
+                cluster_info = sky.status(
+                    refresh=True, cluster_names=settings.cluster_name
                 )
+                if cluster_info:
+                    logger.info(
+                        f"Found existing cluster {settings.cluster_name}. Reusing..."
+                    )
+                    launch_new_cluster = False
+
+                else:
+                    logger.info(
+                        f"Cluster {settings.cluster_name} not found. Launching a new one..."
+                    )
+                    cluster_name = settings.cluster_name
             else:
-                # Find existing cluster
-                for i in sky.status(refresh=True):
-                    if isinstance(
-                        i["handle"].launched_resources.cloud, type(self.cloud)
-                    ):
-                        cluster_name = i["handle"].cluster_name
-                        logger.info(
-                            f"Found existing cluster {cluster_name}. Reusing..."
-                        )
                 cluster_name = self.sanitize_cluster_name(
                     f"{orchestrator_run_name}"
                 )
-                # Launch the cluster
+                logger.info(
+                    f"No cluster name provided. Launching a new cluster with name {cluster_name}..."
+                )
+
+            if launch_new_cluster:
                 sky.launch(
                     task,
                     cluster_name,
@@ -349,7 +355,26 @@ class SkypilotBaseOrchestrator(ContainerizedOrchestrator):
                     idle_minutes_to_autostop=idle_minutes_to_autostop,
                     down=down,
                     stream_logs=settings.stream_logs,
+                    backend=None,
                     detach_setup=True,
+                    detach_run=detach_run,
+                )
+            else:
+                # Make sure the cluster is up -
+                # If the cluster is already up, this will not do anything
+                sky.start(
+                    settings.cluster_name,
+                    down=down,
+                    idle_minutes_to_autostop=idle_minutes_to_autostop,
+                    retry_until_up=settings.retry_until_up,
+                )
+                sky.exec(
+                    task,
+                    settings.cluster_name,
+                    down=down,
+                    stream_logs=settings.stream_logs,
+                    backend=None,
+                    detach_run=detach_run,
                 )
 
         except Exception as e:

zenml/integrations/skypilot_aws/__init__.py

@@ -32,7 +32,7 @@ class SkypilotAWSIntegration(Integration):
 
     NAME = SKYPILOT_AWS
     # all 0.6.x versions of skypilot[aws] are compatible
-    REQUIREMENTS = ["skypilot[aws]~=0.6.0"]
+    REQUIREMENTS = ["skypilot[aws]~=0.8.0"]
     APT_PACKAGES = ["openssh-client", "rsync"]
 
     @classmethod
@@ -48,5 +48,3 @@ class SkypilotAWSIntegration(Integration):
 
         return [SkypilotAWSOrchestratorFlavor]
 
-
-SkypilotAWSIntegration.check_installation()

zenml/integrations/skypilot_azure/__init__.py

@@ -31,7 +31,7 @@ class SkypilotAzureIntegration(Integration):
     """Definition of Skypilot (Azure) Integration for ZenML."""
 
     NAME = SKYPILOT_AZURE
-    REQUIREMENTS = ["skypilot[azure]>=0.6.1"]
+    REQUIREMENTS = ["skypilot[azure]~=0.8.0"]
     APT_PACKAGES = ["openssh-client", "rsync"]
 
     @classmethod
@@ -47,4 +47,3 @@ class SkypilotAzureIntegration(Integration):
 
         return [SkypilotAzureOrchestratorFlavor]
 
-SkypilotAzureIntegration.check_installation()

zenml/integrations/skypilot_gcp/__init__.py

@@ -31,7 +31,7 @@ class SkypilotGCPIntegration(Integration):
     """Definition of Skypilot (GCP) Integration for ZenML."""
 
     NAME = SKYPILOT_GCP
-    REQUIREMENTS = ["skypilot[gcp]~=0.6.0"]
+    REQUIREMENTS = ["skypilot[gcp]~=0.8.0"]
     APT_PACKAGES = ["openssh-client", "rsync"]
 
     @classmethod
@@ -46,5 +46,3 @@ class SkypilotGCPIntegration(Integration):
         )
 
         return [SkypilotGCPOrchestratorFlavor]
-
-SkypilotGCPIntegration.check_installation()

zenml/integrations/skypilot_kubernetes/__init__.py

@@ -32,7 +32,7 @@ class SkypilotKubernetesIntegration(Integration):
 
     NAME = SKYPILOT_KUBERNETES
     # all 0.6.x versions of skypilot[kubernetes] are compatible
-    REQUIREMENTS = ["skypilot[kubernetes]~=0.6.1"]
+    REQUIREMENTS = ["skypilot[kubernetes]~=0.8.0"]
     APT_PACKAGES = ["openssh-client", "rsync"]
 
     @classmethod
@@ -48,5 +48,3 @@ class SkypilotKubernetesIntegration(Integration):
 
         return [SkypilotKubernetesOrchestratorFlavor]
 
-
-SkypilotKubernetesIntegration.check_installation()

zenml/integrations/skypilot_lambda/__init__.py

@@ -31,7 +31,7 @@ class SkypilotLambdaIntegration(Integration):
     """Definition of Skypilot Lambda Integration for ZenML."""
 
     NAME = SKYPILOT_LAMBDA
-    REQUIREMENTS = ["skypilot[lambda]~=0.6.0"]
+    REQUIREMENTS = ["skypilot[lambda]~=0.8.0"]
 
     @classmethod
     def flavors(cls) -> List[Type[Flavor]]:
@@ -46,5 +46,3 @@ class SkypilotLambdaIntegration(Integration):
 
         return [SkypilotLambdaOrchestratorFlavor]
 
-
-SkypilotLambdaIntegration.check_installation()

zenml/integrations/slack/__init__.py

@@ -44,4 +44,3 @@ class SlackIntegration(Integration):
         return [SlackAlerterFlavor]
 
 
-SlackIntegration.check_installation()

zenml/integrations/spark/__init__.py

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

zenml/integrations/tekton/__init__.py

@@ -46,4 +46,3 @@ class TektonIntegration(Integration):
         return [TektonOrchestratorFlavor]
 
 
-TektonIntegration.check_installation()

zenml/integrations/tensorboard/__init__.py

@@ -42,5 +42,3 @@ class TensorBoardIntegration(Integration):
         """Activates the integration."""
         from zenml.integrations.tensorboard import services  # noqa
 
-
-TensorBoardIntegration.check_installation()

zenml/integrations/tensorboard/services/tensorboard_service.py

@@ -19,6 +19,7 @@ from typing import Any, Dict, Union
 from tensorboard import default, program  # type: ignore [import-untyped]
 
 from zenml.logger import get_logger
+from zenml.models.v2.misc.service import ServiceType
 from zenml.services import (
     HTTPEndpointHealthMonitor,
     HTTPEndpointHealthMonitorConfig,
@@ -27,7 +28,6 @@ from zenml.services import (
     LocalDaemonServiceEndpoint,
     LocalDaemonServiceEndpointConfig,
     ServiceEndpointProtocol,
-    ServiceType,
 )
 
 logger = get_logger(__name__)

zenml/integrations/tensorflow/__init__.py

@@ -64,5 +64,3 @@ class TensorflowIntegration(Integration):
             requirements.append("typing-extensions>=4.6.1")
         return requirements
 
-
-TensorflowIntegration.check_installation()

zenml/integrations/vllm/__init__.py

@@ -47,4 +47,3 @@ class VLLMIntegration(Integration):
         return [VLLMModelDeployerFlavor]
 
 
-VLLMIntegration.check_installation()

zenml/integrations/vllm/services/vllm_deployment.py

@@ -19,6 +19,7 @@ from typing import Any, List, Optional, Union
 
 from zenml.constants import DEFAULT_LOCAL_SERVICE_IP_ADDRESS
 from zenml.logger import get_logger
+from zenml.models.v2.misc.service import ServiceType
 from zenml.services import (
     HTTPEndpointHealthMonitor,
     HTTPEndpointHealthMonitorConfig,
@@ -27,7 +28,6 @@ from zenml.services import (
     LocalDaemonServiceEndpoint,
     LocalDaemonServiceEndpointConfig,
     ServiceEndpointProtocol,
-    ServiceType,
 )
 from zenml.services.service import BaseDeploymentService
 

zenml/integrations/wandb/__init__.py

@@ -47,4 +47,3 @@ class WandbIntegration(Integration):
         return [WandbExperimentTrackerFlavor]
 
 
-WandbIntegration.check_installation()

zenml/integrations/whylogs/__init__.py

@@ -65,4 +65,3 @@ class WhylogsIntegration(Integration):
             PandasIntegration.get_requirements(target_os=target_os)
 
 
-WhylogsIntegration.check_installation()

zenml/integrations/xgboost/__init__.py

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

zenml/login/credentials.py

@@ -20,11 +20,11 @@ from uuid import UUID
 
 from pydantic import BaseModel, ConfigDict
 
+from zenml.enums import ServiceState
 from zenml.login.pro.constants import ZENML_PRO_API_URL, ZENML_PRO_URL
 from zenml.login.pro.workspace.models import WorkspaceRead, WorkspaceStatus
 from zenml.models import ServerModel
 from zenml.models.v2.misc.server_models import ServerDeploymentType
-from zenml.services.service_status import ServiceState
 from zenml.utils.enum_utils import StrEnum
 from zenml.utils.string_utils import get_human_readable_time
 from zenml.utils.time_utils import to_local_tz, utc_now

zenml/materializers/__init__.py

@@ -33,6 +33,7 @@ from zenml.materializers.pydantic_materializer import PydanticMaterializer
 from zenml.materializers.service_materializer import ServiceMaterializer
 from zenml.materializers.uuid_materializer import UUIDMaterializer
 
+
 __all__ = [
     "BuiltInContainerMaterializer",
     "BuiltInMaterializer",

zenml/model_deployers/base_model_deployer.py

@@ -30,9 +30,9 @@ from uuid import UUID
 from zenml.client import Client
 from zenml.enums import StackComponentType
 from zenml.logger import get_logger
+from zenml.models.v2.misc.service import ServiceType
 from zenml.services import BaseService, ServiceConfig
 from zenml.services.service import BaseDeploymentService
-from zenml.services.service_type import ServiceType
 from zenml.stack import StackComponent
 from zenml.stack.flavor import Flavor
 from zenml.stack.stack_component import StackComponentConfig