fix: improve worker launch logic

crates/sail-common/src/config/application.rs

@@ -209,6 +209,7 @@ pub struct ClusterConfig {
     pub worker_heartbeat_interval_secs: u64,
     pub worker_heartbeat_timeout_secs: u64,
     pub worker_launch_timeout_secs: u64,
+    pub worker_launch_retry_strategy: RetryStrategy,
     pub worker_task_slots: usize,
     pub task_launch_timeout_secs: u64,
     pub task_stream_buffer: usize,

crates/sail-common/src/config/application.yaml

@@ -171,6 +171,50 @@
   default: "120"
   description: The timeout in seconds for launching a worker.
 
+- key: cluster.worker_launch_retry_strategy.type
+  type: string
+  default: "exponential_backoff"
+  description: |
+    The retry strategy for failed worker launches.
+    Valid values are `fixed` and `exponential_backoff`.
+  experimental: true
+
+- key: cluster.worker_launch_retry_strategy.fixed.max_count
+  type: number
+  default: "3"
+  description: The maximum number of worker launch retries using a fixed delay.
+  experimental: true
+
+- key: cluster.worker_launch_retry_strategy.fixed.delay_secs
+  type: number
+  default: "5"
+  description: The delay in seconds between worker launch retries using a fixed delay.
+  experimental: true
+
+- key: cluster.worker_launch_retry_strategy.exponential_backoff.max_count
+  type: number
+  default: "5"
+  description: The maximum number of worker launch retries using exponential backoff.
+  experimental: true
+
+- key: cluster.worker_launch_retry_strategy.exponential_backoff.initial_delay_secs
+  type: number
+  default: "1"
+  description: The initial delay in seconds before retrying a worker launch.
+  experimental: true
+
+- key: cluster.worker_launch_retry_strategy.exponential_backoff.max_delay_secs
+  type: number
+  default: "30"
+  description: The maximum delay in seconds before retrying a worker launch.
+  experimental: true
+
+- key: cluster.worker_launch_retry_strategy.exponential_backoff.factor
+  type: number
+  default: "2"
+  description: The factor by which the worker launch retry delay increases.
+  experimental: true
+
 - key: cluster.worker_task_slots
   type: number
   default: "8"

crates/sail-common/src/utils/retry.rs

@@ -21,19 +21,55 @@ pub enum RetryStrategy {
     },
 }
 
