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

zenml/image_builders/local_image_builder.py

@@ -17,8 +17,6 @@ import shutil
 import tempfile
 from typing import TYPE_CHECKING, Any, Dict, Optional, Type, cast
 
-from docker.client import DockerClient
-
 from zenml.image_builders import (
     BaseImageBuilder,
     BaseImageBuilderConfig,
@@ -70,15 +68,25 @@ class LocalImageBuilder(BaseImageBuilder):
             )
 
         if not docker_utils.check_docker():
+            # For 3., this is not supported by the python docker library
+            # https://github.com/docker/docker-py/issues/3146
             raise RuntimeError(
-                "Unable to connect to the Docker daemon. There are two "
+                "Unable to connect to the Docker daemon. There are three "
                 "common causes for this:\n"
                 "1) The Docker daemon isn't running.\n"
                 "2) The Docker client isn't configured correctly. The client "
                 "loads its configuration from the following file: "
                 "$HOME/.docker/config.json. If your configuration file is in a "
-                "different location, set it using the `DOCKER_CONFIG` "
-                "environment variable."
+                "different location, set the `DOCKER_CONFIG` environment "
+                "variable to the directory that contains your `config.json` "
+                "file.\n"
+                "3) If your Docker CLI is working fine but you ran into this "
+                "issue, you might be using a non-default Docker context which "
+                "is not supported by the Docker python library. To verify "
+                "this, run `docker context ls` and check which context has a "
+                "`*` next to it. If this is not the `default` context, copy "
+                "the `DOCKER ENDPOINT` value of that context and set the "
+                "`DOCKER_HOST` environment variable to that value."
             )
 
     def build(
@@ -106,7 +114,7 @@ class LocalImageBuilder(BaseImageBuilder):
             # authenticated to access additional registries
             docker_client = container_registry.docker_client
         else:
-            docker_client = DockerClient.from_env()
+            docker_client = docker_utils._try_get_docker_client_from_env()
 
         with tempfile.TemporaryFile(mode="w+b") as f:
             build_context.write_archive(f)

zenml/integrations/__init__.py

@@ -17,12 +17,14 @@ The ZenML integrations module contains sub-modules for each integration that we
 support. This includes orchestrators like Apache Airflow, visualization tools
 like the ``facets`` library, as well as deep learning libraries like PyTorch.
 """
-import sys
-
 from zenml.integrations.airflow import AirflowIntegration  # noqa
+from zenml.integrations.argilla import ArgillaIntegration  # noqa
 from zenml.integrations.aws import AWSIntegration  # noqa
 from zenml.integrations.azure import AzureIntegration  # noqa
 from zenml.integrations.bentoml import BentoMLIntegration  # noqa
+from zenml.integrations.bitbucket import BitbucketIntegration  # noqa
+from zenml.integrations.databricks import DatabricksIntegration  # noqa
+from zenml.integrations.comet import CometIntegration  # noqa
 from zenml.integrations.deepchecks import DeepchecksIntegration  # noqa
 from zenml.integrations.discord import DiscordIntegration  # noqa
 from zenml.integrations.evidently import EvidentlyIntegration  # noqa
@@ -34,21 +36,28 @@ from zenml.integrations.gitlab import GitLabIntegration  # noqa
 from zenml.integrations.great_expectations import (  # noqa
     GreatExpectationsIntegration,
 )
+from zenml.integrations.lightning import LightningIntegration  # noqa
 from zenml.integrations.huggingface import HuggingfaceIntegration  # noqa
+from zenml.integrations.hyperai import HyperAIIntegration  # noqa
 from zenml.integrations.kaniko import KanikoIntegration  # noqa
-from zenml.integrations.kserve import KServeIntegration  # noqa
 from zenml.integrations.kubeflow import KubeflowIntegration  # noqa
 from zenml.integrations.kubernetes import KubernetesIntegration  # noqa
 from zenml.integrations.label_studio import LabelStudioIntegration  # noqa
 from zenml.integrations.langchain import LangchainIntegration  # noqa
 from zenml.integrations.lightgbm import LightGBMIntegration  # noqa
+
 # from zenml.integrations.llama_index import LlamaIndexIntegration  # noqa
 from zenml.integrations.mlflow import MlflowIntegration  # noqa
+from zenml.integrations.modal import ModalIntegration  # noqa
 from zenml.integrations.neptune import NeptuneIntegration  # noqa
 from zenml.integrations.neural_prophet import NeuralProphetIntegration  # noqa
+from zenml.integrations.numpy import NumpyIntegration  # noqa
 from zenml.integrations.openai import OpenAIIntegration  # noqa
+from zenml.integrations.pandas import PandasIntegration  # noqa
+from zenml.integrations.pigeon import PigeonIntegration  # noqa
 from zenml.integrations.pillow import PillowIntegration  # noqa
-from zenml.integrations.polars import PolarsIntegration
+from zenml.integrations.polars import PolarsIntegration  # noqa
+from zenml.integrations.prodigy import ProdigyIntegration  # noqa
 from zenml.integrations.pycaret import PyCaretIntegration  # noqa
 from zenml.integrations.pytorch import PytorchIntegration  # noqa
 from zenml.integrations.pytorch_lightning import (  # noqa
@@ -61,6 +70,8 @@ from zenml.integrations.sklearn import SklearnIntegration  # noqa
 from zenml.integrations.skypilot_aws import SkypilotAWSIntegration  # noqa
 from zenml.integrations.skypilot_gcp import SkypilotGCPIntegration  # noqa
 from zenml.integrations.skypilot_azure import SkypilotAzureIntegration  # noqa
+from zenml.integrations.skypilot_lambda import SkypilotLambdaIntegration  # noqa
+from zenml.integrations.skypilot_kubernetes import SkypilotKubernetesIntegration  # noqa
 from zenml.integrations.slack import SlackIntegration  # noqa
 from zenml.integrations.spark import SparkIntegration  # noqa
 from zenml.integrations.tekton import TektonIntegration  # noqa
@@ -69,3 +80,4 @@ from zenml.integrations.tensorflow import TensorflowIntegration  # noqa
 from zenml.integrations.wandb import WandbIntegration  # noqa
 from zenml.integrations.whylogs import WhylogsIntegration  # noqa
 from zenml.integrations.xgboost import XgboostIntegration  # noqa
+from zenml.integrations.vllm import VLLMIntegration  # noqa

zenml/integrations/airflow/__init__.py

@@ -13,11 +13,9 @@
 #  permissions and limitations under the License.
 """Airflow integration for ZenML.
 
-The Airflow integration sub-module powers an alternative to the local
-orchestrator. You can enable it by registering the Airflow orchestrator with
-the CLI tool, then bootstrap using the ``zenml orchestrator up`` command.
+The Airflow integration powers an alternative orchestrator.
 """
-from typing import List, Optional, Type
+from typing import List, Type
 
 from zenml.integrations.constants import AIRFLOW
 from zenml.integrations.integration import Integration
@@ -30,7 +28,7 @@ class AirflowIntegration(Integration):
     """Definition of Airflow Integration for ZenML."""
 
     NAME = AIRFLOW
-    REQUIREMENTS = ["apache-airflow~=2.4.0"]
+    REQUIREMENTS = []
 
     @classmethod
     def flavors(cls) -> List[Type[Flavor]]:

zenml/integrations/airflow/flavors/airflow_orchestrator_flavor.py

@@ -13,9 +13,9 @@
 #  permissions and limitations under the License.
 """Airflow orchestrator flavor."""
 
-from typing import TYPE_CHECKING, Any, Dict, List, Optional, Type, Union
+from typing import TYPE_CHECKING, Any, Dict, List, Optional, Type
 
-from pydantic import validator
+from pydantic import field_validator
 
 from zenml.config.base_settings import BaseSettings
 from zenml.integrations.airflow import AIRFLOW_ORCHESTRATOR_FLAVOR
@@ -87,10 +87,9 @@ class AirflowOrchestratorSettings(BaseSettings):
 
     custom_dag_generator: Optional[str] = None
 
-    @validator("operator", always=True)
-    def _convert_operator(
-        cls, value: Optional[Union[str, OperatorType]]
-    ) -> Optional[str]:
+    @field_validator("operator", mode="before")
+    @classmethod
+    def _convert_operator(cls, value: Any) -> Any:
         """Converts operator types to source strings.
 
         Args:
@@ -108,7 +107,7 @@ class AirflowOrchestratorSettings(BaseSettings):
             return value
 
 
-class AirflowOrchestratorConfig(  # type: ignore[misc] # https://github.com/pydantic/pydantic/issues/4173
+class AirflowOrchestratorConfig(
     BaseOrchestratorConfig, AirflowOrchestratorSettings
 ):
     """Configuration for the Airflow orchestrator.
@@ -120,6 +119,15 @@ class AirflowOrchestratorConfig(  # type: ignore[misc] # https://github.com/pyda
 
     local: bool = True
 
+    @property
+    def is_schedulable(self) -> bool:
+        """Whether the orchestrator is schedulable or not.
+
+        Returns:
+            Whether the orchestrator is schedulable or not.
+        """
+        return True
+
 
 class AirflowOrchestratorFlavor(BaseOrchestratorFlavor):
     """Flavor for the Airflow orchestrator."""

zenml/integrations/airflow/orchestrators/airflow_orchestrator.py

@@ -16,8 +16,6 @@
 import datetime
 import importlib
 import os
-import platform
-import time
 import zipfile
 from typing import (
     TYPE_CHECKING,
@@ -29,12 +27,8 @@ from typing import (
     Type,
     cast,
 )
-from uuid import UUID
 
 from zenml.config.global_config import GlobalConfiguration
-from zenml.constants import (
-    METADATA_ORCHESTRATOR_URL,
-)
 from zenml.entrypoints import StepEntrypointConfiguration
 from zenml.enums import StackComponentType
 from zenml.integrations.airflow.flavors.airflow_orchestrator_flavor import (
@@ -44,11 +38,10 @@ from zenml.integrations.airflow.flavors.airflow_orchestrator_flavor import (
 )
 from zenml.io import fileio
 from zenml.logger import get_logger
-from zenml.metadata.metadata_types import Uri
 from zenml.orchestrators import ContainerizedOrchestrator
 from zenml.orchestrators.utils import get_orchestrator_run_name
 from zenml.stack import StackValidator
-from zenml.utils import daemon, io_utils
+from zenml.utils import io_utils
 
 if TYPE_CHECKING:
     from zenml.config import ResourceSettings
@@ -57,9 +50,7 @@ if TYPE_CHECKING:
         DagConfiguration,
         TaskConfiguration,
     )
-    from zenml.metadata.metadata_types import MetadataType
-    from zenml.models import PipelineDeploymentResponse
-    from zenml.pipelines import Schedule
+    from zenml.models import PipelineDeploymentResponse, ScheduleResponse
     from zenml.stack import Stack
 
 logger = get_logger(__name__)
@@ -107,18 +98,18 @@ class AirflowOrchestrator(ContainerizedOrchestrator):
     """Orchestrator responsible for running pipelines using Airflow."""
 
     def __init__(self, **values: Any):
-        """Sets environment variables to configure airflow.
+        """Initialize the orchestrator.
 
         Args:
             **values: Values to set in the orchestrator.
         """
         super().__init__(**values)
-        self.airflow_home = os.path.join(
+        self.dags_directory = os.path.join(
             io_utils.get_global_config_directory(),
             "airflow",
             str(self.id),
+            "dags",
         )
-        self._set_env()
 
     @property
     def config(self) -> AirflowOrchestratorConfig:
@@ -138,19 +129,6 @@ class AirflowOrchestrator(ContainerizedOrchestrator):
         """
         return AirflowOrchestratorSettings
 
-    @property
-    def dags_directory(self) -> str:
-        """Returns path to the airflow dags directory.
-
-        Returns:
-            Path to the airflow dags directory.
-        """
-        return os.path.join(self.airflow_home, "dags")
-
-    def _set_env(self) -> None:
-        """Sets environment variables to configure airflow."""
-        os.environ["AIRFLOW_HOME"] = self.airflow_home
-
     @property
     def validator(self) -> Optional["StackValidator"]:
         """Validates the stack.
@@ -346,21 +324,22 @@ class AirflowOrchestrator(ContainerizedOrchestrator):
         """
         io_utils.create_dir_recursive_if_not_exists(output_dir)
 
-        if self.config.local and output_dir != self.dags_directory:
+        if self.config.local and output_dir == self.dags_directory:
             logger.warning(
-                "You're using a local Airflow orchestrator but specified a "
-                "custom DAG output directory `%s`. This DAG will not be found "
-                "by the local Airflow server until you copy it in the DAGs "
-                "directory `%s`.",
+                "You're using a local Airflow orchestrator but have not "
+                "specified a custom DAG output directory. Unless you've "
+                "configured your Airflow server to look for DAGs in this "
+                "directory (%s), this DAG will not be found automatically "
+                "by your local Airflow server.",
                 output_dir,
-                self.dags_directory,
             )
 
         def _write_zip(path: str) -> None:
             with zipfile.ZipFile(path, mode="w") as z:
                 z.write(dag_generator_values.file, arcname="dag.py")
                 z.writestr(
-                    dag_generator_values.config_file_name, dag_config.json()
+                    dag_generator_values.config_file_name,
+                    dag_config.model_dump_json(),
                 )
 
             logger.info("Writing DAG definition to `%s`.", path)
@@ -411,7 +390,7 @@ class AirflowOrchestrator(ContainerizedOrchestrator):
 
     @staticmethod
     def _translate_schedule(
-        schedule: Optional["Schedule"] = None,
+        schedule: Optional["ScheduleResponse"] = None,
     ) -> Dict[str, Any]:
         """Convert ZenML schedule into Airflow schedule.
 
@@ -447,224 +426,8 @@ class AirflowOrchestrator(ContainerizedOrchestrator):
 
         return {
             "schedule": "@once",
-            # set the a start time in the past and disable catchup so airflow
+            # set a start time in the past and disable catchup so airflow
             # runs the dag immediately
             "start_date": datetime.datetime.utcnow() - datetime.timedelta(7),
             "catchup": False,
         }
-
-    #####################
-    #   Local Airflow   #
-    #####################
-
-    @property
-    def pid_file(self) -> str:
-        """Returns path to the daemon PID file.
-
-        Returns:
-            Path to the daemon PID file.
-        """
-        return os.path.join(self.airflow_home, "airflow_daemon.pid")
-
-    @property
-    def log_file(self) -> str:
-        """Returns path to the airflow log file.
-
-        Returns:
-            str: Path to the airflow log file.
-        """
-        return os.path.join(self.airflow_home, "airflow_orchestrator.log")
-
-    @property
-    def password_file(self) -> str:
-        """Returns path to the webserver password file.
-
-        Returns:
-            Path to the webserver password file.
-        """
-        return os.path.join(self.airflow_home, "standalone_admin_password.txt")
-
-    @property
-    def is_running(self) -> bool:
-        """Returns whether the orchestrator is "running".
-
-        In the non-local case, this is always True. Otherwise checks if the
-        local Airflow server is running.
-
-        Returns:
-            If the orchestrator is running.
-
-        Raises:
-            RuntimeError: If port 8080 is occupied.
-        """
-        if not self.config.local:
-            return True
-
-        from airflow.cli.commands.standalone_command import (  # type: ignore
-            StandaloneCommand,
-        )
-        from airflow.jobs.triggerer_job import TriggererJob
-
-        daemon_running = daemon.check_if_daemon_is_running(self.pid_file)
-
-        command = StandaloneCommand()
-        webserver_port_open = command.port_open(8080)
-
-        if not daemon_running:
-            if webserver_port_open:
-                raise RuntimeError(
-                    "The airflow daemon does not seem to be running but "
-                    "local port 8080 is occupied. Make sure the port is "
-                    "available and try again."
-                )
-
-            # exit early so we don't check non-existing airflow databases
-            return False
-
-        # we can't use StandaloneCommand().is_ready() here as the
-        # Airflow SequentialExecutor apparently does not send a heartbeat
-        # while running a task which would result in this returning `False`
-        # even if Airflow is running.
-        airflow_running = webserver_port_open and command.job_running(
-            TriggererJob
-        )
-        return airflow_running
-
-    @property
-    def is_provisioned(self) -> bool:
-        """Returns whether the airflow daemon is currently running.
-
-        Returns:
-            True if the airflow daemon is running, False otherwise.
-        """
-        return self.is_running
-
-    def provision(self) -> None:
-        """Ensures that Airflow is running."""
-        if not self.config.local:
-            return
-
-        if self.is_running:
-            logger.info("Airflow is already running.")
-            self._log_webserver_credentials()
-            return
-
-        self._check_local_server_requirements()
-
-        if not fileio.exists(self.dags_directory):
-            io_utils.create_dir_recursive_if_not_exists(self.dags_directory)
-
-        from airflow.cli.commands.standalone_command import StandaloneCommand
-
-        self._set_server_env()
-        try:
-            command = StandaloneCommand()
-            daemon.run_as_daemon(
-                command.run,
-                pid_file=self.pid_file,
-                log_file=self.log_file,
-            )
-            while not self.is_running:
-                # Wait until the daemon started all the relevant airflow
-                # processes
-                time.sleep(0.1)
-            self._log_webserver_credentials()
-        except Exception as e:
-            logger.error(e)
-            logger.error(
-                "An error occurred while starting the Airflow daemon. If you "
-                "want to start it manually, use the commands described in the "
-                "official Airflow quickstart guide for running Airflow locally."
-            )
-            self.deprovision()
-        finally:
-            logger.warning(
-                "Airflow provisioning using `zenml stack up` is "
-                "deprecated. Please follow the new Airflow quickstart guide "
-                "to run Airflow locally."
-            )
-
-    def deprovision(self) -> None:
-        """Stops the airflow daemon if necessary and tears down resources."""
-        if not self.config.local:
-            return
-
-        if self.is_running:
-            daemon.stop_daemon(self.pid_file)
-
-        fileio.rmtree(self.airflow_home)
-        logger.info("Airflow spun down.")
-
-    def _set_server_env(self) -> None:
-        """Sets environment variables for the local Airflow server process."""
-        os.environ["AIRFLOW__CORE__DAGS_FOLDER"] = self.dags_directory
-        os.environ["AIRFLOW__CORE__LOAD_EXAMPLES"] = "false"
-        # check the DAG folder every 10 seconds for new files
-        os.environ["AIRFLOW__SCHEDULER__DAG_DIR_LIST_INTERVAL"] = "10"
-
-        if platform.system() == "Darwin":
-            # Prevent crashes during forking on MacOS
-            # https://github.com/apache/airflow/issues/28487
-            os.environ["no_proxy"] = "*"
-
-    @staticmethod
-    def _check_local_server_requirements() -> None:
-        """Checks that all packages for a local Airflow server are installed.
-
-        When running a local Airflow server, we require the
-        `apache-airflow-providers-docker` to run steps locally in Docker
-        containers in addition to the basic integration requirements.
-
-        Raises:
-            RuntimeError: If the `apache-airflow-providers-docker` is not
-                installed in the active Python environment.
-        """
-        try:
-            from airflow.providers.docker.operators.docker import (  # noqa
-                DockerOperator,
-            )
-        except ImportError:
-            raise RuntimeError(
-                "Unable to import Airflow `DockerOperator` in the active "
-                "Python environment. Spinning up a local Airflow server to "
-                "run ZenML pipelines requires the `DockerOperator` to be "
-                "available. Please run "
-                "`pip install apache-airflow-providers-docker` to install it "
-                "and try again."
-            )
-
-    def _log_webserver_credentials(self) -> None:
-        """Logs URL and credentials to log in to the airflow webserver.
-
-        Raises:
-            FileNotFoundError: If the password file does not exist.
-        """
-        if fileio.exists(self.password_file):
-            with open(self.password_file) as file:
-                password = file.read().strip()
-        else:
-            raise FileNotFoundError(
-                f"Can't find password file '{self.password_file}'"
-            )
-        logger.info(
-            "To inspect your DAGs, login to `http://localhost:8080` "
-            "with username: `admin` password: `%s`",
-            password,
-        )
-
-    def get_pipeline_run_metadata(
-        self, run_id: UUID
-    ) -> Dict[str, "MetadataType"]:
-        """Get general component-specific metadata for a pipeline run.
-
-        Args:
-            run_id: The ID of the pipeline run.
-
-        Returns:
-            A dictionary of metadata.
-        """
-        if self.config.local:
-            return {
-                METADATA_ORCHESTRATOR_URL: Uri("http://localhost:8080"),
-            }
-        return {}

zenml/integrations/airflow/orchestrators/dag_generator.py

@@ -19,7 +19,7 @@ import os
 import zipfile
 from typing import Any, Dict, List, Optional, Type, Union
 
-from pydantic import BaseModel
+from pydantic import BaseModel, Field
 
 ENV_ZENML_AIRFLOW_RUN_ID = "ZENML_AIRFLOW_RUN_ID"
 ENV_ZENML_LOCAL_STORES_PATH = "ZENML_LOCAL_STORES_PATH"
@@ -51,7 +51,9 @@ class DagConfiguration(BaseModel):
 
     local_stores_path: Optional[str] = None
 
-    schedule: Union[datetime.timedelta, str]
+    schedule: Union[datetime.timedelta, str] = Field(
+        union_mode="left_to_right"
+    )
     start_date: datetime.datetime
     end_date: Optional[datetime.datetime] = None
     catchup: bool = False
@@ -199,7 +201,7 @@ else:
     import airflow
 
     config_str = archive.read(CONFIG_FILENAME)
-    dag_config = DagConfiguration.parse_raw(config_str)
+    dag_config = DagConfiguration.model_validate_json(config_str)
 
     step_name_to_airflow_operator = {}
 

zenml/integrations/argilla/__init__.py

@@ -0,0 +1,46 @@
+#  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.
+"""Initialization of the Argilla integration."""
+from typing import List, Type
+
+from zenml.integrations.constants import ARGILLA
+from zenml.integrations.integration import Integration
+from zenml.stack import Flavor
+
+ARGILLA_ANNOTATOR_FLAVOR = "argilla"
+
+
+class ArgillaIntegration(Integration):
+    """Definition of Argilla integration for ZenML."""
+
+    NAME = ARGILLA
+    REQUIREMENTS = [
+        "argilla>=2.0.0",
+    ]
+
+    @classmethod
+    def flavors(cls) -> List[Type[Flavor]]:
+        """Declare the stack component flavors for the Argilla integration.
+
+        Returns:
+            List of stack component flavors for this integration.
+        """
+        from zenml.integrations.argilla.flavors import (
+            ArgillaAnnotatorFlavor,
+        )
+
+        return [ArgillaAnnotatorFlavor]
+
+
+ArgillaIntegration.check_installation()

zenml/integrations/argilla/annotators/__init__.py

@@ -0,0 +1,20 @@
+#  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.
+"""Initialization of the Argilla annotators submodule."""
+
+from zenml.integrations.argilla.annotators.argilla_annotator import (
+    ArgillaAnnotator,
+)
+
+__all__ = ["ArgillaAnnotator"]