Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Remove ExecutorReservation and change the task assignment philosophy from executor first to task first #823

Merged
merged 2 commits into from
Jul 13, 2023
Merged

Remove ExecutorReservation and change the task assignment philosophy from executor first to task first #823

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
fa1cad0
Remove ExecutorReservation and change the task assignment philosophy …
kyotoYaho Jun 28, 2023
8d574fd
Add PENDING_JOBS_METRIC and RUNNING_JOBS_METRIC to replace INFLIGHT_T…
kyotoYaho Jun 28, 2023
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
271 changes: 67 additions & 204 deletions ballista/scheduler/src/cluster/kv.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -17,13 +17,14 @@

use crate::cluster::storage::{KeyValueStore, Keyspace, Lock, Operation, WatchEvent};
use crate::cluster::{
reserve_slots_bias, reserve_slots_round_robin, ClusterState, ExecutorHeartbeatStream,
JobState, JobStateEvent, JobStateEventStream, JobStatus, TaskDistribution,
bind_task_bias, bind_task_round_robin, BoundTask, ClusterState,
ExecutorHeartbeatStream, ExecutorSlot, JobState, JobStateEvent, JobStateEventStream,
JobStatus, TaskDistribution,
};
use crate::scheduler_server::{timestamp_secs, SessionBuilder};
use crate::state::execution_graph::ExecutionGraph;
use crate::state::executor_manager::ExecutorReservation;
use crate::state::session_manager::create_datafusion_context;
use crate::state::task_manager::JobInfoCache;
use crate::state::{decode_into, decode_protobuf};
use async_trait::async_trait;
use ballista_core::config::BallistaConfig;
Expand Down Expand Up @@ -174,12 +175,12 @@ impl<S: KeyValueStore, T: 'static + AsLogicalPlan, U: 'static + AsExecutionPlan>
Ok(())
}

async fn reserve_slots(
async fn bind_schedulable_tasks(
&self,
num_slots: u32,
distribution: TaskDistribution,
active_jobs: Arc<HashMap<String, JobInfoCache>>,
executors: Option<HashSet<String>>,
) -> Result<Vec<ExecutorReservation>> {
) -> Result<Vec<BoundTask>> {
let lock = self.store.lock(Keyspace::Slots, "global").await?;

with_lock(lock, async {
Expand All @@ -192,7 +193,7 @@ impl<S: KeyValueStore, T: 'static + AsLogicalPlan, U: 'static + AsExecutionPlan>
))
})?;

let mut available_slots: Vec<&mut AvailableTaskSlots> = slots
let available_slots: Vec<&mut AvailableTaskSlots> = slots
.task_slots
.iter_mut()
.filter_map(|data| {
Expand All @@ -205,82 +206,33 @@ impl<S: KeyValueStore, T: 'static + AsLogicalPlan, U: 'static + AsExecutionPlan>
})
.collect();

available_slots.sort_by(|a, b| Ord::cmp(&b.slots, &a.slots));

let reservations = match distribution {
TaskDistribution::Bias => reserve_slots_bias(available_slots, num_slots),
let bound_tasks = match distribution {
TaskDistribution::Bias => {
bind_task_bias(available_slots, active_jobs, |_| false).await
}
TaskDistribution::RoundRobin => {
reserve_slots_round_robin(available_slots, num_slots)
bind_task_round_robin(available_slots, active_jobs, |_| false).await
}
};

if !reservations.is_empty() {
if !bound_tasks.is_empty() {
self.store
.put(Keyspace::Slots, "all".to_owned(), slots.encode_to_vec())
.await?
}

Ok(reservations)
Ok(bound_tasks)
})
.await
}

