zenml-nightly 0.58.2.dev20240617__py3-none-any.whl → 0.58.2.dev20240619__py3-none-any.whl

zenml/container_registries/base_container_registry.py

@@ -16,7 +16,7 @@
 import re
 from typing import TYPE_CHECKING, Optional, Tuple, Type, cast
 
-from pydantic import validator
+from pydantic import field_validator
 
 from zenml.constants import DOCKER_REGISTRY_RESOURCE_TYPE
 from zenml.enums import StackComponentType
@@ -42,7 +42,8 @@ class BaseContainerRegistryConfig(AuthenticationConfigMixin):
 
     uri: str
 
-    @validator("uri")
+    @field_validator("uri")
+    @classmethod
     def strip_trailing_slash(cls, uri: str) -> str:
         """Removes trailing slashes from the URI.
 

zenml/container_registries/default_container_registry.py

@@ -35,7 +35,7 @@ class DefaultContainerRegistryFlavor(BaseContainerRegistryFlavor):
 
     @property
     def docs_url(self) -> Optional[str]:
-        """A url to point at docs explaining this flavor.
+        """A URL to point at docs explaining this flavor.
 
         Returns:
             A flavor docs url.
@@ -44,7 +44,7 @@ class DefaultContainerRegistryFlavor(BaseContainerRegistryFlavor):
 
     @property
     def sdk_docs_url(self) -> Optional[str]:
-        """A url to point at docs explaining this flavor.
+        """A URL to point at docs explaining this flavor.
 
         Returns:
             A flavor docs url.
@@ -53,7 +53,7 @@ class DefaultContainerRegistryFlavor(BaseContainerRegistryFlavor):
 
     @property
     def logo_url(self) -> str:
-        """A url to represent the flavor in the dashboard.
+        """A URL to represent the flavor in the dashboard.
 
         Returns:
             The flavor logo.

zenml/event_hub/base_event_hub.py

@@ -118,7 +118,7 @@ class BaseEventHub(ABC):
             action_callback: The action to trigger.
         """
         request = TriggerExecutionRequest(
-            trigger=trigger.id, event_metadata=event.dict()
+            trigger=trigger.id, event_metadata=event.model_dump()
         )
 
         action_config = trigger.get_metadata().action

zenml/event_sources/base_event_source.py

@@ -13,7 +13,6 @@
 #  permissions and limitations under the License.
 """Base implementation for event sources."""
 
-import json
 from abc import ABC, abstractmethod
 from typing import (
     Any,
@@ -210,7 +209,7 @@ class BaseEventSourceHandler(BasePlugin, ABC):
             event_source=event_source, config=config
         )
         # Serialize the configuration back into the request
-        event_source.configuration = config.dict(exclude_none=True)
+        event_source.configuration = config.model_dump(exclude_none=True)
         # Create the event source in the database
         event_source_response = self.zen_store.create_event_source(
             event_source=event_source
@@ -238,7 +237,9 @@ class BaseEventSourceHandler(BasePlugin, ABC):
             raise
 
         # Serialize the configuration back into the response
-        event_source_response.set_configuration(config.dict(exclude_none=True))
+        event_source_response.set_configuration(
+            config.model_dump(exclude_none=True)
+        )
 
         # Return the response to the user
         return event_source_response
@@ -281,7 +282,7 @@ class BaseEventSourceHandler(BasePlugin, ABC):
             config_update=config_update,
         )
         # Serialize the configuration update back into the update request
-        event_source_update.configuration = config_update.dict(
+        event_source_update.configuration = config_update.model_dump(
             exclude_none=True
         )
 
@@ -320,7 +321,7 @@ class BaseEventSourceHandler(BasePlugin, ABC):
 
         # Serialize the configuration back into the response
         event_source_response.set_configuration(
-            response_config.dict(exclude_none=True)
+            response_config.model_dump(exclude_none=True)
         )
         # Return the response to the user
         return event_source_response
@@ -387,7 +388,9 @@ class BaseEventSourceHandler(BasePlugin, ABC):
                 event_source=event_source, config=config
             )
             # Serialize the configuration back into the response
-            event_source.set_configuration(config.dict(exclude_none=True))
+            event_source.set_configuration(
+                config.model_dump(exclude_none=True)
+            )
 
         # Return the response to the user
         return event_source
@@ -647,11 +650,7 @@ class BaseEventSourceFlavor(BasePluginFlavor, ABC):
         Returns:
             The config schema.
         """
-        config_schema: Dict[str, Any] = json.loads(
-            cls.EVENT_SOURCE_CONFIG_CLASS.schema_json()
-        )
-
-        return config_schema
+        return cls.EVENT_SOURCE_CONFIG_CLASS.model_json_schema()
 
     @classmethod
     def get_event_source_config_schema(cls) -> Dict[str, Any]:
@@ -660,11 +659,7 @@ class BaseEventSourceFlavor(BasePluginFlavor, ABC):
         Returns:
             The config schema.
         """
