Bazel Torch: What Does Success Actually Look Like?

TL;DR: Everyone keeps asking “when will torch work in Bazel?” but nobody stops to define what “work” means. After weeks of fighting overrides and resolution pipelines, I stepped back and mapped the whole landscape: what torch actually is (from a build system’s perspective), what kinds of tests exist, which ones Bazel should own, and what “done” looks like at each level. Turns out the answer isn’t “run everything in Bazel.” It’s “make the build graph correct and run each test in the right place.”

This post is the big-picture companion to the series. Part 1 covered what happens when GPU shows up. Part 2 was the scoreboard. Part 3 went deep on the resolution pipeline. The TIL and the override comedy covered specific mechanics. This post zooms out.

What Even Is Torch? (The CI Engineer’s Version)

If you’re an ML researcher, torch is “the thing you import before doing anything.” If you’re the person maintaining the build system, torch is a 2GB headache with seventeen variants.

Here’s what torch actually is, from the perspective of someone who has to make it resolve, build, test, and cache correctly:

It’s not a Python package. Or rather, it is a Python package the way a car is “a thing with wheels.” Technically true, missing the point. Torch is a massive C++/CUDA library that happens to have Python bindings. When you import torch, Python loads a pile of .so shared libraries — not just .py files. The actual computation happens in compiled C++ and CUDA kernels.

It’s enormous. CPU-only torch is ~800MB. With CUDA support, you’re looking at 2GB+. For comparison, numpy is ~15MB, pydantic is ~2MB. Torch is literally 100x the size of a normal dependency. This matters for CI cache sizes, Docker layer strategies, download times, and disk quotas.

It has build variants. This is the real killer. There isn’t one torch wheel — there are many:

See that +cu124 in the filename? That’s not a version suffix — it’s a completely different binary linked against a completely different CUDA toolkit. You can’t mix them. You can’t “just pip install torch” and expect it to work everywhere.

This is why requirements.txt doesn’t cut it. The moment your CI has multiple environments (macOS dev laptops, Linux CPU runners, Linux GPU runners with different CUDA versions), you need a resolution strategy, not a package list.

The Researcher’s Workflow (and Why I Need to Understand It)

I don’t train models. I’m DevEx — I make the build go fast and the CI go green. But I need to understand what researchers actually do with torch, because that determines what the build system needs to support.

A researcher’s typical day:

1. Write model code. This means import torch on line 1 of basically every file. Model architectures, loss functions, data pipelines, evaluation scripts — they all touch torch.
2. Test locally. Quick sanity checks on a dev machine. “Does this forward pass produce the right tensor shape?” Usually CPU is fine for this.
3. Launch training. The real work happens on GPU clusters — multi-node, multi-GPU, hours to days of compute. This is where NCCL, distributed training, and CUDA kernels actually matter.
4. Iterate. Change architecture, retrain, compare metrics, repeat.

The key insight: most of their code could be tested on CPU. Tensor shapes, model forward passes, data preprocessing, config parsing — none of this needs a GPU. But researchers don’t think in terms of “CPU vs GPU tests.” They just import torch and it works. The separation of “what needs a GPU” vs “what doesn’t” is a CI concern, not a research concern.

Which means it’s my concern.

The Test Taxonomy (The Part I Wish Someone Had Explained To Me)

Not all torch tests are created equal. Once I mapped out what kinds of tests actually exist in our codebase, the Bazel strategy became much clearer.

That top row — pure unit tests — is the easy win. These don’t touch torch at all. They run on any machine, they’re fast, they’re hermetic. Bazel eats these for breakfast. We already have them working.

The second row is the interesting one. Torch CPU tests. These need a real torch installation, but they don’t need a GPU. They test things like “does my model produce the right output shape?” or “does my data augmentation pipeline preserve tensor dimensions?” These are the bulk of our test suite, and they’re stuck in a gap: too torch-dependent for plain Bazel runners, not GPU-dependent enough to justify a GPU runner.

The bottom three rows? Those stay where they are. Docker containers on GPU runners, launched by pytest. Bazel doesn’t need to own these, and honestly, trying to make Bazel orchestrate GPU tests would create more problems than it solves.

Where Each Test Type Runs Today vs Where It Should Run

Here’s the current state and the target state. The “gap” column is where the work is.

See that gap? ~40% of our tests could run on cheap CPU runners in seconds instead of waiting for GPU runner availability and Docker container spin-up. But we can’t do that today because there’s no clean way to get CPU torch into a Bazel test sandbox on Linux.

The Five Levels of “Done”

