From PR to Test Runner: How a Monorepo Decides What to Test

Every PR in our monorepo triggers test selection. You push code, some tests run, results come back. Maybe 40 tests, maybe 400. Nobody asks how the number was chosen. It’s like asking how your thermostat decides when to turn on the heat — it works, so you don’t think about it.

I thought about it. And what I found was not a thermostat. It’s two thermostats, wired to different sensors, running simultaneously, while a third system watches from across the room and pretends the other two don’t exist.

This is a companion piece to the dep graph post. That one was “we looked at the structure.” This one is “we looked at how tests flow through that structure.” Same monorepo, different question. Instead of “what depends on what?” — “what runs where, and who decides?”

The question nobody asks

A PR changes three files in mai_config. CI runs 247 tests across 12 packages. Why those 247? Why not 246? Why not all 15,000?

Somewhere in the pipeline, something looked at your diff, consulted a graph, made decisions about which packages are affected, figured out which of those tests need GPUs, allocated them across runners, and reported back. All of this happens in the time it takes you to switch to Slack and complain about CI being slow.

The answer turns out to involve a workspace graph built from ruff import analysis, a YAML file someone maintains by hand, a bin-packing algorithm from the 1960s, and an AST walker that reads your test decorators without running Python. Plus a newer system running alongside the old one in shadow mode, validating results but not yet in control.

Let’s trace the whole path.

Step 1: What changed?

Every test workflow starts with the same question: which packages are affected by this PR?

The find-changed GitHub Action calls ci.cli find_changed, which does something more involved than you’d expect. It doesn’t just look at which files were modified. It builds the full workspace graph — all 171 packages, all their internal dependencies — by combining three data sources:

ruff import analysis — scans Python source files for import statements
uv.lock — the lockfile that resolves transitive dependencies
AST scanning — parses source code for additional dependency signals

Then it diffs the graph against the base SHA. The output: a list of packages whose transitive dependency closure intersects with the changed files. Not “which packages contain changed files” — “which packages could be affected by the changed files.” Big difference. Change mai_config and you might affect 40 downstream packages, even though you only touched one directory.

I covered the shape of this graph in the dep graph post. The short version: mai_config has enormous fan-in. Touch it, and the shockwave propagates through half the repo. The workspace graph is what quantifies that shockwave.

This step is shared. Both CPU and GPU test workflows consume its output. After this, the paths diverge.

Step 2a: CPU tests (the simple path)

CPU tests are boring. I mean that as a compliment.

The workflow reads cpu_tests.yml, which maps packages to their test paths and pytest flags. For each affected package, it spawns one GitHub Actions job that runs uv run python -um pytest on a standard VM. One package, one job, no drama.

There’s no sharding. There’s no bin-packing. There’s no duration estimation. The CPU test runner gets the package name, runs its tests, and reports back. If mai_config changed and 12 packages are affected, you get 12 parallel jobs. Simple.

The only interesting thing is that this used to be slow — 13.8 minutes — and Yipu got it down to 5.6 through better parallelization and runner allocation. But the architecture is straightforward. One job per package. Done.

GPU tests are where it gets weird.

Step 2b: GPU tests (the interesting path)

GPU tests have a problem that CPU tests don’t: scarcity. CPU runners are cheap VMs you can spin up by the dozen. GPU runners are 8xH100 nodes that cost real money — about $168 per run across 12 shards, as I covered in the sharding post. You can’t just spawn one job per package. You need to pack tests efficiently across a limited number of runners.

The GPU pipeline has three stages, each more involved than the last.

Stage 1: Find the test files

ci.cli tests takes the affected packages from find_changed and maps them to individual test files. But it doesn’t use the workspace graph for this — it uses FileTreeDiffChecker, which does file-level diffing. Different mechanism, finer granularity. The workspace graph says “this package is affected.” The file-level diff says “these specific test files need to run.”

Stage 2: Collect test metadata

ci.cli plan_tests runs pytest --collect-only — in parallel, one worker per package. This doesn’t execute any tests. It asks pytest to discover them: enumerate every test function, every parametrized variant, every marker. The output is a manifest of what exists.

This step is surprisingly expensive. pytest --collect-only has to import the test files, which means importing their dependencies, which for GPU packages means importing torch, which means initializing CUDA contexts. “Just discovering what tests exist” can take 30-60 seconds per package. Running it in parallel across packages is the only thing that keeps this bearable.

Stage 3: Shard

Now the planner has a list of tests and needs to assign them to runners. This is the bin-packing problem, and the algorithm is a classic: Longest Processing Time first (LPT).

Here’s how it works:

Calculate shard count. num_shards = ceil(total_weight / desired_time_per_shard). The target is 900 seconds per shard for PR CI. Clamp to max 12.
Sort tests by weight, heaviest first. test_mha_kernels.py at 364.8 seconds goes first. The 2-second config tests go last.
For each test, find the least-loaded shard. Add a UV_SYNC_COST penalty (240 seconds) if the shard hasn’t seen this package yet — because the first test from a new package triggers a uv sync that subsequent tests from the same package skip.
Assign.

