def test_worker_max_downloads_per_worker():
g = TaskGraph()
a = g.new_task("a", duration=0, outputs=[1, 1, 1, 1])
b = g.new_task("b", duration=0)
b.add_inputs(a.outputs)
s = fixed_scheduler([
(0, a, 0),
(1, b, 0),
])
assert do_sched_test(g, [1, 1], s, SimpleNetModel()) == 2
assert do_sched_test(
g, [Worker(), Worker(max_downloads_per_worker=1)], s,
SimpleNetModel()) == 4
assert do_sched_test(
g, [Worker(), Worker(max_downloads_per_worker=2)], s,
SimpleNetModel()) == 2
assert do_sched_test(
g, [Worker(), Worker(max_downloads_per_worker=3)], s,
SimpleNetModel()) == 2
assert do_sched_test(
g, [Worker(), Worker(max_downloads_per_worker=4)], s,
SimpleNetModel()) == 1
assert do_sched_test(
g, [Worker(), Worker(max_downloads_per_worker=5)], s,
SimpleNetModel()) == 1
def test_trace_task_execution():
tg = TaskGraph()
a = tg.new_task(output_size=0, duration=2)
b = tg.new_task(output_size=0, duration=3)
b.add_input(a)
c = tg.new_task(duration=4)
c.add_input(b)
simulator = do_sched_test(tg, [1, 1],
fixed_scheduler([(0, a, 0), (1, b, 0),
(0, c, 0)]),
trace=True,
return_simulator=True)
start_events = [
e for e in simulator.trace_events
if isinstance(e, TaskStartTraceEvent)
]
assert start_events == [
TaskStartTraceEvent(0, simulator.workers[0], a),
TaskStartTraceEvent(2, simulator.workers[1], b),
TaskStartTraceEvent(5, simulator.workers[0], c)
]
end_events = [
e for e in simulator.trace_events if isinstance(e, TaskEndTraceEvent)
]
assert end_events == [
TaskEndTraceEvent(2, simulator.workers[0], a),
TaskEndTraceEvent(5, simulator.workers[1], b),
TaskEndTraceEvent(9, simulator.workers[0], c)
]
def test_trace_task_fetch():
tg = TaskGraph()
a = tg.new_task(output_size=5, duration=2)
b = tg.new_task(output_size=3, duration=3)
b.add_input(a)
c = tg.new_task(duration=4)
c.add_input(b)
simulator = do_sched_test(tg, [1, 1],
fixed_scheduler([(0, a, 0), (1, b, 0),
(0, c, 0)]),
netmodel=SimpleNetModel(1),
trace=True,
return_simulator=True)
workers = simulator.workers
fetch_start_events = [
e for e in simulator.trace_events
if isinstance(e, FetchStartTraceEvent)
]
assert fetch_start_events == [
FetchStartTraceEvent(2, workers[1], workers[0], a.output),
FetchStartTraceEvent(10, workers[0], workers[1], b.output),
]
fetch_end_events = [
e for e in simulator.trace_events if isinstance(e, FetchEndTraceEvent)
]
assert fetch_end_events == [
FetchEndTraceEvent(7, workers[1], workers[0], a.output),
FetchEndTraceEvent(13, workers[0], workers[1], b.output),
]
def test_random_levels():
graph = TaskGraph()
random_levels([3, 10, 5, 1], [0, 3, 2, 3], lambda: graph.new_task(output_size=1))
check_graph(graph)
assert graph.task_count == 19
assert len(list(graph.arcs)) == 43
def test_worker_running_tasks():
test_graph = TaskGraph()
test_graph.new_task("X", duration=10)
a = test_graph.new_task("A", duration=1, output_size=1)
b = test_graph.new_task("B", duration=8, output_size=1)
b.add_input(a)
remaining_times = []
class Scheduler(SchedulerBase):
def __init__(self):
self.scheduled = False
def schedule(self, new_ready, new_finished):
workers = self.simulator.workers
remaining_times.append([[
t.remaining_time(self.simulator.env.now)
for t in w.running_tasks.values()
] for w in workers])
if not self.scheduled:
tasks = self.simulator.task_graph.tasks
self.scheduled = True
return [
TaskAssignment(workers[0], tasks[0]),
TaskAssignment(workers[1], tasks[1]),
TaskAssignment(workers[1], tasks[2])
]
else:
return ()
scheduler = Scheduler()
do_sched_test(test_graph, 2, scheduler)
assert remaining_times == [[[], []], [[9], []], [[1], []], [[], []]]
def test_worker_priority_block():
g = TaskGraph()
a = g.new_task("a", duration=1)
b = g.new_task("b", duration=1, cpus=3)
c = g.new_task("c", duration=1)
s = fixed_scheduler([(0, a, 3), (0, b, 2), (0, c, 1)])
w = [Worker(cpus=3)]
simulator = do_sched_test(g, w, s, SimpleNetModel(), return_simulator=True)
runtime_state = simulator.runtime_state
assert runtime_state.task_info(a).end_time == pytest.approx(1)
assert runtime_state.task_info(b).end_time == pytest.approx(2)
assert runtime_state.task_info(c).end_time == pytest.approx(1)
s = fixed_scheduler([(0, a, 3), (0, b, 2, 2), (0, c, 1)])
w = [Worker(cpus=3)]
simulator = do_sched_test(g, w, s, SimpleNetModel(), return_simulator=True)
runtime_state = simulator.runtime_state
assert runtime_state.task_info(a).end_time == pytest.approx(1)
assert runtime_state.task_info(b).end_time == pytest.approx(2)
assert runtime_state.task_info(c).end_time == pytest.approx(3)
def test_scheduling_time():
test_graph = TaskGraph()
a = test_graph.new_task("A", duration=3, output_size=1)
b = test_graph.new_task("B", duration=1, output_size=1)
c = test_graph.new_task("C", duration=1, output_size=1)
d = test_graph.new_task("D", duration=1, output_size=1)
b.add_input(a)
c.add_input(b)
d.add_input(c)
times = []
class Scheduler(SchedulerBase):
def schedule(self, new_ready, new_finished):
times.append(self.simulator.env.now)
workers = self.simulator.workers
return [
TaskAssignment(workers[0], t)
for t in new_ready
]
scheduler = Scheduler()
simulator = do_sched_test(
test_graph, 1, scheduler,
SimpleNetModel(bandwidth=2),
scheduling_time=2, return_simulator=True)
runtime_state = simulator.runtime_state
assert times == [0, 5, 8, 11, 14]
assert runtime_state.task_info(a).end_time == 5
assert runtime_state.task_info(b).end_time == 8
assert runtime_state.task_info(c).end_time == 11
assert runtime_state.task_info(d).end_time == 14
def test_simulator_no_events():
task_graph = TaskGraph()
task_graph.new_task("A", duration=1)
scheduler = DoNothingScheduler()
with pytest.raises(RuntimeError):
do_sched_test(task_graph, 1, scheduler)
def test_simulator_cpus1():
test_graph = TaskGraph()
test_graph.new_task("A", duration=1, cpus=1)
test_graph.new_task("B", duration=1, cpus=2)
scheduler = AllOnOneScheduler()
assert do_sched_test(test_graph, [2], scheduler) == 2
scheduler = AllOnOneScheduler()
assert do_sched_test(test_graph, [3], scheduler) == 1
def test_more_outputs_from_same_source():
test_graph = TaskGraph()
a = test_graph.new_task("A", duration=1, outputs=[1, 1, 1])
b = test_graph.new_task("B", duration=1)
b.add_input(a.outputs[0])
b.add_input(a.outputs[2])
s = fixed_scheduler([
(0, a, 0),
(0, b, 0),
])
assert do_sched_test(test_graph, [1], s) == 2
def test_worker_estimate_earliest_time_offset_now():
tg = TaskGraph()
t0 = tg.new_task(expected_duration=3, cpus=1)
t1 = tg.new_task(expected_duration=5, cpus=1)
t2 = tg.new_task(expected_duration=3, cpus=2)
worker = Worker(cpus=2)
worker.assignments = [
TaskAssignment(worker, t0),
TaskAssignment(worker, t1)
]
worker.running_tasks[t0] = RunningTask(t0, 0)
worker.running_tasks[t1] = RunningTask(t1, 0)
assert worker_estimate_earliest_time(worker, t2, 2) == 3
def plan_reverse_cherry1():
"""
a1/10/1 a2/10/1
\ /
\ /
a3
""" # noqa
task_graph = TaskGraph()
a1 = task_graph.new_task("a1", 10, 1)
a2 = task_graph.new_task("a2", 10, 1)
a3 = task_graph.new_task("a3", 1)
a3.add_input(a1)
a3.add_input(a2)
return task_graph
def test_worker_execute_priorities():
SIZE = 20
g = TaskGraph()
b = [g.new_task("b{}".format(i), duration=1) for i in range(SIZE)]
r = random.Random(42)
priorities = list(range(SIZE))
r.shuffle(priorities)
s = fixed_scheduler([(0, t, p) for t, p in zip(b, priorities)])
simulator = do_sched_test(g, [Worker(cpus=2)], s, return_simulator=True)
runtime_state = simulator.runtime_state
for t, p in zip(b, priorities):
assert runtime_state.task_info(t).end_time == pytest.approx(
(SIZE - p - 1) // 2 + 1)
def test_compute_b_level_multiple_outputs():
tg = TaskGraph()
a = tg.new_task(outputs=[2, 4], expected_duration=0)
b = tg.new_task(outputs=[5], expected_duration=0)
c = tg.new_task(outputs=[2], expected_duration=0)
d = tg.new_task(expected_duration=0)
b.add_input(a.outputs[0])
c.add_input(a.outputs[1])
d.add_inputs((b, c))
blevel = compute_b_level_duration_size(None, tg)
assert blevel[a] == 7
assert blevel[b] == 5
assert blevel[c] == 2
assert blevel[d] == 0
def test_worker_estimate_earliest_time():
tg = TaskGraph()
t0 = tg.new_task(expected_duration=3, cpus=2)
t1 = tg.new_task(expected_duration=5, cpus=1)
t2 = tg.new_task(expected_duration=4, cpus=3)
t4 = tg.new_task(expected_duration=4, cpus=2)
t5 = tg.new_task(expected_duration=4, cpus=2)
worker = Worker(cpus=4)
worker.assignments = [
TaskAssignment(worker, t0),
TaskAssignment(worker, t1),
TaskAssignment(worker, t2),
TaskAssignment(worker, t4)
]
worker.running_tasks[t0] = RunningTask(t0, 0)
worker.running_tasks[t1] = RunningTask(t1, 0)
assert worker_estimate_earliest_time(worker, t5, 0) == 7
def json_deserialize(data):
tasks = json.loads(data)
graph = TaskGraph()
id_to_task = {}
for t in tasks:
task = graph.new_task(duration=t["d"],
expected_duration=t["e_d"],
cpus=t["cpus"],
outputs=[o["s"] for o in t["outputs"]])
for (index, output) in enumerate(task.outputs):
output.expected_size = t["outputs"][index]["e_s"]
id_to_task[len(id_to_task)] = task
for (i, t) in enumerate(tasks):
for (parent, output_index) in t["inputs"]:
parent = id_to_task[parent]
id_to_task[i].add_input(parent.outputs[output_index])
return graph
def test_worker_max_downloads_global():
g = TaskGraph()
a1, a2, a3, a4 = [
g.new_task("a{}".format(i), duration=0, output_size=1)
for i in range(4)
]
b = g.new_task("b", duration=0)
b.add_inputs([a1, a2, a3, a4])
s = fixed_scheduler([
(0, a1, 0),
(1, a2, 0),
(2, a3, 0), # worker is 2!
(2, a4, 0), # worker is also 2!
(4, b, 0),
])
def make_workers(max_downloads, max_downloads_per_worker=2):
return [
Worker(),
Worker(),
Worker(),
Worker(),
Worker(max_downloads=max_downloads,
max_downloads_per_worker=max_downloads_per_worker)
]
assert do_sched_test(g, make_workers(1), s,
SimpleNetModel()) == pytest.approx(4)
assert do_sched_test(g, make_workers(2), s,
SimpleNetModel()) == pytest.approx(2)
assert do_sched_test(g, make_workers(3), s,
SimpleNetModel()) == pytest.approx(2)
assert do_sched_test(g, make_workers(3), s,
SimpleNetModel()) == pytest.approx(2)
assert do_sched_test(g, make_workers(4), s,
SimpleNetModel()) == pytest.approx(1)
assert do_sched_test(g, make_workers(4, 1), s,
SimpleNetModel()) == pytest.approx(2)
assert do_sched_test(g, make_workers(3, 1), s,
SimpleNetModel()) == pytest.approx(2)
def test_is_descendant():
graph = TaskGraph()
n1, n2, n3, n4, n5 = [
graph.new_task(output_size=1 if i != 5 else None) for i in range(5)
]
n1.add_input(n2)
n1.add_input(n4)
n2.add_input(n3)
n2.add_input(n4)
assert not n1.is_predecessor_of(n4)
assert n4.is_predecessor_of(n1)
assert not n2.is_predecessor_of(n4)
assert n4.is_predecessor_of(n2)
assert n2.is_predecessor_of(n1)
assert not n1.is_predecessor_of(n5)
assert not n5.is_predecessor_of(n1)
assert not n1.is_predecessor_of(n1)
def test_task_graph_merge(plan1):
task_graph = TaskGraph.merge([plan1, plan1, plan1, plan1])
assert task_graph.task_count == 4 * plan1.task_count
task_graph.validate()
for i, t in enumerate(task_graph.tasks):
assert t.id == i
assert len(t.inputs) == len(plan1.tasks[i % plan1.task_count].inputs)
assert t.duration == plan1.tasks[i % plan1.task_count].duration
assert id(t) != id(plan1.tasks[i % plan1.task_count])
def plan1():
"""
a1/1 a2/3/3
| |
a3/1 | a4/1/6
|\ / /|
o o |/ |
| a5/1 a6/6 a7/2
| \ | /
| \ | /
\--- a8/1
""" # noqa
task_graph = TaskGraph()
a1, a2, a3, a4, a5, a6, a7, a8 = [
task_graph.new_task(
"a{}".format(i + 1),
duration=duration,
expected_duration=duration,
outputs=[TaskOutput(size, size) for size in outputs])
for i, (duration, outputs) in enumerate([
(2, [1]), # a1
(3, [3]), # a2
(2, [1, 1]), # a3
(1, [6]), # a4
(1, [1]), # a5
(6, [1]), # a6
(1, [2]), # a7
(1, []) # a8
])
]
a3.add_input(a1)
a5.add_inputs([a3.outputs[0], a2, a4])
a6.add_input(a4)
a8.add_inputs([a5, a6, a7, a3.outputs[1]])
task_graph.validate()
return task_graph
def test_task_zerocost():
test_graph = TaskGraph()
a = test_graph.new_task("A", duration=1, output_size=100)
b = test_graph.new_task("B", duration=1, output_size=50)
c = test_graph.new_task("C", duration=8)
c.add_inputs((a, b))
d = test_graph.new_task("D", duration=1, outputs=[0])
e = test_graph.new_task("E", duration=1, outputs=[0])
e.add_input(d)
class Scheduler(StaticScheduler):
def static_schedule(self):
workers = self.simulator.workers
tasks = self.simulator.task_graph.tasks
return [
TaskAssignment(workers[0], tasks[0]),
TaskAssignment(workers[1], tasks[1]),
TaskAssignment(workers[2], tasks[3]),
TaskAssignment(workers[0], tasks[4]),
TaskAssignment(workers[2], tasks[2])
]
scheduler = Scheduler()
do_sched_test(test_graph, 3, scheduler, SimpleNetModel(bandwidth=2))
def test_worker_download_priorities1():
SIZE = 20
g = TaskGraph()
a = g.new_task("a", duration=0, outputs=[1] * SIZE)
b = [g.new_task("b{}".format(i), duration=0) for i in range(SIZE)]
for i, t in enumerate(b):
t.add_input(a.outputs[i])
r = random.Random(42)
priorities = list(range(SIZE))
r.shuffle(priorities)
s = fixed_scheduler([(0, a, 0)] + [(1, t, p)
for t, p in zip(b, priorities)])
w = [Worker(), Worker(max_downloads=2, max_downloads_per_worker=2)]
simulator = do_sched_test(g, w, s, SimpleNetModel(), return_simulator=True)
runtime_state = simulator.runtime_state
for t, p in zip(b, priorities):
assert runtime_state.task_info(t).end_time == pytest.approx(
(SIZE - p - 1) // 2 + 1)
def test_simulator_cpus3():
test_graph = TaskGraph()
test_graph.new_task("A", duration=3, cpus=1)
test_graph.new_task("B", duration=1, cpus=2)
test_graph.new_task("C", duration=1, cpus=1)
test_graph.new_task("D", duration=1, cpus=3)
test_graph.new_task("E", duration=1, cpus=1)
test_graph.new_task("F", duration=1, cpus=1)
scheduler = AllOnOneScheduler()
assert do_sched_test(test_graph, [3], scheduler) == 4
scheduler = AllOnOneScheduler()
assert do_sched_test(test_graph, [4], scheduler) == 3
scheduler = AllOnOneScheduler()
assert do_sched_test(test_graph, [5], scheduler) == 3
def test_worker_freecpus():
test_graph = TaskGraph()
test_graph.new_task("A", duration=10, cpus=2, output_size=1)
test_graph.new_task("B", duration=8, cpus=3, output_size=1)
c = test_graph.new_task("C", duration=1, cpus=1, output_size=1)
d = test_graph.new_task("D", duration=3, cpus=3, output_size=1)
d.add_input(c)
free_cpus = []
class Scheduler(SchedulerBase):
def schedule(self, new_ready, new_finished):
worker = self.simulator.workers[0]
free_cpus.append(worker.free_cpus)
return [TaskAssignment(worker, t) for t in new_ready]
scheduler = Scheduler()
do_sched_test(test_graph, [10], scheduler)
assert free_cpus == [10, 5, 5, 8, 10]
def test_worker_download_priorities2():
g = TaskGraph()
a = g.new_task("a", duration=0, outputs=[2, 2])
b = g.new_task("b", duration=4, output_size=2)
d = g.new_task("d", duration=1)
a2 = g.new_task("a2", duration=1)
a2.add_input(a.outputs[0])
b2 = g.new_task("b", duration=1, output_size=1)
b2.add_input(a.outputs[1])
b2.add_input(b)
s = fixed_scheduler([(0, a), (0, b), (1, d, 3), (1, a2, 1), (1, b2, 2)])
w = [Worker(cpus=3), Worker(cpus=1, max_downloads=1)]
simulator = do_sched_test(g, w, s, SimpleNetModel(), return_simulator=True)
assert simulator.runtime_state.task_info(a2).end_time == pytest.approx(3)
assert simulator.runtime_state.task_info(b2).end_time == pytest.approx(7)
def test_simulator_empty_task_graph():
task_graph = TaskGraph()
scheduler = DoNothingScheduler()
assert do_sched_test(task_graph, 1, scheduler) == 0
def test_random_dependencies():
graph = TaskGraph()
random_dependencies(10, 0.2, lambda: graph.new_task(output_size=1))
assert graph.task_count == 10
check_graph(graph)
