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

zenml/utils/typing_utils.py

@@ -0,0 +1,243 @@
+#  Copyright (c) ZenML GmbH 2021. 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.
+"""Carried over version of some functions from the pydantic v1 typing module.
+
+Check out the latest version here:
+https://github.com/pydantic/pydantic/blob/v1.10.14/pydantic/typing.py
+"""
+
+import sys
+import typing
+from typing import Any, Optional, Set, Tuple, Type, Union, cast
+from typing import get_args as _typing_get_args
+from typing import get_origin as _typing_get_origin
+
+from typing_extensions import Annotated, Literal
+
+# Annotated[...] is implemented by returning an instance of one of these
+# classes, depending on python/typing_extensions version.
+AnnotatedTypeNames = {"AnnotatedMeta", "_AnnotatedAlias"}
+
+# None types
+NONE_TYPES: Tuple[Any, Any, Any] = (None, None.__class__, Literal[None])
+
+# Literal types
+LITERAL_TYPES: Set[Any] = {Literal}
+if hasattr(typing, "Literal"):
+    LITERAL_TYPES.add(typing.Literal)
+
+# ----- is_none_type -----
+
+if sys.version_info[:2] == (3, 8):
+
+    def is_none_type(type_: Any) -> bool:
+        """Checks if the provided type is none type.
+
+        Args:
+            type_: type to check.
+
+        Returns:
+            boolean indicating whether the type is none type.
+        """
+        for none_type in NONE_TYPES:
+            if type_ is none_type:
+                return True
+        # With python 3.8, specifically 3.8.10, Literal "is" checks are very
+        # flakey can change on very subtle changes like use of types in other
+        # modules, hopefully this check avoids that issue.
+        if is_literal_type(type_):  # pragma: no cover
+            return all_literal_values(type_) == (None,)
+        return False
+else:
+
+    def is_none_type(type_: Any) -> bool:
+        """Checks if the provided type is a none type.
+
+        Args:
+            type_: type to check.
+
+        Returns:
+            boolean indicating whether the type is a none type.
+        """
+        return type_ in NONE_TYPES
+
+# ----- is_union -----
+
+if sys.version_info < (3, 10):
+
+    def is_union(type_: Optional[Type[Any]]) -> bool:
+        """Checks if the provided type is a union type.
+
+        Args:
+            type_: type to check.
+
+        Returns:
+            boolean indicating whether the type is union type.
+        """
+        return type_ is Union  # type: ignore[comparison-overlap]
+
+
+else:
+
+    def is_union(type_: Optional[Type[Any]]) -> bool:
+        """Checks if the provided type is a union type.
+
+        Args:
+            type_: type to check.
+
+        Returns:
+            boolean indicating whether the type is union type.
+        """
+        import types
+
+        return type_ is Union or type_ is types.UnionType  # type: ignore[comparison-overlap]
+
+
+# ----- literal -----
+
+
+def is_literal_type(type_: Type[Any]) -> bool:
+    """Checks if the provided type is a literal type.
+
+    Args:
+        type_: type to check.
+
+    Returns:
+        boolean indicating whether the type is union type.
+    """
+    return Literal is not None and get_origin(type_) in LITERAL_TYPES
+
+
+def literal_values(type_: Type[Any]) -> Tuple[Any, ...]:
+    """Fetches the literal values defined in a type.
+
+    Args:
+        type_: type to check.
+
+    Returns:
+        tuple of the literal values.
+    """
+    return get_args(type_)
+
+
+def all_literal_values(type_: Type[Any]) -> Tuple[Any, ...]:
+    """Fetches the literal values defined in a type in a recursive manner.
+
+    This method is used to retrieve all Literal values as Literal can be
+    used recursively (see https://www.python.org/dev/peps/pep-0586)
+    e.g. `Literal[Literal[Literal[1, 2, 3], "foo"], 5, None]`
+
+    Args:
+        type_: type to check.
+
+    Returns:
+        tuple of all the literal values defined in the type.
+    """
+    if not is_literal_type(type_):
+        return (type_,)
+
+    values = literal_values(type_)
+    return tuple(x for value in values for x in all_literal_values(value))
+
+
+# ----- get_origin -----
+
+
+def get_origin(tp: Type[Any]) -> Optional[Type[Any]]:
+    """Fetches the origin of a given type.
+
+    We can't directly use `typing.get_origin` since we need a fallback to
+    support custom generic classes like `ConstrainedList`
+    It should be useless once https://github.com/cython/cython/issues/3537 is
+    solved and https://github.com/pydantic/pydantic/pull/1753 is merged.
+
+    Args:
+        tp: type to check
+
+    Returns:
+        the origin type of the provided type.
+    """
+    if type(tp).__name__ in AnnotatedTypeNames:
+        return cast(Type[Any], Annotated)  # mypy complains about _SpecialForm
+    return _typing_get_origin(tp) or getattr(tp, "__origin__", None)
+
+
+# ----- get_args -----
+
+
+def _generic_get_args(tp: Type[Any]) -> Tuple[Any, ...]:
+    """Generic get args function.
+
+    In python 3.9, `typing.Dict`, `typing.List`, ...
+    do have an empty `__args__` by default (instead of the generic ~T
+    for example). In order to still support `Dict` for example and consider
+    it as `Dict[Any, Any]`, we retrieve the `_nparams` value that tells us
+    how many parameters it needs.
+
+    Args:
+        tp: type to check.
+
+    Returns:
+        Tuple of all the args.
+    """
+    if hasattr(tp, "_nparams"):
+        return (Any,) * tp._nparams  # type: ignore[no-any-return]
+    # Special case for `tuple[()]`, which used to return ((),) with
+    # `typing.Tuple in python 3.10- but now returns () for `tuple` and `Tuple`.
+    try:
+        if tp == Tuple[()] or sys.version_info >= (3, 9) and tp == tuple[()]:
+            return ((),)
+    # there is a TypeError when compiled with cython
+    except TypeError:  # pragma: no cover
+        pass
+    return ()
+
+
+def get_args(tp: Type[Any]) -> Tuple[Any, ...]:
+    """Get type arguments with all substitutions performed.
+
+    For unions, basic simplifications used by Union constructor are performed.
+    Examples::
+        get_args(Dict[str, int]) == (str, int)
+        get_args(int) == ()
+        get_args(Union[int, Union[T, int], str][int]) == (int, str)
+        get_args(Union[int, Tuple[T, int]][str]) == (int, Tuple[str, int])
+        get_args(Callable[[], T][int]) == ([], int)
+
+    Args:
+        tp: the type to check.
+
+    Returns:
+        Tuple of all the args.
+    """
+    if type(tp).__name__ in AnnotatedTypeNames:
+        return tp.__args__ + tp.__metadata__  # type: ignore[no-any-return]
+    # the fallback is needed for the same reasons as `get_origin` (see above)
+    return (
+        _typing_get_args(tp)
+        or getattr(tp, "__args__", ())
+        or _generic_get_args(tp)
+    )
+
+
+def is_optional(tp: Type[Any]) -> bool:
+    """Checks whether a given annotation is typing.Optional.
+
+    Args:
+        tp: the type to check.
+
+    Returns:
+        boolean indicating if the type is typing.Optional.
+    """
+    return is_union(get_origin(tp)) and type(None) in get_args(tp)

