Add instant events for when an awaited task is not ready

Eventual-Inc · Oct 25, 2024 · 69092df · 69092df
1 parent 38f343c
commit 69092df
Show file tree

Hide file tree

Showing 2 changed files with 17 additions and 1 deletion.
diff --git a/daft/runners/ray_runner.py b/daft/runners/ray_runner.py
@@ -844,14 +844,18 @@ def place_in_queue(item):
                                 if num_returns == 0:
                                     break
 
-                                readies, _ = ray.wait(
+                                readies, not_readies = ray.wait(
                                     list(inflight_ref_to_task.keys()),
                                     num_returns=num_returns,
                                     timeout=timeout,
                                     fetch_local=False,
                                 )
                                 num_ready += len(readies)
 
+                                for not_ready in not_readies:
+                                    if not_ready in inflight_ref_to_task:
+                                        runner_tracer.task_not_ready(inflight_ref_to_task[not_ready])
+
                                 for ready in readies:
                                     if ready in inflight_ref_to_task:
                                         task_id = inflight_ref_to_task[ready]

diff --git a/daft/runners/tracer.py b/daft/runners/tracer.py
@@ -282,6 +282,18 @@ def task_dispatched(self, task_id: str, stage_id: int, resource_request: Resourc
             }
         )
 
+    def task_not_ready(self, task_id: str):
+        self._write_event(
+            {
+                "id": task_id,
+                "category": "task",
+                "name": "task_awaited_not_ready",
+                "ph": "n",
+                "pid": 2,
+                "tid": 1,
+            }
+        )
+
     def task_ready(self, task_id: str):
         self._write_event(
             {