diff --git a/python/ray/data/_internal/execution/operators/actor_pool_map_operator.py b/python/ray/data/_internal/execution/operators/actor_pool_map_operator.py index 72dbeff66a1c..7a34b175cf9a 100644 --- a/python/ray/data/_internal/execution/operators/actor_pool_map_operator.py +++ b/python/ray/data/_internal/execution/operators/actor_pool_map_operator.py @@ -933,8 +933,20 @@ def scale(self, req: ActorPoolScalingRequest) -> Optional[int]: @override def refresh_actor_state(self): self._alive_node_to_actor_heap.clear() + dead_actors = [] for actor in self._running_actors: - self._update_running_actor_state(actor) + if self._update_running_actor_state(actor): + dead_actors.append(actor) + for actor in dead_actors: + # Release dead actors so they stop counting towards the pool's + # current size; otherwise a fixed-size pool would never replace + # them and the pipeline could silently stall (see #62746). The + # autoscaler will scale the pool back up to its min size. + logger.warning( + f"{self.map_worker_cls_name} actor {actor} is dead; releasing " + "it from the actor pool so it can be replaced." + ) + self._release_running_actor(actor) @override def on_task_submitted(self, actor: ActorHandle): @@ -1068,7 +1080,12 @@ def select_actors( @override def on_task_completed(self, actor: ActorHandle): """Called when a task completes. Returns the provided actor to the pool.""" - state = self._running_actors[actor] + state = self._running_actors.get(actor) + if state is None: + # The actor was already released from the pool (e.g. it died with + # tasks in flight and was removed by `refresh_actor_state`); + # `_release_running_actor` already reconciled the pool's counters. + return assert state.num_tasks_in_flight > 0 state.num_tasks_in_flight -= 1 self._total_num_tasks_in_flight -= 1 @@ -1127,12 +1144,15 @@ def _create_actor( self._actor_to_logical_id[actor] = logical_actor_id return actor, ready_ref, resource_usage - def _update_running_actor_state(self, actor: ActorHandle): + def _update_running_actor_state(self, actor: ActorHandle) -> bool: """Update running actor state. This is called for every actor in `refresh_actor_state`. Args: actor: The running actor that needs state update. + + Returns: + True if the actor is dead (and should be released from the pool). """ actor_state = actor._get_local_state() @@ -1167,6 +1187,10 @@ def _update_running_actor_state(self, actor: ActorHandle): self._update_rank(actor=actor, state=running_actor_state, died=died) + # Only report definitively dead actors for release; a `None` state means + # the state is unknown (possibly transiently), so keep those in the pool. + return actor_state == _ACTOR_STATE_DEAD + def _update_rank(self, actor: ActorHandle, state: _ActorState, died: bool): """Update the scheduling rank for an actor after a state refresh. diff --git a/python/ray/data/tests/test_actor_pool_map_operator.py b/python/ray/data/tests/test_actor_pool_map_operator.py index 688c018c60f2..3af6671e0d8b 100644 --- a/python/ray/data/tests/test_actor_pool_map_operator.py +++ b/python/ray/data/tests/test_actor_pool_map_operator.py @@ -1357,6 +1357,58 @@ def test_completed_when_downstream_op_has_finished_execution(ray_start_regular_s assert actor_pool_map_op.has_completed() +def test_dead_actor_released_and_replaced_e2e(shutdown_only, restore_data_context): + """A dead actor must be released from the pool and replaced (#62746). + + With a fixed-size pool and `max_errored_blocks = -1`, an actor that dies + mid-pipeline (e.g. `sys.exit(0)` inside the UDF) used to stay in + `_running_actors` forever: the pool's size never dropped, the autoscaler + never created a replacement, and the pipeline silently hung. + """ + import sys as _sys + + ray.shutdown() + ray.init(num_cpus=2) + + ctx = ray.data.DataContext.get_current() + ctx.max_errored_blocks = -1 + + @ray.remote(num_cpus=0) + class DeathFlag: + def __init__(self): + self._died = False + + def should_die(self) -> bool: + died, self._died = self._died, True + return not died + + flag = DeathFlag.remote() + + class DiesOnce: + def __init__(self): + self._counter = 0 + + def __call__(self, batch): + self._counter += 1 + if self._counter == 3 and ray.get(flag.should_die.remote()): + _sys.exit(0) + return batch + + ds = ( + ray.data.range(20, override_num_blocks=20) + .map_batches( + DiesOnce, + batch_size=1, + compute=ray.data.ActorPoolStrategy(size=1), + ) + .materialize() + ) + + # The pipeline completes (instead of hanging), with only the blocks that + # were in flight on the dead actor dropped. + assert ds.count() > 0 + + def test_actor_pool_fault_tolerance_e2e(ray_start_cluster, restore_data_context): """Test that a dataset with actor pools can finish, when all nodes in the cluster are removed and added back."""