async fn reserve_slots_exact(
&self,
num_slots: u32,
distribution: TaskDistribution,
executors: Option<HashSet<String>>,
) -> Result<Vec<ExecutorReservation>> {
let lock = self.store.lock(Keyspace::Slots, "global").await?;

with_lock(lock, async {
let resources = self.store.get(Keyspace::Slots, "all").await?;

let mut slots =
ExecutorTaskSlots::decode(resources.as_slice()).map_err(|err| {
BallistaError::Internal(format!(
"Unexpected value in executor slots state: {err:?}"
))
})?;

let mut available_slots: Vec<&mut AvailableTaskSlots> = slots
.task_slots
.iter_mut()
.filter_map(|data| {
(data.slots > 0
&& executors
.as_ref()
.map(|executors| executors.contains(&data.executor_id))
.unwrap_or(true))
.then_some(data)
})
.collect();

available_slots.sort_by(|a, b| Ord::cmp(&b.slots, &a.slots));

let reservations = match distribution {
TaskDistribution::Bias => reserve_slots_bias(available_slots, num_slots),
TaskDistribution::RoundRobin => {
reserve_slots_round_robin(available_slots, num_slots)
}
};

if reservations.len() == num_slots as usize {
self.store
.put(Keyspace::Slots, "all".to_owned(), slots.encode_to_vec())
.await?;
Ok(reservations)
} else {
Ok(vec![])
}
})
.await
}
async fn unbind_tasks(&self, executor_slots: Vec<ExecutorSlot>) -> Result<()> {
let mut increments = HashMap::new();
for (executor_id, num_slots) in executor_slots {
let v = increments.entry(executor_id).or_insert_with(|| 0);
*v += num_slots;
}

async fn cancel_reservations(
&self,
reservations: Vec<ExecutorReservation>,
) -> Result<()> {
let lock = self.store.lock(Keyspace::Slots, "all").await?;

with_lock(lock, async {
Expand All @@ -293,18 +245,9 @@ impl<S: KeyValueStore, T: 'static + AsLogicalPlan, U: 'static + AsExecutionPlan>
))
})?;

let mut increments = HashMap::new();
for ExecutorReservation { executor_id, .. } in reservations {
if let Some(inc) = increments.get_mut(&executor_id) {
*inc += 1;
} else {
increments.insert(executor_id, 1usize);
}
}

for executor_slots in slots.task_slots.iter_mut() {
if let Some(slots) = increments.get(&executor_slots.executor_id) {
executor_slots.slots += *slots as u32;
executor_slots.slots += *slots;
}
}

Expand All @@ -319,8 +262,7 @@ impl<S: KeyValueStore, T: 'static + AsLogicalPlan, U: 'static + AsExecutionPlan>
&self,
metadata: ExecutorMetadata,
spec: ExecutorData,
reserve: bool,
) -> Result<Vec<ExecutorReservation>> {
) -> Result<()> {
let executor_id = metadata.id.clone();

//TODO this should be in a transaction
Expand All @@ -338,83 +280,40 @@ impl<S: KeyValueStore, T: 'static + AsLogicalPlan, U: 'static + AsExecutionPlan>
})
.await?;

if !reserve {
let available_slots = AvailableTaskSlots {
executor_id,
slots: spec.available_task_slots,
};

let lock = self.store.lock(Keyspace::Slots, "all").await?;
let available_slots = AvailableTaskSlots {
executor_id,
slots: spec.available_task_slots,
};

with_lock(lock, async {
let current_slots = self.store.get(Keyspace::Slots, "all").await?;
let lock = self.store.lock(Keyspace::Slots, "all").await?;

let mut current_slots: ExecutorTaskSlots =
decode_protobuf(current_slots.as_slice())?;
with_lock(lock, async {
let current_slots = self.store.get(Keyspace::Slots, "all").await?;

if let Some((idx, _)) =
current_slots.task_slots.iter().find_position(|slots| {
slots.executor_id == available_slots.executor_id
})
{
current_slots.task_slots[idx] = available_slots;
} else {
current_slots.task_slots.push(available_slots);
}
let mut current_slots: ExecutorTaskSlots =
decode_protobuf(current_slots.as_slice())?;

self.store
.put(
Keyspace::Slots,
"all".to_string(),
current_slots.encode_to_vec(),
)
.await
})
.await?;

Ok(vec![])
} else {
let num_slots = spec.available_task_slots as usize;
let mut reservations: Vec<ExecutorReservation> = vec![];
for _ in 0..num_slots {
reservations.push(ExecutorReservation::new_free(executor_id.clone()));
if let Some((idx, _)) = current_slots
.task_slots
.iter()
.find_position(|slots| slots.executor_id == available_slots.executor_id)
{
current_slots.task_slots[idx] = available_slots;
} else {
current_slots.task_slots.push(available_slots);
}

let available_slots = AvailableTaskSlots {
executor_id,
slots: 0,
};

let lock = self.store.lock(Keyspace::Slots, "all").await?;

with_lock(lock, async {
let current_slots = self.store.get(Keyspace::Slots, "all").await?;

let mut current_slots: ExecutorTaskSlots =
decode_protobuf(current_slots.as_slice())?;

if let Some((idx, _)) =
current_slots.task_slots.iter().find_position(|slots| {
slots.executor_id == available_slots.executor_id
})
{
current_slots.task_slots[idx] = available_slots;
} else {
current_slots.task_slots.push(available_slots);
}

self.store
.put(
Keyspace::Slots,
"all".to_string(),
current_slots.encode_to_vec(),
)
.await
})
.await?;
self.store
.put(
Keyspace::Slots,
"all".to_string(),
current_slots.encode_to_vec(),
)
.await
})
.await?;

Ok(reservations)
}
Ok(())
}

