add timeout and fix lock in autoscale

Author: huleilei

.github/workflows/pr-test-suite.yml

@@ -362,7 +362,7 @@ jobs:
 
   integration-test-io:
     runs-on: ubuntu-latest
-    timeout-minutes: 30
+    timeout-minutes: 60
     needs: [skipcheck, integration-test-build]
     if: ${{ needs.skipcheck.outputs.skip == 'false' }}
     env:
@@ -455,7 +455,7 @@ jobs:
   # Same as integration-test-io but runs the tests that require credentials, only on `main`
   integration-test-io-credentialed:
     runs-on: ubuntu-latest
-    timeout-minutes: 30
+    timeout-minutes: 60
     needs: [skipcheck, integration-test-build]
     if: ${{ needs.skipcheck.outputs.skip == 'false' && github.ref == 'refs/heads/main' }}
     env:

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

@@ -205,83 +205,111 @@ impl WorkerManager for RayWorkerManager {
                 )
             });
 
-        let mut state = self
-            .state
-            .lock()
-            .expect("Failed to lock RayWorkerManagerState");
-
-        // If no desired bundles, clear outstanding autoscaler requests instead of scale-up.
-        if bundles.is_empty()
-            || (requested_num_cpus <= 0.0
-                && requested_num_gpus <= 0.0
-                && requested_memory_bytes == 0)
-        {
-            Python::attach(|py| -> DaftResult<()> {
-                let flotilla_module = py.import(pyo3::intern!(py, "daft.runners.flotilla"))?;
-                flotilla_module.call_method0(pyo3::intern!(py, "clear_autoscaling_requests"))?;
-                Ok(())
-            })?;
-            state.max_resources_requested = ResourceRequest::default();
-            return Ok(());
+        enum AutoscaleAction {
+            Noop,
+            Clear,
+            ScaleUp {
+                python_bundles: Vec<HashMap<&'static str, i64>>,
+            },
         }
 
-        let (cluster_num_cpus, cluster_num_gpus, cluster_memory_bytes) = state
-            .ray_workers
-            .values()
-            .fold((0.0, 0.0, 0), |acc, worker| {
-                (
-                    acc.0 + worker.total_num_cpus(),
-                    acc.1 + worker.total_num_gpus(),
-                    acc.2 + worker.total_memory_bytes(),
-                )
-            });
-
-        let resource_request_greater_than_current_capacity = requested_num_cpus > cluster_num_cpus
-            || requested_num_gpus > cluster_num_gpus
-            || requested_memory_bytes > cluster_memory_bytes;
-
-        let resource_request_greater_than_max_requested = requested_num_cpus
-            > state.max_resources_requested.num_cpus().unwrap_or(0.0)
-            || requested_num_gpus > state.max_resources_requested.num_gpus().unwrap_or(0.0)
-            || requested_memory_bytes > state.max_resources_requested.memory_bytes().unwrap_or(0);
-
-        let cluster_is_zero_capacity =
-            cluster_num_cpus <= 0.0 && cluster_num_gpus <= 0.0 && cluster_memory_bytes == 0;
-        let should_bootstrap = cluster_is_zero_capacity
-            && (requested_num_cpus > 0.0 || requested_num_gpus > 0.0 || requested_memory_bytes > 0);
+        let mut action = AutoscaleAction::Noop;
 
-        // Only autoscale if we need more capacity AND this is greater than we've seen before
-        if (resource_request_greater_than_current_capacity
-            && resource_request_greater_than_max_requested)
-            || should_bootstrap
         {
-            // On scale-up demand, allow previously blacklisted workers to be reused immediately.
-            state.pending_release_blacklist.clear();
-            state.last_refresh = None;
-            state.max_resources_requested = ResourceRequest::try_new_internal(
-                Some(requested_num_cpus),
-                Some(requested_num_gpus),
-                Some(requested_memory_bytes),
-            )?;
-            let python_bundles = bundles
-                .iter()
-                .map(|bundle| {
-                    let mut dict = HashMap::new();
-                    dict.insert("CPU", bundle.num_cpus().ceil() as i64);
-                    dict.insert("GPU", bundle.num_gpus().ceil() as i64);
-                    dict.insert("memory", bundle.memory_bytes() as i64);
-                    dict
-                })
-                .collect::<Vec<_>>();
+            let mut state = self
+                .state
+                .lock()
+                .expect("Failed to lock RayWorkerManagerState");
 
-            Python::attach(|py| -> DaftResult<()> {
-                let flotilla_module = py.import(pyo3::intern!(py, "daft.runners.flotilla"))?;
-                flotilla_module
-                    .call_method1(pyo3::intern!(py, "try_autoscale"), (python_bundles,))?;
+            // If no desired bundles, clear outstanding autoscaler requests instead of scale-up.
+            if bundles.is_empty()
+                || (requested_num_cpus <= 0.0
+                    && requested_num_gpus <= 0.0
+                    && requested_memory_bytes == 0)
+            {
+                state.max_resources_requested = ResourceRequest::default();
+                action = AutoscaleAction::Clear;
+            } else {
+                let (cluster_num_cpus, cluster_num_gpus, cluster_memory_bytes) = state
+                    .ray_workers
+                    .values()
+                    .fold((0.0, 0.0, 0), |acc, worker| {
+                        (
+                            acc.0 + worker.total_num_cpus(),
+                            acc.1 + worker.total_num_gpus(),
+                            acc.2 + worker.total_memory_bytes(),
+                        )
+                    });
+
+                let resource_request_greater_than_current_capacity = requested_num_cpus
+                    > cluster_num_cpus
+                    || requested_num_gpus > cluster_num_gpus
+                    || requested_memory_bytes > cluster_memory_bytes;
+
+                let resource_request_greater_than_max_requested = requested_num_cpus
+                    > state.max_resources_requested.num_cpus().unwrap_or(0.0)
+                    || requested_num_gpus > state.max_resources_requested.num_gpus().unwrap_or(0.0)
+                    || requested_memory_bytes
+                        > state.max_resources_requested.memory_bytes().unwrap_or(0);
+
+                let cluster_is_zero_capacity =
+                    cluster_num_cpus <= 0.0 && cluster_num_gpus <= 0.0 && cluster_memory_bytes == 0;
+                let should_bootstrap = cluster_is_zero_capacity
+                    && (requested_num_cpus > 0.0
+                        || requested_num_gpus > 0.0
+                        || requested_memory_bytes > 0);
+
+                // Only autoscale if we need more capacity AND this is greater than we've seen before
+                if (resource_request_greater_than_current_capacity
+                    && resource_request_greater_than_max_requested)
+                    || should_bootstrap
+                {
+                    // On scale-up demand, allow previously blacklisted workers to be reused immediately.
+                    state.pending_release_blacklist.clear();
+                    state.last_refresh = None;
+                    state.max_resources_requested = ResourceRequest::try_new_internal(
+                        Some(requested_num_cpus),
+                        Some(requested_num_gpus),
+                        Some(requested_memory_bytes),
+                    )?;
+
+                    let python_bundles = bundles
+                        .iter()
+                        .map(|bundle| {
+                            let mut dict = HashMap::new();
+                            dict.insert("CPU", bundle.num_cpus().ceil() as i64);
+                            dict.insert("GPU", bundle.num_gpus().ceil() as i64);
+                            dict.insert("memory", bundle.memory_bytes() as i64);
+                            dict
+                        })
+                        .collect::<Vec<_>>();
+
+                    action = AutoscaleAction::ScaleUp { python_bundles };
+                }
+            }
+        }
+
+        match action {
+            AutoscaleAction::Noop => Ok(()),
+            AutoscaleAction::Clear => {
+                Python::attach(|py| -> DaftResult<()> {
+                    let flotilla_module = py.import(pyo3::intern!(py, "daft.runners.flotilla"))?;
+                    flotilla_module
+                        .call_method0(pyo3::intern!(py, "clear_autoscaling_requests"))?;
+                    Ok(())
+                })?;
                 Ok(())
-            })?;
+            }
+            AutoscaleAction::ScaleUp { python_bundles } => {
+                Python::attach(|py| -> DaftResult<()> {
+                    let flotilla_module = py.import(pyo3::intern!(py, "daft.runners.flotilla"))?;
+                    flotilla_module
+                        .call_method1(pyo3::intern!(py, "try_autoscale"), (python_bundles,))?;
+                    Ok(())
+                })?;
+                Ok(())
+            }
         }
-        Ok(())
     }
 
     fn retire_idle_ray_workers(