feat(dashboard): enhance dashboard UI and fix Ray runner state reporting

Author: Jay-ju

daft/context.py

@@ -77,6 +77,21 @@ def _notify_optimization_start(self, query_id: str) -> None:
     def _notify_optimization_end(self, query_id: str, optimized_plan: str) -> None:
         self._ctx.notify_optimization_end(query_id, optimized_plan)
 
+    def _notify_exec_start(self, query_id: str, physical_plan: str) -> None:
+        self._ctx.notify_exec_start(query_id, physical_plan)
+
+    def _notify_exec_end(self, query_id: str) -> None:
+        self._ctx.notify_exec_end(query_id)
+
+    def _notify_exec_operator_start(self, query_id: str, node_id: int) -> None:
+        self._ctx.notify_exec_operator_start(query_id, node_id)
+
+    def _notify_exec_operator_end(self, query_id: str, node_id: int) -> None:
+        self._ctx.notify_exec_operator_end(query_id, node_id)
+
+    def _notify_exec_emit_stats(self, query_id: str, node_id: int, stats: dict[str, int]) -> None:
+        self._ctx.notify_exec_emit_stats(query_id, node_id, stats)
+
     def _notify_result_out(self, query_id: str, result: PartitionT) -> None:
         from daft.recordbatch.micropartition import MicroPartition
 

daft/daft/__init__.pyi

@@ -1997,6 +1997,7 @@ class DistributedPhysicalPlan:
     def num_partitions(self) -> int: ...
     def repr_ascii(self, simple: bool) -> str: ...
     def repr_mermaid(self, options: MermaidOptions) -> str: ...
+    def repr_json(self) -> str: ...
 
 class DistributedPhysicalPlanRunner:
     def __init__(self) -> None: ...
@@ -2197,8 +2198,18 @@ class QueryEndState(Enum):
 class PyQueryMetadata:
     output_schema: PySchema
     unoptimized_plan: str
+    runner: str
+    ray_dashboard_url: str | None
+    entrypoint: str | None
 
-    def __init__(self, output_schema: PySchema, unoptimized_plan: str) -> None: ...
+    def __init__(
+        self,
+        output_schema: PySchema,
+        unoptimized_plan: str,
+        runner: str,
+        ray_dashboard_url: str | None = None,
+        entrypoint: str | None = None,
+    ) -> None: ...
 
 class PyQueryResult:
     end_state: QueryEndState
@@ -2225,6 +2236,11 @@ class PyDaftContext:
     def notify_result_out(self, query_id: str, result: PartitionT) -> None: ...
     def notify_optimization_start(self, query_id: str) -> None: ...
     def notify_optimization_end(self, query_id: str, optimized_plan: str) -> None: ...
+    def notify_exec_start(self, query_id: str, physical_plan: str) -> None: ...
+    def notify_exec_end(self, query_id: str) -> None: ...
+    def notify_exec_operator_start(self, query_id: str, node_id: int) -> None: ...
+    def notify_exec_operator_end(self, query_id: str, node_id: int) -> None: ...
+    def notify_exec_emit_stats(self, query_id: str, node_id: int, stats: dict[str, int]) -> None: ...
 
 def set_runner_ray(
     address: str | None = None,

daft/runners/flotilla.py

@@ -237,7 +237,9 @@ def try_autoscale(bundles: list[dict[str, int]]) -> None:
     num_cpus=0,
 )
 class RemoteFlotillaRunner:
-    def __init__(self) -> None:
+    def __init__(self, dashboard_url: str | None = None) -> None:
+        if dashboard_url:
+            os.environ["DAFT_DASHBOARD_URL"] = dashboard_url
         self.curr_plans: dict[str, DistributedPhysicalPlan] = {}
         self.curr_result_gens: dict[str, AsyncIterator[RayPartitionRef]] = {}
         self.plan_runner = DistributedPhysicalPlanRunner()
