zenml-nightly 0.58.2.dev20240618__py3-none-any.whl → 0.58.2.dev20240620__py3-none-any.whl

zenml/cli/web_login.py

@@ -101,7 +101,7 @@ def web_login(url: str, verify_ssl: Union[str, bool]) -> str:
                 "Content-Type": "application/x-www-form-urlencoded",
                 "User-Agent": user_agent_header.encode(),
             },
-            data=auth_request.dict(),
+            data=auth_request.model_dump(),
             verify=verify_ssl,
             timeout=DEFAULT_HTTP_TIMEOUT,
         )
@@ -152,7 +152,7 @@ def web_login(url: str, verify_ssl: Union[str, bool]) -> str:
         response = requests.post(
             login_url,
             headers={"Content-Type": "application/x-www-form-urlencoded"},
-            data=token_request.dict(),
+            data=token_request.model_dump(),
             verify=verify_ssl,
             timeout=DEFAULT_HTTP_TIMEOUT,
         )

zenml/client.py

@@ -17,6 +17,7 @@ import functools
 import json
 import os
 from abc import ABCMeta
+from collections import Counter
 from datetime import datetime
 from functools import partial
 from pathlib import Path
@@ -36,13 +37,14 @@ from typing import (
 )
 from uuid import UUID, uuid4
 
-from pydantic import SecretStr
+from pydantic import ConfigDict, SecretStr
 
 from zenml.client_lazy_loader import (
     client_lazy_loader,
     evaluate_all_lazy_load_args_in_client_methods,
 )
 from zenml.config.global_config import GlobalConfiguration
+from zenml.config.pipeline_run_configuration import PipelineRunConfiguration
 from zenml.config.source import Source
 from zenml.constants import (
     ENV_ZENML_ACTIVE_STACK_ID,
@@ -247,18 +249,14 @@ class ClientConfiguration(FileSyncModel):
         self.active_stack_id = stack.id
         self._active_stack = stack
 
-    class Config:
-        """Pydantic configuration class."""
-
+    model_config = ConfigDict(
         # Validate attributes when assigning them. We need to set this in order
         # to have a mix of mutable and immutable attributes
-        validate_assignment = True
+        validate_assignment=True,
         # Allow extra attributes from configs of previous ZenML versions to
         # permit downgrading
-        extra = "allow"
-        # all attributes with leading underscore are private and therefore
-        # are mutable and not included in serialization
-        underscore_attrs_are_private = True
+        extra="allow",
+    )
 
 
 class ClientMetaClass(ABCMeta):
@@ -1304,7 +1302,7 @@ class Client(metaclass=ClientMetaClass):
         )
 
         # Create the update model