-struct ExponentialBackoffDelay {
-    delay: Duration,
-    max_delay: Duration,
-    factor: u32,
+#[derive(Debug, Clone, Copy, PartialEq, Eq)]
+pub struct RetryStep {
+    /// The retry number. Retries are numbered from one; the initial attempt is zero.
+    pub retry: usize,
+    pub delay: Duration,
 }
 
-impl Iterator for ExponentialBackoffDelay {
-    type Item = Duration;
+#[derive(Debug, Clone)]
+pub struct RetrySchedule {
+    next: usize,
+    remaining: usize,
+    kind: RetryScheduleKind,
+}
+
+#[derive(Debug, Clone)]
+enum RetryScheduleKind {
+    Fixed {
+        delay: Duration,
+    },
+    ExponentialBackoff {
+        delay: Duration,
+        max_delay: Duration,
+        factor: u32,
+    },
+}
+
+impl Iterator for RetrySchedule {
+    type Item = RetryStep;
 
     fn next(&mut self) -> Option<Self::Item> {
-        let delay = self.delay;
-        self.delay = std::cmp::min(delay * self.factor, self.max_delay);
-        Some(delay)
+        if self.remaining == 0 {
+            return None;
+        }
+        let retry = self.next;
+        self.next += 1;
+        self.remaining -= 1;
+        let delay = match &mut self.kind {
+            RetryScheduleKind::Fixed { delay } => *delay,
+            RetryScheduleKind::ExponentialBackoff {
+                delay,
+                max_delay,
+                factor,
+            } => {
+                let current = *delay;
+                *delay = std::cmp::min(delay.saturating_mul(*factor), *max_delay);
+                current
+            }
+        };
+        Some(RetryStep { retry, delay })
     }
 }
 
@@ -45,7 +81,7 @@ impl RetryStrategy {
         T: Send + 'static,
         E: std::fmt::Display + Send + 'static,
     {
-        let mut delay = self.delay();
+        let mut retries = self.retries();
         let mut attempt = 0;
         loop {
             let span = Span::enter_with_local_parent("RetryStrategy::run")
@@ -55,33 +91,41 @@ impl RetryStrategy {
                 x @ Ok(_) => return x,
                 Err(e) => {
                     warn!("retryable operation failed: {e}");
-                    if let Some(delay) = delay.next() {
-                        tokio::time::sleep(delay).await;
+                    if let Some(step) = retries.next() {
+                        tokio::time::sleep(step.delay).await;
+                        attempt = step.retry;
                     } else {
                         return Err(e);
                     }
                 }
             }
-            attempt += 1;
         }
     }
 
-    fn delay(&self) -> Box<dyn Iterator<Item = Duration> + Send> {
+    /// Returns a finite schedule containing only retries after the initial attempt.
+    ///
+    /// The first item has retry number one. If `max_count` is zero, the schedule is empty.
+    pub fn retries(&self) -> RetrySchedule {
         match self {
             Self::ExponentialBackoff {
                 max_count,
                 initial_delay,
                 max_delay,
                 factor,
-            } => Box::new(
-                ExponentialBackoffDelay {
+            } => RetrySchedule {
+                next: 1,
+                remaining: *max_count,
+                kind: RetryScheduleKind::ExponentialBackoff {
                     delay: *initial_delay,
                     max_delay: *max_delay,
                     factor: *factor,
-                }
-                .take(*max_count),
-            ),
-            Self::Fixed { max_count, delay } => Box::new(std::iter::repeat_n(*delay, *max_count)),
+                },
+            },
+            Self::Fixed { max_count, delay } => RetrySchedule {
+                next: 1,
+                remaining: *max_count,
+                kind: RetryScheduleKind::Fixed { delay: *delay },
+            },
         }
     }
 }
@@ -112,3 +156,93 @@ impl From<&config::RetryStrategy> for RetryStrategy {
         }
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use std::sync::Arc;
+    use std::sync::atomic::{AtomicUsize, Ordering};
+    use std::time::Duration;
+
+    use super::{RetryStep, RetryStrategy};
+
+    #[test]
+    fn fixed_schedule_contains_one_based_retries() {
+        let strategy = RetryStrategy::Fixed {
+            max_count: 3,
+            delay: Duration::from_secs(5),
+        };
+
+        assert_eq!(
+            strategy.retries().collect::<Vec<_>>(),
+            vec![
+                RetryStep {
+                    retry: 1,
+                    delay: Duration::from_secs(5),
+                },
+                RetryStep {
+                    retry: 2,
+                    delay: Duration::from_secs(5),
+                },
+                RetryStep {
+                    retry: 3,
+                    delay: Duration::from_secs(5),
+                },
+            ]
+        );
+    }
+
+    #[test]
+    fn zero_max_count_has_no_retries() {
+        let strategy = RetryStrategy::Fixed {
+            max_count: 0,
+            delay: Duration::from_secs(5),
+        };
+
+        assert_eq!(strategy.retries().next(), None);
+    }
+
+    #[test]
+    fn exponential_backoff_schedule_is_capped() {
+        let strategy = RetryStrategy::ExponentialBackoff {
+            max_count: 4,
+            initial_delay: Duration::from_secs(2),
+            max_delay: Duration::from_secs(5),
+            factor: 2,
+        };
+
+        let retries = strategy
+            .retries()
+            .map(|step| (step.retry, step.delay))
+            .collect::<Vec<_>>();
+        assert_eq!(
+            retries,
+            vec![
+                (1, Duration::from_secs(2)),
+                (2, Duration::from_secs(4)),
+                (3, Duration::from_secs(5)),
+                (4, Duration::from_secs(5)),
+            ]
+        );
+    }
+
+    #[tokio::test]
+    async fn run_performs_initial_attempt_and_scheduled_retries() {
+        let strategy = RetryStrategy::Fixed {
+            max_count: 2,
+            delay: Duration::ZERO,
+        };
+        let calls = Arc::new(AtomicUsize::new(0));
+        let result: Result<(), &str> = strategy
+            .run({
+                let calls = Arc::clone(&calls);
+                move || {
+                    calls.fetch_add(1, Ordering::Relaxed);
+                    async { Err("failed") }
+                }
+            })
+            .await;
+
+        assert_eq!(result, Err("failed"));
+        assert_eq!(calls.load(Ordering::Relaxed), 3);
+    }
+}

crates/sail-execution/src/driver/actor/core.rs

@@ -12,6 +12,7 @@ use sail_common::actor::{Actor, ActorAction, ActorContext};
 use crate::driver::job_scheduler::{JobScheduler, JobSchedulerOptions};
 use crate::driver::task_assigner::{TaskAssigner, TaskAssignerOptions};
 use crate::driver::worker_pool::{WorkerPool, WorkerPoolOptions};
+use crate::driver::worker_scaler::{WorkerScaler, WorkerScalerOptions};
 use crate::driver::{DriverActor, DriverComponents, DriverMessage, DriverOptions};
 use crate::shuffle::{ShuffleBackendKind, celeborn_application_id};
 use crate::stream::celeborn::CelebornStreamManager;
@@ -43,13 +44,16 @@ impl Actor for DriverActor {
         );
         let job_scheduler = JobScheduler::new(JobSchedulerOptions::from(&options), event_reporter);
         let task_assigner = TaskAssigner::new(TaskAssignerOptions::from(&options));
+        let worker_scaler = WorkerScaler::new(WorkerScalerOptions::from(&options));
         Self {
             options,
             worker_pool,
             job_scheduler,
             task_assigner,
+            worker_scaler,
             task_runner: None,
             extensions: Default::default(),
+            activated: false,
             task_sequences: HashMap::new(),
             shutdown_notifier: None,
         }
@@ -129,7 +133,7 @@ impl Actor for DriverActor {
         message: DriverMessage,
     ) -> ActorAction {
         match message {
-            DriverMessage::Activate => self.handle_activate(ctx),
+            DriverMessage::Activate { result } => self.handle_activate(ctx, result),
             DriverMessage::RegisterWorker {
                 worker_id,
                 host,
@@ -146,6 +150,12 @@ impl Actor for DriverActor {
             DriverMessage::ProbePendingWorker { worker_id } => {
                 self.handle_probe_pending_worker(ctx, worker_id)
             }
+            DriverMessage::WorkerFailedToStart { worker_id, message } => {
+                self.handle_worker_failed_to_start(ctx, worker_id, message)
+            }
+            DriverMessage::RetryWorkerDemand { request } => {
+                self.handle_retry_worker_demand(ctx, request)
+            }
             DriverMessage::ProbeIdleWorker { worker_id, instant } => {
                 self.handle_probe_idle_worker(ctx, worker_id, instant)
             }