From 2346aa90663c10e2c2125efac09651be44160c57 Mon Sep 17 00:00:00 2001
From: Erik Bernhardsson <mail@erikbern.com>
Date: Wed, 31 Jul 2024 14:22:48 -0400
Subject: [PATCH] Fix issue with global output manager tree (#2062)

* Preventure test leakage by resetting singleton

* Fix issue by making _tree an instance property
---
 modal/_output.py | 26 +++++++++++++++-----------
 test/conftest.py |  6 ++++++
 2 files changed, 21 insertions(+), 11 deletions(-)

diff --git a/modal/_output.py b/modal/_output.py
index 0c9bb4e4d..e16b95c37 100644
--- a/modal/_output.py
+++ b/modal/_output.py
@@ -154,7 +154,6 @@ def finalize(self):
 
 class OutputManager:
     _instance: ClassVar[Optional["OutputManager"]] = None
-    _tree: ClassVar[Optional[Tree]] = None
 
     _console: Console
     _task_states: Dict[str, int]
@@ -168,6 +167,7 @@ class OutputManager:
     _app_page_url: Optional[str]
     _show_image_logs: bool
     _status_spinner_live: Optional[Live]
+    _tree: Optional[Tree]
 
     def __init__(
         self,
@@ -188,6 +188,7 @@ def __init__(
         self._app_page_url = None
         self._show_image_logs = False
         self._status_spinner_live = None
+        self._tree = None
 
     @classmethod
     def disable(cls):
@@ -389,23 +390,26 @@ def show_status_spinner(self):
     @classmethod
     @contextlib.contextmanager
     def make_tree(cls):
-        # Note: If the output isn't enabled, don't actually show the tree.
-        cls._tree = Tree(step_progress("Creating objects..."), guide_style="gray50")
-
         if output_mgr := OutputManager.get():
-            with output_mgr.make_live(cls._tree):
-                yield
-            cls._tree.label = step_completed("Created objects.")
-            output_mgr.print(output_mgr._tree)
+            tree = output_mgr._tree = Tree(step_progress("Creating objects..."), guide_style="gray50")
+            with output_mgr.make_live(tree):
+                try:
+                    yield
+                finally:
+                    output_mgr._tree = None
+            tree.label = step_completed("Created objects.")
+            output_mgr.print(tree)
         else:
             yield
 
     @classmethod
     def add_status_row(cls) -> "StatusRow":
         # Return a status row to be used for object creation.
-        # If output isn't enabled, the status row might be invisible.
-        assert cls._tree, "Output manager has no tree yet"
-        return StatusRow(cls._tree)
+        # If output isn't enabled, just create a hidden tree
+        if cls._instance and cls._instance._tree:
+            return StatusRow(cls._instance._tree)
+        else:
+            return StatusRow(Tree(""))
 
 
 class ProgressHandler:
diff --git a/test/conftest.py b/test/conftest.py
index 4bc10e3e0..c5cf8cf7e 100644
--- a/test/conftest.py
+++ b/test/conftest.py
@@ -30,6 +30,7 @@
 import modal._serialization
 from modal import __version__, config
 from modal._container_io_manager import _ContainerIOManager
+from modal._output import OutputManager
 from modal._serialization import serialize_data_format
 from modal._utils.async_utils import asyncify, synchronize_api
 from modal._utils.grpc_testing import patch_mock_servicer
@@ -1551,6 +1552,11 @@ async def reset_default_client():
     Client.set_env_client(None)
 
 
+@pytest_asyncio.fixture(scope="function", autouse=True)
+async def reset_output_manager():
+    OutputManager._instance = None
+
+
 @pytest.fixture(name="mock_dir", scope="session")
 def mock_dir_factory():
     """Sets up a temp dir with content as specified in a nested dict