-        update_model = StackUpdate(  # type: ignore[call-arg]
+        update_model = StackUpdate(
             workspace=self.active_workspace.id,
             user=self.active_user.id,
             stack_spec_path=stack_spec_file,
@@ -1596,7 +1594,7 @@ class Client(metaclass=ClientMetaClass):
         service_request = ServiceRequest(
             name=config.service_name,
             service_type=service_type,
-            config=config.dict(),
+            config=config.model_dump(),
             workspace=self.active_workspace.id,
             user=self.active_user.id,
             model_version_id=model_version_id,
@@ -2030,7 +2028,7 @@ class Client(metaclass=ClientMetaClass):
             allow_name_prefix_match=False,
         )
 
-        update_model = ComponentUpdate(  # type: ignore[call-arg]
+        update_model = ComponentUpdate(
             workspace=self.active_workspace.id,
             user=self.active_user.id,
             component_spec_path=component_spec_path,
@@ -2448,6 +2446,171 @@ class Client(metaclass=ClientMetaClass):
             )
             self.zen_store.delete_pipeline(pipeline_id=pipeline.id)
 
+    @_fail_for_sql_zen_store
+    def trigger_pipeline(
+        self,
+        pipeline_name_or_id: Union[str, UUID, None] = None,
+        pipeline_version: Optional[str] = None,
+        run_configuration: Optional[PipelineRunConfiguration] = None,
+        config_path: Optional[str] = None,
+        deployment_id: Optional[UUID] = None,
+        build_id: Optional[UUID] = None,
+        stack_name_or_id: Union[str, UUID, None] = None,
+        synchronous: bool = False,
+    ) -> PipelineRunResponse:
+        """Trigger a pipeline from the server.
+
+        Usage examples:
+        * Run the latest runnable build for the latest version of a pipeline:
+        ```python
+        Client().trigger_pipeline(pipeline_name_or_id=<NAME>)
+        ```
+        * Run the latest runnable build for a specific version of a pipeline:
+        ```python
+        Client().trigger_pipeline(
+            pipeline_name_or_id=<NAME>,
+            pipeline_version=<VERSION>
+        )
+        ```
+        * Run a specific pipeline version on a specific stack:
+        ```python
+        Client().trigger_pipeline(
+            pipeline_name_or_id=<ID>,
+            stack_name_or_id=<ID>
+        )
+        ```
+        * Run a specific deployment:
+        ```python
+        Client().trigger_pipeline(deployment_id=<ID>)
+        ```
+        * Run a specific build:
+        ```python
+        Client().trigger_pipeline(build_id=<ID>)
+        ```
+
+        Args:
+            pipeline_name_or_id: Name or ID of the pipeline. If not given,
+                either the build or deployment that should be run needs to be
+                specified.
+            pipeline_version: Version of the pipeline. This is only used if a
+                pipeline name is given.
+            run_configuration: Configuration for the run. Either this or a
+                path to a config file can be specified.
+            config_path: Path to a YAML configuration file. This file will be
+                parsed as a `PipelineRunConfiguration` object. Either this or
+                the configuration in code can be specified.
+            deployment_id: ID of the deployment to run. Either this or a build
+                to run can be specified.
+            build_id: ID of the build to run. Either this or a deployment to
+                run can be specified.
+            stack_name_or_id: Name or ID of the stack on which to run the
+                pipeline. If not specified, this method will try to find a
+                runnable build on any stack.
+            synchronous: If `True`, this method will wait until the triggered
+                run is finished.
+
+        Raises:
+            RuntimeError: If triggering the pipeline failed.
+
+        Returns:
+            Model of the pipeline run.
+        """
+        from zenml.new.pipelines.run_utils import (
+            validate_run_config_is_runnable_from_server,
+            validate_stack_is_runnable_from_server,
+            wait_for_pipeline_run_to_finish,
+        )
+
+        if Counter([build_id, deployment_id, pipeline_name_or_id])[None] != 2:
+            raise RuntimeError(
+                "You need to specify exactly one of pipeline, build or "
+                "deployment to trigger."
+            )
+
+        if run_configuration and config_path:
+            raise RuntimeError(
+                "Only config path or runtime configuration can be specified."
+            )
+
+        if config_path:
+            run_configuration = PipelineRunConfiguration.from_yaml(config_path)
+
+        if run_configuration:
+            validate_run_config_is_runnable_from_server(run_configuration)
+
+        if deployment_id:
+            if stack_name_or_id:
+                logger.warning(
+                    "Deployment ID and stack specified, ignoring the stack and "
+                    "using stack from deployment instead."
+                )
+
+            run = self.zen_store.run_deployment(
+                deployment_id=deployment_id,
+                run_configuration=run_configuration,
+            )
+        elif build_id:
+            if stack_name_or_id:
+                logger.warning(
+                    "Build ID and stack specified, ignoring the stack and "
+                    "using stack from build instead."
+                )
+
+            run = self.zen_store.run_build(
+                build_id=build_id, run_configuration=run_configuration
+            )
+        else:
+            assert pipeline_name_or_id
+            pipeline = self.get_pipeline(
+                name_id_or_prefix=pipeline_name_or_id, version=pipeline_version
+            )
+
+            stack = None
+            if stack_name_or_id:
+                stack = self.get_stack(
+                    stack_name_or_id, allow_name_prefix_match=False
+                )
+                validate_stack_is_runnable_from_server(
+                    zen_store=self.zen_store, stack=stack
+                )
+
+            builds = depaginate(
+                partial(
+                    self.list_builds,
+                    pipeline_id=pipeline.id,
+                    stack_id=stack.id if stack else None,
+                )
+            )
+
+            for build in builds:
+                if not build.template_deployment_id:
+                    continue
+
+                if not build.stack:
+                    continue
+
+                try:
+                    validate_stack_is_runnable_from_server(
+                        zen_store=self.zen_store, stack=build.stack
+                    )
+                except ValueError:
+                    continue
+
+                run = self.zen_store.run_build(
+                    build_id=build.id, run_configuration=run_configuration
+                )
+                break
+            else:
+                raise RuntimeError(
+                    "Unable to find a runnable build for the given stack and "
+                    "pipeline."
+                )
+
+        if synchronous:
+            run = wait_for_pipeline_run_to_finish(run_id=run.id)
+
+        return run
+
     # -------------------------------- Builds ----------------------------------
 
     def get_build(
@@ -2989,13 +3152,13 @@ class Client(metaclass=ClientMetaClass):
 
     def get_deployment(
         self,
-        id_or_prefix: str,
+        id_or_prefix: Union[str, UUID],
         hydrate: bool = True,
     ) -> PipelineDeploymentResponse:
         """Get a deployment by id or prefix.
 
         Args:
-            id_or_prefix: The id or id prefix of the build.
+            id_or_prefix: The id or id prefix of the deployment.
             hydrate: Flag deciding whether to hydrate the output model(s)
                 by including metadata fields in the response.
 
@@ -3011,10 +3174,12 @@ class Client(metaclass=ClientMetaClass):
 
         # First interpret as full UUID
         if is_valid_uuid(id_or_prefix):
-            return self.zen_store.get_deployment(
-                UUID(id_or_prefix),
-                hydrate=hydrate,
+            id_ = (
+                UUID(id_or_prefix)
+                if isinstance(id_or_prefix, str)
+                else id_or_prefix
             )
+            return self.zen_store.get_deployment(id_, hydrate=hydrate)
 
         entity = self.list_deployments(
             id=f"startswith:{id_or_prefix}",
@@ -4270,7 +4435,7 @@ class Client(metaclass=ClientMetaClass):
             hydrate=True,
         )
 
-        secret_update = SecretUpdate(name=new_name or secret.name)  # type: ignore[call-arg]
+        secret_update = SecretUpdate(name=new_name or secret.name)
 
         if new_scope:
             secret_update.scope = new_scope
@@ -4589,7 +4754,7 @@ class Client(metaclass=ClientMetaClass):
         repo = self.get_code_repository(
             name_id_or_prefix=name_id_or_prefix, allow_name_prefix_match=False
         )
-        update = CodeRepositoryUpdate(  # type: ignore[call-arg]
+        update = CodeRepositoryUpdate(
             name=name, description=description, logo_url=logo_url
         )
         return self.zen_store.update_code_repository(
@@ -5099,8 +5264,6 @@ class Client(metaclass=ClientMetaClass):
             expires_at=expires_at,
             expires_skew_tolerance=expires_skew_tolerance,
             expiration_seconds=expiration_seconds,
-            user=self.active_user.id,
-            workspace=self.active_workspace.id,
         )
 
         # Validate and configure the resources
@@ -5143,21 +5306,32 @@ class Client(metaclass=ClientMetaClass):
 
         if verify:
             # Prefer to verify the connector config server-side if the
-            # implementation if available there, because it ensures
+            # implementation, if available there, because it ensures
             # that the connector can be shared with other users or used
             # from other machines and because some auth methods rely on the
             # server-side authentication environment
+
+            # Convert the update model to a request model for validation
+            connector_request_dict = connector_update.model_dump()
+            connector_request_dict.update(
+                user=self.active_user.id,
+                workspace=self.active_workspace.id,
+            )
+            connector_request = ServiceConnectorRequest.model_validate(
+                connector_request_dict
+            )
+
             if connector.remote:
                 connector_resources = (
                     self.zen_store.verify_service_connector_config(
-                        connector_update,
+                        service_connector=connector_request,
                         list_resources=list_resources,
                     )
                 )
             else:
                 connector_instance = (
                     service_connector_registry.instantiate_connector(
-                        model=connector_update
+                        model=connector_request,
                     )
                 )
                 connector_resources = connector_instance.verify(

zenml/client_lazy_loader.py

@@ -16,7 +16,7 @@
 import contextlib
 from typing import TYPE_CHECKING, Any, Callable, Dict, List, Optional, Type
 
-from pydantic import BaseModel, PrivateAttr
+from pydantic import BaseModel
 
 if TYPE_CHECKING:
     from zenml.client import Client
@@ -35,8 +35,7 @@ class ClientLazyLoader(BaseModel):
 
     method_name: str
     call_chain: List[_CallStep] = []
-
-    _exclude_next_call: bool = PrivateAttr(False)
+    exclude_next_call: bool = False
 
     def __getattr__(self, name: str) -> "ClientLazyLoader":
         """Get attribute not defined in ClientLazyLoader.
@@ -54,7 +53,7 @@ class ClientLazyLoader(BaseModel):
         if name != "__deepcopy__":
             self_.call_chain.append(_CallStep(attribute_name=name))
         else:
-            self_._exclude_next_call = True
+            self_.exclude_next_call = True
         return self_
 
     def __call__(self, *args: Any, **kwargs: Any) -> "ClientLazyLoader":
@@ -68,11 +67,11 @@ class ClientLazyLoader(BaseModel):
             self
         """
         # workaround to protect from infinitely looping over in deepcopy called in invocations
-        if not self._exclude_next_call:
+        if not self.exclude_next_call:
             self.call_chain.append(
                 _CallStep(is_call=True, call_args=args, call_kwargs=kwargs)
             )
-        self._exclude_next_call = False
+        self.exclude_next_call = False
         return self
 
     def __getitem__(self, item: Any) -> "ClientLazyLoader":
@@ -99,18 +98,25 @@ class ClientLazyLoader(BaseModel):
             self: "ClientLazyLoader", self_: Any, call_chain_: List[_CallStep]
         ) -> Any:
             next_step = call_chain_.pop(0)
-            if next_step.is_call:
-                self_ = self_(*next_step.call_args, **next_step.call_kwargs)
-            elif next_step.selector:
-                self_ = self_[next_step.selector]
-            elif next_step.attribute_name:
-                self_ = getattr(self_, next_step.attribute_name)
-            else:
+            try:
+                if next_step.is_call:
+                    self_ = self_(
+                        *next_step.call_args, **next_step.call_kwargs
+                    )
+                elif next_step.selector:
+                    self_ = self_[next_step.selector]
+                elif next_step.attribute_name:
+                    self_ = getattr(self_, next_step.attribute_name)
+                else:
+                    raise ValueError(
+                        "Invalid call chain. Reach out to the ZenML team."
+                    )
+            except Exception as e:
                 raise RuntimeError(
                     f"Failed to evaluate lazy load chain `{self.method_name}` "
                     f"+ `{self.call_chain}`. Reach out to the ZenML team via "
                     "Slack or GitHub to check further."
-                )
+                ) from e
             return self_
 
         self_ = getattr(Client(), self.method_name)

zenml/config/base_settings.py

@@ -16,7 +16,7 @@
 from enum import IntFlag, auto
 from typing import Any, ClassVar, Dict, Union
 
-from pydantic import Extra
+from pydantic import ConfigDict
 
 from zenml.config.secret_reference_mixin import SecretReferenceMixin
 
@@ -46,11 +46,10 @@ class BaseSettings(SecretReferenceMixin):
         ConfigurationLevel.PIPELINE | ConfigurationLevel.STEP
     )
 
-    class Config:
-        """Pydantic configuration class."""
-
+    model_config = ConfigDict(
         # public attributes are immutable
-        allow_mutation = False
+        frozen=True,
         # allow extra attributes so this class can be used to parse dicts
         # of arbitrary subclasses
-        extra = Extra.allow
+        extra="allow",
+    )

zenml/config/build_configuration.py

@@ -60,7 +60,7 @@ class BuildConfiguration(BaseModel):
             The checksum.
         """
         hash_ = hashlib.md5()  # nosec
-        hash_.update(self.settings.json().encode())
+        hash_.update(self.settings.model_dump_json().encode())
         if self.entrypoint:
             hash_.update(self.entrypoint.encode())
 

zenml/config/compiler.py

@@ -286,9 +286,9 @@ class Compiler:
         """
         assert stack_component.settings_class
         # Exclude additional config attributes that aren't part of the settings
-        field_names = set(stack_component.settings_class.__fields__)
-        default_settings = stack_component.settings_class.parse_obj(
-            stack_component.config.dict(
+        field_names = set(stack_component.settings_class.model_fields.keys())
+        default_settings = stack_component.settings_class.model_validate(
+            stack_component.config.model_dump(
                 include=field_names, exclude_unset=True, exclude_defaults=True
             )
         )

zenml/config/docker_settings.py

@@ -16,11 +16,13 @@
 from enum import Enum
 from typing import Any, Dict, List, Optional, Union
 
-from pydantic import Extra, root_validator
+from pydantic import Field, model_validator
+from pydantic_settings import SettingsConfigDict
 
 from zenml.config.base_settings import BaseSettings
 from zenml.logger import get_logger
 from zenml.utils import deprecation_utils
+from zenml.utils.pydantic_utils import before_validator_handler
 
 logger = get_logger(__name__)
 
@@ -70,7 +72,7 @@ class DockerSettings(BaseSettings):
     --------------
     * No `dockerfile` specified: If any of the options regarding
     requirements, environment variables or copying files require us to build an
-    image, ZenML will build this image. Otherwise the `parent_image` will be
+    image, ZenML will build this image. Otherwise, the `parent_image` will be
     used to run the pipeline.
     * `dockerfile` specified: ZenML will first build an image based on the
     specified Dockerfile. If any of the options regarding
@@ -190,8 +192,10 @@ class DockerSettings(BaseSettings):
     python_package_installer_args: Dict[str, Any] = {}
     replicate_local_python_environment: Optional[
         Union[List[str], PythonEnvironmentExportMethod]
-    ] = None
-    requirements: Union[None, str, List[str]] = None
+    ] = Field(default=None, union_mode="left_to_right")
+    requirements: Union[None, str, List[str]] = Field(
+        default=None, union_mode="left_to_right"
+    )
     required_integrations: List[str] = []
     required_hub_plugins: List[str] = []
     install_stack_requirements: bool = True
@@ -208,41 +212,40 @@ class DockerSettings(BaseSettings):
         "copy_files", "copy_global_config"
     )
 
-    @root_validator(pre=True)
-    def _migrate_copy_files(cls, values: Dict[str, Any]) -> Dict[str, Any]:
+    @model_validator(mode="before")
+    @classmethod
+    @before_validator_handler
+    def _migrate_copy_files(cls, data: Dict[str, Any]) -> Dict[str, Any]:
         """Migrates the value from the old copy_files attribute.
 
         Args:
-            values: The settings values.
+            data: The settings values.
 
         Returns:
             The migrated settings values.
         """
-        copy_files = values.get("copy_files", None)
+        copy_files = data.get("copy_files", None)
 
         if copy_files is None:
-            return values
+            return data
 
-        if values.get("source_files", None):
+        if data.get("source_files", None):
             # Ignore the copy files value in favor of the new source files
             logger.warning(
                 "Both `copy_files` and `source_files` specified for the "
                 "DockerSettings, ignoring the `copy_files` value."
             )
         elif copy_files is True:
-            values["source_files"] = SourceFileMode.INCLUDE
+            data["source_files"] = SourceFileMode.INCLUDE
         elif copy_files is False:
-            values["source_files"] = SourceFileMode.IGNORE
+            data["source_files"] = SourceFileMode.IGNORE
 
-        return values
+        return data
 
-    @root_validator
-    def _validate_skip_build(cls, values: Dict[str, Any]) -> Dict[str, Any]:
+    @model_validator(mode="after")
+    def _validate_skip_build(self) -> "DockerSettings":
         """Ensures that a parent image is passed when trying to skip the build.
 
-        Args:
-            values: The settings values.
-
         Returns:
             The validated settings values.
 
@@ -250,10 +253,7 @@ class DockerSettings(BaseSettings):
             ValueError: If the build should be skipped but no parent image
                 was specified.
         """
-        skip_build = values.get("skip_build", False)
-        parent_image = values.get("parent_image")
-
-        if skip_build and not parent_image:
+        if self.skip_build and not self.parent_image:
             raise ValueError(
                 "Docker settings that specify `skip_build=True` must always "
                 "contain a `parent_image`. This parent image will be used "
@@ -261,12 +261,11 @@ class DockerSettings(BaseSettings):
                 "Docker builds on top of it."
             )
 
-        return values
-
-    class Config:
-        """Pydantic configuration class."""
+        return self
 
+    model_config = SettingsConfigDict(
         # public attributes are immutable
-        allow_mutation = False
+        frozen=True,
         # prevent extra attributes during model initialization
-        extra = Extra.forbid
+        extra="forbid",
+    )