Bazel Torch: Test Classification, the Real CI Problem

2026/02/24

BUILDmigrationtorch

TL;DR: I wrote a whole post about what “done” looks like for torch in Bazel and framed Level 3 — CPU torch tests on Linux CI — as a Bazel problem. It’s not. It’s a test classification problem that exists today, with or without Bazel. ~40% of our tests import torch but only need CPU. They’re stuck on GPU runners because nobody’s ever separated “imports torch” from “needs a GPU.” Fix classification, and routing is easy — whether you route with Bazel tags or pytest markers.

This post is a companion to the success levels post. That one defined the gap. This one is the “okay, how do we actually close it” playbook.

The Correction

In the success post, I laid out a test taxonomy:

Test Type Needs GPU? % of tests
Pure unit tests No ~20%
Torch CPU tests No ~40%
Torch GPU tests Yes ~25%
Integration tests Yes ~10%
Benchmarks Yes ~5%

And I identified “Torch CPU tests” as THE GAP — tests that could run on cheap CPU runners but are stuck waiting for GPU runner availability. Then I spent the rest of the post talking about Bazel override mechanisms, lockfile strategies, and hermetic Python.

Here’s what I got wrong: the problem isn’t Bazel can’t resolve torch on Linux. The problem is we can’t tell which tests need a GPU and which don’t. That distinction doesn’t exist anywhere — not in Bazel, not in pytest, not in CI config. It’s not a tooling problem. It’s a classification problem.

If I could snap my fingers and know, for every test file, “CPU-only” or “needs GPU” — routing them to the right runners would be a solved problem in any tool. Bazel tags, pytest markers, GHA matrix splits — pick your favorite. The hard part is getting the labels. Everything downstream is plumbing.

Why This Hasn’t Been Done Already

It seems like it should be obvious, right? Just… figure out which tests need a GPU?

Three reasons it’s harder than it sounds:

1. import torch ≠ “needs a GPU”

In our codebase, import torch is as common as import os. Researchers put it on line 1 of everything. A test that does torch.tensor([1, 2, 3]) doesn’t need a GPU. A test that does torch.cuda.synchronize() does. But they both have import torch in their imports.

Today’s CI treats them identically: “this file imports torch → send it to a GPU runner.” Which means a test that checks tensor shapes waits in the same queue as a test that trains a model for three hours.

2. Transitive dependencies hide the real usage

Even looking at a test file’s direct code isn’t enough. Consider:

# test_model_config.py
from model_arch import build_model

def test_model_has_correct_layers():
    model = build_model(config)
    assert len(model.layers) == 12

Does this test need a GPU? Depends on what build_model does. If it calls model.to("cuda") internally — yes. If it constructs the model on CPU by default — no. You can’t tell from the test file alone. You need to trace through build_model → ModelArch.__init__ → LayerFactory.create → … and see if anything in that call graph touches CUDA.

That’s a transitive analysis problem. It’s the same shape as the mai_config blocker from the dependency resolution post — one import, fifty blocked packages. Except now it’s one function call, unknown GPU dependency.

3. Nobody owns this distinction

Researchers write tests. Researchers don’t think about CPU vs GPU runners. That’s an infrastructure concern. But infrastructure (me) doesn’t know which torch operations need CUDA and which don’t. The knowledge lives in researchers’ heads, expressed as code patterns they use without thinking about it.

There’s no natural owner for “classify all tests by hardware requirement.” It falls between roles.

Until now.

The Three Layers

Getting CPU torch tests onto CPU runners breaks down into three layers. They’re ordered by difficulty, and — crucially — the hardest one has to come first.

Layer 1: Classification (Hard) — “Which tests actually need a GPU?”

This is the whole ballgame. Everything else is plumbing.

What we need: A label on every test that says cpu or gpu. Not based on “does it import torch” but based on “does it actually invoke CUDA operations at runtime.”

Three approaches, used together:

Approach A: Heuristic Static Analysis (catches ~80%)

Scan test files and their transitive imports for GPU-indicating patterns:

GPU_PATTERNS = [
    r"\.cuda\(",            # tensor.cuda(), model.cuda()
    r"\.to\(['\"]cuda",     # .to("cuda"), .to("cuda:0")
    r"torch\.cuda\.",       # torch.cuda.is_available(), torch.cuda.synchronize()
    r"nccl",                # distributed training
    r"@pytest\.mark\.gpu",  # already marked (if any)
    r"device=['\"]cuda",    # device="cuda" in function args
]

CPU_SAFE_PATTERNS = [
    r"torch\.tensor\(",     # basic tensor creation
    r"torch\.zeros\(",
    r"torch\.ones\(",
    r"\.shape",             # shape inspection
    r"\.dtype",
    r"torch\.no_grad\(",    # inference mode
    r"torch\.testing\.",    # test utilities
]