-        config_schema: Dict[str, Any] = json.loads(
-            cls.EVENT_FILTER_CONFIG_CLASS.schema_json()
-        )
-
-        return config_schema
+        return cls.EVENT_FILTER_CONFIG_CLASS.model_json_schema()
 
     @classmethod
     def get_flavor_response_model(

zenml/exceptions.py

@@ -245,6 +245,10 @@ class IllegalOperationError(ZenMLBaseException):
     """Raised when an illegal operation is attempted."""
 
 
+class MethodNotAllowedError(ZenMLBaseException):
+    """Raised when the server does not allow a request method."""
+
+
 class SettingsResolvingError(ZenMLBaseException):
     """Raised when resolving settings failed."""
 

zenml/integrations/airflow/__init__.py

@@ -13,9 +13,7 @@
 #  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, Type
 
@@ -30,9 +28,7 @@ class AirflowIntegration(Integration):
     """Definition of Airflow Integration for ZenML."""
 
     NAME = AIRFLOW
-    # remove pendulum version requirement once Airflow supports
-    # pendulum>-3.0.0
-    REQUIREMENTS = ["apache-airflow~=2.4.0", "pendulum<3.0.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.

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,7 +50,6 @@ if TYPE_CHECKING:
         DagConfiguration,
         TaskConfiguration,
     )
-    from zenml.metadata.metadata_types import MetadataType
     from zenml.models import PipelineDeploymentResponse
     from zenml.pipelines import Schedule
     from zenml.stack import Stack
@@ -107,18 +99,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 +130,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 +325,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)
@@ -447,224 +427,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<3.8.0' "
-                "'apache-airflow~=2.4.0'` 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/flavors/argilla_annotator_flavor.py

@@ -15,7 +15,7 @@
 
 from typing import TYPE_CHECKING, Optional, Type
 
-from pydantic import validator
+from pydantic import field_validator
 
 from zenml.annotators.base_annotator import (
     BaseAnnotatorConfig,
@@ -50,13 +50,14 @@ class ArgillaAnnotatorSettings(BaseSettings):
     """
 
     instance_url: str = DEFAULT_LOCAL_INSTANCE_URL
-    api_key: Optional[str] = SecretField()
+    api_key: Optional[str] = SecretField(default=None)
     workspace: Optional[str] = "admin"
     port: Optional[int]
     extra_headers: Optional[str] = None
     httpx_extra_kwargs: Optional[str] = None
 
-    @validator("instance_url")
+    @field_validator("instance_url")
+    @classmethod
     def ensure_instance_url_ends_without_slash(cls, instance_url: str) -> str:
         """Pydantic validator to ensure instance URL ends without a slash.
 
@@ -69,7 +70,7 @@ class ArgillaAnnotatorSettings(BaseSettings):
         return instance_url.rstrip("/")
 
 
-class ArgillaAnnotatorConfig(  # type: ignore[misc] # https://github.com/pydantic/pydantic/issues/4173
+class ArgillaAnnotatorConfig(
     BaseAnnotatorConfig,
     ArgillaAnnotatorSettings,
     AuthenticationConfigMixin,

zenml/integrations/aws/__init__.py

@@ -39,7 +39,7 @@ class AWSIntegration(Integration):
 
     NAME = AWS
     REQUIREMENTS = [
-        "sagemaker==2.117.0",
+        "sagemaker>=2.117.0",
         "kubernetes",
         "aws-profile-manager",
     ]

zenml/integrations/aws/flavors/aws_container_registry_flavor.py

@@ -15,7 +15,7 @@
 
 from typing import TYPE_CHECKING, Optional, Type
 
-from pydantic import validator
+from pydantic import field_validator
 
 from zenml.constants import DOCKER_REGISTRY_RESOURCE_TYPE
 from zenml.container_registries.base_container_registry import (
@@ -37,7 +37,8 @@ if TYPE_CHECKING:
 class AWSContainerRegistryConfig(BaseContainerRegistryConfig):
     """Configuration for AWS Container Registry."""
 
-    @validator("uri")
+    @field_validator("uri")
+    @classmethod
     def validate_aws_uri(cls, uri: str) -> str:
         """Validates that the URI is in the correct format.
 

zenml/integrations/aws/flavors/sagemaker_orchestrator_flavor.py

@@ -15,6 +15,8 @@
 
 from typing import TYPE_CHECKING, Any, Dict, Optional, Type, Union
 
+from pydantic import Field
+
 from zenml.config.base_settings import BaseSettings
 from zenml.integrations.aws import (
     AWS_RESOURCE_TYPE,
@@ -80,13 +82,17 @@ class SagemakerOrchestratorSettings(BaseSettings):
 
     processor_args: Dict[str, Any] = {}
     input_data_s3_mode: str = "File"
-    input_data_s3_uri: Optional[Union[str, Dict[str, str]]] = None
+    input_data_s3_uri: Optional[Union[str, Dict[str, str]]] = Field(
+        default=None, union_mode="left_to_right"
+    )
 
     output_data_s3_mode: str = "EndOfJob"
-    output_data_s3_uri: Optional[Union[str, Dict[str, str]]] = None
+    output_data_s3_uri: Optional[Union[str, Dict[str, str]]] = Field(
+        default=None, union_mode="left_to_right"
+    )
 
 
-class SagemakerOrchestratorConfig(  # type: ignore[misc] # https://github.com/pydantic/pydantic/issues/4173
+class SagemakerOrchestratorConfig(
     BaseOrchestratorConfig, SagemakerOrchestratorSettings
 ):
     """Config for the Sagemaker orchestrator.
@@ -124,8 +130,8 @@ class SagemakerOrchestratorConfig(  # type: ignore[misc] # https://github.com/py
 
     synchronous: bool = True
     execution_role: str
-    aws_access_key_id: Optional[str] = SecretField()
-    aws_secret_access_key: Optional[str] = SecretField()
+    aws_access_key_id: Optional[str] = SecretField(default=None)
+    aws_secret_access_key: Optional[str] = SecretField(default=None)
     aws_profile: Optional[str] = None
     aws_auth_role_arn: Optional[str] = None
     region: Optional[str] = None