Thanks to visit codestin.com
Credit goes to github.com

Skip to content

Commit 626b226

Browse files
spiralispirali
committed
Rust tests fixed
1 parent 9f90f95 commit 626b226

10 files changed

Lines changed: 238 additions & 482 deletions

File tree

‎crates/pyhq/src/client/job.rs‎

Lines changed: 0 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -160,9 +160,7 @@ fn build_tasks(
160160
.into_iter()
161161
.map(|id| id.into())
162162
.collect(),
163-
data_deps: Vec::new(),
164163
task_desc: build_task_desc(task, submit_dir)?,
165-
data_flags: TaskDataFlags::empty(),
166164
})
167165
})
168166
.collect()

‎crates/tako/benches/benchmarks/worker.rs‎

Lines changed: 0 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -66,12 +66,10 @@ fn create_worker_task(id: u32, resource_rq_id: ResourceRqId) -> Task {
6666
instance_id: Default::default(),
6767
priority: Default::default(),
6868
node_list: vec![],
69-
data_deps: vec![],
7069
entry: None,
7170
},
7271
ComputeTaskSharedData {
7372
time_limit: None,
74-
data_flags: TaskDataFlags::empty(),
7573
body: Default::default(),
7674
},
7775
TaskState::Waiting(0),

‎crates/tako/src/internal/scheduler/query.rs‎