This is the mental model I’ve been using to track progress. Each level builds on the previous one.

Level 1: @pypi//torch resolves on macOS

Status: Done.

The smoke test passes. bazel test //bazel/tests:test_torch_smoke runs on macOS, imports torch, does tensor math, passes. The uv.lock has macOS wheels, aspect_rules_py generates the @pypi//torch target, gazelle wires it into BUILD files. Clean.

This was the easy part because macOS torch is just a normal PyPI wheel. No CUDA variants, no Docker, no system-installed packages. It’s the one platform where torch behaves like a regular pip package.

Level 2: bazel build //... doesn’t crash on Linux

Status: The current fight.

Right now, if you run bazel build //... on a Linux CI runner, any target that depends on @pypi//torch fails because the select() statement in the generated torch target has no Linux match. The lockfile only has macOS wheels (because of the sys_platform == 'darwin' gate in pyproject.toml), so Bazel looks at Linux, finds nothing, and errors out.

62% of our packages depend on torch. So 62% of our packages can’t build on Linux in Bazel. This is… not great.

The override mechanism I explored in The Override That Overrode Everything was an attempt to fix this. It works — but it’s global, which means you can’t have different behavior per platform without getting clever.

Level 3: CPU torch tests run via Bazel on Linux CI

Status: Next.

This is the real prize. Get CPU-only torch available inside Bazel’s hermetic Python on Linux, so that the ~40% of tests that only need CPU torch can run as bazel test targets on cheap Linux runners.

Fast, incremental, cached, remote-cacheable. All the things Bazel is good at.

Level 4: Torch-using packages fully migrated to Bazel

Status: Future.

All torch-dependent packages have BUILD files, all CPU-testable targets are tagged correctly, bazel test //... runs everything that can run without a GPU. The build graph is complete and accurate.

Level 5: GPU tests via Bazel

Status: Dream (and probably not worth it).

Running GPU tests inside Bazel’s sandbox with real CUDA. Theoretically possible with remote execution on GPU machines. Practically? Docker + pytest on GPU runners works fine. The ROI of making Bazel orchestrate GPU scheduling, CUDA driver compatibility, and NCCL networking is… low.

Sometimes the right answer is “don’t.”

The Architecture Options for Levels 2-3

Here’s where it gets interesting. There are multiple ways to get torch into a Bazel build on Linux. Each has real tradeoffs.

Option A: Add Linux CPU Wheels to uv.lock

Remove the sys_platform == 'darwin' gate from pyproject.toml. Let uv resolve Linux CPU wheels alongside macOS wheels. Now @pypi//torch has a Linux match and select() works everywhere.

# Before
"torch>=2.7.0; sys_platform == 'darwin'"

# After
"torch>=2.7.0"

Pros:

• Simplest conceptual model — torch is just a pip package everywhere
• @pypi//torch works on both platforms, no overrides needed
• Native Bazel hermetic Python, no system-packages hacks

Cons:

• The lockfile now includes ~800MB CPU Linux wheels (metadata + hashes)
• Every uv lock resolution has to consider Linux torch and all its transitive deps
• CI runners need to download 800MB of torch on first cache miss
• You’re resolving CPU-only torch, but your GPU CI needs CUDA torch — different binary, different deps

Verdict: Clean but expensive. Works if you can stomach the lockfile size and cache implications.

Option B: http_archive a CPU Torch Wheel

Download a specific CPU torch wheel as a Bazel external dependency. Wrap it as a py_library. Use uv.override_requirement() to point @pypi//torch at it.

# MODULE.bazel
http_archive(
    name = "torch_cpu_linux",
    url = "https://download.pytorch.org/whl/cpu/torch-2.7.0+cpu-cp312-cp312-manylinux_2_17_x86_64.whl",
    sha256 = "...",
)

uv.override_requirement(
    hub_name = "pypi",
    requirement = "torch",
    target = "@torch_cpu_linux//:torch",
)

Pros:

• Exact wheel, exact hash, fully hermetic
• Doesn’t pollute uv.lock with Linux wheels
• You control exactly which torch variant Bazel uses

Cons:

• The override is global — it replaces @pypi//torch everywhere, including macOS (unless you use select() in the override target, which gets complicated)
• You’re now maintaining torch version pins in two places (uv.lock for macOS, MODULE.bazel for Linux)
• All of torch’s transitive deps need to be handled too (and torch has… many)
• This is exactly the rabbit hole I went down in the override story

Verdict: Precise but fragile. You end up fighting the override’s global scope.

Option C: Run Bazel Inside Docker

