zenml-nightly 0.83.0.dev20250619__py3-none-any.whl → 0.83.0.dev20250622__py3-none-any.whl

zenml/VERSION

@@ -1 +1 @@
-0.83.0.dev20250619
+0.83.0.dev20250622

zenml/__init__.py

@@ -14,6 +14,7 @@
 """Initialization for ZenML."""
 
 import os
+from typing import Any
 
 ROOT_DIR = os.path.dirname(os.path.abspath(__file__))
 
@@ -26,6 +27,16 @@ from zenml.logger import init_logging  # noqa
 
 init_logging()
 
+def __getattr__(name: str) -> Any:
+    # We allow directly accessing the entrypoint module as `zenml.entrypoint`
+    # as this is needed for some orchestrators. Instead of directly importing
+    # the entrypoint module here, we import it dynamically. This avoids a
+    # warning when running the `zenml.entrypoints.entrypoint` module directly.
+    if name == "entrypoint":
+        from zenml.entrypoints import entrypoint
+        return entrypoint
+
+    raise AttributeError(f"module '{__name__}' has no attribute '{name}'")
 
 # Need to import zenml.models before zenml.config to avoid circular imports
 from zenml.models import *  # noqa: F401
@@ -50,7 +61,7 @@ from zenml.steps import step, get_step_context
 from zenml.steps.utils import log_step_metadata
 from zenml.utils.metadata_utils import log_metadata
 from zenml.utils.tag_utils import Tag, add_tags, remove_tags
