Revert "[core][compiled graphs] Fix test_torch_tensor_dag_gpu CI fail…

…ure (ray-project#48204)"

This reverts commit 23bb654.

python/ray/_private/ray_experimental_perf.py

@@ -181,7 +181,6 @@ async def _exec_async():
     results += timeit(
         "[unstable] compiled single-actor DAG calls", lambda: _exec(compiled_dag)
     )
-    compiled_dag.teardown()
     del a
 
     # Single-actor asyncio DAG calls
@@ -195,10 +194,6 @@ async def _exec_async():
             "[unstable] compiled single-actor asyncio DAG calls",
         )
     )
-    # TODO: Need to explicitly tear down DAGs with enable_asyncio=True because
-    # these DAGs create a background thread that can segfault if the CoreWorker
-    # is torn down first.
-    compiled_dag.teardown()
     del a
 
     # Scatter-gather DAG calls
@@ -216,7 +211,6 @@ async def _exec_async():
         f"[unstable] compiled scatter-gather DAG calls, n={n_cpu} actors",
         lambda: _exec(compiled_dag),
     )
-    compiled_dag.teardown()
 
     # Scatter-gather asyncio DAG calls
 
@@ -229,10 +223,6 @@ async def _exec_async():
             f"[unstable] compiled scatter-gather asyncio DAG calls, n={n_cpu} actors",
         )
     )
-    # TODO: Need to explicitly tear down DAGs with enable_asyncio=True because
-    # these DAGs create a background thread that can segfault if the CoreWorker
-    # is torn down first.
-    compiled_dag.teardown()
 
     # Chain DAG calls
 
@@ -250,7 +240,6 @@ async def _exec_async():
         f"[unstable] compiled chain DAG calls, n={n_cpu} actors",
         lambda: _exec(compiled_dag),
     )
-    compiled_dag.teardown()
 
     # Chain asyncio DAG calls
 
@@ -263,10 +252,6 @@ async def _exec_async():
     results += loop.run_until_complete(
         exec_async(f"[unstable] compiled chain asyncio DAG calls, n={n_cpu} actors")
     )
-    # TODO: Need to explicitly tear down DAGs with enable_asyncio=True because
-    # these DAGs create a background thread that can segfault if the CoreWorker
-    # is torn down first.
-    compiled_dag.teardown()
 
     # Multiple args with small payloads
 
@@ -289,7 +274,6 @@ async def _exec_async():
         f"n={n_actors} actors",
         lambda: _exec(compiled_dag, num_args=n_actors, payload_size=payload_size),
     )
-    compiled_dag.teardown()
 
     # Multiple args with medium payloads
 
@@ -307,7 +291,6 @@ async def _exec_async():
         f"n={n_actors} actors",
         lambda: _exec(compiled_dag, num_args=n_actors, payload_size=payload_size),
     )
-    compiled_dag.teardown()
 
     # Multiple args with large payloads
 
@@ -325,7 +308,6 @@ async def _exec_async():
         f"n={n_actors} actors",
         lambda: _exec(compiled_dag, num_args=n_actors, payload_size=payload_size),
     )
-    compiled_dag.teardown()
 
     # Worst case for multiple arguments: a single actor takes all the arguments
     # with small payloads.
@@ -346,7 +328,6 @@ async def _exec_async():
         "n=1 actors",
         lambda: _exec(compiled_dag, num_args=n_args, payload_size=payload_size),
     )
-    compiled_dag.teardown()
 
     ray.shutdown()
 

python/ray/_private/worker.py

