def test_convert_gpu_scheds(): sched = [(make_job(), 1, "A"), (make_job(), 1, "B"), (make_job(), 2, "A")] machines = {"A": {'type': 'cpu'}, "B": {'type': 'gpu'}} scheds = group_sched_by_machine(sched) gpu_scheds = convert_gpu_scheds(scheds, machines) assert gpu_scheds["A"] == scheds["A"] assert gpu_scheds["B"] != scheds["B"] assert isinstance(gpu_scheds["B"][0][1], theano.sandbox.cuda.GpuOp)
def test_group_sched_by_machine(): sched = [(make_job(), 1, "A"), (make_job(), 1, "B"), (make_job(), 2, "A")] scheds = group_sched_by_machine(sched) assert set(scheds.keys()) == set('AB') assert all(isinstance(v, tuple) for vs in scheds.values() for v in vs)