fix: Account for unschedulable udf actors (#4987)

## Changes Made

Raises error in situations where actors for actor UDFs (i.e. concurrency
udfs) cannot be scheduled due to resource constraints.

This is implemented via timeout (default 60s) on the readiness check for
the actors. Additionally, each swordfish task that needs to run actor
UDFs will receive the full list of actor handles, and will check for
which actors are ready + local. This is to account for the fact that the
UDF actors can be spawned incrementally, and when killed, can spawn on
different nodes.

## Related Issues

Closes #3934

## 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

daft/context.py

@@ -204,6 +204,7 @@ def set_execution_config(
     native_parquet_writer: bool | None = None,
     use_legacy_ray_runner: bool | None = None,
     min_cpu_per_task: float | None = None,
+    actor_udf_ready_timeout: int | None = None,
 ) -> DaftContext:
     """Globally sets various configuration parameters which control various aspects of Daft execution.
 
@@ -253,6 +254,7 @@ def set_execution_config(
         native_parquet_writer: Whether to use the native parquet writer vs the pyarrow parquet writer. Defaults to `True`.
         use_legacy_ray_runner: Whether to use the legacy ray runner. Defaults to `False`.
         min_cpu_per_task: Minimum CPU per task in the Ray runner. Defaults to 0.5.
+        actor_udf_ready_timeout: Timeout for UDF actors to be ready. Defaults to 60 seconds.
     """
     # Replace values in the DaftExecutionConfig with user-specified overrides
     ctx = get_context()
@@ -289,6 +291,7 @@ def set_execution_config(
             native_parquet_writer=native_parquet_writer,
             use_legacy_ray_runner=use_legacy_ray_runner,
             min_cpu_per_task=min_cpu_per_task,
+            actor_udf_ready_timeout=actor_udf_ready_timeout,
         )
 
         ctx._ctx._daft_execution_config = new_daft_execution_config

daft/daft/__init__.pyi

@@ -1960,6 +1960,7 @@ class PyDaftExecutionConfig:
         native_parquet_writer: bool | None = None,
         use_legacy_ray_runner: bool | None = None,
         min_cpu_per_task: float | None = None,
+        actor_udf_ready_timeout: int | None = None,
     ) -> PyDaftExecutionConfig: ...
     @property
     def scan_tasks_min_size_bytes(self) -> int: ...
@@ -2012,6 +2013,8 @@ class PyDaftExecutionConfig:
     @property
     def min_cpu_per_task(self) -> float: ...
     @property
+    def actor_udf_ready_timeout(self) -> int: ...
+    @property
     def scantask_max_parallel(self) -> int: ...
 
 class PyDaftPlanningConfig:

daft/execution/ray_actor_pool_udf.py

@@ -7,6 +7,8 @@
 from daft.recordbatch.micropartition import MicroPartition
 
 if TYPE_CHECKING:
+    from ray.actor import ActorHandle as RayActorHandle
+
     from daft.daft import PyExpr, PyMicroPartition
 
 try:
@@ -36,30 +38,50 @@ def eval_input(self, input: PyMicroPartition) -> PyMicroPartition:
 
 
 class UDFActorHandle:
-    def __init__(self, node_id: str, actor_ref: ray.ObjectRef) -> None:
-        self.node_id = node_id
+    def __init__(self, actor_ref: RayActorHandle) -> None:
         self.actor = actor_ref
 
+    def actor_id(self) -> str:
+        return self.actor._actor_id.hex()
+
     async def eval_input(self, input: PyMicroPartition) -> PyMicroPartition:
         return await self.actor.eval_input.remote(input)
 
-    def is_on_current_node(self) -> bool:
-        return self.node_id == ray.get_runtime_context().get_node_id()
-
     def teardown(self) -> None:
         ray.kill(self.actor)
 
 
+def get_ready_actors_by_location(
+    actor_handles: list[UDFActorHandle],
+) -> tuple[list[UDFActorHandle], list[UDFActorHandle]]:
+    from ray._private.state import actors
+
+    current_node_id = ray.get_runtime_context().get_node_id()
+
+    local_actors = []
+    remote_actors = []
+    for actor_handle in actor_handles:
+        actor_id = actor_handle.actor_id()
+        actor_state = actors(actor_id)
+        if actor_state["Address"]["NodeID"] == current_node_id:
+            local_actors.append(actor_handle)
+        else:
+            remote_actors.append(actor_handle)
+
+    return local_actors, remote_actors
+
+
 async def start_udf_actors(
     projection: list[PyExpr],
     num_actors: int,
     num_gpus_per_actor: float,
     num_cpus_per_actor: float,
     memory_per_actor: float,
+    timeout: int,
 ) -> list[UDFActorHandle]:
     expr_projection = ExpressionsProjection([Expression._from_pyexpr(expr) for expr in projection])
 