Walk the import graph. If a test file — or anything it imports — matches a GPU pattern, flag it as gpu. If it only matches CPU-safe patterns, flag it as cpu. If it’s ambiguous, flag it as unknown.

This won’t be perfect. Static analysis can’t follow dynamic dispatch, conditional branches (if cuda_available: ...), or patterns you haven’t thought of yet. But it doesn’t need to be perfect. It needs to be good enough to separate the obvious cases.

Expected yield: ~80% of torch-importing tests fall into clearly-CPU or clearly-GPU buckets. The remaining ~20% are ambiguous.

Approach B: Manual Annotation (handles the ambiguous ~20%)

For tests that static analysis can’t classify, ask the humans:

@pytest.mark.cpu
def test_model_forward_pass():
    """Forward pass is CPU-safe — model uses CPU by default."""
    ...

@pytest.mark.gpu
def test_gradient_accumulation_multi_gpu():
    """Needs multi-GPU for gradient sync testing."""
    ...

This is the annotation tax. Someone (probably me, armed with the static analysis output) goes through the unknown bucket and makes a call. It’s not fun. It’s also not that big — we’re talking about ~20% of ~40% of tests, so roughly 8% of the total test count. Manageable in a focused sprint.

The key insight: you only need to annotate once. After the initial classification, new tests get a marker as part of the normal code review process. “Your test imports torch — is it CPU or GPU?” becomes a checkbox in the review template.

Approach C: Runtime Detection (validates everything)

The safety net. Run “classified-as-CPU” tests on a CPU-only runner. If they fail with a CUDA error, reclassify:

# Pseudo-workflow
- run: pytest -m cpu --tb=short
  on: cpu-runner
  continue-on-error: true

- if: failure()
  run: |
    # Parse failures, find CUDA-related errors
    # Reclassify those tests as gpu
    # Open PR to update markers

This catches misclassifications from both static analysis and human annotation. It’s a feedback loop: run on CPU → fail → reclassify → try again. After a few iterations, the classification converges.

The beautiful part: false negatives (test classified as CPU but needs GPU) are safe — they fail loudly and get reclassified. False positives (test classified as GPU but only needs CPU) are wasteful but not broken. The error mode is “wasted GPU runner time,” not “broken tests.”

Putting It Together

Step 1: Static analysis scan → ~80% classified
        ├── clearly CPU → mark @pytest.mark.cpu
        ├── clearly GPU → mark @pytest.mark.gpu
        └── ambiguous → needs human review

Step 2: Human review of ambiguous bucket → remaining ~20%
        └── all tests now have cpu/gpu markers

Step 3: Runtime validation loop
        ├── run cpu-marked tests on CPU runner
        ├── CUDA failures → reclassify as gpu
        └── iterate until stable

After convergence, you have a clean, validated classification of every torch-importing test. This is the foundation everything else builds on.

Layer 2: Torch Installation (Medium) — “CPU runners need CPU torch”

Today, CPU runners might not have torch installed at all. Torch-importing tests live on GPU runners because that’s where torch exists. To run torch CPU tests on CPU runners, you need a CPU-only torch build available there.

This is a solved problem — just not solved for us yet.

Option A: Pre-install in the runner image

# cpu-test-runner.Dockerfile
FROM ubuntu:22.04
RUN pip install torch --index-url https://download.pytorch.org/whl/cpu
# CPU-only torch, ~800MB, no CUDA deps

Add CPU-only torch to the base image for CPU test runners. One-time setup, works with any test framework.

Option B: Install at test time

# In GHA workflow
- name: Install CPU torch
  run: pip install torch --index-url https://download.pytorch.org/whl/cpu
  if: matrix.runner == 'cpu'

Simpler to set up, slower per-run (800MB download + install), but avoids maintaining a custom runner image.

Option C: Bazel-managed (the fancy version)

If you’re running tests through Bazel, torch becomes a @pypi// dependency resolved from the CPU-only index. This is what the success post spent pages discussing — Options A through D in the architecture section. It’s the cleanest long-term answer but has more moving parts.

What I’d actually do: Start with Option A or B. Get CPU torch on a runner, run classified-CPU tests there, prove the classification works. Then worry about making it elegant with Bazel. Proving the concept matters more than proving the architecture.

Layer 3: Routing (Easy) — “Send tests to the right runner”

Once you have labels (Layer 1) and runners that can handle each label (Layer 2), routing is just configuration.

With pytest + GHA matrix:

jobs:
  cpu-tests:
    runs-on: [self-hosted, cpu, big]
    steps:
      - run: pytest -m "cpu or not gpu" --ignore=benchmarks/

  gpu-tests:
    runs-on: [self-hosted, gpu, a100]
    steps:
      - run: pytest -m "gpu" --ignore=benchmarks/

With Bazel tags:

# cpu_tests test_suite
test_suite(
    name = "cpu_tests",
    tags = ["requires-torch-cpu"],
)

# gpu_tests test_suite
test_suite(
    name = "gpu_tests",
    tags = ["requires-gpu"],
)

# GHA workflow
jobs:
  cpu:
    runs-on: cpu-runner
    steps:
      - run: bazel test //:cpu_tests

  gpu:
    runs-on: gpu-runner
    steps:
      - run: bazel test //:gpu_tests

With Bazel + tag filtering (no test_suite needed):

# Run all tests except GPU ones on CPU runner
bazel test //... --test_tag_filters=-requires-gpu

# Run only GPU tests on GPU runner
bazel test //... --test_tag_filters=requires-gpu

See? Once you have labels, routing is a solved problem. The tool doesn’t matter. Pytest markers, Bazel tags, CI matrix dimensions — they’re all the same idea: “run tests with label X on runner type Y.” The question was never “how do we route tests?” It was “how do we get the labels?”

The Concrete Playbook

Here’s what “Level 3: CPU torch tests on Linux CI” actually looks like, broken into doable chunks:

Phase 1: Build the Classifier (1-2 weeks)

  1. Write the static analysis scanner. Input: test file path. Output: cpu, gpu, or unknown. Start with the pattern lists above, iterate as you discover new patterns.

  2. Run it on the entire test suite. Generate a report:

    Total torch-importing tests: 847
Classified CPU:              612 (72%)
Classified GPU:              148 (18%)
Unknown:                      87 (10%)

  3. Manual review the unknowns. This is the sprint. Go through 87 tests (or whatever the number is), read the code, make a call. Tag each one.

  4. Add pytest markers to all classified tests. This is a bulk code change — @pytest.mark.cpu or @pytest.mark.gpu on each test function or class. Could be automated from the classification output.

Phase 2: Set Up the Runner (1 week)

  1. Build or configure a CPU test runner with torch. Simplest version: add pip install torch --index-url https://download.pytorch.org/whl/cpu to the runner setup. Verify torch imports and basic tensor ops work.

  2. Create a CI job that runs pytest -m cpu on the CPU runner. Start with a subset — pick one package with high CPU-test density, run it, watch it work. Then expand.

Phase 3: Validate and Iterate (1-2 weeks)

  1. Run all CPU-classified tests on the CPU runner. Some will fail — CUDA errors from misclassifications, missing test dependencies, fixture issues. Fix each one.

  2. Set up the reclassification feedback loop. CUDA failures → PR to update markers. After 2-3 iterations, classification stabilizes.

  3. Remove CPU-classified tests from the GPU runner job. This is the payoff — GPU runners now only run tests that actually need GPUs. Queue times drop.

Phase 4: Make It Stick (ongoing)

  1. Add classification to the code review checklist. New torch-importing tests get a marker. CI enforces it — a lint check that says “this test imports torch but has no cpu/gpu marker.”

  2. Optionally: migrate CPU tests to Bazel. With classification done and CPU torch available, Bazel migration of these tests becomes straightforward — it’s just adding BUILD files and tags. But this is optimization, not the core win.

What Changes When You Frame It This Way

The success post made Level 3 feel like a Bazel infrastructure problem. “Get CPU torch into Bazel’s hermetic Python on Linux” — that sounds like a months-long adventure through override mechanisms and lockfile strategies.

Reframed as a classification problem, Level 3 is:

  1. Label your tests. (Hard but bounded — it’s an analysis + annotation project, not an infrastructure change.)
  2. Put torch on a CPU runner. (One Dockerfile line or one pip install.)
  3. Route by label. (Standard CI configuration.)

No Bazel override mechanisms. No lockfile surgery. No fights with hermetic Python. You can do the entire thing with pytest markers and GHA matrix splits today, and migrate to Bazel tags later if you want the caching benefits.

The bottleneck was never “which build system.” The bottleneck was “we don’t know which tests need a GPU.” Fix that, and the rest is plumbing.

The Honest Scope

A few things I’m glossing over that deserve acknowledgment:

Where This Fits in the Series

The series arc, updated:

The levels of done still stand. The correction is about how to reach Level 3. It’s not a Bazel-first path — it’s a classification-first path that works with any test runner, and makes Bazel migration easier when you get to it.

Next up: actually doing it. I have a scanner to write.

Previous in series: Torch in Bazel: What Does Success Actually Look Like? | How Bazel Resolves Python Deps