def default_strategies():
return {
"never": OptimizationStrategy(),
"remove": Remove(),
"replace": Replace(),
}
def default_strategies():
return {
'never': OptimizationStrategy(),
'remove': Remove(),
'replace': Replace(),
}
class TestOptimize(unittest.TestCase):
strategies = {
'never': OptimizationStrategy(),
'remove': Remove(),
'replace': Replace(),
}
def make_task(self,
label,
optimization=None,
task_def=None,
optimized=None,
task_id=None,
dependencies=None):
task_def = task_def or {'sample': 'task-def'}
task = Task(kind='test', label=label, attributes={}, task=task_def)
task.optimization = optimization
task.task_id = task_id
if dependencies is not None:
task.task['dependencies'] = sorted(dependencies)
return task
def make_graph(self, *tasks_and_edges):
tasks = {t.label: t for t in tasks_and_edges if isinstance(t, Task)}
edges = {e for e in tasks_and_edges if not isinstance(e, Task)}
return TaskGraph(tasks, graph.Graph(set(tasks), edges))
def make_opt_graph(self, *tasks_and_edges):
tasks = {t.task_id: t for t in tasks_and_edges if isinstance(t, Task)}
edges = {e for e in tasks_and_edges if not isinstance(e, Task)}
return TaskGraph(tasks, graph.Graph(set(tasks), edges))
def make_triangle(self, **opts):
"""
Make a "triangle" graph like this:
t1 <-------- t3
`---- t2 --'
"""
return self.make_graph(self.make_task('t1', opts.get('t1')),
self.make_task('t2', opts.get('t2')),
self.make_task('t3', opts.get('t3')),
('t3', 't2', 'dep'), ('t3', 't1', 'dep2'),
('t2', 't1', 'dep'))
def assert_remove_tasks(self, graph, exp_removed, do_not_optimize=set()):
optimize.registry = self.strategies
got_removed = optimize.remove_tasks(
target_task_graph=graph,
optimizations=optimize._get_optimizations(graph, self.strategies),
params={},
do_not_optimize=do_not_optimize)
self.assertEqual(got_removed, exp_removed)
def test_remove_tasks_never(self):
"A graph full of optimization=never has nothing removed"
graph = self.make_triangle()
self.assert_remove_tasks(graph, set())
def test_remove_tasks_all(self):
"A graph full of optimization=remove has removes everything"
graph = self.make_triangle(t1={'remove': None},
t2={'remove': None},
t3={'remove': None})
self.assert_remove_tasks(graph, {'t1', 't2', 't3'})
def test_remove_tasks_blocked(self):
"Removable tasks that are depended on by non-removable tasks are not removed"
graph = self.make_triangle(t1={'remove': None}, t3={'remove': None})
self.assert_remove_tasks(graph, {'t3'})
def test_remove_tasks_do_not_optimize(self):
"Removable tasks that are marked do_not_optimize are not removed"
graph = self.make_triangle(
t1={'remove': None},
t2={'remove': None}, # but do_not_optimize
t3={'remove': None})
self.assert_remove_tasks(graph, {'t3'}, do_not_optimize={'t2'})
def assert_replace_tasks(self,
graph,
exp_replaced,
exp_removed=set(),
exp_label_to_taskid={},
do_not_optimize=None,
label_to_taskid=None,
removed_tasks=None,
existing_tasks=None):
do_not_optimize = do_not_optimize or set()
label_to_taskid = label_to_taskid or {}
removed_tasks = removed_tasks or set()
existing_tasks = existing_tasks or {}
got_replaced = optimize.replace_tasks(
target_task_graph=graph,
optimizations=optimize._get_optimizations(graph, self.strategies),
params={},
do_not_optimize=do_not_optimize,
label_to_taskid=label_to_taskid,
removed_tasks=removed_tasks,
existing_tasks=existing_tasks)
self.assertEqual(got_replaced, exp_replaced)
self.assertEqual(removed_tasks, exp_removed)
self.assertEqual(label_to_taskid, exp_label_to_taskid)
def test_replace_tasks_never(self):
"No tasks are replaced when strategy is 'never'"
graph = self.make_triangle()
self.assert_replace_tasks(graph, set())
def test_replace_tasks_all(self):
"All replacable tasks are replaced when strategy is 'replace'"
graph = self.make_triangle(t1={'replace': 'e1'},
t2={'replace': 'e2'},
t3={'replace': 'e3'})
self.assert_replace_tasks(graph,
exp_replaced={'t1', 't2', 't3'},
exp_label_to_taskid={
't1': 'e1',
't2': 'e2',
't3': 'e3'
})
def test_replace_tasks_blocked(self):
"A task cannot be replaced if it depends on one that was not replaced"
graph = self.make_triangle(t1={'replace': 'e1'}, t3={'replace': 'e3'})
self.assert_replace_tasks(graph,
exp_replaced={'t1'},
exp_label_to_taskid={'t1': 'e1'})
def test_replace_tasks_do_not_optimize(self):
"A task cannot be replaced if it depends on one that was not replaced"
graph = self.make_triangle(
t1={'replace': 'e1'},
t2={'replace': 'xxx'}, # but do_not_optimize
t3={'replace': 'e3'})
self.assert_replace_tasks(graph,
exp_replaced={'t1'},
exp_label_to_taskid={'t1': 'e1'},
do_not_optimize={'t2'})
def test_replace_tasks_removed(self):
"A task can be replaced with nothing"
graph = self.make_triangle(t1={'replace': 'e1'},
t2={'replace': True},
t3={'replace': True})
self.assert_replace_tasks(graph,
exp_replaced={'t1'},
exp_removed={'t2', 't3'},
exp_label_to_taskid={'t1': 'e1'})
def assert_subgraph(self, graph, removed_tasks, replaced_tasks,
label_to_taskid, exp_subgraph, exp_label_to_taskid):
self.maxDiff = None
optimize.slugid = ('tid{}'.format(i) for i in range(1, 10)).next
try:
got_subgraph = optimize.get_subgraph(graph, removed_tasks,
replaced_tasks,
label_to_taskid)
finally:
optimize.slugid = slugid
self.assertEqual(got_subgraph.graph, exp_subgraph.graph)
self.assertEqual(got_subgraph.tasks, exp_subgraph.tasks)
self.assertEqual(label_to_taskid, exp_label_to_taskid)
def test_get_subgraph_no_change(self):
"get_subgraph returns a similarly-shaped subgraph when nothing is removed"
graph = self.make_triangle()
self.assert_subgraph(
graph, set(), set(), {},
self.make_opt_graph(
self.make_task('t1', task_id='tid1', dependencies={}),
self.make_task('t2', task_id='tid2', dependencies={'tid1'}),
self.make_task('t3',
task_id='tid3',
dependencies={'tid1',
'tid2'}), ('tid3', 'tid2', 'dep'),
('tid3', 'tid1', 'dep2'), ('tid2', 'tid1', 'dep')), {
't1': 'tid1',
't2': 'tid2',
't3': 'tid3'
})
def test_get_subgraph_removed(self):
"get_subgraph returns a smaller subgraph when tasks are removed"
graph = self.make_triangle()
self.assert_subgraph(
graph, {'t2', 't3'}, set(), {},
self.make_opt_graph(
self.make_task('t1', task_id='tid1', dependencies={})),
{'t1': 'tid1'})
def test_get_subgraph_replaced(self):
"get_subgraph returns a smaller subgraph when tasks are replaced"
graph = self.make_triangle()
self.assert_subgraph(
graph, set(), {'t1', 't2'}, {
't1': 'e1',
't2': 'e2'
},
self.make_opt_graph(
self.make_task('t3', task_id='tid1',
dependencies={'e1', 'e2'})), {
't1': 'e1',
't2': 'e2',
't3': 'tid1'
})
def test_get_subgraph_removed_dep(self):
"get_subgraph raises an Exception when a task depends on a removed task"
graph = self.make_triangle()
with self.assertRaises(Exception):
optimize.get_subgraph(graph, {'t2'}, set(), {})
IndexSearch,
OptimizationStrategy,
SkipUnlessChanged,
SkipUnlessSchedules,
)
class RandomOptimizer(OptimizationStrategy):
probability = 0.5 # probability task is optimized away
def should_remove_task(self, task, params, _):
if random.random() < self.probability:
return True
return False
STRATEGIES = {
'never':
OptimizationStrategy(), # "never" is the default behavior
'index-search':
IndexSearch(),
'seta':
RandomOptimizer(),
'skip-unless-changed':
SkipUnlessChanged(),
'skip-unless-schedules':
SkipUnlessSchedules(),
'skip-unless-schedules-or-seta':
Either(SkipUnlessSchedules(), RandomOptimizer()),
}