@@ -1874,6 +1874,11 @@ def shutdown(_exiting_interpreter: bool = False):
             and false otherwise. If we are exiting the interpreter, we will
             wait a little while to print any extra error messages.
     """
+    # Make sure to clean up compiled dag node if exists.
+    from ray.dag.compiled_dag_node import _shutdown_all_compiled_dags
+
+    _shutdown_all_compiled_dags()
+
     if _exiting_interpreter and global_worker.mode == SCRIPT_MODE:
         # This is a duration to sleep before shutting down everything in order
         # to make sure that log messages finish printing.

python/ray/dag/compiled_dag_node.py

@@ -1,3 +1,4 @@
+import weakref
 import asyncio
 from collections import defaultdict
 from dataclasses import dataclass, asdict
@@ -8,6 +9,7 @@
 import uuid
 import traceback
 
+import ray.exceptions
 from ray.experimental.channel.cached_channel import CachedChannel
 from ray.experimental.channel.gpu_communicator import GPUCommunicator
 import ray
@@ -52,6 +54,21 @@
 
 logger = logging.getLogger(__name__)
 
+# Keep tracking of every compiled dag created during the lifetime of
+# this process. It tracks them as weakref meaning when the compiled dag
+# is GC'ed, it is automatically removed from here. It is used to teardown
+# compiled dags at interpret shutdown time.
+_compiled_dags = weakref.WeakValueDictionary()
+
+
+# Relying on __del__ doesn't work well upon shutdown because
+# the destructor order is not guaranteed. We call this function
+# upon `ray.worker.shutdown` which is registered to atexit handler
+# so that teardown is properly called before objects are destructed.
+def _shutdown_all_compiled_dags():
+    for _, compiled_dag in _compiled_dags.items():
+        compiled_dag.teardown()
+
 
 @DeveloperAPI
 def do_allocate_channel(
@@ -1648,7 +1665,7 @@ def _is_same_actor(idx1: int, idx2: int) -> bool:
         return False
 
     def _monitor_failures(self):
-        outer = self
+        outer = weakref.proxy(self)
 
         class Monitor(threading.Thread):
             def __init__(self):
@@ -1657,6 +1674,8 @@ def __init__(self):
                 # Lock to make sure that we only perform teardown for this DAG
                 # once.
                 self.in_teardown_lock = threading.Lock()
+                self.name = "CompiledGraphMonitorThread"
+                self._teardown_done = False
 
             def wait_teardown(self):
                 for actor, ref in outer.worker_task_refs.items():
@@ -1686,6 +1705,9 @@ def wait_teardown(self):
             def teardown(self, wait: bool):
                 do_teardown = False
                 with self.in_teardown_lock:
+                    if self._teardown_done:
+                        return
+
                     if not self.in_teardown:
                         do_teardown = True
                         self.in_teardown = True
@@ -1709,9 +1731,11 @@ def teardown(self, wait: bool):
                 ]
                 for cancel_ref in cancel_refs:
                     try:
-                        # TODO(swang): Suppress exceptions from actors trying to
-                        # read closed channels when DAG is being torn down.
                         ray.get(cancel_ref, timeout=30)
+                    except ray.exceptions.RayChannelError:
+                        # Channel error happens when a channel is closed
+                        # or timed out. In this case, do not log.
+                        pass
                     except Exception:
                         logger.exception("Error cancelling worker task")
                         pass
@@ -1724,6 +1748,9 @@ def teardown(self, wait: bool):
                     self.wait_teardown()
                     logger.info("Teardown complete")
 
+                with self.in_teardown_lock:
+                    self._teardown_done = True
+
             def run(self):
                 try:
                     ray.get(list(outer.worker_task_refs.values()))
@@ -2011,11 +2038,7 @@ def teardown(self):
     def __del__(self):
         monitor = getattr(self, "_monitor", None)
         if monitor is not None:
-            # Teardown asynchronously.
-            # NOTE(swang): Somehow, this can get called after the CoreWorker
-            # has already been destructed, so it is not safe to block in
-            # ray.get.
-            monitor.teardown(wait=False)
+            monitor.teardown(wait=True)
 
 
 @DeveloperAPI
@@ -2044,4 +2067,6 @@ def _build_compiled_dag(node):
     root = dag._find_root()
     root.traverse_and_apply(_build_compiled_dag)
     compiled_dag._get_or_compile()
+    global _compiled_dags
+    _compiled_dags[compiled_dag.get_id()] = compiled_dag
     return compiled_dag