zenml-nightly 0.84.0.dev20250729__py3-none-any.whl → 0.84.0.dev20250730__py3-none-any.whl

zenml/VERSION

@@ -1 +1 @@
-0.84.0.dev20250729
+0.84.0.dev20250730

zenml/client.py

@@ -7339,19 +7339,24 @@ class Client(metaclass=ClientMetaClass):
     def create_service_account(
         self,
         name: str,
+        full_name: Optional[str] = None,
         description: str = "",
     ) -> ServiceAccountResponse:
         """Create a new service account.
 
         Args:
             name: The name of the service account.
+            full_name: The display name of the service account.
             description: The description of the service account.
 
         Returns:
             The created service account.
         """
         service_account = ServiceAccountRequest(
-            name=name, description=description, active=True
+            name=name,
+            full_name=full_name or "",
+            description=description,
+            active=True,
         )
         created_service_account = self.zen_store.create_service_account(
             service_account=service_account

zenml/constants.py

@@ -174,6 +174,12 @@ ENV_ZENML_ENABLE_IMPLICIT_AUTH_METHODS = "ZENML_ENABLE_IMPLICIT_AUTH_METHODS"
 ENV_ZENML_DISABLE_PIPELINE_LOGS_STORAGE = "ZENML_DISABLE_PIPELINE_LOGS_STORAGE"
 ENV_ZENML_DISABLE_STEP_LOGS_STORAGE = "ZENML_DISABLE_STEP_LOGS_STORAGE"
 ENV_ZENML_DISABLE_STEP_NAMES_IN_LOGS = "ZENML_DISABLE_STEP_NAMES_IN_LOGS"
+ENV_ZENML_LOGS_STORAGE_MAX_QUEUE_SIZE = "ZENML_LOGS_STORAGE_MAX_QUEUE_SIZE"
+ENV_ZENML_LOGS_STORAGE_QUEUE_TIMEOUT = "ZENML_LOGS_STORAGE_QUEUE_TIMEOUT"
+
+ENV_ZENML_LOGS_WRITE_INTERVAL_SECONDS = "ZENML_LOGS_WRITE_INTERVAL_SECONDS"
+ENV_ZENML_LOGS_MERGE_INTERVAL_SECONDS = "ZENML_LOGS_MERGE_INTERVAL_SECONDS"
+
 ENV_ZENML_CUSTOM_SOURCE_ROOT = "ZENML_CUSTOM_SOURCE_ROOT"
 ENV_ZENML_PIPELINE_RUN_API_TOKEN_EXPIRATION = (
     "ZENML_PIPELINE_API_TOKEN_EXPIRATION"
@@ -504,3 +510,22 @@ STACK_DEPLOYMENT_API_TOKEN_EXPIRATION = 60 * 6  # 6 hours
 ZENML_PIPELINE_RUN_API_TOKEN_EXPIRATION = handle_int_env_var(
     ENV_ZENML_PIPELINE_RUN_API_TOKEN_EXPIRATION, default=0
 )
+
+# Logs storage constants
+# Maximum number of log batches to queue.
+LOGS_STORAGE_MAX_QUEUE_SIZE = handle_int_env_var(
+    ENV_ZENML_LOGS_STORAGE_MAX_QUEUE_SIZE, default=100000
+)
+# Queue timeout controls log dropping vs application blocking:
+# - Positive value (e.g., 30): Wait N seconds, then drop logs if queue full
+# - Negative value (e.g., -1): Block indefinitely until queue has space (never drop logs)
+LOGS_STORAGE_QUEUE_TIMEOUT = handle_int_env_var(
+    ENV_ZENML_LOGS_STORAGE_QUEUE_TIMEOUT, default=-1
+)
+# Log batch-write and merge windows in seconds.
+LOGS_WRITE_INTERVAL_SECONDS = handle_int_env_var(
+    ENV_ZENML_LOGS_WRITE_INTERVAL_SECONDS, default=2
+)
+LOGS_MERGE_INTERVAL_SECONDS = handle_int_env_var(
+    ENV_ZENML_LOGS_MERGE_INTERVAL_SECONDS, default=10 * 60
+)

zenml/integrations/gcp/__init__.py

@@ -43,9 +43,12 @@ class GcpIntegration(Integration):
     """Definition of Google Cloud Platform integration for ZenML."""
 
     NAME = GCP
+    # Adding the gcsfs<=2024.12.0 for now to solve the issue with the 
+    # increased number of list calls. This is a temporary fix and should be
+    # removed once the issue is resolved.
     REQUIREMENTS = [
         "kfp>=2.6.0",
-        "gcsfs!=2025.5.0,!=2025.5.0.post1",
+        "gcsfs!=2025.5.0,!=2025.5.0.post1,<=2024.12.0",
         "google-cloud-secret-manager",
         "google-cloud-container>=2.21.0",
         "google-cloud-artifact-registry>=1.11.3",

zenml/logging/__init__.py

@@ -11,14 +11,3 @@
 #  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.
-
-"""Logging utilities."""
-
-# How many seconds to wait before uploading logs to the artifact store
-STEP_LOGS_STORAGE_INTERVAL_SECONDS: int = 15
-
-# How many messages to buffer before uploading logs to the artifact store
-STEP_LOGS_STORAGE_MAX_MESSAGES: int = 100
-
-# How often to merge logs into a single file
-STEP_LOGS_STORAGE_MERGE_INTERVAL_SECONDS: int = 10 * 60

zenml/logging/step_logging.py

@@ -13,10 +13,13 @@
 #  permissions and limitations under the License.
 """ZenML logging handler."""
 
+import asyncio
 import logging
 import os
+import queue
 import re
 import sys
+import threading
 import time
 from contextlib import nullcontext
 from contextvars import ContextVar
@@ -35,15 +38,14 @@ from zenml.client import Client
 from zenml.constants import (
     ENV_ZENML_DISABLE_PIPELINE_LOGS_STORAGE,
     ENV_ZENML_DISABLE_STEP_NAMES_IN_LOGS,
+    LOGS_MERGE_INTERVAL_SECONDS,
+    LOGS_STORAGE_MAX_QUEUE_SIZE,
+    LOGS_STORAGE_QUEUE_TIMEOUT,
+    LOGS_WRITE_INTERVAL_SECONDS,
     handle_bool_env_var,
 )
 from zenml.exceptions import DoesNotExistException
 from zenml.logger import get_logger
-from zenml.logging import (
-    STEP_LOGS_STORAGE_INTERVAL_SECONDS,
-    STEP_LOGS_STORAGE_MAX_MESSAGES,
-    STEP_LOGS_STORAGE_MERGE_INTERVAL_SECONDS,
-)
 from zenml.models import (
     LogsRequest,
     PipelineDeploymentResponse,
@@ -58,6 +60,8 @@ logger = get_logger(__name__)
 
 redirected: ContextVar[bool] = ContextVar("redirected", default=False)
 
+ansi_escape = re.compile(r"\x1B(?:[@-Z\\-_]|\[[0-?]*[ -/]*[@-~])")
+
 LOGS_EXTENSION = ".log"
 PIPELINE_RUN_LOGS_FOLDER = "pipeline_runs"
 
@@ -71,7 +75,6 @@ def remove_ansi_escape_codes(text: str) -> str:
     Returns:
         the version of the input string where the escape codes are removed.
     """
-    ansi_escape = re.compile(r"\x1B(?:[@-Z\\-_]|\[[0-?]*[ -/]*[@-~])")
     return ansi_escape.sub("", text)
 
 
@@ -232,93 +235,206 @@ def fetch_logs(
 
 
 class PipelineLogsStorage:
-    """Helper class which buffers and stores logs to a given URI."""
+    """Helper class which buffers and stores logs to a given URI using a background thread."""
 
     def __init__(
         self,
         logs_uri: str,
         artifact_store: "BaseArtifactStore",
-        max_messages: int = STEP_LOGS_STORAGE_MAX_MESSAGES,
-        time_interval: int = STEP_LOGS_STORAGE_INTERVAL_SECONDS,
-        merge_files_interval: int = STEP_LOGS_STORAGE_MERGE_INTERVAL_SECONDS,
+        max_queue_size: int = LOGS_STORAGE_MAX_QUEUE_SIZE,
+        queue_timeout: int = LOGS_STORAGE_QUEUE_TIMEOUT,
+        write_interval: int = LOGS_WRITE_INTERVAL_SECONDS,
+        merge_files_interval: int = LOGS_MERGE_INTERVAL_SECONDS,
     ) -> None:
         """Initialization.
 
         Args:
             logs_uri: the URI of the log file or folder.
             artifact_store: Artifact Store from the current step context
-            max_messages: the maximum number of messages to save in the buffer.
-            time_interval: the amount of seconds before the buffer gets saved
-                automatically.
+            max_queue_size: maximum number of individual messages to queue.
+            queue_timeout: timeout in seconds for putting items in queue when full.
+                - Positive value: Wait N seconds, then drop logs if queue still full
+                - Negative value: Block indefinitely until queue has space (never drop logs)
+            write_interval: the amount of seconds before the created files
+                get written to the artifact store.
             merge_files_interval: the amount of seconds before the created files
                 get merged into a single file.
         """
         # Parameters
         self.logs_uri = logs_uri
-        self.max_messages = max_messages
-        self.time_interval = time_interval
+        self.max_queue_size = max_queue_size
+        self.queue_timeout = queue_timeout
+        self.write_interval = write_interval
         self.merge_files_interval = merge_files_interval
 
         # State
-        self.buffer: List[str] = []
-        self.disabled_buffer: List[str] = []
-        self.last_save_time = time.time()
-        self.disabled = False
         self.artifact_store = artifact_store
 
         # Immutable filesystems state
         self.last_merge_time = time.time()
 
+        # Queue and log storage thread for async processing
+        self.log_queue: queue.Queue[str] = queue.Queue(maxsize=max_queue_size)
+        self.log_storage_thread: Optional[threading.Thread] = None
+        self.shutdown_event = threading.Event()
+        self.merge_event = threading.Event()
+
+        # Start the log storage thread
+        self._start_log_storage_thread()
+
+    def _start_log_storage_thread(self) -> None:
+        """Start the log storage thread for processing log queue."""
+        if (
+            self.log_storage_thread is None
+            or not self.log_storage_thread.is_alive()
+        ):
+            self.log_storage_thread = threading.Thread(
+                target=self._log_storage_worker,
+                name="LogsStorage-Worker",
+            )
+            self.log_storage_thread.start()
+
+    def _process_log_queue(self, force_merge: bool = False) -> None:
+        """Write and merge logs to the artifact store using time-based batching.
+
+        Args:
+            force_merge: Whether to force merge the logs.
+        """
+        try:
+            messages = []
+
+            # Get first message (blocking with timeout)
+            try:
+                first_message = self.log_queue.get(timeout=1)
+                messages.append(first_message)
+            except queue.Empty:
+                return
+
+            # Get any remaining messages without waiting (drain quickly)
+            while True:
+                try:
+                    additional_message = self.log_queue.get_nowait()
+                    messages.append(additional_message)
+                except queue.Empty:
+                    break
+
+            # Write the messages to the artifact store
+            if messages:
+                self.write_buffer(messages)
+
+            # Merge the log files if needed
+            if (
+                self._is_merge_needed
+                or self.merge_event.is_set()
+                or force_merge
+            ):
+                self.merge_event.clear()
+
+                self.merge_log_files(merge_all_files=force_merge)
+
+        except Exception as e:
+            logger.error("Error in log storage thread: %s", e)
+        finally:
+            # Always mark all queue tasks as done
+            for _ in messages:
+                self.log_queue.task_done()
+
+            time.sleep(self.write_interval)
+
+    def _log_storage_worker(self) -> None:
+        """Log storage thread worker that processes the log queue."""
+        # Process the log queue until shutdown is requested
+        while not self.shutdown_event.is_set():
+            self._process_log_queue()
+
+        # Shutdown requested - drain remaining queue items and merge log files
+        self._process_log_queue(force_merge=True)
+
+    def _shutdown_log_storage_thread(self, timeout: int = 5) -> None:
+        """Shutdown the log storage thread gracefully.
+
+        Args:
+            timeout: Maximum time to wait for thread shutdown.
+        """
+        if self.log_storage_thread and self.log_storage_thread.is_alive():
+            # Then signal the worker to begin graceful shutdown
+            self.shutdown_event.set()
+
+            # Wait for thread to finish (it will drain the queue automatically)
+            self.log_storage_thread.join(timeout=timeout)
+
     def write(self, text: str) -> None:
-        """Main write method.
+        """Main write method that sends individual messages directly to queue.
 
         Args:
             text: the incoming string.
         """
+        # Skip empty lines
         if text == "\n":
             return
 
-        if not self.disabled:
-            # Add timestamp to the message when it's received
+        # If the current thread is the log storage thread, do nothing
+        # to prevent recursion when the storage thread itself generates logs
+        if (
+            self.log_storage_thread
+            and threading.current_thread() == self.log_storage_thread
+        ):
+            return
+
+        # If the current thread is the fsspec IO thread, do nothing
+        if self._is_fsspec_io_thread:
+            return
+
+        try:
+            # Format the message with timestamp
             timestamp = utc_now().strftime("%Y-%m-%d %H:%M:%S")
             formatted_message = (
                 f"[{timestamp} UTC] {remove_ansi_escape_codes(text)}"
             )
-            self.buffer.append(formatted_message.rstrip())
-            self.save_to_file()
+            formatted_message = formatted_message.rstrip()
+
+            # Send individual message directly to queue
+            if not self.shutdown_event.is_set():
+                try:
+                    if self.queue_timeout < 0:
+                        # Negative timeout = block indefinitely until queue has space
+                        # Guarantees no log loss but may hang application
+                        self.log_queue.put(formatted_message)
+                    else:
+                        # Positive timeout = wait specified time then drop logs
+                        # Prevents application hanging but may lose logs
+                        self.log_queue.put(
+                            formatted_message, timeout=self.queue_timeout
+                        )
+                except queue.Full:
+                    # This only happens with positive timeout
+                    # Queue is full - just skip this message to avoid blocking
+                    # Better to drop logs than hang the application
+                    pass
+
+        except Exception:
+            # Silently ignore errors to prevent recursion
+            pass
 
     @property
-    def _is_write_needed(self) -> bool:
-        """Checks whether the buffer needs to be written to disk.
+    def _is_merge_needed(self) -> bool:
+        """Checks whether the log files need to be merged.
 
         Returns:
-            whether the buffer needs to be written to disk.
+            whether the log files need to be merged.
         """
         return (
-            len(self.buffer) >= self.max_messages
-            or time.time() - self.last_save_time >= self.time_interval
+            self.artifact_store.config.IS_IMMUTABLE_FILESYSTEM
+            and time.time() - self.last_merge_time > self.merge_files_interval
         )
 
-    def _get_timestamped_filename(self, suffix: str = "") -> str:
-        """Returns a timestamped filename.
-
-        Args:
-            suffix: optional suffix for the file name
+    @property
+    def _is_fsspec_io_thread(self) -> bool:
+        """Checks if the current thread is the fsspec IO thread.
 
         Returns:
-            The timestamped filename.
+            whether the current thread is the fsspec IO thread.
         """
-        return f"{time.time()}{suffix}{LOGS_EXTENSION}"
-
-    def save_to_file(self, force: bool = False) -> None:
-        """Method to save the buffer to the given URI.
-
-        Args:
-            force: whether to force a save even if the write conditions not met.
-        """
-        import asyncio
-        import threading
-
         # Most artifact stores are based on fsspec, which converts between
         # sync and async operations by using a separate AIO thread.
         # It may happen that the fsspec call itself will log something,
@@ -329,82 +445,79 @@ class PipelineLogsStorage:
         # To avoid this, we simply check if we're running in the fsspec AIO
         # thread and skip the save if that's the case.
         try:
-            if (
+            return (
                 asyncio.events.get_running_loop() is not None
                 and threading.current_thread().name == "fsspecIO"
-            ):
-                return
+            )
         except RuntimeError:
             # No running loop
-            pass
+            return False
 
-        if not self.disabled and (self._is_write_needed or force):
-            # IMPORTANT: keep this as the first code line in this method! The
-            # code that follows might still emit logging messages, which will
-            # end up triggering this method again, causing an infinite loop.
-            self.disabled = True
+    def _get_timestamped_filename(self, suffix: str = "") -> str:
+        """Returns a timestamped filename.
 
-            try:
-                # The configured logging handler uses a lock to ensure that
-                # logs generated by different threads are not interleaved.
-                # Given that most artifact stores are based on fsspec, which
-                # use a separate thread for async operations, it may happen that
-                # the fsspec library itself will log something, which will end
-                # up in a deadlock.
-                # To avoid this, we temporarily disable the lock in the logging
-                # handler while writing to the file.
-                logging_handler = logging.getLogger().handlers[0]
-                logging_lock = logging_handler.lock
-                logging_handler.lock = None
-
-                if self.buffer:
-                    if self.artifact_store.config.IS_IMMUTABLE_FILESYSTEM:
-                        _logs_uri = self._get_timestamped_filename()
-                        with self.artifact_store.open(
-                            os.path.join(
-                                self.logs_uri,
-                                _logs_uri,
-                            ),
-                            "w",
-                        ) as file:
-                            for message in self.buffer:
-                                file.write(f"{message}\n")
-                    else:
-                        with self.artifact_store.open(
-                            self.logs_uri, "a"
-                        ) as file:
-                            for message in self.buffer:
-                                file.write(f"{message}\n")
-                        self.artifact_store._remove_previous_file_versions(
-                            self.logs_uri
-                        )
+        Args:
+            suffix: optional suffix for the file name
 
-            except (OSError, IOError) as e:
-                # This exception can be raised if there are issues with the
-                # underlying system calls, such as reaching the maximum number
-                # of open files, permission issues, file corruption, or other
-                # I/O errors.
-                logger.error(f"Error while trying to write logs: {e}")
-            finally:
-                # Restore the original logging handler lock
-                logging_handler.lock = logging_lock
+        Returns:
+            The timestamped filename.
+        """
+        return f"{time.time()}{suffix}{LOGS_EXTENSION}"
 
-                self.buffer = []
-                self.last_save_time = time.time()
+    def write_buffer(self, buffer_to_write: List[str]) -> None:
+        """Write the given buffer to file. This runs in the log storage thread.
 
-                self.disabled = False
-        # merge created files on a given interval (defaults to 10 minutes)
-        # only runs on Immutable Filesystems
-        if (
-            self.artifact_store.config.IS_IMMUTABLE_FILESYSTEM
-            and time.time() - self.last_merge_time > self.merge_files_interval
-        ):
-            try:
-                self.merge_log_files()
-            except (OSError, IOError) as e:
-                logger.error(f"Error while trying to roll up logs: {e}")
-            finally:
-                self.last_merge_time = time.time()
+        Args:
+            buffer_to_write: The buffer contents to write to file.
+        """
+        if not buffer_to_write:
+            return
+
+        # The configured logging handler uses a lock to ensure that
+        # logs generated by different threads are not interleaved.
+        # Given that most artifact stores are based on fsspec, which
+        # use a separate thread for async operations, it may happen that
+        # the fsspec library itself will log something, which will end
+        # up in a deadlock.
+        # To avoid this, we temporarily disable the lock in the logging
+        # handler while writing to the file.
+        logging_handler = None
+        logging_lock = None
+        try:
+            # Only try to access logging handler if it exists
+            root_logger = logging.getLogger()
+            if root_logger.handlers:
+                logging_handler = root_logger.handlers[0]
+                logging_lock = getattr(logging_handler, "lock", None)
+                if logging_lock:
+                    logging_handler.lock = None
+
+            # If the artifact store is immutable, write the buffer to a new file
+            if self.artifact_store.config.IS_IMMUTABLE_FILESYSTEM:
+                _logs_uri = self._get_timestamped_filename()
+                with self.artifact_store.open(
+                    os.path.join(
+                        self.logs_uri,
+                        _logs_uri,
+                    ),
+                    "w",
+                ) as file:
+                    for message in buffer_to_write:
+                        file.write(f"{message}\n")
+
+            # If the artifact store is mutable, append the buffer to the existing file
+            else:
+                with self.artifact_store.open(self.logs_uri, "a") as file:
+                    for message in buffer_to_write:
+                        file.write(f"{message}\n")
+                self.artifact_store._remove_previous_file_versions(
+                    self.logs_uri
+                )
+
+        finally:
+            # Re-enable the logging lock
+            if logging_handler and logging_lock:
+                logging_handler.lock = logging_lock
 
     def merge_log_files(self, merge_all_files: bool = False) -> None:
         """Merges all log files into one in the given URI.
@@ -414,12 +527,15 @@ class PipelineLogsStorage:
         Args:
             merge_all_files: whether to merge all files or only raw files
         """
+        # If the artifact store is immutable, merge the log files
         if self.artifact_store.config.IS_IMMUTABLE_FILESYSTEM:
             merged_file_suffix = "_merged"
             files_ = self.artifact_store.listdir(self.logs_uri)
             if not merge_all_files:
                 # already merged files will not be merged again
-                files_ = [f for f in files_ if merged_file_suffix not in f]
+                files_ = [
+                    f for f in files_ if merged_file_suffix not in str(f)
+                ]
             file_name_ = self._get_timestamped_filename(
                 suffix=merged_file_suffix
             )
@@ -453,6 +569,13 @@ class PipelineLogsStorage:
                             os.path.join(self.logs_uri, str(file))
                         )
 
+            # Update the last merge time
+            self.last_merge_time = time.time()
+
+    def send_merge_event(self) -> None:
+        """Send a merge event to the log storage thread."""
+        self.merge_event.set()
+
 
 class PipelineLogsStorageContext:
     """Context manager which patches stdout and stderr during pipeline run execution."""
@@ -474,6 +597,7 @@ class PipelineLogsStorageContext:
             logs_uri=logs_uri, artifact_store=artifact_store
         )
         self.prepend_step_name = prepend_step_name
+        self._original_methods_saved = False
 
     def __enter__(self) -> "PipelineLogsStorageContext":
         """Enter condition of the context manager.
@@ -484,17 +608,16 @@ class PipelineLogsStorageContext:
         Returns:
             self
         """
+        # Save original write methods (flush methods are not wrapped)
         self.stdout_write = getattr(sys.stdout, "write")
-        self.stdout_flush = getattr(sys.stdout, "flush")
-
         self.stderr_write = getattr(sys.stderr, "write")
-        self.stderr_flush = getattr(sys.stderr, "flush")
+        self._original_methods_saved = True
 
+        # Wrap stdout/stderr write methods only
+        # Note: We don't wrap flush() as it's unnecessary overhead and the
+        # background thread handles log flushing based on time/buffer size
         setattr(sys.stdout, "write", self._wrap_write(self.stdout_write))
-        setattr(sys.stdout, "flush", self._wrap_flush(self.stdout_flush))
-
         setattr(sys.stderr, "write", self._wrap_write(self.stderr_write))
-        setattr(sys.stderr, "flush", self._wrap_flush(self.stderr_flush))
 
         redirected.set(True)
         return self
@@ -514,20 +637,20 @@ class PipelineLogsStorageContext:
 
         Restores the `write` method of both stderr and stdout.
         """
-        self.storage.save_to_file(force=True)
-
-        setattr(sys.stdout, "write", self.stdout_write)
-        setattr(sys.stdout, "flush", self.stdout_flush)
-
-        setattr(sys.stderr, "write", self.stderr_write)
-        setattr(sys.stderr, "flush", self.stderr_flush)
-
-        redirected.set(False)
+        # Step 1: Restore stdout/stderr FIRST to prevent logging during shutdown
+        try:
+            if self._original_methods_saved:
+                setattr(sys.stdout, "write", self.stdout_write)
+                setattr(sys.stderr, "write", self.stderr_write)
+                redirected.set(False)
+        except Exception:
+            pass
 
+        # Step 2: Shutdown thread (it will automatically drain queue and merge files)
         try:
-            self.storage.merge_log_files(merge_all_files=True)
-        except (OSError, IOError) as e:
-            logger.warning(f"Step logs roll-up failed: {e}")
+            self.storage._shutdown_log_storage_thread()
+        except Exception:
+            pass
 
     def _wrap_write(self, method: Callable[..., Any]) -> Callable[..., Any]:
         """Wrapper function that utilizes the storage object to store logs.
@@ -577,23 +700,6 @@ class PipelineLogsStorageContext:
 
         return wrapped_write
 
-    def _wrap_flush(self, method: Callable[..., Any]) -> Callable[..., Any]:
-        """Wrapper function that flushes the buffer of the storage object.
-
-        Args:
-            method: the original flush method
-
-        Returns:
-            the wrapped flush method.
-        """
-
-        def wrapped_flush(*args: Any, **kwargs: Any) -> Any:
-            output = method(*args, **kwargs)
-            self.storage.save_to_file()
-            return output
-
-        return wrapped_flush
-
 
 def setup_orchestrator_logging(
     run_id: str, deployment: "PipelineDeploymentResponse"

zenml/models/__init__.py

@@ -283,6 +283,8 @@ from zenml.models.v2.core.service_account import (
     ServiceAccountResponseMetadata,
     ServiceAccountUpdate,
     ServiceAccountRequest,
+    ServiceAccountInternalRequest,
+    ServiceAccountInternalUpdate,
     ServiceAccountResponse,
 )
 from zenml.models.v2.core.service_connector import (
@@ -740,6 +742,8 @@ __all__ = [
     "ServiceAccountResponseMetadata",
     "ServiceAccountUpdate",
     "ServiceAccountRequest",
+    "ServiceAccountInternalRequest",
+    "ServiceAccountInternalUpdate",
     "ServiceAccountResponse",
     "ServiceConnectorConfiguration",
     "ServiceConnectorRequest",