Skip to content

CU9 — Branches run at the same time

g = Graph.somatize(
Source()
>> ((Encoder() >> Bottleneck()) | (Other() >> Other2()))
>> Join()
)
g.plan() # Sequence([Execute, Wave([Sequence, Sequence]), Execute])
g.forward(x) # both branches, on two threads, start to finish

Status: closed. 88 tests in Rust, 86 in Python.

Plan::Parallel was added in CU3 and removed in CU4 because it broke on the diamond: its branches overlapped — both claimed the join node — and it executed twice. It was said then that it would come back “when it means something it does not mean today: spreading across threads”. This is that day, and the question left to answer is what goes inside each branch.

Two answers were tried and the first was rejected with a counterexample:

  • By topological level (Kahn by levels). Each wave is an antichain, and its members are lone steps. It is correct and no node can be duplicated. But with a >> (b >> b2 >> b3 | c >> c2) >> d you get Seq([a, Wave([b,c]), Wave([b2,c2]), b3, d]): lockstep. b2 does not start until c finishes even though it does not depend on it, and c2 finishes and sits watching while b3 runs alone. Worse for what is coming: torch’s device is thread-local, so a branch that hops threads on every wave cannot set it once.

  • By branch, which is what stayed. Seq([a, Wave([Seq([b,b2,b3]), Seq([c,c2])]), d]). One thread per branch, start to finish, and a single join.

The missing piece: decompose, do not flatten

Section titled “The missing piece: decompose, do not flatten”

compile stops walking the topological order and recovers the tree. And it has to recover it from the graph, not from the expression: decision 6 of CU5 says that the same graph built with node()/edge() in a loop gives the same plan, and a loop has no tree. The DSL expression is the oracle, not the source.

Four cases, and the order matters:

caseyields
no nodesEmpty
one nodeExecute
the subgraph splits into connected componentsWave, one branch per component
there is a series cutSequence of the two sides
no cutflat sequence: it is not series-parallel

A series cut (A, B) is what a >> does: the crossing edges run from all the sinks of A to all the sources of B, and from nowhere else. Both halves of the check are needed — without the first, an edge leaving an interior node passes as good.

Testing the prefixes of a topological order sufficed, and it is provable: in a serial composition every node of A reaches a sink of A, every sink of A has an edge to every source of B, and every node of B is reachable from a source of B. Therefore every node of A precedes every node of B in any topological order. There is no need to enumerate subsets.

Before the series cut, cutting at a barrier node was tried — one such that ancestors(x) ∪ {x} ∪ descendants(x) was everything. It only gets it right when the join is a single node. The counterexample, which is in the tests:

(a >> a2 | b) >> (c | d)
with barrier → Seq([ Wave([a, b]), a2, Wave([c, d]) ]) ← splits the branch
correct → Seq([ Wave([Seq([a,a2]), b]), Wave([c,d]) ])
  1. Wave(Vec<Plan>), not Vec<Execute>. A branch is a whole plan. The restrictive shape was simpler to execute, but cannot express a branch of several nodes, which is exactly the case.
  2. A wave means “they are launched at the same time”, not “they are independent”. That was CU4’s lesson: a variant that only describes structure buys nothing.
  3. The branches are connected components, so they are disjoint by construction and no node can appear in two. The bug that killed Parallel cannot come back: there is a test that checks it over a battery of topologies.
  4. std::thread::scope, with no dependencies. It lends out &Catalog and &Driver without wrapping them. The bounds that allow it — Node: Send + Sync, Driver: Send + Sync — have been there since CU2 for another reason: PyO3 requires Send on a pyclass. Rayon would have been the obvious answer and the worst one.
  5. Each branch copies what was produced and returns its own; the parent merges on joining. Copying is cheap because a Value clones by Arc, and in exchange there is not a single lock.
  6. The error is the first declared branch’s, not that of whichever failed earlier on the clock. If two branches break at once, which arrives first is a race and the message cannot depend on it.
  7. A panic inside a branch is not swallowed: it propagates with resume_unwind after scope has waited on the others.
  8. A linear chain compiles to the previous plan, identically. It is the regression that matters most: everything closed from CU2 to CU8 is a chain.
  9. What is not series-parallel is walked in sequence, as before. It is neither a failure nor a warning: it is what there was.

