From 3b91f53fbd45d38187cb181c47eba3410313c452 Mon Sep 17 00:00:00 2001 From: Stephanie Wang Date: Fri, 26 Jul 2024 16:31:01 -0700 Subject: [PATCH] fix Signed-off-by: Stephanie Wang --- python/ray/dag/compiled_dag_node.py | 4 +++- python/ray/experimental/compiled_dag_ref.py | 2 +- 2 files changed, 4 insertions(+), 2 deletions(-) diff --git a/python/ray/dag/compiled_dag_node.py b/python/ray/dag/compiled_dag_node.py index 18a2e272b546..b1f3e51989a1 100644 --- a/python/ray/dag/compiled_dag_node.py +++ b/python/ray/dag/compiled_dag_node.py @@ -1431,7 +1431,9 @@ async def execute_async( fut = asyncio.Future() await self._fut_queue.put(fut) - return CompiledDAGFuture(self, self._execution_index, fut) + fut = CompiledDAGFuture(self, self._execution_index, fut) + self._execution_index += 1 + return fut def teardown(self): """Teardown and cancel all actor tasks for this DAG. After this diff --git a/python/ray/experimental/compiled_dag_ref.py b/python/ray/experimental/compiled_dag_ref.py index 93f1397b60eb..e224250dfa77 100644 --- a/python/ray/experimental/compiled_dag_ref.py +++ b/python/ray/experimental/compiled_dag_ref.py @@ -92,7 +92,7 @@ def get(self, timeout: Optional[float] = None): @PublicAPI(stability="alpha") -class CompiledDAGFuture(CompiledDAGRef): +class CompiledDAGFuture: """ A reference to a compiled DAG execution result, when executed with asyncio. This differs from CompiledDAGRef in that `await` must be called on the