Cleaner

Author: freider

modal/_output.py

@@ -70,6 +70,11 @@ class OutputManager(Protocol):
     Using a protocol allows code to work with any output manager without checking for None.
     """
 
+    @property
+    def is_enabled(self) -> bool:
+        """Whether rich output is enabled."""
+        ...
+
     @property
     def _stdout(self) -> Any:
         """The stdout stream for PTY shell output."""
@@ -162,6 +167,10 @@ class DisabledOutputManager:
     without checking if the output manager exists.
     """
 
+    @property
+    def is_enabled(self) -> bool:
+        return False
+
     @property
     def _stdout(self) -> Any:
         return sys.stdout

modal/_rich_output.py

@@ -165,6 +165,10 @@ class RichOutputManager:
     _show_timestamps: bool
     _object_tree: Tree | None
 
+    @property
+    def is_enabled(self) -> bool:
+        return True
+
     def __init__(
         self,
         *,

modal/image.py

@@ -58,7 +58,7 @@
 from .gpu import GPU_T, parse_gpu_config
 from .mount import _Mount, python_standalone_mount_name
 from .network_file_system import _NetworkFileSystem
-from .output import _get_output_manager, _is_output_enabled
+from .output import _get_output_manager
 from .secret import _Secret
 from .volume import _Volume
 
@@ -685,7 +685,7 @@ async def _load(self: _Image, resolver: Resolver, load_context: LoadContext, exi
                     raise RemoteError(f"Image build for {image_id} failed with the exception:\n{result.exception}")
                 else:
                     msg = f"Image build for {image_id} failed. See build logs for more details."
-                    if not _is_output_enabled():
+                    if not _get_output_manager().is_enabled:
                         msg += " (Hint: Use `modal.enable_output()` to see logs from the process building the Image.)"
                     raise RemoteError(msg)
             elif result.status == api_pb2.GenericResult.GENERIC_STATUS_TERMINATED:

modal/output.py

@@ -8,21 +8,15 @@
 
 import contextlib
 from collections.abc import Generator
-from typing import TYPE_CHECKING, Union
-
-if TYPE_CHECKING:
-    from ._output import DisabledOutputManager
-    from ._rich_output import RichOutputManager
 
+from ._output import _DISABLED_OUTPUT_MANAGER, OutputManager
 
 # Module-level state for output management
-_current_output_manager: "Union[RichOutputManager, None]" = None
+_current_output_manager: OutputManager = _DISABLED_OUTPUT_MANAGER
 
 
 @contextlib.contextmanager
-def enable_output(
-    show_progress: bool = True, show_timestamps: bool = False
-) -> Generator["RichOutputManager | None", None, None]:
+def enable_output(show_progress: bool = True, show_timestamps: bool = False) -> Generator[OutputManager, None, None]:
     """Context manager that enable output when using the Python SDK.
 
     This will print to stdout and stderr things such as
@@ -39,6 +33,7 @@ def enable_output(
     ```
     """
     global _current_output_manager
+    previous_output_manager = _current_output_manager
 
     if show_progress:
         from ._rich_output import RichOutputManager
@@ -47,10 +42,10 @@ def enable_output(
     try:
         yield _current_output_manager
     finally:
-        _current_output_manager = None
+        _current_output_manager = previous_output_manager
 
 
-def _get_output_manager() -> "Union[RichOutputManager, DisabledOutputManager]":
+def _get_output_manager() -> OutputManager:
     """Get the current output manager.
 
     Returns a RichOutputManager when output is enabled, otherwise returns
@@ -59,22 +54,7 @@ def _get_output_manager() -> "Union[RichOutputManager, DisabledOutputManager]":
     This allows code to call output methods without checking if output is enabled,
     simplifying the calling code.
     """
-    if _current_output_manager is not None:
-        return _current_output_manager
-
-    # Return the singleton disabled output manager
-    from ._output import _DISABLED_OUTPUT_MANAGER
-
-    return _DISABLED_OUTPUT_MANAGER
-
-
-def _is_output_enabled() -> bool:
-    """Check if rich output is enabled.
-
-    This is useful for code that needs to conditionally perform operations
-    based on whether output is truly enabled (e.g., starting a logs loop).
-    """
-    return _current_output_manager is not None
+    return _current_output_manager
 
 
 def _disable_output_manager() -> None:
@@ -84,4 +64,4 @@ def _disable_output_manager() -> None:
     to _get_output_manager() return a DisabledOutputManager.
     """
     global _current_output_manager
-    _current_output_manager = None
+    _current_output_manager = _DISABLED_OUTPUT_MANAGER

modal/runner.py

@@ -35,7 +35,7 @@
 from .config import config, logger
 from .environments import _get_environment_cached
 from .exception import ConnectionError, InteractiveTimeoutError, InvalidError, RemoteError, _CliUserExecutionError
-from .output import _get_output_manager, _is_output_enabled, enable_output
+from .output import _get_output_manager, enable_output
 from .running_app import RunningApp, running_app_from_layout
 from .sandbox import _Sandbox
 from .secret import _Secret
@@ -283,8 +283,9 @@ async def _run_app(
             app.set_description(__main__.__name__)
 
     app_state = api_pb2.APP_STATE_DETACHED if detach else api_pb2.APP_STATE_EPHEMERAL
+    output_mgr = _get_output_manager()
 
-    if interactive and not _is_output_enabled():
+    if interactive and not output_mgr.is_enabled:
         msg = "Interactive mode requires output to be enabled. (Use the the `modal.enable_output()` context manager.)"
         raise InvalidError(msg)
 
@@ -311,8 +312,7 @@ def heartbeat():
         heartbeat_loop = tc.infinite_loop(heartbeat, sleep=HEARTBEAT_INTERVAL, log_exception=not detach)
         logs_loop: Optional[asyncio.Task] = None
 
-        output_mgr = _get_output_manager()
-        if _is_output_enabled():
+        if output_mgr.is_enabled:
             from modal._output import get_app_logs_loop
 
             with output_mgr.make_live(output_mgr.step_progress("Initializing...")):
@@ -393,7 +393,7 @@ def heartbeat():
         except ConnectionError as e:
             # If we lose connection to the server after a detached App has started running, it will continue
             # I think we can only exit "nicely" if we are able to print output though, otherwise we should raise
-            if detach and _is_output_enabled():
+            if detach and output_mgr.is_enabled:
                 output_mgr.print(":white_exclamation_mark: Connection lost!")
                 output_mgr.print(detached_disconnect_msg)
                 return
@@ -413,8 +413,9 @@ def heartbeat():
                 logger.warning("Timed out waiting for final app logs.")
 
     # Print completion message if output is still enabled (it may have been disabled during PTY mode)
-    if _is_output_enabled():
-        output_mgr = _get_output_manager()
+    # Re-fetch the output manager in case it was disabled
+    output_mgr = _get_output_manager()
+    if output_mgr.is_enabled:
         output_mgr.print(
             output_mgr.step_completed(
                 f"App completed. [grey70]View run at [underline]{running_app.app_page_url}[/underline][/grey70]"

modal/serving.py

@@ -14,7 +14,7 @@
 from .cli.import_refs import ImportRef, import_app_from_ref
 from .client import _Client
 from .config import config
-from .output import _get_output_manager, _is_output_enabled, enable_output
+from .output import _get_output_manager, enable_output
 from .runner import _run_app, serve_update
 
 if TYPE_CHECKING:
@@ -35,7 +35,7 @@ async def _restart_serve(
 ) -> SpawnProcess:
     ctx = multiprocessing.get_context("spawn")  # Needed to reload the interpreter
     is_ready = ctx.Event()
-    show_progress = _is_output_enabled()
+    show_progress = _get_output_manager().is_enabled
     p = ctx.Process(target=_run_serve, args=(import_ref, existing_app_id, is_ready, environment_name, show_progress))
     p.start()
     await asyncify(is_ready.wait)(timeout)