feat: Flotilla scheduler and dispatcher actors (#4375)

## Changes Made

Implement scheduler + dispatcher actors for flotilla. Including unit
tests for both.

## Related Issues

<!-- Link to related GitHub issues, e.g., "Closes #123" -->

## Checklist

- [ ] Documented in API Docs (if applicable)
- [ ] Documented in User Guide (if applicable)
- [ ] If adding a new documentation page, doc is added to
`docs/mkdocs.yml` navigation
- [ ] Documentation builds and is formatted properly (tag @/ccmao1130
for docs review)

Author: colin-ho

Cargo.lock

@@ -17,6 +17,7 @@ rand = {workspace = true}
 tokio = {workspace = true}
 tokio-stream = {workspace = true}
 tokio-util = {workspace = true}
+tracing = {workspace = true}
 uuid.workspace = true
 
 [features]

src/daft-distributed/Cargo.toml

@@ -1,68 +1,76 @@
-use std::{any::Any, collections::HashMap, sync::Arc};
+use std::{any::Any, collections::HashMap, future::Future, sync::Arc};
 
 use common_error::DaftResult;
 use common_partitioning::{Partition, PartitionRef};
 use daft_local_plan::PyLocalPhysicalPlan;
 use pyo3::{pyclass, pymethods, FromPyObject, PyObject, PyResult, Python};
 
-use crate::scheduling::task::{SwordfishTask, TaskResultHandle};
+use crate::scheduling::{
+    task::{SwordfishTask, TaskResultHandle},
+    worker::WorkerId,
+};
 
 /// TaskHandle that wraps a Python RaySwordfishTaskHandle
 #[allow(dead_code)]
 pub(crate) struct RayTaskResultHandle {
-    /// The handle to the task
-    handle: Option<PyObject>,
+    /// The handle to the RaySwordfishTaskHandle
+    handle: PyObject,
+    /// The coroutine to await the result of the task
+    coroutine: Option<PyObject>,
     /// The task locals, i.e. the asyncio event loop
     task_locals: Option<pyo3_async_runtimes::TaskLocals>,
+    /// The worker id
+    worker_id: WorkerId,
 }
 
 impl RayTaskResultHandle {
     /// Create a new TaskHandle from a Python RaySwordfishTaskHandle
     #[allow(dead_code)]
-    pub fn new(handle: PyObject, task_locals: pyo3_async_runtimes::TaskLocals) -> Self {
+    pub fn new(
+        handle: PyObject,
+        coroutine: PyObject,
+        task_locals: pyo3_async_runtimes::TaskLocals,
+        worker_id: WorkerId,
+    ) -> Self {
         Self {
-            handle: Some(handle),
+            handle,
+            coroutine: Some(coroutine),
             task_locals: Some(task_locals),
+            worker_id,
         }
     }
 }
 
 impl TaskResultHandle for RayTaskResultHandle {
     /// Get the result of the task, awaiting if necessary
-    async fn get_result(&mut self) -> DaftResult<PartitionRef> {
-        let handle = self
-            .handle
-            .take()
-            .expect("Task handle should be present during get_result");
-        let coroutine = Python::with_gil(|py| {
-            let coroutine = handle
-                .call_method0(py, pyo3::intern!(py, "get_result"))?
-                .into_bound(py);
-            pyo3_async_runtimes::tokio::into_future(coroutine)
-        })?;
+    fn get_result(&mut self) -> impl Future<Output = DaftResult<PartitionRef>> + Send + 'static {
+        // Create a rust future that will await the coroutine
+        let coroutine = self.coroutine.take().unwrap();
+        let task_locals = self.task_locals.take().unwrap();
 
-        // await the rust future in the scope of the asyncio event loop
-        let task_locals = self
-            .task_locals
-            .take()
-            .expect("Task locals should be present during get_result");
-        let materialized_result = pyo3_async_runtimes::tokio::scope(task_locals, coroutine).await?;
+        let await_coroutine = async move {
+            let result = Python::with_gil(|py| {
+                pyo3_async_runtimes::tokio::into_future(coroutine.into_bound(py))
+            })?
+            .await?;
+            DaftResult::Ok(result)
+        };
 
-        let ray_part_ref =
-            Python::with_gil(|py| materialized_result.extract::<RayPartitionRef>(py))?;
-        Ok(Arc::new(ray_part_ref))
+        async move {
+            let materialized_result =
+                pyo3_async_runtimes::tokio::scope(task_locals, await_coroutine).await?;
+            let ray_part_ref =
+                Python::with_gil(|py| materialized_result.extract::<RayPartitionRef>(py))?;
+            let partition_ref = Arc::new(ray_part_ref) as PartitionRef;
+            Ok(partition_ref)
+        }
     }
-}
 
-impl Drop for RayTaskResultHandle {
-    fn drop(&mut self) {
-        if let Some(handle) = self.handle.take() {
-            Python::with_gil(|py| {
-                handle
-                    .call_method0(py, "cancel")
-                    .expect("Failed to cancel ray task")
-            });
-        }
+    fn cancel_callback(&mut self) -> DaftResult<()> {
+        Python::with_gil(|py| {
+            self.handle.call_method0(py, "cancel")?;
+            Ok(())
+        })
     }
 }
 

src/daft-distributed/src/python/ray/task.rs

@@ -8,7 +8,8 @@ use pyo3::prelude::*;
 
 use super::{task::RayTaskResultHandle, RaySwordfishTask};
 use crate::scheduling::{
-    task::{SwordfishTask, TaskDetails, TaskId},
+    scheduler::SchedulableTask,
+    task::{SwordfishTask, Task, TaskDetails, TaskId, TaskResultHandleAwaiter},
     worker::{Worker, WorkerId},
 };
 
@@ -44,49 +45,54 @@ impl RaySwordfishWorker {
 
 #[allow(dead_code)]
 impl RaySwordfishWorker {
-    pub fn mark_task_finished(&self, task_id: TaskId) {
+    pub fn mark_task_finished(&self, task_id: &TaskId) {
         self.active_task_details
             .lock()
             .expect("Active task ids should be present")
-            .remove(&task_id);
+            .remove(task_id);
     }
 
     pub fn submit_tasks(
         &self,
-        tasks: Vec<SwordfishTask>,
+        tasks: Vec<SchedulableTask<SwordfishTask>>,
         py: Python<'_>,
         task_locals: &pyo3_async_runtimes::TaskLocals,
-    ) -> DaftResult<Vec<RayTaskResultHandle>> {
-        let (task_details, ray_swordfish_tasks) = tasks
-            .into_iter()
-            .map(|task| (TaskDetails::from(&task), RaySwordfishTask::new(task)))
-            .unzip::<_, _, Vec<_>, Vec<_>>();
-
-        let py_task_handles = self
-            .actor_handle
-            .call_method1(
+    ) -> DaftResult<Vec<TaskResultHandleAwaiter<RayTaskResultHandle>>> {
+        let mut task_handles = Vec::with_capacity(tasks.len());
+        for task in tasks {
+            let (task, result_tx, cancel_token) = task.into_inner();
+            let task_id = task.task_id().clone();
+            let task_details = TaskDetails::from(&task);
+
+            let ray_swordfish_task = RaySwordfishTask::new(task);
+            let py_task_handle = self.actor_handle.call_method1(
                 py,
-                pyo3::intern!(py, "submit_tasks"),
-                (ray_swordfish_tasks,),
-            )?
-            .extract::<Vec<PyObject>>(py)?;
-
-        let task_handles = py_task_handles
-            .into_iter()
-            .map(|py_task_handle| {
-                let task_locals = task_locals.clone_ref(py);
-                RayTaskResultHandle::new(py_task_handle, task_locals)
-            })
-            .collect::<Vec<_>>();
-
-        self.active_task_details
-            .lock()
-            .expect("Active task ids should be present")
-            .extend(
-                task_details
-                    .into_iter()
-                    .map(|details| (details.id.clone(), details)),
+                pyo3::intern!(py, "submit_task"),
+                (ray_swordfish_task,),
+            )?;
+            let coroutine = py_task_handle.call_method0(py, pyo3::intern!(py, "get_result"))?;
+
+            self.active_task_details
+                .lock()
+                .expect("Active task details should be present")
+                .insert(task_id.clone(), task_details);
+
+            let task_locals = task_locals.clone_ref(py);
+            let ray_task_result_handle = RayTaskResultHandle::new(
+                py_task_handle,
+                coroutine,
+                task_locals,
+                self.worker_id.clone(),
+            );
+            let task_result_handle_awaiter = TaskResultHandleAwaiter::new(
+                task_id,
+                self.worker_id.clone(),
+                ray_task_result_handle,
+                result_tx,
+                cancel_token,
             );
+            task_handles.push(task_result_handle_awaiter);
+        }
 
         Ok(task_handles)
     }

src/daft-distributed/src/python/ray/worker.rs

@@ -5,7 +5,8 @@ use pyo3::prelude::*;
 
 use super::{task::RayTaskResultHandle, worker::RaySwordfishWorker};
 use crate::scheduling::{
-    task::{SwordfishTask, TaskId},
+    scheduler::SchedulableTask,
+    task::{SwordfishTask, TaskId, TaskResultHandleAwaiter},
     worker::{Worker, WorkerId, WorkerManager},
 };
 
@@ -43,11 +44,11 @@ impl WorkerManager for RayWorkerManager {
 
     fn submit_tasks_to_workers(
         &self,
-        total_tasks: usize,
-        tasks_per_worker: HashMap<WorkerId, Vec<SwordfishTask>>,
-    ) -> DaftResult<Vec<RayTaskResultHandle>> {
+        tasks_per_worker: HashMap<WorkerId, Vec<SchedulableTask<SwordfishTask>>>,
+    ) -> DaftResult<Vec<TaskResultHandleAwaiter<RayTaskResultHandle>>> {
         Python::with_gil(|py| {
-            let mut task_result_handles = Vec::with_capacity(total_tasks);
+            let mut task_result_handles =
+                Vec::with_capacity(tasks_per_worker.values().map(|v| v.len()).sum());
             for (worker_id, tasks) in tasks_per_worker {
                 let handles = self
                     .ray_workers
@@ -64,9 +65,9 @@ impl WorkerManager for RayWorkerManager {
         &self.ray_workers
     }
 
-    fn mark_task_finished(&self, task_id: TaskId, worker_id: WorkerId) {
+    fn mark_task_finished(&self, task_id: &TaskId, worker_id: &WorkerId) {
         self.ray_workers
-            .get(&worker_id)
+            .get(worker_id)
             .expect("Worker should be present in RayWorkerManager")
             .mark_task_finished(task_id);
     }