add tests for ray runner clean up

daft/runners/ray_runner.py

@@ -417,7 +417,7 @@ def next(self, result_uuid: str) -> ray.ObjectRef | StopIteration:
 
         return result
 
-    def run_plan(
+    def start_plan(
         self,
         plan_scheduler: PhysicalPlanScheduler,
         psets: dict[str, ray.ObjectRef],
@@ -439,6 +439,9 @@ def run_plan(
         t.start()
         self.threads_by_df[result_uuid] = t
 
+    def active_plans(self) -> list[str]:
+        return [r_uuid for r_uuid, is_active in self.active_by_df.items() if is_active]
+
     def stop_plan(self, result_uuid: str) -> None:
         if result_uuid in self.active_by_df:
             # Mark df as non-active
@@ -631,6 +634,12 @@ def __init__(
                 max_task_backlog=max_task_backlog,
             )
 
+    def active_plans(self) -> list[str]:
+        if isinstance(self.ray_context, ray.client_builder.ClientContext):
+            return ray.get(self.scheduler_actor.active_plans.remote())
+        else:
+            return self.scheduler.active_plans()
+
     def run_iter(self, builder: LogicalPlanBuilder, results_buffer_size: int | None = None) -> Iterator[ray.ObjectRef]:
         # Optimize the logical plan.
         builder = builder.optimize()
@@ -646,7 +655,7 @@ def run_iter(self, builder: LogicalPlanBuilder, results_buffer_size: int | None
         result_uuid = str(uuid.uuid4())
         if isinstance(self.ray_context, ray.client_builder.ClientContext):
             ray.get(
-                self.scheduler_actor.run_plan.remote(
+                self.scheduler_actor.start_plan.remote(
                     plan_scheduler=plan_scheduler,
                     psets=psets,
                     result_uuid=result_uuid,
@@ -655,7 +664,7 @@ def run_iter(self, builder: LogicalPlanBuilder, results_buffer_size: int | None
             )
 
         else:
-            self.scheduler.run_plan(
+            self.scheduler.start_plan(
                 plan_scheduler=plan_scheduler,
                 psets=psets,
                 result_uuid=result_uuid,

tests/ray/runner.py

@@ -0,0 +1,48 @@
+from __future__ import annotations
+
+import pytest
+
+import daft
+from daft.context import get_context
+
+
+@pytest.mark.skipif(get_context().runner_config.name != "ray", reason="Needs to run on Ray runner")
+def test_active_plan_clean_up_df_show():
+    path = "tests/assets/parquet-data/mvp.parquet"
+    df = daft.read_parquet([path, path])
+    df.show()
+    runner = get_context().runner()
+    assert len(runner.active_plans()) == 0
+
+
+@pytest.mark.skipif(get_context().runner_config.name != "ray", reason="Needs to run on Ray runner")
+def test_active_plan_single_iter_partitions():
+    path = "tests/assets/parquet-data/mvp.parquet"
+    df = daft.read_parquet([path, path])
+    iter = df.iter_partitions()
+    next(iter)
+    runner = get_context().runner()
+    assert len(runner.active_plans()) == 1
+    del iter
+    assert len(runner.active_plans()) == 0
+
+
+@pytest.mark.skipif(get_context().runner_config.name != "ray", reason="Needs to run on Ray runner")
+def test_active_plan_multiple_iter_partitions():
+    path = "tests/assets/parquet-data/mvp.parquet"
+    df = daft.read_parquet([path, path])
+    iter = df.iter_partitions()
+    next(iter)
+    runner = get_context().runner()
+    assert len(runner.active_plans()) == 1
+
+    df2 = daft.read_parquet([path, path])
+    iter2 = df2.iter_partitions()
+    next(iter2)
+    assert len(runner.active_plans()) == 2
+
+    del iter
+    assert len(runner.active_plans()) == 1
+
+    del iter2
+    assert len(runner.active_plans()) == 0