-from zenml.entrypoints import entrypoint
+
 
 __all__ = [
     "add_tags",
@@ -72,5 +83,4 @@ __all__ = [
     "register_artifact",
     "show",
     "step",
-    "entrypoint",
 ]

zenml/analytics/context.py

@@ -105,8 +105,10 @@ class AnalyticsContext:
 
             # Fetch the `user_id`
             if self.in_server:
-                from zenml.zen_server.auth import get_auth_context
-                from zenml.zen_server.utils import server_config
+                from zenml.zen_server.utils import (
+                    get_auth_context,
+                    server_config,
+                )
 
                 # If the code is running on the server, use the auth context.
                 auth_context = get_auth_context()

zenml/config/server_config.py

@@ -33,6 +33,7 @@ from zenml.constants import (
     DEFAULT_REPORTABLE_RESOURCES,
     DEFAULT_ZENML_JWT_TOKEN_ALGORITHM,
     DEFAULT_ZENML_JWT_TOKEN_LEEWAY,
+    DEFAULT_ZENML_SERVER_AUTH_THREAD_POOL_SIZE,
     DEFAULT_ZENML_SERVER_DEVICE_AUTH_POLLING,
     DEFAULT_ZENML_SERVER_DEVICE_AUTH_TIMEOUT,
     DEFAULT_ZENML_SERVER_FILE_DOWNLOAD_SIZE_LIMIT,
@@ -45,6 +46,7 @@ from zenml.constants import (
     DEFAULT_ZENML_SERVER_MAX_REQUEST_BODY_SIZE_IN_BYTES,
     DEFAULT_ZENML_SERVER_NAME,
     DEFAULT_ZENML_SERVER_PIPELINE_RUN_AUTH_WINDOW,
+    DEFAULT_ZENML_SERVER_REQUEST_CACHE_TIMEOUT,
     DEFAULT_ZENML_SERVER_REQUEST_TIMEOUT,
     DEFAULT_ZENML_SERVER_SECURE_HEADERS_CACHE,
     DEFAULT_ZENML_SERVER_SECURE_HEADERS_CONTENT,
@@ -354,7 +356,10 @@ class ServerConfiguration(BaseModel):
     auto_activate: bool = False
 
     thread_pool_size: int = DEFAULT_ZENML_SERVER_THREAD_POOL_SIZE
-    server_request_timeout: int = DEFAULT_ZENML_SERVER_REQUEST_TIMEOUT
+    auth_thread_pool_size: int = DEFAULT_ZENML_SERVER_AUTH_THREAD_POOL_SIZE
+    request_timeout: int = DEFAULT_ZENML_SERVER_REQUEST_TIMEOUT
+    request_deduplication: bool = True
+    request_cache_timeout: int = DEFAULT_ZENML_SERVER_REQUEST_CACHE_TIMEOUT
 
     max_request_body_size_in_bytes: int = (
         DEFAULT_ZENML_SERVER_MAX_REQUEST_BODY_SIZE_IN_BYTES

zenml/constants.py

@@ -274,6 +274,7 @@ CODE_HASH_PARAMETER_NAME = "step_source"
 # Server settings
 DEFAULT_ZENML_SERVER_NAME = "default"
 DEFAULT_ZENML_SERVER_THREAD_POOL_SIZE = 40
+DEFAULT_ZENML_SERVER_AUTH_THREAD_POOL_SIZE = 5
 DEFAULT_ZENML_JWT_TOKEN_LEEWAY = 10
 DEFAULT_ZENML_JWT_TOKEN_ALGORITHM = "HS256"
 DEFAULT_ZENML_AUTH_SCHEME = AuthScheme.OAUTH2_PASSWORD_BEARER
@@ -283,6 +284,7 @@ DEFAULT_ZENML_SERVER_DEVICE_AUTH_TIMEOUT = 60 * 5  # 5 minutes
 DEFAULT_ZENML_SERVER_DEVICE_AUTH_POLLING = 5  # seconds
 DEFAULT_HTTP_TIMEOUT = 30
 DEFAULT_ZENML_SERVER_REQUEST_TIMEOUT = 20  # seconds
+DEFAULT_ZENML_SERVER_REQUEST_CACHE_TIMEOUT = 300  # seconds
 SERVICE_CONNECTOR_VERIFY_REQUEST_TIMEOUT = 120  # seconds
 ZENML_API_KEY_PREFIX = "ZENKEY_"
 DEFAULT_ZENML_SERVER_PIPELINE_RUN_AUTH_WINDOW = 60 * 48  # 48 hours
@@ -372,6 +374,7 @@ EVENT_FLAVORS = "/event-flavors"
 EVENT_SOURCES = "/event-sources"
 FLAVORS = "/flavors"
 HEALTH = "/health"
+READY = "/ready"
 INFO = "/info"
 LOAD_INFO = "/load-info"
 LOGIN = "/login"

zenml/entrypoints/step_entrypoint_configuration.py

@@ -16,9 +16,11 @@
 import os
 import sys
 from typing import TYPE_CHECKING, Any, List, Set
+from uuid import UUID
 
 from zenml.client import Client
 from zenml.entrypoints.base_entrypoint_configuration import (
+    DEPLOYMENT_ID_OPTION,
     BaseEntrypointConfiguration,
 )
 from zenml.integrations.registry import integration_registry
@@ -147,6 +149,18 @@ class StepEntrypointConfiguration(BaseEntrypointConfiguration):
             kwargs[STEP_NAME_OPTION],
         ]
 
+    def load_deployment(self) -> "PipelineDeploymentResponse":
+        """Loads the deployment.
+
+        Returns:
+            The deployment.
+        """
+        deployment_id = UUID(self.entrypoint_args[DEPLOYMENT_ID_OPTION])
+        step_name = self.entrypoint_args[STEP_NAME_OPTION]
+        return Client().zen_store.get_deployment(
+            deployment_id=deployment_id, step_configuration_filter=[step_name]
+        )
+
     def run(self) -> None:
         """Prepares the environment and runs the configured step."""
         deployment = self.load_deployment()

zenml/models/__init__.py

@@ -80,6 +80,14 @@ from zenml.models.v2.core.api_key import (
     APIKeyInternalUpdate,
     APIKeyRotateRequest,
 )
+from zenml.models.v2.core.api_transaction import (
+    ApiTransactionRequest,
+    ApiTransactionUpdate,
+    ApiTransactionResponse,
+    ApiTransactionResponseBody,
+    ApiTransactionResponseMetadata,
+    ApiTransactionResponseResources,
+)
 from zenml.models.v2.core.artifact import (
     ArtifactFilter,
     ArtifactRequest,
@@ -433,6 +441,7 @@ from zenml.models.v2.misc.info_models import (
 ActionResponseResources.model_rebuild()
 ActionResponseMetadata.model_rebuild()
 APIKeyResponseBody.model_rebuild()
+ApiTransactionResponse.model_rebuild()
 ArtifactResponse.model_rebuild()
 ArtifactResponseBody.model_rebuild()
 ArtifactResponseMetadata.model_rebuild()
@@ -570,6 +579,12 @@ __all__ = [
     "APIKeyInternalResponse",
     "APIKeyInternalUpdate",
     "APIKeyRotateRequest",
+    "ApiTransactionRequest",
+    "ApiTransactionUpdate",
+    "ApiTransactionResponse",
+    "ApiTransactionResponseBody",
+    "ApiTransactionResponseMetadata",
+    "ApiTransactionResponseResources",
     "ArtifactFilter",
     "ArtifactRequest",
     "ArtifactResponse",

zenml/models/v2/core/api_transaction.py

@@ -0,0 +1,193 @@
+#  Copyright (c) ZenML GmbH 2024. All Rights Reserved.
+#
+#  Licensed under the Apache License, Version 2.0 (the "License");
+#  you may not use this file except in compliance with the License.
+#  You may obtain a copy of the License at:
+#
+#       https://www.apache.org/licenses/LICENSE-2.0
+#
+#  Unless required by applicable law or agreed to in writing, software
+#  distributed under the License is distributed on an "AS IS" BASIS,
+#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express
+#  or implied. See the License for the specific language governing
+#  permissions and limitations under the License.
+"""Models representing API transactions."""
+
+from typing import (
+    TYPE_CHECKING,
+    Optional,
+    TypeVar,
+)
+from uuid import UUID
+
+from pydantic import Field, SecretStr
+
+from zenml.models.v2.base.base import BaseUpdate
+from zenml.models.v2.base.scoped import (
+    UserScopedRequest,
+    UserScopedResponse,
+    UserScopedResponseBody,
+    UserScopedResponseMetadata,
+    UserScopedResponseResources,
+)
+from zenml.utils.secret_utils import PlainSerializedSecretStr
+
+if TYPE_CHECKING:
+    from zenml.zen_stores.schemas import BaseSchema
+
+    AnySchema = TypeVar("AnySchema", bound=BaseSchema)
+
+# ------------------ Request Model ------------------
+
+
+class ApiTransactionRequest(UserScopedRequest):
+    """Request model for API transactions."""
+
+    transaction_id: UUID = Field(
+        title="The ID of the transaction.",
+    )
+    method: str = Field(
+        title="The HTTP method of the transaction.",
+    )
+    url: str = Field(
+        title="The URL of the transaction.",
+    )
+
+
+# ------------------ Update Model ------------------
+
+
+class ApiTransactionUpdate(BaseUpdate):
+    """Update model for stack components."""
+
+    result: Optional[PlainSerializedSecretStr] = Field(
+        default=None,
+        title="The response payload.",
+    )
+    cache_time: int = Field(
+        title="The time in seconds that the transaction is kept around after "
+        "completion."
+    )
+
+    def get_result(self) -> Optional[str]:
+        """Get the result of the API transaction.
+
+        Returns:
+            the result of the API transaction.
+        """
+        result = self.result
+        if result is None:
+            return None
+        return result.get_secret_value()
+
+    def set_result(self, result: str) -> None:
+        """Set the result of the API transaction.
+
+        Args:
+            result: the result of the API transaction.
+        """
+        self.result = SecretStr(result)
+
+
+# ------------------ Response Model ------------------
+
+
+class ApiTransactionResponseBody(UserScopedResponseBody):
+    """Response body for API transactions."""
+
+    method: str = Field(
+        title="The HTTP method of the transaction.",
+    )
+    url: str = Field(
+        title="The URL of the transaction.",
+    )
+    completed: bool = Field(
+        title="Whether the transaction is completed.",
+    )
+    result: Optional[PlainSerializedSecretStr] = Field(
+        default=None,
+        title="The response payload.",
+    )
+
+
+class ApiTransactionResponseMetadata(UserScopedResponseMetadata):
+    """Response metadata for API transactions."""
+
+
+class ApiTransactionResponseResources(UserScopedResponseResources):
+    """Response resources for API transactions."""
+
+
+class ApiTransactionResponse(
+    UserScopedResponse[
+        ApiTransactionResponseBody,
+        ApiTransactionResponseMetadata,
+        ApiTransactionResponseResources,
+    ]
+):
+    """Response model for API transactions."""
+
+    def get_hydrated_version(self) -> "ApiTransactionResponse":
+        """Get the hydrated version of this API transaction.
+
+        Returns:
+            an instance of the same entity with the metadata field attached.
+        """
+        return self
+
+    # Body and metadata properties
+
+    @property
+    def method(self) -> str:
+        """The `method` property.
+
+        Returns:
+            the value of the property.
+        """
+        return self.get_body().method
+
+    @property
+    def url(self) -> str:
+        """The `url` property.
+
+        Returns:
+            the value of the property.
+        """
+        return self.get_body().url
+
+    @property
+    def completed(self) -> bool:
+        """The `completed` property.
+
+        Returns:
+            the value of the property.
+        """
+        return self.get_body().completed
+
+    @property
+    def result(self) -> Optional[PlainSerializedSecretStr]:
+        """The `result` property.
+
+        Returns:
+            the value of the property.
+        """
+        return self.get_body().result
+
+    def get_result(self) -> Optional[str]:
+        """Get the result of the API transaction.
+
+        Returns:
+            the result of the API transaction.
+        """
+        result = self.result
+        if result is None:
+            return None
+        return result.get_secret_value()
+
+    def set_result(self, result: str) -> None:
+        """Set the result of the API transaction.
+
+        Args:
+            result: the result of the API transaction.
+        """
+        self.get_body().result = SecretStr(result)

zenml/models/v2/core/pipeline_build.py

@@ -200,6 +200,10 @@ class PipelineBuildResponseBody(ProjectScopedResponseBody):
 class PipelineBuildResponseMetadata(ProjectScopedResponseMetadata):
     """Response metadata for pipeline builds."""
 
+    __zenml_skip_dehydration__: ClassVar[List[str]] = [
+        "images",
+    ]
+
     pipeline: Optional["PipelineResponse"] = Field(
         default=None, title="The pipeline that was used for this build."
     )

zenml/models/v2/core/pipeline_deployment.py

@@ -13,7 +13,7 @@
 #  permissions and limitations under the License.
 """Models representing pipeline deployments."""
 
-from typing import Any, Dict, Optional, Union
+from typing import Any, ClassVar, Dict, List, Optional, Union
 from uuid import UUID
 
 from pydantic import Field
@@ -130,6 +130,13 @@ class PipelineDeploymentResponseBody(ProjectScopedResponseBody):
 class PipelineDeploymentResponseMetadata(ProjectScopedResponseMetadata):
     """Response metadata for pipeline deployments."""
 
+    __zenml_skip_dehydration__: ClassVar[List[str]] = [
+        "pipeline_configuration",
+        "step_configurations",
+        "client_environment",
+        "pipeline_spec",
+    ]
+
     run_name_template: str = Field(
         title="The run name template for runs created using this deployment.",
     )

zenml/models/v2/core/pipeline_run.py

@@ -190,6 +190,13 @@ class PipelineRunResponseBody(ProjectScopedResponseBody):
 class PipelineRunResponseMetadata(ProjectScopedResponseMetadata):
     """Response metadata for pipeline runs."""
 
+    __zenml_skip_dehydration__: ClassVar[List[str]] = [
+        "run_metadata",
+        "config",
+        "client_environment",
+        "orchestrator_environment",
+    ]
+
     run_metadata: Dict[str, MetadataType] = Field(
         default={},
         title="Metadata associated with this pipeline run.",

zenml/models/v2/core/step_run.py

@@ -199,6 +199,12 @@ class StepRunResponseBody(ProjectScopedResponseBody):
 class StepRunResponseMetadata(ProjectScopedResponseMetadata):
     """Response metadata for step runs."""
 
+    __zenml_skip_dehydration__: ClassVar[List[str]] = [
+        "config",
+        "spec",
+        "metadata",
+    ]
+
     # Configuration
     config: "StepConfiguration" = Field(title="The configuration of the step.")
     spec: "StepSpec" = Field(title="The spec of the step.")

zenml/orchestrators/input_utils.py

@@ -65,17 +65,40 @@ def resolve_step_inputs(
     steps_to_fetch.difference_update(step_runs.keys())
 
     if steps_to_fetch:
-        step_runs.update(
-            {
-                run_step.name: run_step
-                for run_step in pagination_utils.depaginate(
-                    Client().list_run_steps,
-                    pipeline_run_id=pipeline_run.id,
-                    project=pipeline_run.project_id,
-                    name="oneof:" + json.dumps(list(steps_to_fetch)),
-                )
-            }
-        )
+        # The list of steps might be too big to fit in the default max URL
+        # length of 8KB supported by most servers. So we need to split it into
+        # smaller chunks.
+        steps_list = list(steps_to_fetch)
+        chunks = []
+        current_chunk = []
+        current_length = 0
+        # stay under 6KB for good measure.
+        max_chunk_length = 6000
+
+        for step_name in steps_list:
+            current_chunk.append(step_name)
+            current_length += len(step_name) + 5  # 5 is for the JSON encoding
+
+            if current_length > max_chunk_length:
+                chunks.append(current_chunk)
+                current_chunk = []
+                current_length = 0
+
+        if current_chunk:
+            chunks.append(current_chunk)
+
+        for chunk in chunks:
+            step_runs.update(
+                {
+                    run_step.name: run_step
+                    for run_step in pagination_utils.depaginate(
+                        Client().list_run_steps,
+                        pipeline_run_id=pipeline_run.id,
+                        project=pipeline_run.project_id,
+                        name="oneof:" + json.dumps(chunk),
+                    )
+                }
+            )
 
     input_artifacts: Dict[str, StepRunInputResponse] = {}
     for name, input_ in step.spec.inputs.items():

zenml/utils/json_utils.py

@@ -113,7 +113,7 @@ def pydantic_encoder(obj: Any) -> Any:
     from pydantic import BaseModel
 
     if isinstance(obj, BaseModel):
-        return obj.model_dump()
+        return obj.model_dump(mode="json")
     elif is_dataclass(obj):
         return asdict(obj)
 

zenml/zen_server/auth.py

@@ -13,13 +13,16 @@
 #  permissions and limitations under the License.
 """Authentication module for ZenML server."""
 
-from contextvars import ContextVar
+import functools
 from datetime import datetime, timedelta
-from typing import Callable, Optional, Tuple, Union
+from functools import wraps
+from typing import Any, Awaitable, Callable, Optional, Tuple, Union
 from urllib.parse import urlencode, urlparse
 from uuid import UUID, uuid4
 
+import anyio.to_thread
 import requests
+from anyio import CapacityLimiter
 from fastapi import Depends, Response
 from fastapi.security import (
     HTTPBasic,
@@ -73,40 +76,13 @@ from zenml.zen_server.jwt import JWTToken
 from zenml.zen_server.utils import (
     get_zenml_headers,
     is_same_or_subdomain,
+    request_manager,
     server_config,
     zen_store,
 )
 
 logger = get_logger(__name__)
 
-# create a context variable to store the authentication context
-_auth_context: ContextVar[Optional["AuthContext"]] = ContextVar(
-    "auth_context", default=None
-)
-
-
-def get_auth_context() -> Optional["AuthContext"]:
-    """Returns the current authentication context.
-
-    Returns:
-        The authentication context.
-    """
-    auth_context = _auth_context.get()
-    return auth_context
-
-
-def set_auth_context(auth_context: "AuthContext") -> "AuthContext":
-    """Sets the current authentication context.
-
-    Args:
-        auth_context: The authentication context.
-
-    Returns:
-        The authentication context.
-    """
-    _auth_context.set(auth_context)
-    return auth_context
-
 
 class AuthContext(BaseModel):
     """The authentication context."""
@@ -1128,4 +1104,50 @@ def authentication_provider() -> Callable[..., AuthContext]:
         raise ValueError(f"Unknown authentication scheme: {auth_scheme}")
 
 
-authorize = authentication_provider()
+def get_authorization_provider() -> Callable[..., Awaitable[AuthContext]]:
+    """Fetches the async authorization provider.
+
+    Returns:
+        The async authorization provider.
+    """
+    provider = authentication_provider()
+    # Create a custom thread pool limiter with a limit of 1 thread for all
+    # auth calls
+    thread_limiter = CapacityLimiter(server_config().auth_thread_pool_size)
+
+    @wraps(provider)
+    async def async_authorize_fn(*args: Any, **kwargs: Any) -> AuthContext:
+        from zenml.zen_server.utils import get_system_metrics_log_str
+
+        request_context = request_manager().current_request
+
+        @wraps(provider)
+        def sync_authorize_fn(*args: Any, **kwargs: Any) -> AuthContext:
+            assert request_context is not None
+
+            logger.debug(
+                f"[{request_context.log_request_id}] API STATS - "
+                f"{request_context.log_request} "
+                f"AUTHORIZING "
+                f"{get_system_metrics_log_str(request_context.request)}"
+            )
+
+            try:
+                auth_context = provider(*args, **kwargs)
+                request_context.auth_context = auth_context
+                return auth_context
+            finally:
+                logger.debug(
+                    f"[{request_context.log_request_id}] API STATS - "
+                    f"{request_context.log_request} "
+                    f"AUTHORIZED "
+                    f"{get_system_metrics_log_str(request_context.request)}"
+                )
+
+        func = functools.partial(sync_authorize_fn, *args, **kwargs)
+        return await anyio.to_thread.run_sync(func, limiter=thread_limiter)
+
+    return async_authorize_fn
+
+
+authorize = get_authorization_provider()

zenml/zen_server/cloud_utils.py

@@ -1,7 +1,6 @@
 """Utils concerning anything concerning the cloud control plane backend."""
 
 import logging
-import threading
 import time
 from datetime import datetime, timedelta
 from threading import RLock
@@ -17,7 +16,11 @@ from zenml.exceptions import (
 )
 from zenml.logger import get_logger
 from zenml.utils.time_utils import utc_now
-from zenml.zen_server.utils import get_zenml_headers, server_config
+from zenml.zen_server.utils import (
+    get_system_metrics_log_str,
+    get_zenml_headers,
+    server_config,
+)
 
 logger = get_logger(__name__)
 
@@ -60,13 +63,21 @@ class ZenMLCloudConnection:
         Returns:
             The response.
         """
+        from zenml.zen_server.utils import get_current_request_context
+
         url = self._config.api_url + endpoint
 
+        log_request_id = "N/A"
+        try:
+            request_context = get_current_request_context()
+            log_request_id = request_context.log_request_id
+        except RuntimeError:
+            pass
+
         if logger.isEnabledFor(logging.DEBUG):
-            # Get the request ID from the current thread object
-            request_id = threading.current_thread().name
             logger.debug(
-                f"[{request_id}] RBAC STATS - {method} {endpoint} started"
+                f"[{log_request_id}] RBAC STATS - {method} "
+                f"{endpoint} started {get_system_metrics_log_str()}"
             )
             start_time = time.time()
 
@@ -107,8 +118,9 @@ class ZenMLCloudConnection:
             if logger.isEnabledFor(logging.DEBUG):
                 duration = (time.time() - start_time) * 1000
                 logger.debug(
-                    f"[{request_id}] RBAC STATS - {status_code} {method} "
-                    f"{endpoint} completed in {duration:.2f}ms"
+                    f"[{log_request_id}] RBAC STATS - "
+                    f"{status_code} {method} {endpoint} completed in "
+                    f"{duration:.2f}ms {get_system_metrics_log_str()}"
                 )
 
         return response