async fn save_executor_metadata(&self, metadata: ExecutorMetadata) -> Result<()> {
Expand Down Expand Up @@ -502,18 +401,17 @@ impl<S: KeyValueStore, T: 'static + AsLogicalPlan, U: 'static + AsExecutionPlan>
impl<S: KeyValueStore, T: 'static + AsLogicalPlan, U: 'static + AsExecutionPlan> JobState
for KeyValueState<S, T, U>
{
async fn accept_job(
&self,
job_id: &str,
job_name: &str,
queued_at: u64,
) -> Result<()> {
fn accept_job(&self, job_id: &str, job_name: &str, queued_at: u64) -> Result<()> {
self.queued_jobs
.insert(job_id.to_string(), (job_name.to_string(), queued_at));

Ok(())
}

fn pending_job_number(&self) -> usize {
self.queued_jobs.len()
}

async fn submit_job(&self, job_id: String, graph: &ExecutionGraph) -> Result<()> {
if self.queued_jobs.get(&job_id).is_some() {
let status = graph.status();
Expand Down Expand Up @@ -774,62 +672,17 @@ async fn with_lock<Out, F: Future<Output = Out>>(mut lock: Box<dyn Lock>, op: F)

#[cfg(test)]
mod test {

use crate::cluster::kv::KeyValueState;
use crate::cluster::storage::sled::SledClient;
use crate::cluster::test::{
test_executor_registration, test_fuzz_reservations, test_job_lifecycle,
test_job_planning_failure, test_reservation,
};
use crate::cluster::TaskDistribution;
use crate::cluster::test_util::{test_job_lifecycle, test_job_planning_failure};
use crate::test_utils::{
test_aggregation_plan, test_join_plan, test_two_aggregations_plan,
};
use ballista_core::error::Result;
use ballista_core::serde::BallistaCodec;
use ballista_core::utils::default_session_builder;

#[cfg(feature = "sled")]
fn make_sled_state() -> Result<KeyValueState<SledClient>> {
Ok(KeyValueState::new(
"",
SledClient::try_new_temporary()?,
BallistaCodec::default(),
default_session_builder,
))
}

#[cfg(feature = "sled")]
#[tokio::test]
async fn test_sled_executor_reservation() -> Result<()> {
test_executor_registration(make_sled_state()?).await
}

#[cfg(feature = "sled")]
#[tokio::test]
async fn test_sled_reserve() -> Result<()> {
test_reservation(make_sled_state()?, TaskDistribution::Bias).await?;
test_reservation(make_sled_state()?, TaskDistribution::RoundRobin).await?;

Ok(())
}

#[cfg(feature = "sled")]
#[tokio::test]
async fn test_sled_fuzz_reserve() -> Result<()> {
test_fuzz_reservations(make_sled_state()?, 10, TaskDistribution::Bias, 10, 10)
.await?;
test_fuzz_reservations(
make_sled_state()?,
10,
TaskDistribution::RoundRobin,
10,
10,
)
.await?;

Ok(())
}

#[cfg(feature = "sled")]
#[tokio::test]
async fn test_sled_job_lifecycle() -> Result<()> {
Expand All @@ -854,4 +707,14 @@ mod test {

Ok(())
}

#[cfg(feature = "sled")]
fn make_sled_state() -> Result<KeyValueState<SledClient>> {
Ok(KeyValueState::new(
"",
SledClient::try_new_temporary()?,
BallistaCodec::default(),
default_session_builder,
))
}
}
Loading