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