-    actors = [
+    actors: list[RayActorHandle] = [
         UDFActor.options(  # type: ignore
             scheduling_strategy="SPREAD",
             num_gpus=num_gpus_per_actor,
@@ -68,6 +90,20 @@ async def start_udf_actors(
         ).remote(expr_projection)
         for _ in range(num_actors)
     ]
-    node_ids = await asyncio.gather(*[actor.get_node_id.remote() for actor in actors])
-    handles = [UDFActorHandle(node_id, actor) for actor, node_id in zip(actors, node_ids)]
-    return handles
+
+    # Wait for actors to be ready
+    ready_futures = [asyncio.wrap_future(actor.__ray_ready__.remote().future()) for actor in actors]
+    ready_refs, _ = await asyncio.wait(ready_futures, return_when=asyncio.ALL_COMPLETED, timeout=timeout)
+
+    # Verify that the __ray_ready__ calls were successful
+    await asyncio.gather(*ready_refs)
+
+    if not ready_refs:
+        raise RuntimeError(
+            f"UDF actors failed to start within {timeout} seconds, please increase the actor_udf_ready_timeout config via daft.set_execution_config(actor_udf_ready_timeout=timeout)"
+        )
+
+    # Return the ready actors
+    ready_indices = [ready_futures.index(ref) for ref in ready_refs]
+    ready_actors = [UDFActorHandle(actors[i]) for i in ready_indices]
+    return ready_actors

src/common/daft-config/src/lib.rs

@@ -74,6 +74,7 @@ pub struct DaftExecutionConfig {
     pub native_parquet_writer: bool,
     pub use_legacy_ray_runner: bool,
     pub min_cpu_per_task: f64,
+    pub actor_udf_ready_timeout: usize,
 }
 
 impl Default for DaftExecutionConfig {
@@ -110,6 +111,7 @@ impl Default for DaftExecutionConfig {
             native_parquet_writer: true,
             use_legacy_ray_runner: false,
             min_cpu_per_task: 0.5,
+            actor_udf_ready_timeout: 60,
         }
     }
 }

src/common/daft-config/src/python.rs

@@ -118,6 +118,7 @@ impl PyDaftExecutionConfig {
         native_parquet_writer=None,
         use_legacy_ray_runner=None,
         min_cpu_per_task=None,
+        actor_udf_ready_timeout=None,
     ))]
     fn with_config_values(
         &self,
@@ -150,6 +151,7 @@ impl PyDaftExecutionConfig {
         native_parquet_writer: Option<bool>,
         use_legacy_ray_runner: Option<bool>,
         min_cpu_per_task: Option<f64>,
+        actor_udf_ready_timeout: Option<usize>,
     ) -> PyResult<Self> {
         let mut config = self.config.as_ref().clone();
 
@@ -266,6 +268,10 @@ impl PyDaftExecutionConfig {
             config.min_cpu_per_task = min_cpu_per_task;
         }
 
+        if let Some(actor_udf_ready_timeout) = actor_udf_ready_timeout {
+            config.actor_udf_ready_timeout = actor_udf_ready_timeout;
+        }
+
         Ok(Self {
             config: Arc::new(config),
         })
@@ -396,6 +402,11 @@ impl PyDaftExecutionConfig {
     fn min_cpu_per_task(&self) -> PyResult<f64> {
         Ok(self.config.min_cpu_per_task)
     }
+
+    #[getter]
+    fn actor_udf_ready_timeout(&self) -> PyResult<usize> {
+        Ok(self.config.actor_udf_ready_timeout)
+    }
 }
 
 impl_bincode_py_state_serialization!(PyDaftExecutionConfig);

src/daft-distributed/src/pipeline_node/actor_udf.rs

@@ -14,10 +14,8 @@ use super::{
     PipelineNodeContext, SubmittableTaskStream, TreeDisplay,
 };
 use crate::{
-    scheduling::{
-        scheduler::SubmittableTask,
-        task::{SwordfishTask, Task},
-    },
+    pipeline_node::append_plan_to_existing_task,
+    scheduling::{scheduler::SubmittableTask, task::SwordfishTask},
     stage::{StageConfig, StageExecutionContext},
     utils::{
         channel::{Sender, create_channel},
@@ -37,6 +35,7 @@ impl UDFActors {
     async fn initialize_actors(
         projection: &[BoundExpr],
         udf_properties: &UDFProperties,
+        actor_ready_timeout: usize,
     ) -> DaftResult<Vec<PyObjectWrapper>> {
         let (task_locals, py_exprs) = Python::with_gil(|py| {
             let task_locals = crate::utils::runtime::PYO3_ASYNC_RUNTIME_LOCALS
@@ -77,6 +76,7 @@ impl UDFActors {
                         gpu_request,
                         cpu_request,
                         memory_request,
+                        actor_ready_timeout,
                     ),
                 )?;
                 pyo3_async_runtimes::tokio::into_future(coroutine)
@@ -98,10 +98,12 @@ impl UDFActors {
         Ok(actors)
     }
 
-    async fn get_actors(&mut self) -> DaftResult<Vec<PyObjectWrapper>> {
+    async fn get_actors(&mut self, actor_ready_timeout: usize) -> DaftResult<Vec<PyObjectWrapper>> {
         match self {
             Self::Uninitialized(projection, udf_properties) => {
-                let actors = Self::initialize_actors(projection, udf_properties).await?;
+                let actors =
+                    Self::initialize_actors(projection, udf_properties, actor_ready_timeout)
+                        .await?;
                 *self = Self::Initialized {
                     actors: actors.clone(),
                 };
@@ -130,6 +132,7 @@ pub(crate) struct ActorUDF {
     child: Arc<dyn DistributedPipelineNode>,
     projection: Vec<BoundExpr>,
     udf_properties: UDFProperties,
+    actor_ready_timeout: usize,
 }
 
 impl ActorUDF {
@@ -164,6 +167,7 @@ impl ActorUDF {
             child,
             projection,
             udf_properties,
+            actor_ready_timeout: stage_config.config.actor_udf_ready_timeout,
         })
     }
 
@@ -181,9 +185,9 @@ impl ActorUDF {
 
         let mut running_tasks = JoinSet::new();
         while let Some(task) = input_task_stream.next().await {
-            let actors = udf_actors.get_actors().await?;
+            let actors = udf_actors.get_actors(self.actor_ready_timeout).await?;
 
-            let modified_task = self.append_actor_udf_to_task(task, actors)?;
+            let modified_task = self.append_actor_udf_to_task(task, actors);
             let (submittable_task, notify_token) = modified_task.add_notify_token();
             running_tasks.spawn(notify_token);
             if result_tx.send(submittable_task).await.is_err() {
@@ -202,10 +206,10 @@ impl ActorUDF {
     }
 
     fn append_actor_udf_to_task(
-        &self,
+        self: &Arc<Self>,
         submittable_task: SubmittableTask<SwordfishTask>,
         actors: Vec<PyObjectWrapper>,
-    ) -> DaftResult<SubmittableTask<SwordfishTask>> {
+    ) -> SubmittableTask<SwordfishTask> {
         let memory_request = self
             .udf_properties
             .resource_request
@@ -214,33 +218,22 @@ impl ActorUDF {
             .map(|m| m as u64)
             .unwrap_or(0);
 
-        let mut task_context = submittable_task.task().task_context();
-        if let Some(logical_node_id) = self.context.logical_node_id {
-            task_context.add_logical_node_id(logical_node_id);
-        }
-        let task_plan = submittable_task.task().plan();
-        let actor_pool_project_plan = LocalPhysicalPlan::distributed_actor_pool_project(
-            task_plan,
-            actors,
-            self.udf_properties.batch_size,
-            memory_request,
-            self.config.schema.clone(),
-            StatsState::NotMaterialized,
-        );
-
-        // Set scheduling strategy based on whether we have a valid worker ID
-        let scheduling_strategy = submittable_task.task().strategy().clone();
-        let psets = submittable_task.task().psets().clone();
-
-        let task = submittable_task.with_new_task(SwordfishTask::new(
-            task_context,
-            actor_pool_project_plan,
-            self.config.execution_config.clone(),
-            psets,
-            scheduling_strategy,
-            self.context.to_hashmap(),
-        ));
-        Ok(task)
+        let batch_size = self.udf_properties.batch_size;
+        let schema = self.config.schema.clone();
+        append_plan_to_existing_task(
+            submittable_task,
+            &(self.clone() as Arc<dyn DistributedPipelineNode>),
+            &move |input| {
+                LocalPhysicalPlan::distributed_actor_pool_project(
+                    input,
+                    actors.clone(),
+                    batch_size,
+                    memory_request,
+                    schema.clone(),
+                    StatsState::NotMaterialized,
+                )
+            },
+        )
     }
 
     fn multiline_display(&self) -> Vec<String> {

src/daft-local-execution/src/intermediate_ops/cross_join.rs

@@ -142,7 +142,7 @@ impl IntermediateOperator for CrossJoinOperator {
         ]
     }
 
-    fn make_state(&self) -> DaftResult<Self::State> {
+    async fn make_state(&self) -> DaftResult<Self::State> {
         Ok(CrossJoinState::new(self.state_bridge.clone()))
     }
 }