Use the existing NGC-based Docker image (nvcr.io/nvidia/pytorch:25.03-py3). Install bazelisk in the image. Run bazel test inside the container where torch is already system-installed.

FROM nvcr.io/nvidia/pytorch:25.03-py3
RUN curl -L https://github.com/bazelbuild/bazelisk/releases/... -o /usr/local/bin/bazel
# torch is already here, system-installed

Pros:

• Torch “just works” — it’s in the system Python
• Same image used for GPU tests, no separate torch management
• Researchers don’t need to care about Bazel’s torch story

Cons:

• Bazel uses a hermetic Python toolchain (rules_python downloads Python 3.12.2 into its own sandbox). This Python can’t see system site-packages by default.
• You need to either: teach Bazel’s Python to use --system-site-packages, or bypass Bazel’s hermetic Python and use the system Python (defeating half the point of Bazel)
• Docker image startup adds latency vs bare-metal runners
• Cache sharing between Docker runs is harder

The hermetic Python problem deserves its own paragraph. When Bazel runs a py_test, it uses a Python interpreter downloaded by rules_python — typically Python 3.12.2 in its own isolated directory. This interpreter has its own site-packages that Bazel populates from @pypi// wheels. It cannot see /usr/lib/python3/dist-packages/torch or wherever the system torch lives. Bazel’s entire value proposition is isolation. You’re asking it to be less isolated. It doesn’t want to.

Verdict: Pragmatic but fights Bazel’s philosophy. You gain torch access at the cost of hermetic purity.

Option D: Don’t Run Torch Tests in Bazel

Keep torch tests in pytest. Docker + GPU runners for anything that imports torch. Bazel only owns non-torch packages.

# ci/configs/cpu_tests.yml
# Only non-torch packages here
packages:
  - mai_utils
  - data_pipeline
  - config_lib
  # NOT model_arch, NOT training, NOT evaluation

Pros:

• Zero torch-in-Bazel complexity
• Each tool does what it’s good at: Bazel for hermetic builds, Docker for GPU runtime
• This is what most of the industry does (ask anyone at Google — torch tests don’t go through Blaze)

Cons:

• ~40% of tests (torch CPU tests) stay in the slow Docker pipeline even though they don’t need a GPU
• No incremental testing for torch-dependent code
• The build graph is incomplete — bazel build //... can’t process 62% of the repo

Verdict: The pragmatic default. Not exciting, but reliable.

The Key Insight: Separation of Concerns

After going through all four options (and personally living through Option B’s failure mode), here’s what I think the real answer is:

You don’t need to run everything in Bazel. You need the build graph to be correct and the tests to run in the right place.

The mental model:

The bridge between these worlds is: @pypi//torch resolves (even if it resolves to a stub or CPU-only wheel) so that the build graph is complete. Then test targets are tagged or gated to run in the right environment:

py_test(
    name = "test_model_forward",
    srcs = ["test_model_forward.py"],
    deps = ["//model_arch", "@pypi//torch"],
    tags = ["requires-torch-cpu"],  # CI knows where to run this
)

py_test(
    name = "test_gpu_kernel",
    srcs = ["test_gpu_kernel.py"],
    deps = ["//kernels", "@pypi//torch"],
    target_compatible_with = ["@platforms//os:linux"],
    tags = ["requires-gpu"],  # Stays in Docker pipeline
)

The build graph is complete. Bazel knows the dependency structure. Each test runs where it can actually pass. Nobody’s trying to run CUDA kernels on a GitHub-hosted runner, and nobody’s waiting 10 minutes for a Docker container to spin up just to test tensor shapes.

So Where Are We?

Honest status as of today:

• Level 1 (macOS resolve): Done. Smoke test passes, torch imports, tensors work.
• Level 2 (Linux build): In progress. The override mechanism exists but the global scope is a problem. Need either a platform-aware override or a different approach entirely.
• Level 3 (CPU torch tests): Blocked on Level 2. Can’t run what you can’t build.
• Level 4 (full migration): Future. Depends on Level 3, plus gazelle improvements for the remaining 62%.
• Level 5 (GPU in Bazel): Not planned. Docker wins here. That’s fine.

The next step is probably a hybrid of Options A and D: get Linux CPU wheels into the lockfile for the build graph (accepting the lockfile bloat), tag torch-dependent test targets appropriately, and keep GPU tests in Docker. Not glamorous. But it closes the gap without fighting Bazel’s philosophy.

You don’t need to boil the ocean. You need the graph to be correct and the tests to run in the right place. Everything else is optimization.

Previous in series: The Override That Overrode Everything | How Bazel Resolves Python Deps