zenml-nightly 0.58.2.dev20240615__py3-none-any.whl → 0.58.2.dev20240623__py3-none-any.whl

zenml/cli/utils.py

@@ -290,29 +290,31 @@ def print_pydantic_models(
             if isinstance(model, BaseIdentifiedResponse):
                 include_columns = ["id"]
 
-                if "name" in model.__fields__:
+                if "name" in model.model_fields:
                     include_columns.append("name")
 
                 include_columns.extend(
                     [
                         k
-                        for k in model.__fields__[
-                            "body"
-                        ].type_.__fields__.keys()
-                        if k not in exclude_columns
-                    ]
-                    + [
-                        k
-                        for k in model.__fields__[
-                            "metadata"
-                        ].type_.__fields__.keys()
+                        for k in model.get_body().model_fields.keys()
                         if k not in exclude_columns
                     ]
                 )
 
+                if model.metadata is not None:
+                    include_columns.extend(
+                        [
+                            k
+                            for k in model.get_metadata().model_fields.keys()
+                            if k not in exclude_columns
+                        ]
+                    )
+
             else:
                 include_columns = [
-                    k for k in model.dict().keys() if k not in exclude_columns
+                    k
+                    for k in model.model_dump().keys()
+                    if k not in exclude_columns
                 ]
         else:
             include_columns = columns
@@ -325,7 +327,7 @@ def print_pydantic_models(
             #  we want to attempt to represent them by name, if they contain
             #  such a field, else the id is used
             if isinstance(value, BaseIdentifiedResponse):
-                if "name" in value.__fields__:
+                if "name" in value.model_fields:
                     items[k] = str(getattr(value, "name"))
                 else:
                     items[k] = str(value.id)
@@ -335,7 +337,7 @@ def print_pydantic_models(
             elif isinstance(value, list):
                 for v in value:
                     if isinstance(v, BaseIdentifiedResponse):
-                        if "name" in v.__fields__:
+                        if "name" in v.model_fields:
                             items.setdefault(k, []).append(
                                 str(getattr(v, "name"))
                             )
@@ -425,27 +427,31 @@ def print_pydantic_model(
         if isinstance(model, BaseIdentifiedResponse):
             include_columns = ["id"]
 
-            if "name" in model.__fields__:
+            if "name" in model.model_fields:
                 include_columns.append("name")
 
             include_columns.extend(
                 [
                     k
-                    for k in model.__fields__["body"].type_.__fields__.keys()
-                    if k not in exclude_columns
-                ]
-                + [
-                    k
-                    for k in model.__fields__[
-                        "metadata"
-                    ].type_.__fields__.keys()
+                    for k in model.get_body().model_fields.keys()
                     if k not in exclude_columns
                 ]
             )
 
+            if model.metadata is not None:
+                include_columns.extend(
+                    [
+                        k
+                        for k in model.get_metadata().model_fields.keys()
+                        if k not in exclude_columns
+                    ]
+                )
+
         else:
             include_columns = [
-                k for k in model.dict().keys() if k not in exclude_columns
+                k
+                for k in model.model_dump().keys()
+                if k not in exclude_columns
             ]
     else:
         include_columns = list(columns)
@@ -455,7 +461,7 @@ def print_pydantic_model(
     for k in include_columns:
         value = getattr(model, k)
         if isinstance(value, BaseIdentifiedResponse):
-            if "name" in value.__fields__:
+            if "name" in value.model_fields:
                 items[k] = str(getattr(value, "name"))
             else:
                 items[k] = str(value.id)
@@ -465,7 +471,7 @@ def print_pydantic_model(
         elif isinstance(value, list):
             for v in value:
                 if isinstance(v, BaseIdentifiedResponse):
-                    if "name" in v.__fields__:
+                    if "name" in v.model_fields:
                         items.setdefault(k, []).append(str(getattr(v, "name")))
                     else:
                         items.setdefault(k, []).append(str(v.id))
@@ -1078,16 +1084,6 @@ def install_packages(
         else:
             raise e
 
-    if "label-studio" in packages:
-        warning(
-            "There is a known issue with Label Studio installations "
-            "via zenml. You might find that the Label Studio "
-            "installation breaks the ZenML CLI. In this case, please "
-            "run `pip install 'pydantic<1.11,>=1.9.0'` to fix the "
-            "issue or message us on Slack if you need help with this. "
-            "We are working on a more definitive fix."
-        )
-
 
 def uninstall_package(package: str, use_uv: bool = False) -> None:
     """Uninstalls pypi package from the current environment with pip or uv.
@@ -1280,7 +1276,7 @@ def pretty_print_model_version_table(
             "NAME": model_version.registered_model.name,
             "MODEL_VERSION": model_version.version,
             "VERSION_DESCRIPTION": model_version.description,
-            "METADATA": model_version.metadata.dict()
+            "METADATA": model_version.metadata.model_dump()
             if model_version.metadata
             else {},
         }
@@ -1322,7 +1318,7 @@ def pretty_print_model_version_details(
             if model_version.last_updated_at
             else "N/A"
         ),
-        "METADATA": model_version.metadata.dict()
+        "METADATA": model_version.metadata.model_dump()
         if model_version.metadata
         else {},
         "MODEL_SOURCE_URI": model_version.model_source_uri,
@@ -2256,7 +2252,7 @@ def _scrub_secret(config: StackComponentConfig) -> Dict[str, Any]:
         A configuration with secret values removed.
     """
     config_dict = {}
-    config_fields = dict(config.__class__.__fields__)
+    config_fields = config.__class__.model_fields
     for key, value in config_fields.items():
         if secret_utils.is_secret_field(value):
             config_dict[key] = "********"
@@ -2533,7 +2529,7 @@ def list_options(filter_model: Type[BaseFilter]) -> Callable[[F], F]:
     def inner_decorator(func: F) -> F:
         options = []
         data_type_descriptors = set()
-        for k, v in filter_model.__fields__.items():
+        for k, v in filter_model.model_fields.items():
             if k not in filter_model.CLI_EXCLUDE_FIELDS:
                 options.append(
                     click.option(

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
             )
         )