Followup for #1119

CHANGELOG.md

@@ -2,10 +2,15 @@
 
 ## Unreleased
 
+### New features
+
+* The server scheduler now contains a safety limit for computation, configurable via `--scheduler-time-limit` (default: 5s)
+
 ### Fixes
 
 * Fixed server crash in a specific situation when an unschedulable high-priority task occurs
 
+
 ## v0.26.2
 
 ### Fixes

crates/hyperqueue/src/client/commands/server.rs

@@ -130,6 +130,14 @@ pub struct ServerStartOpts {
     /// USE AT YOUR OWN RISK.
     #[arg(long)]
     disable_worker_authentication_and_encryption: bool,
+
+    #[arg(
+        long,
+        value_parser = parse_hms_or_human_time,
+        default_value = "5s",
+        help = duration_doc!("Maximum time the scheduler's placement solve may run per round.")
+    )]
+    scheduler_time_limit: Duration,
 }
 
 #[derive(Parser)]
@@ -222,6 +230,7 @@ async fn start_server(gsettings: &GlobalSettings, opts: ServerStartOpts) -> anyh
                 }
             }),
         server_uid: access_file.as_ref().map(|a| a.server_uid().to_string()),
+        scheduler_mip_time_limit: opts.scheduler_time_limit,
     };
 
     init_hq_server(gsettings, server_cfg).await

crates/hyperqueue/src/server/bootstrap.rs

@@ -48,6 +48,8 @@ pub struct ServerConfig {
     pub worker_secret_key: Option<Arc<SecretKey>>,
     pub client_secret_key: Option<Arc<SecretKey>>,
     pub server_uid: Option<String>,
+    /// See `tako::server::SchedulerConfig::mip_time_limit`.
+    pub scheduler_mip_time_limit: Duration,
 }
 
 /// This function initializes the HQ server.
@@ -204,6 +206,10 @@ pub async fn initialize_server(
         None,
         state_ref.get().server_info().server_uid.clone(),
         worker_id_initial_value,
+        tako::server::SchedulerConfig {
+            mip_time_limit: server_cfg.scheduler_mip_time_limit,
+            ..Default::default()
+        },
     )?;
     let (autoalloc_service, autoalloc_process) =
         create_autoalloc_service(server_ref.clone(), queue_id_initial_value, events.clone());

crates/hyperqueue/src/tests/server.rs

@@ -57,6 +57,7 @@ where
         worker_secret_key: None,
         client_secret_key: None,
         server_uid: None,
+        scheduler_mip_time_limit: Duration::from_secs(5),
     };
     let (fut, notify, _state, _senders) =
         initialize_server(&gsettings, server_cfg, 1.into(), 1, None)

crates/pyhq/src/cluster/server.rs

