Addressed review comments

Signed-off-by: Srinath Krishnamachari <[email protected]>
ray-project · Oct 18, 2024 · 217ca5b · 217ca5b
1 parent 1b506b5
commit 217ca5b
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 7 deletions.
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
@@ -217,13 +217,10 @@ def _task_done_callback(actor_to_return):
                 # Dipsatch more tasks.
                 self._dispatch_tasks()
 
-            # For some reason, if we don't define a new variable `actor_to_return`,
-            # the following lambda won't capture the correct `actor` variable.
-            actor_to_return = actor
+            from functools import partial
+
             self._submit_data_task(
-                gen,
-                bundle,
-                lambda: _task_done_callback(actor_to_return),
+                gen, bundle, partial(_task_done_callback, actor_to_return=actor)
             )
 
     def _refresh_actor_cls(self):

diff --git a/python/ray/data/tests/test_actor_pool_map_operator.py b/python/ray/data/tests/test_actor_pool_map_operator.py
@@ -649,7 +649,7 @@ async def wait_for_nodes_restarted(self):
     signal_actor = Signal.remote()
 
     # Spin up nodes
-    num_nodes = 2
+    num_nodes = 5
     nodes = []
     for _ in range(num_nodes):
         nodes.append(cluster.add_node(num_cpus=10, num_gpus=1))