zenml/utils/yaml_utils.py

@@ -193,7 +193,7 @@ def is_json_serializable(obj: Any) -> bool:
     Returns:
         Whether the object is JSON serializable using pydantics encoder class.
     """
-    from pydantic.json import pydantic_encoder
+    from zenml.utils.json_utils import pydantic_encoder
 
     try:
         json.dumps(obj, default=pydantic_encoder)

zenml/zen_server/auth.py

@@ -154,7 +154,7 @@ def _fetch_and_verify_api_key(
     # Update the "last used" timestamp of the API key
     store.update_internal_api_key(
         api_key.id,
-        APIKeyInternalUpdate(update_last_login=True),  # type: ignore[call-arg]
+        APIKeyInternalUpdate(update_last_login=True),
     )
 
     return api_key
@@ -544,7 +544,7 @@ def authenticate_external_user(external_access_token: str) -> AuthContext:
 
         if isinstance(payload, dict):
             try:
-                external_user = ExternalUserModel.parse_obj(payload)
+                external_user = ExternalUserModel.model_validate(payload)
             except Exception as e:
                 logger.exception(
                     f"Error parsing user information from external "

zenml/zen_server/cloud_utils.py

@@ -4,7 +4,7 @@ import os
 from typing import Any, Dict, Optional
 
 import requests
-from pydantic import BaseModel, validator
+from pydantic import BaseModel, ConfigDict, field_validator
 from requests.adapters import HTTPAdapter, Retry
 
 from zenml.exceptions import SubscriptionUpgradeRequiredError
@@ -23,7 +23,8 @@ class ZenMLCloudConfiguration(BaseModel):
     oauth2_audience: str
     auth0_domain: str
 
-    @validator("api_url")
+    @field_validator("api_url")
+    @classmethod
     def _strip_trailing_slashes_url(cls, url: str) -> str:
         """Strip any trailing slashes on the API URL.
 