@@ -40,6 +40,7 @@ impl RunningServer {
                 worker_secret_key: None,
                 server_uid: None,
                 journal_flush_period: Duration::from_secs(30),
+                scheduler_mip_time_limit: Duration::from_secs(5),
             };
 
             let main_future = async move {

crates/tako/src/control.rs

@@ -16,7 +16,9 @@ use crate::internal::common::error::DsError;
 use crate::internal::common::resources::ResourceRqId;
 use crate::internal::messages::worker::ToWorkerMessage;
 use crate::internal::scheduler::query::compute_new_worker_query;
-use crate::internal::scheduler::{run_scheduling, scheduler_loop};
+use crate::internal::scheduler::{
+    SchedulerConfig, SchedulerResult, run_scheduling, scheduler_loop,
+};
 use crate::internal::server::client::handle_new_tasks;
 use crate::internal::server::comm::{Comm, CommSenderRef};
 use crate::internal::server::core::{CoreRef, CustomConnectionHandler};
@@ -130,7 +132,24 @@ impl ServerRef {
         let mut core = self.core_ref.get_mut();
         let mut comm = self.comm_ref.get_mut();
         if comm.get_scheduling_flag() {
-            run_scheduling(&mut core, &mut comm, Instant::now())
+            // Try to finish scheduling, at most 3 times
+            if (0..3)
+                .find_map(
+                    |_| match run_scheduling(&mut core, &mut comm, Instant::now()) {
+                        SchedulerResult::Done => Some(true),
+                        SchedulerResult::NoProgress => Some(false),
+                        SchedulerResult::NeedMoreCompute => None,
+                    },
+                )
+                .unwrap_or(false)
+            {
+                comm.reset_scheduling_flag();
+            } else {
+                return Ok(NewWorkerAllocationResponse {
+                    single_node_workers_per_query: Vec::new(),
+                    multi_node_allocations: Vec::new(),
+                });
+            }
         }
         Ok(compute_new_worker_query(&mut core, queries))
     }
@@ -237,6 +256,7 @@ pub fn server_start(
     custom_conn_handler: Option<CustomConnectionHandler>,
     server_uid: String,
     worker_id_initial_value: WorkerId,
+    scheduler_config: SchedulerConfig,
 ) -> crate::Result<(ServerRef, impl Future<Output = crate::Result<()>>)> {
     let listener_port = listener.local_addr()?.port();
 
@@ -251,6 +271,7 @@ pub fn server_start(
         custom_conn_handler,
         server_uid,
         worker_id_initial_value,
+        scheduler_config,
     );
     let connections = crate::internal::server::rpc::connection_initiator(
         listener,

crates/tako/src/internal/scheduler/main.rs

@@ -16,40 +16,69 @@ pub(crate) async fn scheduler_loop(
     let mut last_schedule = Instant::now().checked_sub(minimum_delay * 2).unwrap();
     loop {
         scheduler_wakeup.notified().await;
+        let mut now = Instant::now();
         if !comm_ref.get().get_scheduling_flag() {
-            last_schedule = Instant::now();
+            last_schedule = now;
             continue;
         }
-        let mut now = Instant::now();
         let since_last_schedule = now - last_schedule;
         if minimum_delay > since_last_schedule {
             sleep(minimum_delay - since_last_schedule).await;
             now = Instant::now();
         }
-        let mut comm = comm_ref.get_mut();
-        if !comm.get_scheduling_flag() {
+        if !comm_ref.get_mut().get_scheduling_flag() {
             last_schedule = now;
             continue;
         }
-        let mut core = core_ref.get_mut();
-        run_scheduling(&mut core, &mut comm, now);
+        while matches!(
+            run_scheduling(&mut core_ref.get_mut(), &mut comm_ref.get_mut(), now),
+            SchedulerResult::NeedMoreCompute
+        ) {
+            sleep(minimum_delay).await;
+        }
+        comm_ref.get_mut().reset_scheduling_flag();
         last_schedule = Instant::now();
     }
 }
 
-pub(crate) fn run_scheduling_inner(core: &mut Core, comm: &mut impl Comm, now: Instant) {
+pub(crate) enum SchedulerResult {
+    Done,
+    NeedMoreCompute,
+    NoProgress,
+}
+
+pub(crate) fn run_scheduling_inner(
+    core: &mut Core,
+    comm: &mut impl Comm,
+    now: Instant,
+) -> SchedulerResult {
     let batches = create_task_batches(core, now, None);
     let solution = run_scheduling_solver(core, now, &batches, None);
+    let need_more_compute = if !solution.is_optimal {
+        if solution.is_empty() {
+            log::error!("Scheduler made no progress within given time limit");
+            SchedulerResult::NoProgress
+        } else {
+            log::debug!("Scheduler dispatched a non-optimal placement this round");
+            SchedulerResult::NeedMoreCompute
+        }
+    } else {
+        SchedulerResult::Done
+    };
     let mapping = create_task_mapping(core, solution);
     //mapping.dump();
     mapping.send_messages(core, comm);
+    need_more_compute
 }
 
-pub(crate) fn run_scheduling(core: &mut Core, comm: &mut CommSender, now: Instant) {
+pub(crate) fn run_scheduling(
+    core: &mut Core,
+    comm: &mut CommSender,
+    now: Instant,
+) -> SchedulerResult {
     trace_time!(
         "scheduler",
         "run_scheduling_inner",
         run_scheduling_inner(core, comm, now)
-    );
-    comm.reset_scheduling_flag();
+    )
 }

crates/tako/src/internal/scheduler/mod.rs

@@ -8,8 +8,9 @@ mod state;
 mod taskqueue;
 
 pub(crate) use batches::{TaskBatch, create_task_batches};
-pub(crate) use main::{run_scheduling, scheduler_loop};
+pub(crate) use main::{SchedulerResult, run_scheduling, scheduler_loop};
 pub(crate) use solver::run_scheduling_solver;
+pub use state::SchedulerConfig;
 pub(crate) use state::SchedulerState;
 pub(crate) use taskqueue::TaskQueues;
 
@@ -20,7 +21,7 @@ pub(crate) use main::run_scheduling_inner;
 pub(crate) use mapping::{WorkerTaskMapping, create_task_mapping};
 
 #[cfg(test)]
-pub(crate) use state::SchedulerConfig;
+pub(crate) use solver::SchedulingSolution;
 
 #[cfg(test)]
 pub(crate) use batches::PriorityCut;

crates/tako/src/internal/scheduler/solver.rs

@@ -7,10 +7,30 @@ use crate::resources::{CPU_RESOURCE_ID, ResourceRqId};
 use crate::{Map, ResourceVariantId, Set, WorkerId};
 use thin_vec::ThinVec;
 
-#[derive(Default, Debug)]
+#[derive(Debug)]
 pub(crate) struct SchedulingSolution {
     pub(crate) sn_counts: Map<(ResourceRqId, ResourceVariantId), Map<WorkerId, u32>>,
     pub(crate) mn_workers: Map<(ResourceRqId, ResourceVariantId), Vec<ThinVec<WorkerId>>>,
+    /// `false` if the MILP solve hit its time limit before proving
+    /// optimality (see `SchedulerConfig::mip_time_limit`).
+    pub(crate) is_optimal: bool,
+}
+
+impl Default for SchedulingSolution {
+    fn default() -> Self {
+        SchedulingSolution {
+            sn_counts: Map::new(),
+            mn_workers: Map::new(),
+            is_optimal: true,
+        }
+    }
+}
+
+impl SchedulingSolution {
+    pub fn is_empty(&self) -> bool {
+        self.sn_counts.values().all(|v| v.is_empty())
+            && self.mn_workers.values().all(|v| v.is_empty())
+    }
 }
 
 pub(crate) fn run_scheduling_solver(
@@ -410,9 +430,11 @@ pub(crate) fn run_scheduling_solver(
     }
 
     let mut result = SchedulingSolution::default();
-    let Some((solution, _)) = solver.solve() else {
+    let Some((solution, is_optimal)) = solver.solve_bounded(scheduler_cache.config.mip_time_limit)
+    else {
         return result;
     };
+    result.is_optimal = is_optimal;
 
     for batch in task_batches {
         let resource_rq_id = batch.resource_rq_id;

crates/tako/src/internal/scheduler/state.rs

@@ -1,5 +1,6 @@
 use crate::internal::scheduler::gap::GapCache;
 use crate::{Map, ResourceVariantId, TaskId, WorkerId};
+use std::time::Duration;
 
 pub struct SchedulerConfig {
     /// The number of tasks that are never prefilled (wrt a given resource request)
@@ -9,17 +10,35 @@ pub struct SchedulerConfig {
     /// TODO: Maybe we can choose it dynamically wrt. resources of workers
     ///       but small tens looks reasonable
     pub proactive_filling_max: u32,
+    /// Hard wall-clock cap on a single scheduler MILP solve. An emergency
+    /// backstop, not a tuning knob: on expiry, the best incumbent found so
+    /// far is dispatched and marked non-optimal.
+    pub mip_time_limit: Duration,
 }
 
 impl Default for SchedulerConfig {
     fn default() -> Self {
         SchedulerConfig {
             proactive_filling_reserve: 16,
             proactive_filling_max: 40,
+            mip_time_limit: default_mip_time_limit(),
         }
     }
 }
 
+// Unit tests assert exact placement counts on small instances, so default to
+// a generous limit. Tests exercising the bounded solve set mip_time_limit
+// explicitly via TestEnv::set_scheduler_config.
+#[cfg(test)]
+fn default_mip_time_limit() -> Duration {
+    Duration::from_secs(60)
+}
+
+#[cfg(not(test))]
+fn default_mip_time_limit() -> Duration {
+    Duration::from_secs(5)
+}
+
 #[derive(Default)]
 pub(crate) struct SchedulerState {
     pub gap_cache: GapCache,

crates/tako/src/internal/server/core.rs

@@ -6,7 +6,7 @@ use crate::internal::common::resources::map::{
 };
 use crate::internal::common::resources::{ResourceId, ResourceRequestVariants, ResourceRqId};
 use crate::internal::common::{Set, WrappedRcRefCell};
-use crate::internal::scheduler::{SchedulerState, TaskQueues};
+use crate::internal::scheduler::{SchedulerConfig, SchedulerState, TaskQueues};
 use crate::internal::server::rpc::ConnectionDescriptor;
 use crate::internal::server::task::{Task, TaskRuntimeState};
 use crate::internal::server::taskmap::TaskMap;
@@ -64,13 +64,15 @@ pub(crate) struct Core {
 pub(crate) type CoreRef = WrappedRcRefCell<Core>;
 
 impl CoreRef {
+    #[allow(clippy::too_many_arguments)]
     pub fn new(
         worker_listen_port: u16,
         secret_key: Option<Arc<SecretKey>>,
         idle_timeout: Option<Duration>,
         custom_conn_handler: Option<CustomConnectionHandler>,
         server_uid: String,
         worker_id_initial_value: WorkerId,
+        scheduler_config: SchedulerConfig,
     ) -> Self {
         CoreRef::wrap(Core {
             worker_listen_port,
@@ -79,6 +81,10 @@ impl CoreRef {
             custom_conn_handler,
             server_uid,
             worker_id_counter: worker_id_initial_value.as_num(),
+            scheduler_state: SchedulerState {
+                config: scheduler_config,
+                ..Default::default()
+            },
             ..Default::default()
         })
     }