[data] Remove ray.kill in ActorPoolMapOperator

python/ray/data/_internal/execution/operators/actor_pool_map_operator.py

@@ -640,13 +640,15 @@ def _kill_all_running_actors(self):
 
     def _kill_running_actor(self, actor: ray.actor.ActorHandle):
         """Kill the provided actor and remove it from the pool."""
-        ray.kill(actor)
         del self._num_tasks_in_flight[actor]
+        del self._actor_locations[actor]
+        assert actor not in self._pending_actors
 
     def _kill_pending_actor(self, ready_ref: ray.ObjectRef):
         """Kill the provided pending actor and remove it from the pool."""
         actor = self._pending_actors.pop(ready_ref)
-        ray.kill(actor)
+        assert actor not in self._num_tasks_in_flight
+        assert actor not in self._actor_locations
 
     def _get_location(self, bundle: RefBundle) -> Optional[NodeIdStr]:
         """Ask Ray for the node id of the given bundle.

python/ray/data/tests/test_map.py

@@ -13,9 +13,11 @@
 import pytest
 
 import ray
+from ray._private.test_utils import wait_for_condition
 from ray.data._internal.execution.interfaces.ref_bundle import (
     _ref_bundles_iterator_to_block_refs_list,
 )
+from ray.data._internal.execution.operators.actor_pool_map_operator import _MapWorker
 from ray.data.context import DataContext
 from ray.data.exceptions import UserCodeException
 from ray.data.tests.conftest import *  # noqa
@@ -76,6 +78,19 @@ def test_basic_actors(shutdown_only):
             concurrency=(8, 4),
         )
 
+    # Make sure all actors are dead after dataset execution finishes.
+    def _all_actors_dead():
+        actor_table = ray.state.actors()
+        actors = {
+            id: actor_info
+            for actor_info in actor_table.values()
+            if actor_info["ActorClassName"] == _MapWorker.__name__
+        }
+        assert len(actors) > 0
+        return all(actor_info["State"] == "DEAD" for actor_info in actors.values())
+
+    wait_for_condition(_all_actors_dead)
+
 
 def test_callable_classes(shutdown_only):
     ray.init(num_cpus=2)