zenml-nightly 0.58.1.dev20240608__tar.gz → 0.58.2.dev20240611__tar.gz

{zenml_nightly-0.58.1.dev20240608 → zenml_nightly-0.58.2.dev20240611}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: zenml-nightly
-Version: 0.58.1.dev20240608
+Version: 0.58.2.dev20240611
 Summary: ZenML: Write production-ready ML code.
 Home-page: https://zenml.io
 License: Apache-2.0
@@ -426,7 +426,7 @@ the Apache License Version 2.0.
     <a href="https://github.com/zenml-io/zenml-projects">Projects Showcase</a>
     <br />
     <br />
-    🎉 Version 0.58.1 is out. Check out the release notes
+    🎉 Version 0.58.2 is out. Check out the release notes
     <a href="https://github.com/zenml-io/zenml/releases">here</a>.
     <br />
     🖥️ Download our VS Code Extension <a href="https://marketplace.visualstudio.com/items?itemName=ZenML.zenml-vscode">here</a>.

{zenml_nightly-0.58.1.dev20240608 → zenml_nightly-0.58.2.dev20240611}/README.md

@@ -289,7 +289,7 @@ the Apache License Version 2.0.
     <a href="https://github.com/zenml-io/zenml-projects">Projects Showcase</a>
     <br />
     <br />
-    🎉 Version 0.58.1 is out. Check out the release notes
+    🎉 Version 0.58.2 is out. Check out the release notes
     <a href="https://github.com/zenml-io/zenml/releases">here</a>.
     <br />
     🖥️ Download our VS Code Extension <a href="https://marketplace.visualstudio.com/items?itemName=ZenML.zenml-vscode">here</a>.

{zenml_nightly-0.58.1.dev20240608 → zenml_nightly-0.58.2.dev20240611}/RELEASE_NOTES.md

@@ -1,5 +1,29 @@
 <!-- markdown-link-check-disable -->
 
+# 0.58.2
+
+The 0.58.2 minor release is packed with a set of improvements to the ZenML logging and ZenML Server.
+
+With this release ZenML logging will:
+- Offer pagination of the logs during fetching via REST API
+- Store the full logs history on GCS Artifact Stores
+- Be performant running logging-heavy tasks, like TQDM logging or logging of training in any Deep Learning framework (also TQDM-backed)
+
+## What's Changed
+* update test-migrations.sh with latest versions by @safoinme in https://github.com/zenml-io/zenml/pull/2757
+* Fix overriding expiration date for api tokens by @schustmi in https://github.com/zenml-io/zenml/pull/2753
+* Step logs pagination by @schustmi in https://github.com/zenml-io/zenml/pull/2731
+* Fix broken links (round 2) by @strickvl in https://github.com/zenml-io/zenml/pull/2760
+* Remove default system flag in docker UV by @avishniakov in https://github.com/zenml-io/zenml/pull/2764
+* Another batch of small fixes and expansions by @AlexejPenner in https://github.com/zenml-io/zenml/pull/2762
+* Server scalability improvements by @stefannica in https://github.com/zenml-io/zenml/pull/2752
+* Add option to start parallel kubernetes steps with delay by @schustmi in https://github.com/zenml-io/zenml/pull/2758
+* Move `thread_limiter` to app startup event by @avishniakov in https://github.com/zenml-io/zenml/pull/2765
+* Logging performance improvements and GCP logging fix by @avishniakov in https://github.com/zenml-io/zenml/pull/2755
+
+
+**Full Changelog**: https://github.com/zenml-io/zenml/compare/0.58.1...0.58.2
+
 # 0.58.1
 
 The 0.58.1 release brings a set of minor enhancement and bugfix to the ZenML framework, such as the ability to delete all versions of a pipeline using the Client/CLI, providing greater flexibility and control over pipeline management. Users can now specify Python package installer arguments. Furthermore, a fix has been implemented for the Sentencepiece tokenizer materializer.

{zenml_nightly-0.58.1.dev20240608 → zenml_nightly-0.58.2.dev20240611}/pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "zenml-nightly"
-version = "0.58.1.dev20240608"
+version = "0.58.2.dev20240611"
 packages = [{ include = "zenml", from = "src" }]
 description = "ZenML: Write production-ready ML code."
 authors = ["ZenML GmbH <info@zenml.io>"]

zenml_nightly-0.58.2.dev20240611/src/zenml/VERSION

@@ -0,0 +1 @@
+0.58.2.dev20240611

{zenml_nightly-0.58.1.dev20240608 → zenml_nightly-0.58.2.dev20240611}/src/zenml/artifact_stores/base_artifact_store.py