That’s LPT. It’s the same algorithm you’d use if you had 12 dishwashers and wanted to distribute plates so nobody finishes way before or after the others — give the biggest plate to whoever has the lightest load. It’s a 4/3 approximation to optimal makespan, which is a fancy way of saying it’s provably within 33% of the best possible distribution. For a greedy algorithm you can implement in 20 lines, that’s excellent.

The UV sync cost penalty is the clever bit. Without it, the algorithm might scatter tests from the same package across many shards, and each shard pays the one-time package setup cost independently. With the penalty, it naturally clusters same-package tests together — because the second test from a package is “free” (no sync cost), making its shard look cheaper.

Where the weights come from

This is the part that surprised me.

The weights live in .github/ci/python/gpu_test_durations.json. 284 entries, file-level granularity, ranging from 0.0 to 364.8 seconds. The default for unknown tests is 5.0 seconds.

These are maintained manually.

There’s no automated duration collection pipeline. No CI job that watches test execution times and updates the file. Someone — and I can tell you it’s been Sriram — periodically looks at recent CI runs, computes average durations with some variance smoothing, and updates the JSON file by hand.

I wrote about this in the sharding post already, but it bears repeating because it’s the kind of thing that sounds absurd until you think about why the alternative is hard. Automated weight collection requires: reliable test-level timing from CI (not always available), handling of flaky tests that sometimes take 3 seconds and sometimes 300, dealing with tests that are parametrized differently in different runs, filtering out cold-start outliers from warm-cache norms, and making sure a bad data point doesn’t corrupt the weights and cascade into terrible shard distributions. The manual approach is: look at the numbers, does this seem right, yes, commit. It’s low-tech. It works.

There’s also package_sync_durations.json, which tracks the one-time uv sync overhead per package. This feeds the UV_SYNC_COST penalty in the sharding algorithm. Same deal — manually maintained, updated periodically.

Two systems at once

Here’s where it gets architecturally interesting. The system I just described? That’s Gen 1. There’s a Gen 2 running alongside it right now, in validation mode. Same inputs, different architecture, not yet driving actual shard assignments.

Gen 1: YAML-driven

Gen 1 classifies tests using static config files:

cpu_tests.yml — maps packages to CPU test paths
test_exec_mode.yml — classifies test paths into execution modes (serial, isolated, concurrent, SKIP) by prefix matching
GPU_TEST_PKGS — hardcoded list of packages that get GPU CI

Test discovery is the pytest --collect-only pipeline. Sharding uses TestTarget objects with RunnerType and the weighted LPT algorithm.

The strength: it’s simple, predictable, and battle-tested. The weakness: those YAML files are manually maintained and have to stay in sync with the actual test code. If someone adds a new GPU test to a package that isn’t in GPU_TEST_PKGS, it silently doesn’t get GPU CI. If someone moves a test file and doesn’t update test_exec_mode.yml, it runs in the wrong mode.

Gen 2: Annotation-driven

Gen 2 flips the model. Instead of external YAML files classifying tests, tests classify themselves.

@pytest.mark.runs_on("gpu")
@pytest.mark.run_mode("isolated")
def test_distributed_gradient_sync():
    ...

The test declares its requirements in code — same file, same review, same PR. No separate YAML to keep in sync.

The planner reads these annotations via a new AST walker (ci/src/ci/plan/ast.py) that parses Python source files without executing them. It builds an AnnotatedObject tree with decorator-to-annotation mapping, handling inheritance at three levels: file → class → method. A pytestmark at file level applies to everything in the file. A marker on a class applies to all its methods. A marker on a method overrides its class.

It also handles our custom @mesh_test decorator, which is a compound annotation — it implies runs_on("gpu") plus specific resource requirements.

The output is a Ticket — a (test, ShardClass, RunMode) triple. Shard classes map to runner types (gpu, cpu, big, redis). Run modes control isolation (isolated, serial, concurrent). The config fallback file test_exec_modes.yml still exists for tests that haven’t been annotated yet, but the direction is clear: annotations replace config.

Gen 2 has five sharding algorithms: none, package, weighted, weighted-random, and random. Gen 1 only had weighted. The package mode assigns all tests from a package to the same shard — no UV sync penalty needed because there’s no cross-package mixing. The random modes exist for testing whether the weighted approach actually matters (spoiler: it does).

Running in parallel

The collect_and_plan job in the GPU workflow already runs Gen 2 alongside Gen 1. Both systems see the same test files. Both produce shard assignments. Gen 1’s assignments drive the actual test execution. Gen 2’s assignments get logged for comparison.

This is the responsible way to deploy a new planner: run it in shadow mode, compare its decisions against the production system, find the cases where they disagree, and figure out who’s right. Once the disagreement rate drops low enough, you flip the switch.

The annotation bet