@@ -51,12 +52,11 @@ class ZenMLCloudConfiguration(BaseModel):
 
         return ZenMLCloudConfiguration(**env_config)
 
-    class Config:
-        """Pydantic configuration class."""
-
+    model_config = ConfigDict(
         # Allow extra attributes from configs of previous ZenML versions to
         # permit downgrading
-        extra = "allow"
+        extra="allow"
+    )
 
 
 class ZenMLCloudSession:

zenml/zen_server/deploy/base_provider.py

@@ -78,7 +78,7 @@ class BaseServerProvider(ABC):
         if isinstance(config, cls.CONFIG_TYPE):
             return config
         try:
-            return cls.CONFIG_TYPE(**config.dict())
+            return cls.CONFIG_TYPE(**config.model_dump())
         except ValidationError as e:
             raise ServerDeploymentConfigurationError(
                 f"Invalid configuration for provider {cls.TYPE.value}: {e}"

zenml/zen_server/deploy/deployment.py

@@ -15,7 +15,7 @@
 
 from typing import Optional
 
-from pydantic import BaseModel
+from pydantic import BaseModel, ConfigDict
 
 from zenml.enums import ServerProviderType
 from zenml.services.service_status import ServiceState
@@ -34,16 +34,14 @@ class ServerDeploymentConfig(BaseModel):
 
     name: str
     provider: ServerProviderType
-
-    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 to be set in the base class. The concrete
         # classes are responsible for validating the attributes.
-        extra = "allow"
+        extra="allow",
+    )
 
 
 class ServerDeploymentStatus(BaseModel):