@@ -171,6 +171,7 @@ class BaseArtifactStoreConfig(StackComponentConfig):
     path: str
 
     SUPPORTED_SCHEMES: ClassVar[Set[str]]
+    IS_IMMUTABLE_FILESYSTEM: ClassVar[bool] = False
 
     @root_validator(skip_on_failure=True)
     def _ensure_artifact_store(cls, values: Dict[str, Any]) -> Any:

{zenml_nightly-0.58.1.dev20240608 → zenml_nightly-0.58.2.dev20240611}/src/zenml/config/server_config.py

@@ -39,6 +39,7 @@ from zenml.constants import (
     DEFAULT_ZENML_SERVER_SECURE_HEADERS_REFERRER,
     DEFAULT_ZENML_SERVER_SECURE_HEADERS_XFO,
     DEFAULT_ZENML_SERVER_SECURE_HEADERS_XXP,
+    DEFAULT_ZENML_SERVER_THREAD_POOL_SIZE,
     DEFAULT_ZENML_SERVER_USE_LEGACY_DASHBOARD,
     ENV_ZENML_SERVER_PREFIX,
 )
@@ -301,6 +302,8 @@ class ServerConfiguration(BaseModel):
     display_updates: bool = True
     auto_activate: bool = False
 
+    thread_pool_size: int = DEFAULT_ZENML_SERVER_THREAD_POOL_SIZE
+
     _deployment_id: Optional[UUID] = None
 
     @root_validator(pre=True)

{zenml_nightly-0.58.1.dev20240608 → zenml_nightly-0.58.2.dev20240611}/src/zenml/constants.py

@@ -230,6 +230,7 @@ STEP_SOURCE_PARAMETER_NAME = "step_source"
 
 # Server settings
 DEFAULT_ZENML_SERVER_NAME = "default"
+DEFAULT_ZENML_SERVER_THREAD_POOL_SIZE = 40
 DEFAULT_ZENML_JWT_TOKEN_LEEWAY = 10
 DEFAULT_ZENML_JWT_TOKEN_ALGORITHM = "HS256"
 DEFAULT_ZENML_AUTH_SCHEME = AuthScheme.OAUTH2_PASSWORD_BEARER

{zenml_nightly-0.58.1.dev20240608 → zenml_nightly-0.58.2.dev20240611}/src/zenml/integrations/airflow/__init__.py

@@ -17,7 +17,7 @@ The Airflow integration sub-module powers an alternative to the local
 orchestrator. You can enable it by registering the Airflow orchestrator with
 the CLI tool, then bootstrap using the ``zenml orchestrator up`` command.
 """
-from typing import List, Optional, Type
+from typing import List, Type
 
 from zenml.integrations.constants import AIRFLOW
 from zenml.integrations.integration import Integration

{zenml_nightly-0.58.1.dev20240608 → zenml_nightly-0.58.2.dev20240611}/src/zenml/integrations/gcp/artifact_stores/gcp_artifact_store.py

@@ -35,9 +35,11 @@ from zenml.integrations.gcp.flavors.gcp_artifact_store_flavor import (
     GCPArtifactStoreConfig,
 )
 from zenml.io.fileio import convert_to_str
+from zenml.logger import get_logger
 from zenml.secret.schemas import GCPSecretSchema
 from zenml.stack.authentication_mixin import AuthenticationMixin
 
+logger = get_logger(__name__)
 PathType = Union[bytes, str]
 
 
@@ -109,6 +111,10 @@ class GCPArtifactStore(BaseArtifactStore, AuthenticationMixin):
         Returns:
             A file-like object that can be used to read or write to the file.
         """
