def test_fmap():
inc = lambda x: x + 1
assert fmap(inc, {1: 2, 3: 4}) == {1: 3, 3: 5}
def distribute(inputs,
outputs,
input_shapes,
machines,
commtime,
comptime,
makespan=100):
known_shapes = shape_of_variables(inputs, outputs, input_shapes)
variables = theano.gof.graph.variables(inputs, outputs)
dag, dinputs, doutputs = dicdag.theano.theano_graph_to_dag(inputs, outputs)
vars = set.union(set(dinputs), set(doutputs), set(dag.keys()),
{v
for value in dag.values() for v in value['args']})
assert len(vars) == len(map(str, vars))
unidag = dicdag.unidag.dag_to_unidag(dag)
# TODO: This should be an input
is_gpu = lambda m: machines[m]['type'] == 'gpu'
can_start_on = lambda v, m: not is_gpu(m)
can_end_on = lambda v, m: not is_gpu(m)
def dag_commtime(job, a, b):
inputs, op, output = job
return commtime(output, a, b)
def dag_comptime(job, a):
if job[1] == dicdag.index:
return 0
return comptime(make_apply(*job), a)
# Compute Schedule
dags, sched, makespan = tompkins.schedule(unidag, machines, dag_comptime,
dag_commtime, lambda j: 0,
lambda j, a: 1, makespan)
cleaner_dags = fmap(replace_send_recvs, dags)
full_dags = fmap(dicdag.unidag.unidag_to_dag, cleaner_dags)
check_send_recv(full_dags)
merge_dags = merge_gpu_dags(full_dags, machines)
check_send_recv(merge_dags)
rankfile = {machine: i for i, machine in enumerate(sorted(merge_dags))}
theano_graphs = {
machine: dag_to_theano_graph(
dag, make_ith_output(rankfile, tagof, known_shapes, machine))
for machine, dag in merge_dags.items()
}
# Check that all inputs and outputs are inputs/outputs to the computation
# or mpi
def valid_inp(x):
return x.name in map(str, inputs)
def valid_out(x):
return x.name in map(str, outputs) or 'mpi_token' in str(x)
assert all(valid_inp(x) for g in theano_graphs.values() for x in g[0])
assert all(valid_out(x) for g in theano_graphs.values() for x in g[1])
if not all(count == 2 for count in tagof.counts.values()):
print "issue with tag counts"
for x in tagof.counts:
print x
for x in tagof.cache:
print x
scheds = tompkins_to_theano_scheds(sched, machines)
return theano_graphs, scheds, rankfile, makespan
def distribute(inputs, outputs, input_shapes, machines, commtime, comptime, makespan=100):
known_shapes = shape_of_variables(inputs, outputs, input_shapes)
variables = theano.gof.graph.variables(inputs, outputs)
dag, dinputs, doutputs = dicdag.theano.theano_graph_to_dag(inputs, outputs)
vars = set.union(set(dinputs), set(doutputs), set(dag.keys()),
{v for value in dag.values() for v in value['args']})
assert len(vars) == len(map(str, vars))
unidag = dicdag.unidag.dag_to_unidag(dag)
# TODO: This should be an input
is_gpu = lambda m: machines[m]['type'] == 'gpu'
can_start_on = lambda v, m: not is_gpu(m)
can_end_on = lambda v, m: not is_gpu(m)
def dag_commtime(job, a, b):
inputs, op, output = job
return commtime(output, a, b)
def dag_comptime(job, a):
if job[1]==dicdag.index:
return 0
return comptime(make_apply(*job), a)
# Compute Schedule
dags, sched, makespan = tompkins.schedule(
unidag, machines, dag_comptime, dag_commtime,
lambda j:0, lambda j,a:1, makespan)
cleaner_dags = fmap(replace_send_recvs, dags)
full_dags = fmap(dicdag.unidag.unidag_to_dag, cleaner_dags)
check_send_recv(full_dags)
merge_dags = merge_gpu_dags(full_dags, machines)
check_send_recv(merge_dags)
rankfile = {machine: i for i, machine in enumerate(sorted(merge_dags))}
theano_graphs = {machine: dag_to_theano_graph(dag,
make_ith_output(rankfile, tagof, known_shapes, machine))
for machine, dag in merge_dags.items()}
# Check that all inputs and outputs are inputs/outputs to the computation
# or mpi
def valid_inp(x):
return x.name in map(str, inputs)
def valid_out(x):
return x.name in map(str, outputs) or 'mpi_token' in str(x)
assert all(valid_inp(x) for g in theano_graphs.values()
for x in g[0])
assert all(valid_out(x) for g in theano_graphs.values()
for x in g[1])
if not all(count == 2 for count in tagof.counts.values()):
print "issue with tag counts"
for x in tagof.counts: print x
for x in tagof.cache: print x
scheds = tompkins_to_theano_scheds(sched, machines)
return theano_graphs, scheds, rankfile, makespan
def test_fmap():
inc = lambda x : x + 1
assert fmap(inc, {1: 2, 3: 4}) == {1: 3, 3: 5}