The image of the DSL is exactly the series-parallel graphs. >> composes serially by connecting all terminals to all heads, | composes in parallel with disjoint union, and there is no third operation.

The minimal pattern that is not series-parallel is the “N” — a→c, a→d, b→d — and it cannot be written with >> and |. Getting to it requires node()/edge(). So the line explains itself in one sentence: if you wrote it with the DSL, it parallelizes; if not, it parallelizes where it can. That there are DAGs without a tree is a theorem, not a gap in the algorithm — see Valdes, Tarjan and Lawler, “The recognition of series parallel digraphs”, SIAM J. Comput. 11(2), 1982.

Graph.forward held the GIL while the engine ran. The moment a wave spawns threads that call a Python object’s forward, those threads block asking for it and the whole process freezes — not even a join(timeout=…) on the main thread returns, because it needs the GIL too. The fix is one line, py.allow_threads, and its test has to live in another process: a hang like that cannot be caught by anything inside.

What allow_threads does not fix, and must not be confused with the above: two pure Python nodes in the same wave interleave, but do not overlap. The wave puts both in flight — which is why a rendezvous between them resolves, and the test that proves it has no driver in the way — but the GIL means only one runs at any instant, so there is no time to gain. Time is gained when the work releases the GIL: torch in its dispatch, waiting on a network driver, I/O.

And a distinction that costs dearly if lost: the busy-driver test — two_branches_can_keep_the_driver_busy_at_the_same_time — is not what proves the branches are concurrent. The rendezvous tests already prove that, with no driver. What that test measures is that the shared driver is not the bottleneck: it is lent as &dyn Driver to both threads and serves two requests at once.

Decomposition (soma-core/tests/unit/plan.rs)

  • empty, one node, and a linear chain identical to the previous one
  • output fan, input fan and diamond, each with its wave
  • a >> (b >> b2 >> b3 | c >> c2) >> d gives a wave of two sequences
  • (a >> a2 | b) >> (c | d) does not split the branch — the barrier counterexample
  • (a | b) >> (c | d) is two waves, not one of four
  • a wave inside another wave’s branch
  • two unrelated graphs, each long, are two branches
  • the N is walked in sequence, and does not spoil the parallelism beside it

Invariants, over a battery of ten topologies

  • no node executes twice or is left out
  • the order the plan dictates respects every edge
  • every step declares exactly its predecessors in the graph
  • a wave’s branches share no node
  • the same graph always compiles the same

The oracle (soma-core/tests/unit/build.rs)

  • seven DSL expressions, and their plan is the tree that was written

Execution (soma-core/tests/unit/execution.rs)

  • the real execution order respects the edges, with threads in the way
  • a whole branch runs on the same thread, and two branches on different threads
  • two and three branches really do run at once — without sleeping: they agree to meet, and were they sequential the first would exhaust the deadline
  • the diamond’s result is the same spread out as in a row
  • what a branch produces inside reaches whoever reads it
  • two failing branches always give the first declared one’s error
  • a panic inside a branch is not swallowed
  • two branches can keep the driver busy at the same time
  • a wave that is the whole plan returns the map of its leaves

Python (soma-python/tests/test_waves.py)

  • the engine releases the GIL — in another process, with a deadline
  • two and three branches really run at once, without a driver: the rendezvous resolves, therefore both are in flight
  • two pure Python nodes give the right result even though the GIL interleaves them — the price, said plainly
  • the DSL with branches gives the same plan as node()/edge()
  • threads, real order, failures and the N, as in Rust

The device. Plan::Execute still does not say where. This slice is what enables it — one branch per thread is what makes pinning a device per branch mean something — but .on("cuda:1") is the next use case.

Micro-batches. Overlapping inside a branch, not across branches, is another problem and another variant.

Spreading a wave across processes. It needs transport, and Opaque does not cross a wire. That is the other next use case.