Record task metadata the GCS has not published yet

raydar/task_tracker/task_tracker.py

@@ -1,6 +1,7 @@
 import asyncio
 import itertools
 import logging
+import time
 from collections.abc import Iterable
 from typing import Literal
 
@@ -34,14 +35,17 @@ class AsyncMetadataTrackerCallback:
     def __init__(self, name: str, namespace: str):
         self.actor = ray.get_actor(name, namespace)
 
-    def process(self, obj_refs: Iterable[ray.ObjectRef]) -> None:
+    def process(self, obj_refs: Iterable[ray.ObjectRef], settle_timeout: float = 60.0, settle_interval: float = 0.5) -> None:
         """Processes an interable collection of ray.ObjectRefs.
 
         Iterates through the collection, finds completed references, and returns those references to the
         self.actor attribute via its .callback remote function.
 
         Args:
             obj_refs: An iterable collection of (possibly) in-progress ray object references
+            settle_timeout: How long to keep asking the tracker to re-resolve tasks whose
+                metadata the GCS has not published yet.
+            settle_interval: How long to wait between those attempts.
         """
         active_tasks = set(obj_refs)
         while len(active_tasks) > 0:
@@ -55,6 +59,16 @@ def process(self, obj_refs: Iterable[ray.ObjectRef]) -> None:
             if len(finished_tasks) > 0:
                 self.actor.callback.remote(finished_tasks)
 
+        # ray.wait reports an object ready before the GCS has published the task's final
+        # state, so tasks the tracker could not resolve are held in its pending list. Only
+        # a later callback revisits that list, and the loop above has just run out of
+        # tasks to send, so the last batch would otherwise never be recorded. Calling back
+        # with nothing re-resolves the pending list on its own.
+        deadline = time.monotonic() + settle_timeout
+        while time.monotonic() < deadline and ray.get(self.actor.has_pending_tasks.remote()):
+            time.sleep(settle_interval)
+            ray.get(self.actor.callback.remote([]))
+
     def exit(self) -> None:
         """Terminate this actor"""
         ray.actor.exit_actor()
@@ -165,6 +179,10 @@ def drain(self) -> dict | None:
     def get_dashboard_mode(self) -> str | None:
         return self.dashboard_mode
 
+    def has_pending_tasks(self) -> bool:
+        """Whether any task is still waiting on the GCS to publish its final state."""
+        return bool(self.pending_tasks)
+
     def callback(self, tasks: Iterable[ray.ObjectRef]) -> None:
         """A remote function used by this actor's processor actor attribute. Will be called by a separate actor
         with a collection of ray object references once those ObjectReferences are not in the "RUNNING" or

raydar/tests/test_schema.py

@@ -0,0 +1,28 @@
+"""The schema must match what Ray actually reports.
+
+These run without a Ray cluster: the defect they guard is a type declaration,
+and asserting it directly is both deterministic and immediate.
+"""
+
+import polars as pl
+import pytest
+
+from raydar.task_tracker.schema import schema
+
+# Ray reports ids as hex strings, e.g. TaskState.actor_id.
+RAY_ID = "22c21d18d3081db73e11271a01000000"
+
+ID_COLUMNS = ("task_id", "actor_id", "job_id", "node_id", "worker_id", "parent_task_id", "placement_group_id")
+
+
+@pytest.mark.parametrize("column", ID_COLUMNS)
+def test_id_columns_hold_ray_hex_ids(column):
+    # Declaring these numeric made get_df raise on any actor task and rendered
+    # every id as 0.0 in the dashboard.
+    frame = pl.DataFrame({column: [RAY_ID]}, schema_overrides={column: schema[column]})
+    assert frame[column][0] == RAY_ID
+
+
+@pytest.mark.parametrize("column", ID_COLUMNS)
+def test_id_columns_are_declared_as_strings(column):
+    assert schema[column] == pl.Utf8

raydar/tests/test_task_tracker.py

@@ -14,13 +14,6 @@ def do_some_work():
     return True
 
 
-@ray.remote
-class SomeActor:
-    def do_some_work(self):
-        time.sleep(0.1)
-        return True
-
-
 def wait_for(fetch, ready, timeout=120, interval=0.5):
     """Poll `fetch` until `ready` accepts the value, then return it."""
     deadline = time.time() + timeout
@@ -73,21 +66,16 @@ def test_dashboard_is_off_by_default(self, trackers):
         task_tracker = trackers()
         assert task_tracker.dashboard_url is None
 
-    def test_actor_task_ids_survive_as_strings(self, trackers):
-        # Ray reports actor_id as a hex string. Declaring it numeric made get_df
-        # raise and rendered every id as 0.0 in the dashboard.
+    def test_a_single_task_is_still_recorded(self, trackers):
+        # One task completes in one wait round, so the tracker gets exactly one
+        # callback. The GCS has not published the task's state that early, so
+        # nothing was recorded until the processor learned to wait for it.
         task_tracker = trackers(dashboard="local")
-        actor = SomeActor.remote()
-        refs = [actor.do_some_work.remote() for _ in range(3)]
-        task_tracker.process(refs)
-        ray.get(refs)
+        task_tracker.process([do_some_work.remote()])
 
-        df = wait_for(task_tracker.get_df, lambda d: not d.is_empty())
-        assert not df.is_empty(), "tracker recorded no finished actor tasks"
-
-        actor_ids = [a for a in df["actor_id"].to_list() if a]
-        assert actor_ids, "actor_id was not recorded"
-        assert all(isinstance(a, str) and int(a, 16) for a in actor_ids)
+        df = wait_for(task_tracker.get_df, lambda d: not d.is_empty(), timeout=90)
+        assert not df.is_empty(), "the only task the tracker was given went unrecorded"
+        assert df[["name", "state"]].row(0) == ("do_some_work", "FINISHED")
 
     def test_dashboard_options_reach_the_dashboard(self, trackers):
         layout = {"sizes": [1], "viewers": {}}