@@ -83,7 +81,7 @@ class ServerDeployment(BaseModel):
     """
 
     config: ServerDeploymentConfig
-    status: Optional[ServerDeploymentStatus]
+    status: Optional[ServerDeploymentStatus] = None
 
     @property
     def is_running(self) -> bool:

zenml/zen_server/deploy/docker/docker_zen_server.py

@@ -16,6 +16,8 @@
 import os
 from typing import Dict, List, Optional, Tuple, cast
 
+from pydantic import ConfigDict
+
 import zenml
 from zenml.config.global_config import GlobalConfiguration
 from zenml.config.store_config import StoreConfiguration
@@ -70,10 +72,7 @@ class DockerServerDeploymentConfig(ServerDeploymentConfig):
     store: Optional[StoreConfiguration] = None
     use_legacy_dashboard: bool = DEFAULT_ZENML_SERVER_USE_LEGACY_DASHBOARD
 
-    class Config:
-        """Pydantic configuration."""
-
-        extra = "forbid"
+    model_config = ConfigDict(extra="forbid")
 
 
 class DockerZenServerConfig(ContainerServiceConfig):

zenml/zen_server/deploy/local/local_provider.py

@@ -62,7 +62,6 @@ class LocalServerProvider(BaseServerProvider):
         try:
             # Make sure the ZenML Server dependencies are installed
             import fastapi  # noqa
-            import fastapi_utils  # noqa
             import jwt  # noqa
             import multipart  # noqa
             import uvicorn  # noqa

zenml/zen_server/deploy/local/local_zen_server.py

@@ -17,6 +17,8 @@ import ipaddress
 import os
 from typing import Dict, List, Optional, Tuple, Union, cast
 
+from pydantic import ConfigDict, Field
+
 from zenml.client import Client
 from zenml.config.global_config import GlobalConfiguration
 from zenml.config.store_config import StoreConfiguration
@@ -60,17 +62,15 @@ class LocalServerDeploymentConfig(ServerDeploymentConfig):
     """
 
     port: int = 8237
