[core][accelerated DAGs] Make DAG teardown blocking and fix bug durin…

…g close (ray-project#45099)

## Why are these changes needed?

Contains a few fixes related to DAG teardown:
- Removes an unnecessary `.close()` call that would error if the DAG has
a single output (instead of a MultiOutputNode)
- Makes `dag.teardown()` blocking to ensure that actors can be reused
after the teardown call returns.
- Makes DAG teardown in `__del__` asynchronous. if synchronous, this can
hang the driver upon shutdown. I'm not exactly sure why but I believe
this happens if the CoreWorker is shut down before `dag.teardown()` is
complete.

---------
Signed-off-by: Stephanie Wang <[email protected]>

python/ray/dag/compiled_dag_node.py

@@ -295,7 +295,7 @@ def __init__(
 
  # ObjectRef for each worker's task. The task is an infinite loop that
  # repeatedly executes the method specified in the DAG.
- self.worker_task_refs: List["ray.ObjectRef"] = []
+ self.worker_task_refs: Dict["ray.actor.ActorHandle", "ray.ObjectRef"] = {}
  # Set of actors present in the DAG.
  self.actor_refs = set()
 
@@ -525,14 +525,14 @@ def _get_node_id(self):
 
  # Assign the task with the correct input and output buffers.
  worker_fn = task.dag_node._get_remote_method("__ray_call__")
- self.worker_task_refs.append(
- worker_fn.options(concurrency_group="_ray_system").remote(
-  do_exec_compiled_task,
-  resolved_args,
-  task.dag_node.get_method_name(),
-  output_wrapper_fn=task.output_wrapper_fn,
-  has_type_hints=bool(self._type_hints),
- )
+ self.worker_task_refs[
+ task.dag_node._get_actor_handle()
+ ] = worker_fn.options(concurrency_group="_ray_system").remote(
+ do_exec_compiled_task,
+ resolved_args,
+ task.dag_node.get_method_name(),
+ output_wrapper_fn=task.output_wrapper_fn,
+ has_type_hints=bool(self._type_hints),
  )
 
  # Wrapper function for inputs provided to dag.execute().
@@ -586,9 +586,37 @@ def __init__(self):
  super().__init__(daemon=True)
  self.in_teardown = False
 
- def teardown(self):
+ def wait_teardown(self):
+ for actor, ref in outer.worker_task_refs.items():
+ timeout = False
+ try:
+ ray.get(ref, timeout=10)
+ except ray.exceptions.GetTimeoutError:
+ logger.warn(
+ f"Compiled DAG actor {actor} is still running 10s "
+ "after teardown(). Teardown may hang."
+ )
+ timeout = True
+ except Exception:
+ # We just want to check that the task has finished so
+ # we don't care if the actor task ended in an
+ # exception.
+ pass
+
+ if not timeout:
+ continue
+
+ try:
+ ray.get(ref)
+ except Exception:
+ pass
+
+ def teardown(self, wait: bool):
  if self.in_teardown:
+ if wait:
+ self.wait_teardown()
  return
+
  logger.info("Tearing down compiled DAG")
 
  outer._dag_submitter.close()
@@ -604,24 +632,20 @@ def teardown(self):
  except Exception:
  logger.exception("Error cancelling worker task")
  pass
- logger.info("Waiting for worker tasks to exit")
- for ref in outer.worker_task_refs:
- try:
- ray.get(ref)
- except Exception:
- pass
- logger.info("Teardown complete")
+
+ if wait:
+ logger.info("Waiting for worker tasks to exit")
+ self.wait_teardown()
+ logger.info("Teardown complete")
 
  def run(self):
  try:
- ray.get(outer.worker_task_refs)
+ ray.get(list(outer.worker_task_refs.values()))
  except Exception as e:
  logger.debug(f"Handling exception from worker tasks: {e}")
  if self.in_teardown:
  return
- for output_channel in outer.dag_output_channels:
- output_channel.close()
- self.teardown()
+ self.teardown(wait=True)
 
  monitor = Monitor()
  monitor.start()
@@ -701,13 +725,21 @@ async def execute_async(
  return AwaitableDAGOutput(fut, self._dag_output_fetcher)
 
  def teardown(self):
- """Teardown and cancel all worker tasks for this DAG."""
+ """Teardown and cancel all actor tasks for this DAG. After this
+ function returns, the actors should be available to execute new tasks
+ or compile a new DAG."""
  monitor = getattr(self, "_monitor", None)
  if monitor is not None:
- monitor.teardown()
+ monitor.teardown(wait=True)
 
  def __del__(self):
- self.teardown()
+ 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)
 
 
 @DeveloperAPI

python/ray/dag/tests/test_accelerated_dag.py

@@ -213,7 +213,10 @@ def f(x):
 
 
 def test_dag_fault_tolerance(ray_start_regular_shared):
- actors = [Actor.remote(0, fail_after=100, sys_exit=False) for _ in range(4)]
+ actors = [
+ Actor.remote(0, fail_after=100 if i == 0 else None, sys_exit=False)
+ for i in range(4)
+ ]
  with InputNode() as i:
  out = [a.inc.bind(i) for a in actors]
  dag = MultiOutputNode(out)
@@ -235,9 +238,30 @@ def test_dag_fault_tolerance(ray_start_regular_shared):
 
  compiled_dag.teardown()
 
+ # All actors are still alive.
+ ray.get([actor.echo.remote("hello") for actor in actors])
+
+ # Remaining actors can be reused.
+ actors.pop(0)
+ with InputNode() as i:
+ out = [a.inc.bind(i) for a in actors]
+ dag = MultiOutputNode(out)
+
+ compiled_dag = dag.experimental_compile()
+ for i in range(100):
+ output_channels = compiled_dag.execute(1)
+ # TODO(swang): Replace with fake ObjectRef.
+ output_channels.begin_read()
+ output_channels.end_read()
+
+ compiled_dag.teardown()
+
 
 def test_dag_fault_tolerance_sys_exit(ray_start_regular_shared):
- actors = [Actor.remote(0, fail_after=100, sys_exit=True) for _ in range(4)]
+ actors = [
+ Actor.remote(0, fail_after=100 if i == 0 else None, sys_exit=True)
+ for i in range(4)
+ ]
  with InputNode() as i:
  out = [a.inc.bind(i) for a in actors]
  dag = MultiOutputNode(out)
@@ -257,6 +281,26 @@ def test_dag_fault_tolerance_sys_exit(ray_start_regular_shared):
  output_channels.begin_read()
  output_channels.end_read()
 
+ # Remaining actors are still alive.
+ with pytest.raises(ray.exceptions.RayActorError):
+ ray.get(actors[0].echo.remote("hello"))
+ actors.pop(0)
+ ray.get([actor.echo.remote("hello") for actor in actors])
+
+ # Remaining actors can be reused.
+ with InputNode() as i:
+ out = [a.inc.bind(i) for a in actors]
+ dag = MultiOutputNode(out)
+
+ compiled_dag = dag.experimental_compile()
+ for i in range(100):
+ output_channels = compiled_dag.execute(1)
+ # TODO(swang): Replace with fake ObjectRef.
+ output_channels.begin_read()
+ output_channels.end_read()
+
+ compiled_dag.teardown()
+
 
 def test_dag_teardown_while_running(ray_start_regular_shared):
  a = Actor.remote(0)