@@ -340,19 +342,21 @@ class FlotillaRunner:
 
     def __init__(self) -> None:
         head_node_id = get_head_node_id()
+        dashboard_url = os.environ.get("DAFT_DASHBOARD_URL")
         self.runner = RemoteFlotillaRunner.options(  # type: ignore
             name=get_flotilla_runner_actor_name(),
             namespace=FLOTILLA_RUNNER_NAMESPACE,
             get_if_exists=True,
+            runtime_env={"env_vars": {"DAFT_DASHBOARD_URL": dashboard_url}} if dashboard_url else None,
             scheduling_strategy=(
                 ray.util.scheduling_strategies.NodeAffinitySchedulingStrategy(
                     node_id=head_node_id,
                     soft=False,
                 )
-                if head_node_id is not None
-                else "DEFAULT"
+                if head_node_id
+                else None
             ),
-        ).remote()
+        ).remote(dashboard_url=dashboard_url)
 
     def stream_plan(
         self,

daft/runners/native_runner.py

@@ -96,12 +96,19 @@ def run_iter(
         output_schema = builder.schema()
 
         # Optimize the logical plan.
-        ctx._notify_query_start(query_id, PyQueryMetadata(output_schema._schema, builder.repr_json()))
+        import sys
+
+        entrypoint = "python " + " ".join(sys.argv)
+        ctx._notify_query_start(
+            query_id,
+            PyQueryMetadata(output_schema._schema, builder.repr_json(), "Native (Swordfish)", None, entrypoint),
+        )
         ctx._notify_optimization_start(query_id)
         builder = builder.optimize(ctx.daft_execution_config)
         ctx._notify_optimization_end(query_id, builder.repr_json())
 
         plan = LocalPhysicalPlan.from_logical_plan_builder(builder._builder)
+
         executor = NativeExecutor()
         results_gen = executor.run(
             plan,
@@ -112,8 +119,10 @@ def run_iter(
         )
 
         try:
+            total_rows = 0
             for result in results_gen:
                 ctx._notify_result_out(query_id, result.partition())
+                total_rows += len(result.partition())
                 yield result
         except KeyboardInterrupt as e:
             query_result = PyQueryResult(QueryEndState.Canceled, "Query canceled by the user.")

daft/runners/ray_runner.py

@@ -1,6 +1,9 @@
 from __future__ import annotations
 
+import json
 import logging
+import os
+import sys
 import time
 import uuid
 from collections.abc import Generator, Iterable, Iterator
@@ -45,6 +48,9 @@
     FileFormatConfig,
     FileInfos,
     IOConfig,
+    PyQueryMetadata,
+    PyQueryResult,
+    QueryEndState,
 )
 from daft.datatype import DataType
 from daft.filesystem import glob_path_with_stats
@@ -548,19 +554,84 @@ def run_iter(
         ctx = get_context()
         query_id = str(uuid.uuid4())
         daft_execution_config = ctx.daft_execution_config
+        output_schema = builder.schema()
 
-        # Optimize the logical plan.
-        builder = builder.optimize(daft_execution_config)
-
-        distributed_plan = DistributedPhysicalPlan.from_logical_plan_builder(
-            builder._builder, query_id, daft_execution_config
+        # Notify query start
+        ray_dashboard_url = None
+        try:
+            ray_dashboard_url = ray.worker.get_dashboard_url()
+            if ray_dashboard_url:
+                if not ray_dashboard_url.startswith("http"):
+                    ray_dashboard_url = f"http://{ray_dashboard_url}"
+
+                # Try to append Job ID
+                try:
+                    job_id = ray.get_runtime_context().get_job_id()
+                    if job_id:
+                        ray_dashboard_url = f"{ray_dashboard_url}/#/jobs/{job_id}"
+                except Exception:
+                    pass
+        except Exception:
+            pass
+
+        entrypoint = "python " + " ".join(sys.argv)
+
+        ctx._notify_query_start(
+            query_id,
+            PyQueryMetadata(
+                output_schema._schema, builder.repr_json(), "Ray (Flotilla)", ray_dashboard_url, entrypoint
+            ),
         )
-        if self.flotilla_plan_runner is None:
-            self.flotilla_plan_runner = FlotillaRunner()
+        ctx._notify_optimization_start(query_id)
 
-        yield from self.flotilla_plan_runner.stream_plan(
-            distributed_plan, self._part_set_cache.get_all_partition_sets()
-        )
+        # Log Dashboard URL if configured
+        dashboard_url = os.environ.get("DAFT_DASHBOARD_URL")
+        if dashboard_url:
+            print(f"Daft Dashboard: {dashboard_url}/query/{query_id}")
+
+        try:
+            # Optimize the logical plan.
+            builder = builder.optimize(daft_execution_config)
+            ctx._notify_optimization_end(query_id, builder.repr_json())
+
+            distributed_plan = DistributedPhysicalPlan.from_logical_plan_builder(
+                builder._builder, query_id, daft_execution_config
+            )
+            physical_plan_json = distributed_plan.repr_json()
+            ctx._notify_exec_start(query_id, physical_plan_json)
+
+            if self.flotilla_plan_runner is None:
+                self.flotilla_plan_runner = FlotillaRunner()
+
+            total_rows = 0
+            for result in self.flotilla_plan_runner.stream_plan(
+                distributed_plan, self._part_set_cache.get_all_partition_sets()
+            ):
+                if result.metadata() is not None:
+                    total_rows += result.metadata().num_rows
+                yield result
+
+            # Mark all operators as finished to clean up the Dashboard UI before notify_exec_end
+            try:
+                plan_dict = json.loads(physical_plan_json)
+
+                def notify_end(node: dict[str, Any]) -> None:
+                    if "id" in node:
+                        ctx._notify_exec_operator_end(query_id, node["id"])
+                    if "children" in node:
+                        for child in node["children"]:
+                            notify_end(child)
+
+                notify_end(plan_dict)
+            except Exception:
+                pass
+
+            ctx._notify_exec_end(query_id)
+            ctx._notify_query_end(query_id, PyQueryResult(QueryEndState.Finished, ""))
+
+        except Exception as e:
+            ctx._notify_query_end(query_id, PyQueryResult(QueryEndState.Failed, str(e)))
+            raise
 
     def run_iter_tables(
         self, builder: LogicalPlanBuilder, results_buffer_size: int | None = None

daft/utils.py

@@ -168,10 +168,10 @@ def np_datetime64_to_timestamp(dt: np.datetime64) -> tuple[int, PyTimeUnit | Non
     val: np.int64 = dt.astype(np.int64) * np.int64(count)
 
     if unit in ("Y", "M", "W", "D"):
-        val = np.datetime64(dt, "D").astype(np.int64)  # type: ignore
+        val = np.datetime64(dt, "D").astype(np.int64)
         return val.item(), None
     elif unit in ("h", "m"):
-        val = np.datetime64(dt, "s").astype(np.int64)  # type: ignore
+        val = np.datetime64(dt, "s").astype(np.int64)
         return val.item(), PyTimeUnit.seconds()
     elif unit == "s":
         return val.item(), PyTimeUnit.seconds()
@@ -183,7 +183,7 @@ def np_datetime64_to_timestamp(dt: np.datetime64) -> tuple[int, PyTimeUnit | Non
         return val.item(), PyTimeUnit.nanoseconds()
     else:
         # unit is too small, just convert to nanoseconds
-        val = np.datetime64(dt, "ns").astype(np.int64)  # type: ignore
+        val = np.datetime64(dt, "ns").astype(np.int64)
         return val.item(), PyTimeUnit.nanoseconds()
 
 

src/daft-context/Cargo.toml

@@ -11,11 +11,14 @@ daft-micropartition = {path = "../daft-micropartition", default-features = false
 daft-dashboard = {path = "../daft-dashboard", default-features = false}
 daft-runners = {workspace = true}
 pyo3 = {workspace = true, optional = true}
+tokio = {workspace = true}
+oneshot = "0.1.8"
 log = {workspace = true}
 # For debug subscriber
 dashmap = {workspace = true}
 # Client for submitting to dashboard server
-reqwest = {workspace = true, default-features = false}
+reqwest = {workspace = true, default_features = false}
+uuid = {workspace = true}
 
 [features]
 python = [

src/daft-context/src/lib.rs

@@ -184,6 +184,100 @@ impl DaftContext {
             Ok::<(), DaftError>(())
         })
     }
+
+    pub fn notify_exec_start(&self, query_id: QueryID, physical_plan: String) -> DaftResult<()> {
+        self.with_state(|state| {
+            for subscriber in state.subscribers.values() {
+                subscriber.on_exec_start(query_id.clone(), physical_plan.clone().into())?;
+            }
+            Ok::<(), DaftError>(())
+        })
+    }
+
+    pub fn notify_exec_end(&self, query_id: QueryID) -> DaftResult<()> {
+        let subscribers = self.with_state(|state| {
+            state
+                .subscribers
+                .values()
+                .cloned()
+                .collect::<Vec<Arc<dyn Subscriber>>>()
+        });
+        let rt = common_runtime::get_io_runtime(false);
+        for subscriber in subscribers {
+            let query_id = query_id.clone();
+            let _ = rt.block_within_async_context(async move {
+                if let Err(e) = subscriber.on_exec_end(query_id).await {
+                    log::error!("Failed to notify exec end: {}", e);
+                }
+            });
+        }
+        Ok(())
+    }
+
+    pub fn notify_exec_operator_start(&self, query_id: QueryID, node_id: usize) -> DaftResult<()> {
+        let subscribers = self.with_state(|state| {
+            state
+                .subscribers
+                .values()
+                .cloned()
+                .collect::<Vec<Arc<dyn Subscriber>>>()
+        });
+        let rt = common_runtime::get_io_runtime(false);
+        for subscriber in subscribers {
+            let query_id = query_id.clone();
+            rt.spawn(async move {
+                if let Err(e) = subscriber.on_exec_operator_start(query_id, node_id).await {
+                    log::error!("Failed to notify exec operator start: {}", e);
+                }
+            });
+        }
+        Ok(())
+    }
+
+    pub fn notify_exec_operator_end(&self, query_id: QueryID, node_id: usize) -> DaftResult<()> {
+        let subscribers = self.with_state(|state| {
+            state
+                .subscribers
+                .values()
+                .cloned()
+                .collect::<Vec<Arc<dyn Subscriber>>>()
+        });
+        let rt = common_runtime::get_io_runtime(false);
+        for subscriber in subscribers {
+            let query_id = query_id.clone();
+            rt.spawn(async move {
+                if let Err(e) = subscriber.on_exec_operator_end(query_id, node_id).await {
+                    log::error!("Failed to notify exec operator end: {}", e);
+                }
+            });
+        }
+        Ok(())
+    }
+
+    pub fn notify_exec_emit_stats(
+        &self,
+        query_id: QueryID,
+        stats: Vec<(usize, common_metrics::StatSnapshot)>,
+    ) -> DaftResult<()> {
+        let subscribers = self.with_state(|state| {
+            state
+                .subscribers
+                .values()
+                .cloned()
+                .collect::<Vec<Arc<dyn Subscriber>>>()
+        });
+        let rt = common_runtime::get_io_runtime(false);
+        for subscriber in subscribers {
+            let stats = stats.clone();
+            let query_id = query_id.clone();
+            rt.spawn(async move {
+                if let Err(e) = subscriber.on_exec_emit_stats(query_id, &stats).await {
+                    log::error!("Failed to notify exec emit stats: {}", e);
+                }
+            });
+        }
+        Ok(())
+    }
 }
 
 static DAFT_CONTEXT: OnceLock<DaftContext> = OnceLock::new();