GPU Test Targeting with the Bazel Dep Graph

GPU runners are expensive. Like, “we track minutes the way finance tracks headcount” expensive. And right now, the entire question of whether a PR needs GPU tests is answered by a YAML list that a human maintains:

# run_gpu_tests.yml
GPU_TEST_PKGS: "rocket mai_layers mai_distributed ..."
# ~13 manually maintained entries

Change anything in those packages → spin up a GPU runner. Change anything else → don’t. The mechanism works. The maintenance model doesn’t.

What’s wrong with a list

Nothing, in the same way nothing is wrong with a phone book. It works until it doesn’t, and you don’t find out it doesn’t until someone’s tests silently don’t run.

The failure modes are symmetric:

False positives. Someone touches mai_stats — which is in the GPU list — but the change is to a stats logging function. No tensors. No CUDA. But the list doesn’t know that. It sees the package name, triggers a GPU runner, and burns 15 minutes of GPU time on tests that could’ve run on a CPU.

False negatives. Someone adds a torch-dependent test in dataprocessing. dataprocessing isn’t in the list. Why would it be? It sounds like a data package. The test exists. It needs a GPU. CI never runs it. Nobody notices until it breaks on main three weeks later and the oncall engineer gets to debug it at 11 PM.

Staleness is invisible. There’s no mechanism to verify that the list is correct. No CI check. No lint rule. Someone has to periodically look at the list, look at the code, and compare. In a monorepo with 100+ packages, “someone” is “nobody.”

The list is debt that accrues silently. Every new package with GPU tests is a potential false negative. Every refactored package that drops GPU deps is a permanent false positive. The list was correct on the day it was written. Today? We genuinely don’t know.

The dep graph already knows the answer

The previous post teased this idea. With the full Bazel dep graph covering 100% of the monorepo, you can ask the question directly:

# "Does this change affect any torch-tagged test targets?"
bazel query 'attr(tags, "torch", rdeps(//..., //changed/package/...))'

If the result is empty → no GPU runner. If it’s not empty → spin one up, and you know exactly which tests need it.

No list. No human curation. The graph IS the list, and it updates itself every time someone adds a dependency to a BUILD file.

The key is the torch tag convention. Tests that need a GPU get tags = ["torch"] in their BUILD file. This is already how our Bazel CI distinguishes torch tests — py_venv_test with include_system_site_packages = True uses it. The tag exists. We just haven’t been using it for test selection.

What I built

bazel/tools/gpu_test_targeting.py — 837 lines of Python with four modes:

The core is two Bazel queries:

# 1. Find all torch-tagged test targets in the repo
bazel query 'attr(tags, "torch", kind(".*test", //...))'

# 2. For a set of changed files, find affected torch tests
bazel query 'attr(tags, "torch", rdeps(//..., //changed_package/...))'

Query 1 is the audit — what does Bazel know about GPU tests? Query 2 is the targeting — given these changes, which GPU tests are affected? Compare the results against what the hardcoded list would do, and you get the false positive / false negative delta.

The audit: what the list gets wrong today

First question: is the hardcoded list even correct?

Two findings jump out.

All 13 GPU_TEST_PKGS entries aren’t on Bazel yet. The graph can see them as manual stubs (for dependency traversal), but they don’t have real BUILD files with py_test targets and torch tags. So the graph can’t independently verify whether these packages actually need GPU tests. The hardcoded list is the only authority for these 13 packages — and it might be right! We just can’t confirm it.

Two packages have torch-tagged targets that the list misses entirely. dataprocessing/item_response_theory has a test tagged torch. bazel/tests has the torch smoke test. Neither is in GPU_TEST_PKGS. These are real false negatives — torch tests that never get triggered by the hardcoded selection logic.

This is the pattern that silently breaks CI. Someone writes a test that needs a GPU, puts it in a package nobody thinks of as a “GPU package,” and the test exists in the codebase without ever running in CI.

The simulation: 7 scenarios, two approaches

The audit is a static snapshot. The simulation answers the dynamic question: given a specific code change, what does each approach do?

Read this table column by column and a pattern emerges.

Scenarios 1, 2, 7: hardcoded fires, graph doesn’t. These are packages in the GPU list that haven’t migrated to Bazel. The hardcoded approach catches them. The graph can’t — it doesn’t have torch-tagged test targets for these packages yet. The false positives here reflect the migration gap, not a flaw in graph-based targeting.

Scenarios 3, 5, 6: graph fires, hardcoded doesn’t. These are the interesting ones. A package with a torch-tagged test that’s not in the GPU list. A transitive dependency through lib/mai_nano that reaches a torch test. A direct change to a torch test file. The hardcoded list has no way to catch these. The graph catches all three.

Scenario 4: both agree. Leaf package, no torch relationship. Neither approach triggers GPU tests. This is the happy path, and it’s the majority of real-world changes.

The punchline: the two approaches are complementary. The hardcoded list covers the migration gap (packages the graph can’t see). The graph covers the blind spots (transitive deps and unlisted packages the list doesn’t know about). The union of both is strictly better than either alone.

The shadow CI step

Theory is nice. Data is better. So I added a shadow step to run_gpu_tests.yml:

- name: Shadow GPU test targeting (graph comparison)
  continue-on-error: true
  run: |
    uv run python bazel/tools/gpu_test_targeting.py shadow \
      --base-ref "${{ github.event.pull_request.base.sha }}" \
      --output gpu-targeting-report.json

- name: Upload GPU targeting report
  uses: actions/upload-artifact@v4
  with:
    name: gpu-targeting-shadow-report
    path: gpu-targeting-report.json

The shadow step runs alongside the real GPU test selection. It compares what the hardcoded list chose against what the graph would choose, and uploads the delta as a CI artifact. continue-on-error: true means it never blocks tests. It’s purely observational — a dashcam recording the difference between two decision-making processes.

Over time, this produces a dataset. How often does the list select GPU tests that the graph says are unnecessary? How often does the graph catch torch-dependent tests that the list misses? As more packages migrate to Bazel, the graph gets more complete, and the shadow reports tell us exactly when the crossover happens.

The migration gap, honestly

I could pitch this as “the graph replaces the list.” It doesn’t. Not today.

13 out of 13 GPU_TEST_PKGS entries aren’t on Bazel. The packages most likely to need GPU tests — mai_layers, mai_distributed, mai_trainer, rocket — are the heavy ones. Lots of CUDA deps, lots of native code, lots of reasons they haven’t migrated yet. The graph has manual stubs for dependency traversal, but stubs don’t have py_test targets, and without test targets there’s nothing to tag with torch.

So the honest state of things:

This is a tool that gets more powerful automatically. Every package that migrates to Bazel and gets torch-tagged tests adds signal to the graph. I don’t have to update a YAML file. I don’t have to remember that dataprocessing now has a GPU test. The BUILD file declares it, the tag marks it, and the query finds it.

The trajectory matters more than the current state:

1. Now: Shadow CI collects comparison data on every PR
2. As GPU packages migrate: Graph coverage grows, shadow reports show convergence
3. At ~80% coverage: Switch to graph-primary with hardcoded fallback
4. Eventually: Retire the list entirely

Each step requires zero additional work beyond what the Bazel migration is already doing. The tool is built. The shadow step is running. The only variable is migration velocity.

Why this matters for the team

Let me make the case directly, because this post is also the pitch.

The hardcoded list is maintenance debt that compounds silently. Every new GPU-dependent package is a potential false negative. Every refactored package is a potential false positive. There’s no feedback loop — no CI check that says “hey, this package has torch tests but isn’t in your list.” The failure mode is always silent omission, the kind you discover through a broken main or a wasted GPU budget.

The dep graph eliminates the maintenance entirely. A torch tag in a BUILD file is self-describing. It says “this test needs a GPU” in the same place where you declare all other test properties. bazel query finds it without anyone maintaining a separate list. The BUILD file is already the source of truth for dependencies — making it the source of truth for GPU requirements is just using the same mechanism.

The investment pays off automatically. Here’s the part I really want to emphasize. This tool doesn’t require ongoing work. Every package that gets a BUILD file with torch-tagged tests makes GPU targeting more precise. Not because someone remembered to update a YAML file. Because the graph IS the truth, and the graph updates when the code updates.

If the team is already committed to Bazel migration — and we are — then graph-based GPU targeting is free. It’s a downstream benefit of work we’re already doing. The only incremental cost was writing the tool (done) and adding the shadow step (done). Everything else is riding the migration wave.

And the data will show us when we’re ready. The shadow CI step isn’t just collecting data for fun. It’s building the case for the transition. When the shadow reports consistently show >80% agreement between the graph and the hardcoded list, we switch. When they show 100%, we delete the list.

No guessing. No “I think we’re ready.” The data tells us.

The endgame

The endgame looks like this:

# run_gpu_tests.yml (future)
- name: Determine GPU test targets
  run: |
    TORCH_TARGETS=$(bazel query 'attr(tags, "torch", rdeps(//..., set($CHANGED_TARGETS)))')
    if [ -z "$TORCH_TARGETS" ]; then
      echo "No torch-affected targets. Skipping GPU runner."
      exit 0
    fi
    echo "Running: $TORCH_TARGETS"
    bazel test $TORCH_TARGETS --test_tag_filters=torch

No list. No maintenance. Change a leaf package → zero GPU tests. Change a core library that torch depends on → exactly the affected torch tests. Change a torch test directly → that test runs.

The precision is target-level, not package-level. The hardcoded list says “run everything in mai_layers.” The graph says “run these 3 specific tests in mai_layers whose transitive deps include the file you changed.” When you’re paying per-minute for GPU runners, that precision is money.

But that’s the endgame. Today, we’re at step 1: shadow CI, collecting data, building the case. The tool is built. The plumbing is in place. Now we wait for the migration to catch up — and every BUILD file with a torch tag brings us closer.

This is a deep-dive on one of the use cases from What a Full Dep Graph Actually Unlocks. The dep graph coverage that makes this possible came from The Dep Graph Is the Product. Previous posts in the series: Your Monorepo Has Two Dependency Graphs (and They Disagree), When Homegrown Tools Hit the Ceiling, Graduating Test Selection from find_changed to Bazel.