Lines changed: 10 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -20,13 +20,22 @@ pub(crate) fn compute_new_worker_query(
2020

2121
let fake_worker_id_base = core.worker_counter() + 1;
2222
let mut fake_worker_counter = fake_worker_id_base;
23+
24+
/* Make sure that all resources provided by Worker has an Id */
25+
for query in queries {
26+
for item in &query.descriptor.resources {
27+
dbg!(&item.name);
28+
core.get_or_create_resource_id(&item.name);
29+
}
30+
}
31+
2332
let resource_map = core.resource_map().create_resource_id_map();
2433
let mut fake_workers = Vec::new();
2534
let now = std::time::Instant::now();
35+
2636
queries.iter().for_each(|query| {
2737
for _ in 0..query.max_sn_workers {
2838
let mut resources = query.descriptor.clone();
29-
3039
if query.partial {
3140
// If query is partial, add a fake maximal resources for resource that was not explicitly defined
3241
for name in resource_map.iter_names() {

‎crates/tako/src/internal/scheduler/taskqueue.rs‎

Lines changed: 18 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -125,6 +125,22 @@ impl TaskQueue {
125125
}
126126
}
127127

128+
#[cfg(test)]
129+
pub(crate) fn is_ready(&self, task_id: TaskId, priority: Priority) -> bool {
130+
let Some(item) = self.queue.get(&Reverse(priority)) else {
131+
return false;
132+
};
133+
match item {
134+
OneOrMoreTaskIds::One(t) => *t == task_id,
135+
OneOrMoreTaskIds::More(ts) => ts.contains(&task_id),
136+
}
137+
}
138+
139+
#[cfg(test)]
140+
pub(crate) fn is_prefilled(&self, task_id: TaskId) -> bool {
141+
self.prefill.as_ref().unwrap().1.contains(&task_id)
142+
}
143+
128144
fn check_dispose_prefill(&mut self, priority: Priority, retracted: &mut Vec<TaskId>) {
129145
if self.prefill.as_ref().map_or(false, |(p, _)| *p < priority) {
130146
let (p, ts) = self.prefill.take().unwrap();
@@ -202,7 +218,7 @@ impl TaskQueue {
202218
}
203219

204220
pub fn shrink_to_fit(&mut self) {
205-
todo!()
221+
// Do nothing
206222
}
207223

208224
pub fn size(&self) -> u32 {
@@ -294,7 +310,7 @@ impl TaskQueue {
294310
if let Some(prefill) = &mut self.prefill {
295311
assert_eq!(prefill.0, priority);
296312
for task_id in &result {
297-
prefill.1.remove(task_id);
313+
prefill.1.insert(*task_id);
298314
}
299315
} else {
300316
self.prefill = Some((priority, result.iter().copied().collect()))

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

Lines changed: 12 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -384,6 +384,13 @@ impl Core {
384384
}
385385
assert_eq!(*unfinished_deps, count);
386386
worker_check_sn(self, task.id, None, None);
387+
let priority = task.priority();
388+
assert_eq!(
389+
self.task_queues
390+
.get(task.resource_rq_id)
391+
.is_ready(task_id, priority),
392+
*unfinished_deps == 0
393+
)
387394
}
388395

389396
TaskRuntimeState::Assigned { worker_id, .. }
@@ -409,6 +416,11 @@ impl Core {
409416
TaskRuntimeState::Prefilled { worker_id } => {
410417
fw_check(task);
411418
worker_check_sn(self, task.id, None, Some(*worker_id));
419+
assert!(
420+
self.task_queues
421+
.get(task.resource_rq_id)
422+
.is_prefilled(task_id)
423+
);
412424
}
413425
TaskRuntimeState::RunningMultiNode(ws) => {
414426
assert!(!ws.is_empty());

‎crates/tako/src/internal/server/reactor.rs‎

Lines changed: 46 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -278,8 +278,7 @@ pub(crate) fn on_task_update(
278278
task_id,
279279
rv_id: rv_id,
280280
} => {
281-
task_reject(core, &mut *comm, worker_id, task_id, rv_id);
282-
need_scheduling = true;
281+
need_scheduling |= task_reject(core, &mut *comm, worker_id, task_id, rv_id);
283282
}
284283
WorkerTaskUpdate::EnableRequest {
285284
resource_rq_id: rq_id,
@@ -393,19 +392,23 @@ fn task_reject(
393392
worker_id: WorkerId,
394393
task_id: TaskId,
395394
resource_rq_variant: ResourceVariantId,
396-
) {
395+
) -> bool {
397396
let CoreSplitMut {
398397
task_map,
399398
task_queues,
400399
worker_map,
401400
request_map,
401+
scheduler_state,
402402
..
403403
} = core.split_mut();
404404
let Some(task) = task_map.find_task_mut(task_id) else {
405405
log::debug!("Unknown task rejected id={task_id}");
406-
return;
406+
return false;
407407
};
408408
log::debug!("Task id={task_id} (variant={resource_rq_variant}) rejected on worker={worker_id}");
409+
let worker = worker_map.get_worker_mut(worker_id);
410+
let resource_rq_id = task.resource_rq_id;
411+
worker.block_request(resource_rq_id, resource_rq_variant);
409412
match &task.state {
410413
TaskRuntimeState::Assigned {
411414
worker_id: w_id,
@@ -417,12 +420,35 @@ fn task_reject(
417420
if resource_rq_variant != *rv_id {
418421
log::debug!("Rejection from invalid worker");
419422
}
423+
let rq = request_map.get(resource_rq_id).get(resource_rq_variant);
424+
worker.remove_sn_task(task_id, rq);
420425
}
421-
TaskRuntimeState::Prefilled { .. } => {
422-
todo!()
426+
TaskRuntimeState::Prefilled { worker_id: w_id } => {
427+
if worker_id != *w_id {
428+
log::debug!("Rejection from invalid worker");
429+
}
430+
worker.remove_prefill_task(task_id);
431+
task_queues
432+
.get_mut(resource_rq_id)
433+
.remove_prefilled(task_id);
423434
}
424-
TaskRuntimeState::Retracting { .. } => {
425-
todo!()
435+
TaskRuntimeState::Retracting { worker_id: w_id } => {
436+
if worker_id != *w_id {
437+
log::debug!("Rejection from invalid worker");
438+
return false;
439+
}
440+
if let Some((target_id, rv_id)) = scheduler_state.redirects.remove(&task_id) {
441+
log::debug!("Transfering to {target_id}");
442+
task.state = TaskRuntimeState::Assigned {
443+
worker_id: target_id,
444+
rv_id,
445+
};
446+
comm.send_worker_message(
447+
target_id,
448+
&ComputeTasksBuilder::single_task(task, rv_id, Vec::new()),
449+
);
450+
return false;
451+
}
426452
}
427453
TaskRuntimeState::Waiting { .. }
428454
| TaskRuntimeState::Running { .. }
@@ -432,15 +458,10 @@ fn task_reject(
432458
}
433459
};
434460
task.state = TaskRuntimeState::Waiting { unfinished_deps: 0 };
435-
let resource_rq_id = task.resource_rq_id;
436461
let mut retracted = Vec::new();
437462
task_queues.add_ready_task(&task, &mut retracted);
438463
process_retracted(task_map, worker_map, comm, retracted);
439-
let rq = request_map.get(resource_rq_id).get(resource_rq_variant);
440-
441-
let worker = worker_map.get_worker_mut(worker_id);
442-
worker.remove_sn_task(task_id, rq);
443-
worker.block_request(resource_rq_id, resource_rq_variant);
464+
true
444465
}
445466

446467
fn request_enabled(
@@ -615,6 +636,7 @@ fn task_failed(
615636
worker_map,
616637
request_map,
617638
task_queues,
639+
scheduler_state,
618640
..
619641
} = core.split_mut();
620642
if let Some(task) = task_map.find_task(task_id) {
@@ -649,6 +671,16 @@ fn task_failed(
649671
.get_worker_mut(worker_id)
650672
.remove_prefill_task(task_id);
651673
}
674+
TaskRuntimeState::Retracting { worker_id: w } => {
675+
assert_eq!(worker_id, *w);
676+
try_remove_redirection(
677+
worker_map,
678+
scheduler_state,
679+
request_map,
680+
task_id,
681+
task.resource_rq_id,
682+
);
683+
}
652684
_ => {}
653685
}
654686
}

‎crates/tako/src/internal/tests/test_reactor.rs‎

Lines changed: 42 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -305,10 +305,11 @@ fn test_assignments_and_finish() {
305305
comm.check_need_scheduling();
306306
assert_eq!(comm.client.take_task_finished(1)[0], t2);
307307
comm.emptiness_check();
308-
308+
rt.sanity_check();
309309
on_task_update(rt.core(), &mut comm, ws[0], smallvec![task_finished(t1)]);
310-
310+
assert_eq!(comm.client.take_task_finished(1)[0], t1);
311311
comm.check_need_scheduling();
312+
comm.emptiness_check();
312313

313314
let mut comm = rt.schedule();
314315

@@ -329,8 +330,6 @@ fn test_assignments_and_finish() {
329330
..
330331
}) if tasks[0].id == t4
331332
));
332-
assert_eq!(comm.client.take_task_finished(1)[0], t1);
333-
comm.emptiness_check();
334333
rt.sanity_check();
335334

336335
on_task_update(rt.core(), &mut comm, ws[0], smallvec![task_finished(t3)]);
@@ -934,6 +933,23 @@ fn test_prefill_started() {
934933
rt.sanity_check();
935934
}
936935

936+
#[test]
937+
fn test_prefill_rejected() {
938+
let mut rt = TestEnv::new();
939+
let (w1, _t1, t2) = setup_prefill(&mut rt);
940+
let up = WorkerTaskUpdate::RejectRequest {
941+
task_id: t2,
942+
rv_id: 0.into(),
943+
};
944+
let mut comm = TestComm::new();
945+
on_task_update(rt.core(), &mut comm, w1, smallvec![up]);
946+
comm.check_need_scheduling();
947+
comm.emptiness_check();
948+
assert!(rt.task(t2).is_waiting());
949+
assert!(!rt.worker(w1).blocked_requests.is_empty());
950+
rt.sanity_check();
951+
}
952+
937953
#[test]
938954
fn test_prefill_failed() {
939955
let mut rt = TestEnv::new();
@@ -1103,6 +1119,28 @@ fn test_steal_source_worker_lost() {
11031119
rt.sanity_check();
11041120
}
11051121

1122+
#[test]
1123+
fn test_steal_rejected() {
1124+
let mut rt = TestEnv::new();
1125+
let (w1, w2, t) = setup_retracting(&mut rt);
1126+
let up = WorkerTaskUpdate::RejectRequest {
1127+
task_id: t,
1128+
rv_id: 0.into(),
1129+
};
1130+
let mut comm = TestComm::new();
1131+
on_task_update(rt.core(), &mut comm, w1, smallvec![up]);
1132+
comm.take_worker_msgs(w2, 1);
1133+
comm.emptiness_check();
1134+
match &rt.task(t).state {
1135+
TaskRuntimeState::Assigned { worker_id, .. } => {
1136+
assert_eq!(*worker_id, w2);
1137+
}
1138+
_ => panic!(),
1139+
}
1140+
assert!(!rt.worker(w1).blocked_requests.is_empty());
1141+
rt.sanity_check();
1142+
}
1143+
11061144
#[test]
11071145
fn test_steal_target_worker_lost() {
11081146
let mut rt = TestEnv::new();

0 commit comments

Comments
 (0)