diff --git a/raydar/tests/test_task_tracker.py b/raydar/tests/test_task_tracker.py index b6962b0..cf1aca2 100644 --- a/raydar/tests/test_task_tracker.py +++ b/raydar/tests/test_task_tracker.py @@ -1,3 +1,4 @@ +import logging import time import httpx @@ -30,72 +31,84 @@ def wait_for(fetch, ready, timeout=120, interval=0.5): return value +@pytest.fixture +def trackers(): + """Build trackers and release them afterwards. + + Each RayTaskTracker pins 0.2 of the head node's 1.0 `node:__internal_head__` + budget across its two detached actors, so leaving them alive starves the + fifth tracker in a module and its metadata never arrives. + """ + created = [] + + def make(**kwargs): + tracker = RayTaskTracker(**kwargs) + created.append(tracker) + return tracker + + yield make + + for tracker in created: + try: + tracker.exit() + except Exception: + logging.getLogger(__name__).exception("Failed to release tracker %s", tracker.name) + + @pytest.mark.usefixtures("unittest_ray_cluster") class TestRayTaskTracker: - def test_construction_and_dataframe(self): - task_tracker = RayTaskTracker(dashboard="local") - try: - assert len(task_tracker.namespace.split("-")) == 2 - refs = [do_some_work.remote() for _ in range(10)] - task_tracker.process(refs) - - # Metadata arrives via GCS polling, so wait on the result rather than - # on a fixed sleep, which was slow on a fast box and flaky on a slow one. - df = wait_for(task_tracker.get_df, lambda d: not d.is_empty()) - assert not df.is_empty(), "tracker recorded no finished tasks" - assert df[["name", "state"]].row(0) == ("do_some_work", "FINISHED") - finally: - task_tracker.dashboard.stop() - - def test_dashboard_is_off_by_default(self): - task_tracker = RayTaskTracker() + def test_construction_and_dataframe(self, trackers): + task_tracker = trackers(dashboard="local") + assert len(task_tracker.namespace.split("-")) == 2 + refs = [do_some_work.remote() for _ in range(10)] + task_tracker.process(refs) + + # Metadata arrives via GCS polling, so wait on the result rather than + # on a fixed sleep, which was slow on a fast box and flaky on a slow one. + df = wait_for(task_tracker.get_df, lambda d: not d.is_empty()) + assert not df.is_empty(), "tracker recorded no finished tasks" + assert df[["name", "state"]].row(0) == ("do_some_work", "FINISHED") + + 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): + 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. - task_tracker = RayTaskTracker(dashboard="local") - try: - actor = SomeActor.remote() - refs = [actor.do_some_work.remote() for _ in range(3)] - task_tracker.process(refs) - ray.get(refs) + 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) - df = wait_for(task_tracker.get_df, lambda d: not d.is_empty()) - assert not df.is_empty(), "tracker recorded no finished actor tasks" + 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) - finally: - task_tracker.dashboard.stop() + 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) - def test_dashboard_options_reach_the_dashboard(self): + def test_dashboard_options_reach_the_dashboard(self, trackers): layout = {"sizes": [1], "viewers": {}} - task_tracker = RayTaskTracker(dashboard="local", dashboard_options={"title": "custom", "layout": layout}) - try: - dashboard = task_tracker.dashboard.dashboard - dashboard.apply({"schemas": {"t": {"a": "integer"}}}) - assert dashboard.state.layout == layout - assert "custom" in httpx.get(task_tracker.dashboard_url).text - finally: - task_tracker.dashboard.stop() - - def test_local_dashboard_serves_tables_pulled_from_the_actor(self): - task_tracker = RayTaskTracker(dashboard="local") - try: - assert task_tracker.dashboard_url.startswith("http://127.0.0.1:") - task_tracker.create_table("custom", {"a": "string", "b": "integer"}) - task_tracker.update_table("custom", [{"a": "foo", "b": 1}]) - - tables = task_tracker.dashboard.dashboard.tables - expected = ["custom", "task_tracker_data"] - names = wait_for(lambda: sorted(tables.names()), lambda n: n == expected, timeout=60) - assert names == expected - - # Wait on the row too: names alone pass before the update is applied. - rows = wait_for(lambda: tables._tables["custom"].size(), lambda n: n > 0, timeout=60) - assert rows == 1 - assert httpx.get(task_tracker.dashboard_url).status_code == 200 - finally: - task_tracker.dashboard.stop() + task_tracker = trackers(dashboard="local", dashboard_options={"title": "custom", "layout": layout}) + dashboard = task_tracker.dashboard.dashboard + dashboard.apply({"schemas": {"t": {"a": "integer"}}}) + assert dashboard.state.layout == layout + assert "custom" in httpx.get(task_tracker.dashboard_url).text + + def test_local_dashboard_serves_tables_pulled_from_the_actor(self, trackers): + task_tracker = trackers(dashboard="local") + assert task_tracker.dashboard_url.startswith("http://127.0.0.1:") + task_tracker.create_table("custom", {"a": "string", "b": "integer"}) + task_tracker.update_table("custom", [{"a": "foo", "b": 1}]) + + tables = task_tracker.dashboard.dashboard.tables + expected = ["custom", "task_tracker_data"] + names = wait_for(lambda: sorted(tables.names()), lambda n: n == expected, timeout=60) + assert names == expected + + # Wait on the row too: names alone pass before the update is applied. + rows = wait_for(lambda: tables._tables["custom"].size(), lambda n: n > 0, timeout=60) + assert rows == 1 + assert httpx.get(task_tracker.dashboard_url).status_code == 200