-    ip_address: Union[ipaddress.IPv4Address, ipaddress.IPv6Address] = (
-        ipaddress.IPv4Address(DEFAULT_LOCAL_SERVICE_IP_ADDRESS)
+    ip_address: Union[ipaddress.IPv4Address, ipaddress.IPv6Address] = Field(
+        default=ipaddress.IPv4Address(DEFAULT_LOCAL_SERVICE_IP_ADDRESS),
+        union_mode="left_to_right",
     )
     blocking: bool = False
     store: Optional[StoreConfiguration] = None
     use_legacy_dashboard: bool = DEFAULT_ZENML_SERVER_USE_LEGACY_DASHBOARD
 
-    class Config:
-        """Pydantic configuration."""
-
-        extra = "forbid"
+    model_config = ConfigDict(extra="forbid")
 
 
 class LocalZenServerConfig(LocalDaemonServiceConfig):

zenml/zen_server/deploy/terraform/terraform_zen_server.py

@@ -19,6 +19,8 @@ from pathlib import Path
 from typing import Any, Dict, Optional, cast
 from uuid import UUID
 
+from pydantic import ConfigDict
+
 from zenml.logger import get_logger
 from zenml.services import ServiceType
 from zenml.services.terraform.terraform_service import (
@@ -137,11 +139,7 @@ class TerraformServerDeploymentConfig(ServerDeploymentConfig):
     database_ssl_key: str = ""
     database_ssl_verify_server_cert: bool = True
     analytics_opt_in: bool = True
-
-    class Config:
-        """Pydantic configuration."""
-
-        extra = "allow"
+    model_config = ConfigDict(extra="allow")
 
 
 class TerraformZenServerConfig(TerraformServiceConfig):
@@ -199,7 +197,7 @@ class TerraformZenServer(TerraformService):
         # filter keys that are not modeled as terraform deployment vars
         vars = {
             k: str(v) if isinstance(v, UUID) else v
-            for k, v in self.config.server.dict().items()
+            for k, v in self.config.server.model_dump().items()
             if k not in filter_vars
         }
         assert self.status.runtime_path

zenml/zen_server/exceptions.py

@@ -24,6 +24,7 @@ from zenml.exceptions import (
     DuplicateRunNameError,
     EntityExistsError,
     IllegalOperationError,
+    MethodNotAllowedError,
     SecretExistsError,
     StackComponentExistsError,
     StackExistsError,
@@ -39,7 +40,7 @@ if TYPE_CHECKING:
 class ErrorModel(BaseModel):
     """Base class for error responses."""
 
-    detail: Any
+    detail: Optional[Any] = None
 
 
 error_response = dict(model=ErrorModel)
@@ -93,6 +94,8 @@ REST_API_EXCEPTIONS: List[Tuple[Type[Exception], int]] = [
     (RuntimeError, 500),
     # 501 Not Implemented,
     (NotImplementedError, 501),
+    # 405 Method Not Allowed
+    (MethodNotAllowedError, 405),
 ]
 
 

zenml/zen_server/feature_gate/zenml_cloud_feature_gate.py

@@ -109,7 +109,7 @@ class ZenMLCloudFeatureGateInterface(FeatureGateInterface, ZenMLCloudSession):
                 "tenant_id": str(server_config.external_server_id),
                 "resource_id": str(resource_id),
             },
-        ).dict()
+        ).model_dump()
         response = self._post(endpoint=USAGE_EVENT_ENDPOINT, data=data)
         if response.status_code != 200:
             logger.error(

zenml/zen_server/pipeline_deployment/utils.py

@@ -9,6 +9,7 @@ from uuid import UUID
 from fastapi import BackgroundTasks
 from packaging import version
 
+from zenml.config.base_settings import BaseSettings
 from zenml.config.pipeline_configurations import PipelineConfiguration
 from zenml.config.pipeline_run_configuration import PipelineRunConfiguration
 from zenml.config.step_configurations import Step, StepConfiguration
@@ -29,10 +30,11 @@ from zenml.models import (
 )
 from zenml.new.pipelines.run_utils import (
     create_placeholder_run,
+    validate_run_config_is_runnable_from_server,
+    validate_stack_is_runnable_from_server,
 )
-from zenml.orchestrators import BaseOrchestratorConfig
 from zenml.stack.flavor import Flavor
-from zenml.utils import dict_utils, pydantic_utils
+from zenml.utils import dict_utils, pydantic_utils, settings_utils
 from zenml.zen_server.auth import AuthContext
 from zenml.zen_server.pipeline_deployment.runner_entrypoint_configuration import (
     RunnerEntrypointConfiguration,
@@ -73,21 +75,32 @@ def run_pipeline(
     if not stack:
         raise ValueError("Unable to run deployment without associated stack.")
 
-    validate_stack(stack)
+    validate_stack_is_runnable_from_server(zen_store=zen_store(), stack=stack)
+    if run_config:
+        validate_run_config_is_runnable_from_server(run_config)
 
     deployment_request = apply_run_config(
         deployment=deployment,
         run_config=run_config or PipelineRunConfiguration(),
         user_id=auth_context.user.id,
     )
+
+    ensure_async_orchestrator(deployment=deployment_request, stack=stack)
+
     new_deployment = zen_store().create_deployment(deployment_request)
     placeholder_run = create_placeholder_run(deployment=new_deployment)
     assert placeholder_run
 
-    api_token = auth_context.encoded_access_token
-    if not api_token:
-        assert auth_context.access_token
-        api_token = auth_context.access_token.encode()
+    if auth_context.access_token:
+        token = auth_context.access_token
+        token.pipeline_id = deployment_request.pipeline
+
+        # We create a non-expiring token to make sure its active for the entire
+        # duration of the pipeline run
+        api_token = token.encode(expires=None)
+    else:
+        assert auth_context.encoded_access_token
+        api_token = auth_context.encoded_access_token
 
     server_url = server_config().server_url
     if not server_url:
@@ -170,39 +183,35 @@ def run_pipeline(
     return placeholder_run
 
 
-def validate_stack(stack: StackResponse) -> None:
-    """Validate a stack model.
+def ensure_async_orchestrator(
+    deployment: PipelineDeploymentRequest, stack: StackResponse
+) -> None:
+    """Ensures the orchestrator is configured to run async.
 
     Args:
-        stack: The stack to validate.
-
-    Raises:
-        ValueError: If the stack has components of a custom flavor, local
-            components or a synchronous orchestrator.
+        deployment: Deployment request in which the orchestrator
+            configuration should be updated to ensure the orchestrator is
+            running async.
+        stack: The stack on which the deployment will run.
     """
-    for component_list in stack.components.values():
-        assert len(component_list) == 1
-        component = component_list[0]
-        flavors = zen_store().list_flavors(
-            FlavorFilter(name=component.flavor, type=component.type)
-        )
-        assert len(flavors) == 1
-        flavor_model = flavors[0]
-
-        if flavor_model.workspace is not None:
-            raise ValueError("No custom stack component flavors allowed.")
-
-        flavor = Flavor.from_model(flavor_model)
-        component_config = flavor.config_class(**component.configuration)
+    orchestrator = stack.components[StackComponentType.ORCHESTRATOR][0]
+    flavors = zen_store().list_flavors(
+        FlavorFilter(name=orchestrator.flavor, type=orchestrator.type)
+    )
+    flavor = Flavor.from_model(flavors[0])
 
-        if component_config.is_local:
-            raise ValueError("No local stack components allowed.")
+    if "synchronous" in flavor.config_class.model_fields:
+        key = settings_utils.get_flavor_setting_key(flavor)
 
-        if flavor.type == StackComponentType.ORCHESTRATOR:
-            assert isinstance(component_config, BaseOrchestratorConfig)
+        if settings := deployment.pipeline_configuration.settings.get(key):
+            settings_dict = settings.model_dump()
+        else:
+            settings_dict = {}
 
-            if component_config.is_synchronous:
-                raise ValueError("No synchronous orchestrator allowed.")
+        settings_dict["synchronous"] = False
+        deployment.pipeline_configuration.settings[key] = (
+            BaseSettings.model_validate(settings_dict)
+        )
 
 
 def get_requirements_for_stack(
@@ -326,57 +335,28 @@ def apply_run_config(
         run_config: The run configuration to apply.
         user_id: The ID of the user that wants to run the deployment.
 
-    Raises:
-        ValueError: If the run configuration contains values that can't be
-            updated when running a pipeline deployment.
-
     Returns:
         The updated deployment.
     """
-    pipeline_updates = {}
-
-    if run_config.parameters:
-        raise ValueError(
-            "Can't set parameters when running pipeline via Rest API."
-        )
-
-    if run_config.build:
-        raise ValueError("Can't set build when running pipeline via Rest API.")
-
-    if run_config.schedule:
-        raise ValueError(
-            "Can't set schedule when running pipeline via Rest API."
-        )
-
-    if run_config.settings.get("docker"):
-        raise ValueError(
-            "Can't set DockerSettings when running pipeline via Rest API."
-        )
-
-    pipeline_updates = run_config.dict(
-        exclude_none=True, include=set(PipelineConfiguration.__fields__)
+    pipeline_updates = run_config.model_dump(
+        exclude_none=True, include=set(PipelineConfiguration.model_fields)
     )
 
     pipeline_configuration = pydantic_utils.update_model(
         deployment.pipeline_configuration, update=pipeline_updates
     )
-    pipeline_configuration_dict = pipeline_configuration.dict(
+    pipeline_configuration_dict = pipeline_configuration.model_dump(
         exclude_none=True
     )
     steps = {}
     for invocation_id, step in deployment.step_configurations.items():
         step_config_dict = dict_utils.recursive_update(
             copy.deepcopy(pipeline_configuration_dict),
-            update=step.config.dict(exclude_none=True),
+            update=step.config.model_dump(exclude_none=True),
         )
-        step_config = StepConfiguration.parse_obj(step_config_dict)
+        step_config = StepConfiguration.model_validate(step_config_dict)
 
         if update := run_config.steps.get(invocation_id):
-            if update.settings.get("docker"):
-                raise ValueError(
-                    "Can't set DockerSettings when running pipeline via Rest API."
-                )
-
             step_config = pydantic_utils.update_model(
                 step_config, update=update
             )