Python multiprocessing for tasks depending on each other

I have a simulation that runs a process on a tree structure, where the slow processes down processes are separate from each other. This looked like a thing that would be a prime candidate for multiprocessing to me, so I tried to wrap my head around the python multiprocessing module and its locks, pools etc. and wrote a toy example that walks down a length-annotated tree, where each branch of the tree waits for its parent.

import multiprocessing
import time
import newick

manager = multiprocessing.Manager()
generated_languages = manager.dict()
# The root should be allowed to start computation immediately. The parent of
# the root is None, initialize it with some starting state of the
# simulation which is not false.
generated_languages[None] = True
# Processes will want to write results to file. This can take some time,
# provide a lock around the file access.
io_lock = manager.Lock()

def mp_worker(node):
"""Run the simulation on the branch leading to node.

That process takes time proportional to the branch length.

"""
def parent_done():
"""Check whether the parent node has been generated yet."""
print("Parent {:} generated: {:}".format(
parent, generated_languages.get(parent)))
return generated_languages.get(parent)

name = node.name
# We cannot have duplicate names, and we cannot have unnamed nodes.
assert name not in generated_languages
parent = None if node.ancestor is None else node.ancestor.name

print("Process {:s}\tWaiting for parent {:}".format(name, parent))
start_from = parent_done()
while not start_from:
time.sleep(1)
start_from = parent_done()

print("Process {:}\t‘running the simulation’ (ie., waiting {:}s)".format(
(node.name, node.length))
time.sleep(int(node.length))
result = start_from

generated_languages[name] = result

# Write result to file.
io_lock.acquire()
print("Process {:}\tDONE with result {:}".format(
name, result))
io_lock.release()

return result

p = multiprocessing.Pool(2)

for result in p.imap(mp_worker, t.walk()):
print(result)


Now this works, it generates the expected output (tree nodes in pre-order, with no node process started before its parent is finished) but it does not look like good code to me. I noticed a few problems myself which I don't know how to do cleanly

• I think manager should be used in a context manager, but mp_worker can only get one argument from pool.imap, so I want mp_worker to be a closure on manager, which means it must be valid at definition-time as well as at run-time of mp_worker, which feels wrong.
• The waiting loop should obviously replaced by something else. Should it be one Condition? One Condition for each Node? Something else entirely? How?
• With its current structure, a long branch is likely to block its children when the simulation could be continued somewhere else in the tree. The heuristic way around this is to sort nodes by their distance from the root and iterate them in highest-first-order, which is possible with my structure here. A cleverer way might involve a queue of nodes still waiting for their parent, what would be the best pattern (instead of pool.imap) to implement that?
• The bounding factors for writing to a file will be different from the ones for running the simulation. Writing to a file may take seconds, running the simulation on one branch hours. I assume it is fine to have one blocking process doing both, given these different orders of magnitudes?