zenml-nightly 0.82.1.dev20250521__py3-none-any.whl → 0.82.1.dev20250524__py3-none-any.whl

zenml/VERSION

@@ -1 +1 @@
-0.82.1.dev20250521
+0.82.1.dev20250524

zenml/client.py

@@ -2024,7 +2024,9 @@ class Client(metaclass=ClientMetaClass):
             name=name,
             type=component_type,
             flavor=flavor,
-            configuration=configuration,
+            configuration=validated_config.model_dump(
+                mode="json", exclude_unset=True
+            ),
             labels=labels,
         )
 
@@ -2112,7 +2114,9 @@ class Client(metaclass=ClientMetaClass):
             assert validated_config is not None
             warn_if_config_server_mismatch(validated_config)
 
-            update_model.configuration = existing_configuration
+            update_model.configuration = validated_config.model_dump(
+                mode="json", exclude_unset=True
+            )
 
         if labels is not None:
             existing_labels = component.labels or {}

zenml/config/build_configuration.py

@@ -108,6 +108,13 @@ class BuildConfiguration(BaseModel):
             )
             if digest:
                 hash_.update(digest.encode())
+            elif self.settings.skip_build:
+                # If the user has specified to skip the build, the image being
+                # used for the run is whatever they set for `parent_image`.
+                # In this case, the bevavior depends on the orchestrators image
+                # pull policy, and whether there is any caching involved. This
+                # is out of our control here, so we don't log any warning.
+                pass
             else:
                 logger.warning(
                     "Unable to fetch parent image digest for image `%s`. "

zenml/config/docker_settings.py

@@ -386,6 +386,31 @@ class DockerSettings(BaseSettings):
 
         return self
 
+    @model_validator(mode="before")
+    @classmethod
+    @before_validator_handler
+    def _warn_about_future_default_installer(
+        cls, data: Dict[str, Any]
+    ) -> Dict[str, Any]:
+        """Warns about the future change of default package installer from pip to uv.
+
+        Args:
+            data: The model data.
+
+        Returns:
+            The validated settings values.
+        """
+        if "python_package_installer" not in data:
+            logger.warning(
+                "In a future release, the default Python package installer "
+                "used by ZenML to build container images for your "
+                "containerized pipelines will change from 'pip' to 'uv'. "
+                "To maintain current behavior, you can explicitly set "
+                "`python_package_installer=PythonPackageInstaller.PIP` "
+                "in your DockerSettings."
+            )
+        return data
+
     model_config = ConfigDict(
         # public attributes are immutable
         frozen=True,

zenml/config/server_config.py

@@ -45,6 +45,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_TIMEOUT,
     DEFAULT_ZENML_SERVER_SECURE_HEADERS_CACHE,
     DEFAULT_ZENML_SERVER_SECURE_HEADERS_CONTENT,
     DEFAULT_ZENML_SERVER_SECURE_HEADERS_CSP,
@@ -252,6 +253,11 @@ class ServerConfiguration(BaseModel):
             used.
         file_download_size_limit: The maximum size of the file download in
             bytes. If not specified, the default value of 2GB will be used.
+        thread_pool_size: The size of the thread pool for handling requests. If
+            not specified, the default value of 40 will be used.
+        server_request_timeout: The timeout for server requests in seconds. If
+            not specified, the default value of 20 seconds will be used. This
+            value should be lower than the client's request timeout.
     """
 
     deployment_type: ServerDeploymentType = ServerDeploymentType.OTHER
@@ -348,6 +354,7 @@ 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
 
     max_request_body_size_in_bytes: int = (
         DEFAULT_ZENML_SERVER_MAX_REQUEST_BODY_SIZE_IN_BYTES

zenml/constants.py

@@ -282,6 +282,7 @@ DEFAULT_ZENML_SERVER_MAX_DEVICE_AUTH_ATTEMPTS = 3
 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
 SERVICE_CONNECTOR_VERIFY_REQUEST_TIMEOUT = 120  # seconds
 ZENML_API_KEY_PREFIX = "ZENKEY_"
 DEFAULT_ZENML_SERVER_PIPELINE_RUN_AUTH_WINDOW = 60 * 48  # 48 hours

zenml/enums.py

@@ -314,6 +314,7 @@ class EnvironmentType(StrEnum):
     LIGHTNING_AI_STUDIO = "lightning_ai_studio"
     GITHUB_CODESPACES = "github_codespaces"
     VSCODE_REMOTE_CONTAINER = "vscode_remote_container"
+    ZENML_CODESPACE = "zenml_codespace"
 
 
 class ModelStages(StrEnum):

zenml/environment.py

@@ -79,6 +79,8 @@ def get_environment() -> str:
         return EnvironmentType.GENERIC_CI
     elif Environment.in_github_codespaces():
         return EnvironmentType.GITHUB_CODESPACES
+    elif Environment.in_zenml_codespace():
+        return EnvironmentType.ZENML_CODESPACE
     elif Environment.in_vscode_remote_container():
         return EnvironmentType.VSCODE_REMOTE_CONTAINER
     elif Environment.in_lightning_ai_studio():
@@ -254,6 +256,16 @@ class Environment(metaclass=SingletonMetaClass):
             or "GITHUB_CODESPACES_PORT_FORWARDING_DOMAIN" in os.environ
         )
 
+    @staticmethod
+    def in_zenml_codespace() -> bool:
+        """If the current Python process is running in ZenML Codespaces.
+
+        Returns:
+            `True` if the current Python process is running in ZenML Codespaces,
+            `False` otherwise.
+        """
+        return os.environ.get("ZENML_ENVIRONMENT") == "codespace"
+
     @staticmethod
     def in_vscode_remote_container() -> bool:
         """If the current Python process is running in a VS Code Remote Container.

zenml/integrations/gcp/__init__.py

@@ -45,7 +45,7 @@ class GcpIntegration(Integration):
     NAME = GCP
     REQUIREMENTS = [
         "kfp>=2.6.0",
-        "gcsfs",
+        "gcsfs!=2025.5.0,!=2025.5.0.post1",
         "google-cloud-secret-manager",
         "google-cloud-container>=2.21.0",
         "google-cloud-artifact-registry>=1.11.3",

zenml/integrations/gcp/service_connectors/gcp_service_connector.py

@@ -78,7 +78,7 @@ from zenml.service_connectors.service_connector import (
 from zenml.utils.enum_utils import StrEnum
 from zenml.utils.pydantic_utils import before_validator_handler
 from zenml.utils.secret_utils import PlainSerializedSecretStr
-from zenml.utils.time_utils import utc_now
+from zenml.utils.time_utils import to_utc_timezone, utc_now
 
 logger = get_logger(__name__)
 
@@ -1200,14 +1200,13 @@ class GCPServiceConnector(ServiceConnector):
         elif auth_method == GCPAuthenticationMethods.OAUTH2_TOKEN:
             assert isinstance(cfg, GCPOAuth2TokenConfig)
 
-            expires_at = self.expires_at
-            if expires_at:
-                # Remove the UTC timezone
-                expires_at = expires_at.replace(tzinfo=None)
-
             credentials = gcp_credentials.Credentials(
                 token=cfg.token.get_secret_value(),
-                expiry=expires_at,
+                # Currently GCP expects the expiry to be a timezone-naive
+                # UTC datetime.
+                expiry=to_utc_timezone(self.expires_at).replace(tzinfo=None)
+                if self.expires_at
+                else None,
                 scopes=scopes,
             )
 
@@ -1303,10 +1302,12 @@ class GCPServiceConnector(ServiceConnector):
                 )
 
         if credentials.expiry:
-            # Add the UTC timezone to the expiration time
-            expires_at = credentials.expiry.replace(
-                tzinfo=datetime.timezone.utc
-            )
+            expires_at = credentials.expiry
+
+        if expires_at:
+            # Add the UTC timezone to the expiration time, if it's not already
+            # set
+            expires_at = to_utc_timezone(expires_at)
 
         return credentials, expires_at
 

zenml/integrations/hyperai/orchestrators/hyperai_orchestrator.py

@@ -425,7 +425,7 @@ class HyperAIOrchestrator(ContainerizedOrchestrator):
         with tempfile.NamedTemporaryFile(mode="w", delete=True) as f:
             # Define bash line and command line
             bash_line = "#!/bin/bash\n"
-            command_line = f'cd {directory_name} && echo {ENV_ZENML_HYPERAI_RUN_ID}="{deployment_id}_$(date +\%s)" > .env && docker compose up -d'
+            command_line = rf'cd {directory_name} && echo {ENV_ZENML_HYPERAI_RUN_ID}="{deployment_id}_$(date +\%s)" > .env && docker compose up -d'
 
             # Write script to temporary file
             with f.file as f_:

zenml/integrations/kubernetes/orchestrators/kube_utils.py

@@ -270,18 +270,22 @@ def wait_pod(
         # Stream logs to `zenml.logger.info()`.
         # TODO: can we do this without parsing all logs every time?
         if stream_logs and pod_is_not_pending(resp):
-            response = core_api.read_namespaced_pod_log(
-                name=pod_name,
-                namespace=namespace,
-                _preload_content=False,
-            )
-            raw_data = response.data
-            decoded_log = raw_data.decode("utf-8", errors="replace")
-            logs = decoded_log.splitlines()
-            if len(logs) > logged_lines:
-                for line in logs[logged_lines:]:
-                    logger.info(line)
-                logged_lines = len(logs)
+            try:
+                response = core_api.read_namespaced_pod_log(
+                    name=pod_name,
+                    namespace=namespace,
+                    _preload_content=False,
+                )
+            except ApiException as e:
+                logger.error(f"Error reading pod logs: {e}. Retrying...")
+            else:
+                raw_data = response.data
+                decoded_log = raw_data.decode("utf-8", errors="replace")
+                logs = decoded_log.splitlines()
+                if len(logs) > logged_lines:
+                    for line in logs[logged_lines:]:
+                        logger.info(line)
+                    logged_lines = len(logs)
 
         # Raise an error if the pod failed.
         if pod_failed(resp):

zenml/materializers/built_in_materializer.py

@@ -352,7 +352,9 @@ class BuiltInContainerMaterializer(BaseMaterializer):
                 ):
                     type_ = find_type_by_str(type_str)
                     materializer_class = materializer_registry[type_]
-                    materializer = materializer_class(uri=path_)
+                    materializer = materializer_class(
+                        uri=path_, artifact_store=self.artifact_store
+                    )
                     element = materializer.load(type_)
                     outputs.append(element)
 
@@ -364,7 +366,9 @@ class BuiltInContainerMaterializer(BaseMaterializer):
                     materializer_class = source_utils.load(
                         entry["materializer"]
                     )
-                    materializer = materializer_class(uri=path_)
+                    materializer = materializer_class(
+                        uri=path_, artifact_store=self.artifact_store
+                    )
                     element = materializer.load(type_)
                     outputs.append(element)
 
@@ -427,7 +431,9 @@ class BuiltInContainerMaterializer(BaseMaterializer):
                 self.artifact_store.mkdir(element_path)
                 type_ = type(element)
                 materializer_class = materializer_registry[type_]
-                materializer = materializer_class(uri=element_path)
+                materializer = materializer_class(
+                    uri=element_path, artifact_store=self.artifact_store
+                )
                 materializers.append(materializer)
                 metadata.append(
                     {

zenml/zen_server/cloud_utils.py

@@ -1,5 +1,8 @@
 """Utils concerning anything concerning the cloud control plane backend."""
 
+import logging
+import threading
+import time
 from datetime import datetime, timedelta
 from threading import RLock
 from typing import Any, Dict, Optional
@@ -12,9 +15,12 @@ from zenml.exceptions import (
     IllegalOperationError,
     SubscriptionUpgradeRequiredError,
 )
+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
 
+logger = get_logger(__name__)
+
 _cloud_connection: Optional["ZenMLCloudConnection"] = None
 
 
@@ -56,16 +62,16 @@ class ZenMLCloudConnection:
         """
         url = self._config.api_url + endpoint
 
-        response = self.session.request(
-            method=method,
-            url=url,
-            params=params,
-            json=data,
-            timeout=self._config.http_timeout,
-        )
-        if response.status_code == 401:
-            # Refresh the auth token and try again
-            self._reset_login()
+        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"
+            )
+            start_time = time.time()
+
+        status_code: Optional[int] = None
+        try:
             response = self.session.request(
                 method=method,
                 url=url,
@@ -73,18 +79,36 @@ class ZenMLCloudConnection:
                 json=data,
                 timeout=self._config.http_timeout,
             )
+            if response.status_code == 401:
+                # Refresh the auth token and try again
+                self._reset_login()
+                response = self.session.request(
+                    method=method,
+                    url=url,
+                    params=params,
+                    json=data,
+                    timeout=self._config.http_timeout,
+                )
 
-        try:
-            response.raise_for_status()
-        except requests.HTTPError as e:
-            if response.status_code == 402:
-                raise SubscriptionUpgradeRequiredError(response.json())
-            elif response.status_code == 403:
-                raise IllegalOperationError(response.json())
-            else:
-                raise RuntimeError(
-                    f"Failed while trying to contact the central zenml pro "
-                    f"service: {e}"
+            status_code = response.status_code
+            try:
+                response.raise_for_status()
+            except requests.HTTPError as e:
+                if response.status_code == 402:
+                    raise SubscriptionUpgradeRequiredError(response.json())
+                elif response.status_code == 403:
+                    raise IllegalOperationError(response.json())
+                else:
+                    raise RuntimeError(
+                        f"Failed while trying to contact the central zenml pro "
+                        f"service: {e}"
+                    )
+        finally:
+            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"
                 )
 
         return response

zenml/zen_server/routers/actions_endpoints.py

@@ -44,7 +44,7 @@ from zenml.zen_server.rbac.utils import (
     verify_permission_for_model,
 )
 from zenml.zen_server.utils import (
-    handle_exceptions,
+    async_fastapi_endpoint_wrapper,
     make_dependable,
     plugin_flavor_registry,
     zen_store,
@@ -61,7 +61,7 @@ router = APIRouter(
     "",
     responses={401: error_response, 404: error_response, 422: error_response},
 )
-@handle_exceptions
+@async_fastapi_endpoint_wrapper
 def list_actions(
     action_filter_model: ActionFilter = Depends(make_dependable(ActionFilter)),
     hydrate: bool = False,
@@ -132,7 +132,7 @@ def list_actions(
     "/{action_id}",
     responses={401: error_response, 404: error_response, 422: error_response},
 )
-@handle_exceptions
+@async_fastapi_endpoint_wrapper
 def get_action(
     action_id: UUID,
     hydrate: bool = True,
@@ -179,7 +179,7 @@ def get_action(
     "",
     responses={401: error_response, 409: error_response, 422: error_response},
 )
-@handle_exceptions
+@async_fastapi_endpoint_wrapper
 def create_action(
     action: ActionRequest,
     _: AuthContext = Security(authorize),
@@ -225,7 +225,7 @@ def create_action(
     "/{action_id}",
     responses={401: error_response, 404: error_response, 422: error_response},
 )
-@handle_exceptions
+@async_fastapi_endpoint_wrapper
 def update_action(
     action_id: UUID,
     action_update: ActionUpdate,
@@ -280,7 +280,7 @@ def update_action(
     "/{action_id}",
     responses={401: error_response, 404: error_response, 422: error_response},
 )
-@handle_exceptions
+@async_fastapi_endpoint_wrapper
 def delete_action(
     action_id: UUID,
     force: bool = False,

zenml/zen_server/routers/artifact_endpoint.py

@@ -36,7 +36,7 @@ from zenml.zen_server.rbac.endpoint_utils import (
 )
 from zenml.zen_server.rbac.models import ResourceType
 from zenml.zen_server.utils import (
-    handle_exceptions,
+    async_fastapi_endpoint_wrapper,
     make_dependable,
     zen_store,
 )
@@ -52,7 +52,7 @@ artifact_router = APIRouter(
     "",
     responses={401: error_response, 404: error_response, 422: error_response},
 )
-@handle_exceptions
+@async_fastapi_endpoint_wrapper
 def list_artifacts(
     artifact_filter_model: ArtifactFilter = Depends(
         make_dependable(ArtifactFilter)
@@ -83,7 +83,7 @@ def list_artifacts(
     "",
     responses={401: error_response, 409: error_response, 422: error_response},
 )
-@handle_exceptions
+@async_fastapi_endpoint_wrapper
 def create_artifact(
     artifact: ArtifactRequest,
     _: AuthContext = Security(authorize),
@@ -106,7 +106,7 @@ def create_artifact(
     "/{artifact_id}",
     responses={401: error_response, 404: error_response, 422: error_response},
 )
-@handle_exceptions
+@async_fastapi_endpoint_wrapper
 def get_artifact(
     artifact_id: UUID,
     hydrate: bool = True,
@@ -133,7 +133,7 @@ def get_artifact(
     "/{artifact_id}",
     responses={401: error_response, 404: error_response, 422: error_response},
 )
-@handle_exceptions
+@async_fastapi_endpoint_wrapper
 def update_artifact(
     artifact_id: UUID,
     artifact_update: ArtifactUpdate,
@@ -160,7 +160,7 @@ def update_artifact(
     "/{artifact_id}",
     responses={401: error_response, 404: error_response, 422: error_response},
 )
-@handle_exceptions
+@async_fastapi_endpoint_wrapper
 def delete_artifact(
     artifact_id: UUID,
     _: AuthContext = Security(authorize),

zenml/zen_server/routers/artifact_version_endpoints.py

@@ -66,7 +66,7 @@ from zenml.zen_server.rbac.utils import (
     get_allowed_resource_ids,
 )
 from zenml.zen_server.utils import (
-    handle_exceptions,
+    async_fastapi_endpoint_wrapper,
     make_dependable,
     set_filter_project_scope,
     zen_store,
@@ -83,7 +83,7 @@ artifact_version_router = APIRouter(
     "",
     responses={401: error_response, 404: error_response, 422: error_response},
 )
-@handle_exceptions
+@async_fastapi_endpoint_wrapper
 def list_artifact_versions(
     artifact_version_filter_model: ArtifactVersionFilter = Depends(
         make_dependable(ArtifactVersionFilter)
@@ -127,7 +127,7 @@ def list_artifact_versions(
     "",
     responses={401: error_response, 409: error_response, 422: error_response},
 )
-@handle_exceptions
+@async_fastapi_endpoint_wrapper
 def create_artifact_version(
     artifact_version: ArtifactVersionRequest,
     _: AuthContext = Security(authorize),
@@ -150,7 +150,7 @@ def create_artifact_version(
     BATCH,
     responses={401: error_response, 409: error_response, 422: error_response},
 )
-@handle_exceptions
+@async_fastapi_endpoint_wrapper
 def batch_create_artifact_version(
     artifact_versions: List[ArtifactVersionRequest],
     _: AuthContext = Security(authorize),
@@ -173,7 +173,7 @@ def batch_create_artifact_version(
     "/{artifact_version_id}",
     responses={401: error_response, 404: error_response, 422: error_response},
 )
-@handle_exceptions
+@async_fastapi_endpoint_wrapper
 def get_artifact_version(
     artifact_version_id: UUID,
     hydrate: bool = True,
@@ -200,7 +200,7 @@ def get_artifact_version(
     "/{artifact_version_id}",
     responses={401: error_response, 404: error_response, 422: error_response},
 )
-@handle_exceptions
+@async_fastapi_endpoint_wrapper
 def update_artifact_version(
     artifact_version_id: UUID,
     artifact_version_update: ArtifactVersionUpdate,
@@ -227,7 +227,7 @@ def update_artifact_version(
     "/{artifact_version_id}",
     responses={401: error_response, 404: error_response, 422: error_response},
 )
-@handle_exceptions
+@async_fastapi_endpoint_wrapper
 def delete_artifact_version(
     artifact_version_id: UUID,
     _: AuthContext = Security(authorize),
@@ -248,7 +248,7 @@ def delete_artifact_version(
     "",
     responses={401: error_response, 404: error_response, 422: error_response},
 )
-@handle_exceptions
+@async_fastapi_endpoint_wrapper
 def prune_artifact_versions(
     project_name_or_id: Union[str, UUID],
     only_versions: bool = True,
@@ -275,7 +275,7 @@ def prune_artifact_versions(
     "/{artifact_version_id}" + VISUALIZE,
     responses={401: error_response, 404: error_response, 422: error_response},
 )
-@handle_exceptions
+@async_fastapi_endpoint_wrapper
 def get_artifact_visualization(
     artifact_version_id: UUID,
     index: int = 0,
@@ -303,7 +303,7 @@ def get_artifact_visualization(
     "/{artifact_version_id}" + DOWNLOAD_TOKEN,
     responses={401: error_response, 404: error_response, 422: error_response},
 )
-@handle_exceptions
+@async_fastapi_endpoint_wrapper
 def get_artifact_download_token(
     artifact_version_id: UUID,
     _: AuthContext = Security(authorize),
@@ -334,7 +334,7 @@ def get_artifact_download_token(
     "/{artifact_version_id}" + DATA,
     responses={401: error_response, 404: error_response, 422: error_response},
 )
-@handle_exceptions
+@async_fastapi_endpoint_wrapper
 def download_artifact_data(
     artifact_version_id: UUID, token: str
 ) -> FileResponse:

zenml/zen_server/routers/auth_endpoints.py

@@ -71,8 +71,8 @@ from zenml.zen_server.rbac.utils import (
     verify_permission_for_model,
 )
 from zenml.zen_server.utils import (
+    async_fastapi_endpoint_wrapper,
     get_ip_location,
-    handle_exceptions,
     server_config,
     zen_store,
 )
@@ -222,11 +222,11 @@ class OAuthLoginRequestForm:
     LOGIN,
     response_model=Union[OAuthTokenResponse, OAuthRedirectResponse],
 )
+@async_fastapi_endpoint_wrapper
 @rate_limit_requests(
     day_limit=server_config().login_rate_limit_day,
     minute_limit=server_config().login_rate_limit_minute,
 )
-@handle_exceptions
 def token(
     request: Request,
     response: Response,
@@ -341,6 +341,7 @@ def logout(
     DEVICE_AUTHORIZATION,
     response_model=OAuthDeviceAuthorizationResponse,
 )
+@async_fastapi_endpoint_wrapper
 def device_authorization(
     request: Request,
     client_id: UUID = Form(...),
@@ -471,7 +472,7 @@ def device_authorization(
     API_TOKEN,
     response_model=str,
 )
-@handle_exceptions
+@async_fastapi_endpoint_wrapper
 def api_token(
     token_type: APITokenType = APITokenType.GENERIC,
     expires_in: Optional[int] = None,