diff --git a/raydar/task_tracker/task_tracker.py b/raydar/task_tracker/task_tracker.py index 1e80d1e..74be3ae 100644 --- a/raydar/task_tracker/task_tracker.py +++ b/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,7 +35,7 @@ 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 @@ -42,6 +43,9 @@ def process(self, obj_refs: Iterable[ray.ObjectRef]) -> None: 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 diff --git a/raydar/tests/test_schema.py b/raydar/tests/test_schema.py new file mode 100644 index 0000000..88e2a26 --- /dev/null +++ b/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 diff --git a/raydar/tests/test_task_tracker.py b/raydar/tests/test_task_tracker.py index cf1aca2..03a1819 100644 --- a/raydar/tests/test_task_tracker.py +++ b/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": {}}