+        if mode in ("a", "ab"):
+            logger.warning(
+                "GCS Filesystem is immutable, so append mode will overwrite existing files."
+            )
         return self.filesystem.open(path=path, mode=mode)
 
     def copyfile(

{zenml_nightly-0.58.1.dev20240608 → zenml_nightly-0.58.2.dev20240611}/src/zenml/integrations/gcp/flavors/gcp_artifact_store_flavor.py

@@ -36,6 +36,7 @@ class GCPArtifactStoreConfig(
     """Configuration for GCP Artifact Store."""
 
     SUPPORTED_SCHEMES: ClassVar[Set[str]] = {GCP_PATH_PREFIX}
+    IS_IMMUTABLE_FILESYSTEM: ClassVar[bool] = True
 
 
 class GCPArtifactStoreFlavor(BaseArtifactStoreFlavor):

{zenml_nightly-0.58.1.dev20240608 → zenml_nightly-0.58.2.dev20240611}/src/zenml/integrations/kubernetes/flavors/kubernetes_orchestrator_flavor.py

@@ -82,6 +82,9 @@ class KubernetesOrchestratorConfig(  # type: ignore[misc] # https://github.com/p
             containers).
         skip_local_validations: If `True`, the local validations will be
             skipped.
+        parallel_step_startup_waiting_period: How long to wait in between
+            starting parallel steps. This can be used to distribute server
+            load when running pipelines with a huge amount of parallel steps.
     """
 
     incluster: bool = False
@@ -89,6 +92,7 @@ class KubernetesOrchestratorConfig(  # type: ignore[misc] # https://github.com/p
     kubernetes_namespace: str = "zenml"
     local: bool = False
     skip_local_validations: bool = False
+    parallel_step_startup_waiting_period: Optional[float] = None
 
     @property
     def is_remote(self) -> bool:

{zenml_nightly-0.58.1.dev20240608 → zenml_nightly-0.58.2.dev20240611}/src/zenml/integrations/kubernetes/orchestrators/kubernetes_orchestrator_entrypoint.py

@@ -142,7 +142,14 @@ def main() -> None:
         )
         logger.info(f"Pod of step `{step_name}` completed.")
 
-    ThreadedDagRunner(dag=pipeline_dag, run_fn=run_step_on_kubernetes).run()
+    parallel_node_startup_waiting_period = (
+        orchestrator.config.parallel_step_startup_waiting_period or 0.0
+    )
+    ThreadedDagRunner(
+        dag=pipeline_dag,
+        run_fn=run_step_on_kubernetes,
+        parallel_node_startup_waiting_period=parallel_node_startup_waiting_period,
+    ).run()
 
     logger.info("Orchestration pod completed.")
 

{zenml_nightly-0.58.1.dev20240608 → zenml_nightly-0.58.2.dev20240611}/src/zenml/logging/__init__.py

@@ -12,8 +12,13 @@
 #  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_nightly-0.58.2.dev20240611/src/zenml/logging/step_logging.py

@@ -0,0 +1,487 @@
+#  Copyright (c) ZenML GmbH 2023. 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.
+"""ZenML logging handler."""
+
+import os
+import re
+import sys
+import time
+from contextvars import ContextVar
+from types import TracebackType
+from typing import Any, Callable, List, Optional, Type, Union
+from uuid import UUID, uuid4
+
+from zenml.artifact_stores import BaseArtifactStore
+from zenml.artifacts.utils import (
+    _load_artifact_store,
+    _load_file_from_artifact_store,
+)
+from zenml.client import Client
+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.zen_stores.base_zen_store import BaseZenStore
+
+# Get the logger
+logger = get_logger(__name__)
+
+redirected: ContextVar[bool] = ContextVar("redirected", default=False)
+
+LOGS_EXTENSION = ".log"
+
+
+def remove_ansi_escape_codes(text: str) -> str:
+    """Auxiliary function to remove ANSI escape codes from a given string.
+
+    Args:
+        text: the input string
+
+    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)
+
+
+def prepare_logs_uri(
+    artifact_store: "BaseArtifactStore",
+    step_name: str,
+    log_key: Optional[str] = None,
+) -> str:
+    """Generates and prepares a URI for the log file or folder for a step.
+
+    Args:
+        artifact_store: The artifact store on which the artifact will be stored.
+        step_name: Name of the step.
+        log_key: The unique identification key of the log file.
+
+    Returns:
+        The URI of the log storage (file or folder).
+    """
+    if log_key is None:
+        log_key = str(uuid4())
+
+    logs_base_uri = os.path.join(
+        artifact_store.path,
+        step_name,
+        "logs",
+    )
+
+    # Create the dir
+    if not artifact_store.exists(logs_base_uri):
+        artifact_store.makedirs(logs_base_uri)
+
+    # Delete the file if it already exists
+    if artifact_store.config.IS_IMMUTABLE_FILESYSTEM:
+        logs_uri_folder = os.path.join(logs_base_uri, log_key)
+        if artifact_store.exists(logs_uri_folder):
+            logger.warning(
+                f"Logs directory {logs_uri_folder} already exists! Removing old log directory..."
+            )
+            artifact_store.rmtree(logs_uri_folder)
+
+        artifact_store.makedirs(logs_uri_folder)
+        return logs_uri_folder
+    else:
+        logs_uri = os.path.join(logs_base_uri, f"{log_key}.log")
+        if artifact_store.exists(logs_uri):
+            logger.warning(
+                f"Logs file {logs_uri} already exists! Removing old log file..."
+            )
+            artifact_store.remove(logs_uri)
+        return logs_uri
+
+
+def fetch_logs(
+    zen_store: "BaseZenStore",
+    artifact_store_id: Union[str, UUID],
+    logs_uri: str,
+    offset: int = 0,
+    length: int = 1024 * 1024 * 16,  # Default to 16MiB of data
+) -> str:
+    """Fetches the logs from the artifact store.
+
+    Args:
+        zen_store: The store in which the artifact is stored.
+        artifact_store_id: The ID of the artifact store.
+        logs_uri: The URI of the artifact.
+        offset: The offset from which to start reading.
+        length: The amount of bytes that should be read.
+
+    Returns:
+        The logs as a string.
+
+    Raises:
+        DoesNotExistException: If the artifact does not exist in the artifact
+            store.
+    """
+
+    def _read_file(
+        uri: str, offset: int = 0, length: Optional[int] = None
+    ) -> str:
+        return str(
+            _load_file_from_artifact_store(
+                uri,
+                artifact_store=artifact_store,
+                mode="rb",
+                offset=offset,
+                length=length,
+            ).decode()
+        )
+
+    artifact_store = _load_artifact_store(artifact_store_id, zen_store)
+    if not artifact_store.isdir(logs_uri):
+        return _read_file(logs_uri, offset, length)
+    else:
+        files = artifact_store.listdir(logs_uri)
+        if len(files) == 1:
+            return _read_file(
+                os.path.join(logs_uri, str(files[0])), offset, length
+            )
+        else:
+            is_negative_offset = offset < 0
+            files.sort(reverse=is_negative_offset)
+
+            # search for the first file we need to read
+            latest_file_id = 0
+            for i, file in enumerate(files):
+                file_size: int = artifact_store.size(
+                    os.path.join(logs_uri, str(file))
+                )  # type: ignore[assignment]
+
+                if is_negative_offset:
+                    if file_size >= -offset:
+                        latest_file_id = -(i + 1)
+                        break
+                    else:
+                        offset += file_size
+                else:
+                    if file_size > offset:
+                        latest_file_id = i
+                        break
+                    else:
+                        offset -= file_size
+
+            # read the files according to pre-filtering
+            files.sort()
+            ret = []
+            for file in files[latest_file_id:]:
+                ret.append(
+                    _read_file(
+                        os.path.join(logs_uri, str(file)),
+                        offset,
+                        length,
+                    )
+                )
+                offset = 0
+                length -= len(ret[-1])
+                if length <= 0:
+                    # stop further reading, if the whole length is already read
+                    break
+
+            if not ret:
+                raise DoesNotExistException(
+                    f"Folder '{logs_uri}' is empty in artifact store "
+                    f"'{artifact_store.name}'."
+                )
+            return "".join(ret)
+
+
+class StepLogsStorage:
+    """Helper class which buffers and stores logs to a given URI."""
+
+    def __init__(
+        self,
+        logs_uri: str,
+        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,
+    ) -> None:
+        """Initialization.
+
+        Args:
+            logs_uri: the URI of the log file or folder.
+            max_messages: the maximum number of messages to save in the buffer.
+            time_interval: the amount of seconds before the buffer gets saved
+                automatically.
+            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.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: Optional["BaseArtifactStore"] = None
+
+        # Immutable filesystems state
+        self.last_merge_time = time.time()
+        self.log_files_not_merged: List[str] = []
+        self.next_merged_file_name: str = f"{time.time()}{LOGS_EXTENSION}"
+
+    @property
+    def artifact_store(self) -> "BaseArtifactStore":
+        """Returns the active artifact store.
+
+        Returns:
+            The active artifact store.
+        """
+        if self._artifact_store is None:
+            self._artifact_store = Client().active_stack.artifact_store
+        return self._artifact_store
+
+    def write(self, text: str) -> None:
+        """Main write method.
+
+        Args:
+            text: the incoming string.
+        """
+        if text == "\n":
+            return
+
+        if not self.disabled:
+            self.buffer.append(text)
+            self.save_to_file()
+
+    @property
+    def _is_write_needed(self) -> bool:
+        """Checks whether the buffer needs to be written to disk.
+
+        Returns:
+            whether the buffer needs to be written to disk.
+        """
+        return (
+            len(self.buffer) >= self.max_messages
+            or time.time() - self.last_save_time >= self.time_interval
+        )
+
+    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.
+        """
+        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
+
+            try:
+                if self.buffer:
+                    if self.artifact_store.config.IS_IMMUTABLE_FILESYSTEM:
+                        if not self.log_files_not_merged:
+                            self.next_merged_file_name = (
+                                f"{time.time()}{LOGS_EXTENSION}"
+                            )
+                        log_file_ = f"{time.time()}{LOGS_EXTENSION}"
+                        self.log_files_not_merged.append(log_file_)
+                        with self.artifact_store.open(
+                            os.path.join(
+                                self.logs_uri,
+                                log_file_,
+                            ),
+                            "w",
+                        ) as file:
+                            for message in self.buffer:
+                                file.write(
+                                    remove_ansi_escape_codes(message) + "\n"
+                                )
+                    else:
+                        with self.artifact_store.open(
+                            self.logs_uri, "a"
+                        ) as file:
+                            for message in self.buffer:
+                                file.write(
+                                    remove_ansi_escape_codes(message) + "\n"
+                                )
+
+            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:
+                self.buffer = []
+                self.last_save_time = time.time()
+
+                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(
+                    self.next_merged_file_name, self.log_files_not_merged
+                )
+            except (OSError, IOError) as e:
+                logger.error(f"Error while trying to roll up logs: {e}")
+            else:
+                self.log_files_not_merged = []
+            finally:
+                self.last_merge_time = time.time()
+
+    def merge_log_files(
+        self,
+        file_name: Optional[str] = None,
+        files: Optional[List[str]] = None,
+    ) -> None:
+        """Merges all log files into one in the given URI.
+
+        Called on the logging context exit.
+
+        Args:
+            file_name: The name of the merged log file.
+            files: The list of log files to merge.
+        """
+        if self.artifact_store.config.IS_IMMUTABLE_FILESYSTEM:
+            files_ = files or self.artifact_store.listdir(self.logs_uri)
+            file_name_ = file_name or f"full_log{LOGS_EXTENSION}"
+            if len(files_) > 1:
+                files_.sort()
+                logger.debug("Log files count: %s", len(files_))
+
+                try:
+                    # dump all logs to a local file first
+                    with self.artifact_store.open(
+                        os.path.join(self.logs_uri, file_name_), "w"
+                    ) as merged_file:
+                        for file in files_:
+                            merged_file.write(
+                                str(
+                                    _load_file_from_artifact_store(
+                                        os.path.join(self.logs_uri, str(file)),
+                                        artifact_store=self.artifact_store,
+                                        mode="r",
+                                    )
+                                )
+                            )
+                except Exception as e:
+                    logger.warning(f"Failed to merge log files. {e}")
+                else:
+                    # clean up left over files
+                    for file in files_:
+                        self.artifact_store.remove(
+                            os.path.join(self.logs_uri, str(file))
+                        )
+
+
+class StepLogsStorageContext:
+    """Context manager which patches stdout and stderr during step execution."""
+
+    def __init__(self, logs_uri: str) -> None:
+        """Initializes and prepares a storage object.
+
+        Args:
+            logs_uri: the URI of the logs file.
+        """
+        self.storage = StepLogsStorage(logs_uri=logs_uri)
+
+    def __enter__(self) -> "StepLogsStorageContext":
+        """Enter condition of the context manager.
+
+        Wraps the `write` method of both stderr and stdout, so each incoming
+        message gets stored in the step logs storage.
+
+        Returns:
+            self
+        """
+        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")
+
+        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.stdout_write))
+        setattr(sys.stderr, "flush", self._wrap_flush(self.stdout_flush))
+
+        redirected.set(True)
+        return self
+
+    def __exit__(
+        self,
+        exc_type: Optional[Type[BaseException]],
+        exc_val: Optional[BaseException],
+        exc_tb: Optional[TracebackType],
+    ) -> None:
+        """Exit condition of the context manager.
+
+        Args:
+            exc_type: The class of the exception
+            exc_val: The instance of the exception
+            exc_tb: The traceback of the exception
+
+        Restores the `write` method of both stderr and stdout.
+        """
+        self.storage.save_to_file(force=True)
+        self.storage.merge_log_files()
+
+        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)
+
+    def _wrap_write(self, method: Callable[..., Any]) -> Callable[..., Any]:
+        """Wrapper function that utilizes the storage object to store logs.
+
+        Args:
+            method: the original write method
+
+        Returns:
+            the wrapped write method.
+        """
+
+        def wrapped_write(*args: Any, **kwargs: Any) -> Any:
+            output = method(*args, **kwargs)
+            if args:
+                self.storage.write(args[0])
+            return output
+
+        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