The key insight behind Gen 2 is this: the current system has test classification scattered across three YAML files that someone has to manually keep in sync with the code. The new system makes the code the source of truth.

You annotate a test with @pytest.mark.runs_on("gpu") and the planner picks it up automatically. No YAML edit. No separate PR. No “I added a GPU test but forgot to update the config and it silently ran on a CPU runner for three weeks.”

This is the same pattern as Bazel’s BUILD files versus our old find_changed tool. With find_changed, dependency information lived in a separate system that had to mirror the code. With Bazel, the code declares its own dependencies. The Gen 2 planner is doing the same thing for test execution requirements.

I wrote about this convergence in the test classification post — the bottleneck was never “which routing system” but “where do the labels live?” Moving labels into the code is the answer, whether you’re talking about Bazel tags or pytest markers.

Bazel is watching from across the room

Here’s the thing: none of this applies to Bazel tests.

Bazel has its own CI workflow, triggered by file path patterns only — *.bazel, bazel/**, .bazelrc. If your PR doesn’t touch Bazel files, Bazel CI doesn’t run. It has no connection to the workspace graph. It doesn’t know about find_changed. It doesn’t use the planner.

Aspect	Python CI	Bazel CI
Trigger	Always on PR	Only on BUILD/Bazel file changes
Change detection	Workspace graph + import analysis	GitHub Actions `paths:` filter
Test discovery	`pytest --collect-only` or AST scanning	`bazel test //...`
Sharding	Custom Python planner (LPT)	Bazel native sharding
Dependency graph	Reconstructed from imports + lockfile	Declared in BUILD files

If a PR changes both Python source and BUILD files, both pipelines run independently. No coordination. No shared planning. Two separate systems looking at the same code through different lenses.

This is fine for now. Bazel covers maybe 15% of the monorepo. But the endgame is obvious: once Bazel covers everything, the Python CI planner becomes unnecessary. Bazel’s bazel test //... with --test_tag_filters does change detection, test discovery, and sharding in one command. The entire custom pipeline I just described collapses into a build system primitive.

We’re building the road to make Bazel the road. Same story as the dep graph post — the networkx script that plans the Bazel migration gets replaced by Bazel itself. The planner that decides which tests to run gets replaced by Bazel itself. The tools keep building the path to their own retirement.

The full picture

Let me draw the whole thing. This is what happens between “you push code” and “tests run”:

PR opened
    │
    ▼
find-changed action
(workspace graph: ruff imports + uv.lock + AST)
"which packages are affected?"
    │
    ├──────────────────────────────┐
    ▼                              ▼
CPU workflow                   GPU workflow
    │                              │
cpu_tests.yml                  ci.cli tests
(package → path map)           (file-level diff)
    │                              │
1 job per package              ci.cli plan_tests
    │                          (pytest --collect-only, parallel)
pytest on GHA VMs                  │
    │                         ┌────┴────┐
    ▼                         ▼         ▼
  done                    Gen 1      Gen 2
                         (YAML +    (AST annotations +
                         weighted    Ticket model +
                         LPT)       pluggable sharder)
                            │         │ (shadow mode)
                            ▼         │
                      shard assignments
                            │
                            ▼
                      12 GPU shards
                      (H100 nodes via Slurm)
                            │
                            ▼
                          done

         ═══════════════════════════════
         Bazel CI (separate universe)
         Triggered only by BUILD file changes
         bazel test //...
         No connection to anything above

Three systems. Two generations of the planner running simultaneously. One build system watching from the side. All answering the same question: what do we need to test?

What I learned from looking inside

A few things struck me about this system.

The manual parts are the ones that matter most. The weight file is manually maintained. The YAML configs are manually maintained. The Gen 2 annotations are manually added by engineers. Automation handles the graph analysis, the bin-packing, the parallel collection. But the critical inputs — “how long does this test take?” and “what kind of runner does it need?” — come from humans looking at data and making calls. The system is as good as the people feeding it.

Two generations is not an accident. You can’t just replace a production test planner. You shadow it. You run both, compare, find disagreements, fix them. The gap between “this new system works on paper” and “this new system works on our 15,000 tests” is months of validation. Gen 1 and Gen 2 running in parallel isn’t technical debt. It’s engineering discipline.

The workspace graph does more than anyone realizes. The same graph that drives test selection also drives the Bazel migration planner, the dep graph analysis, and the impact radius estimation I wrote about in the dep graph post. It’s the MRI machine that keeps getting used for new diagnoses. The CI team built it for one purpose. We keep finding more.

Bazel will eventually replace all of this. Not next quarter. Maybe not next year. But the architecture is clear. Declared dependencies in BUILD files replace the reconstructed workspace graph. Bazel’s --test_tag_filters replaces the YAML classification files. Bazel’s native sharding replaces the LPT algorithm. Every custom component has a Bazel-native equivalent waiting for the migration to reach it.

Until then — the planner runs, the shards fill, the tests pass or fail, and nobody asks how the number was chosen. Which is exactly how infrastructure should work. Invisible until you go looking.