Why Bazel Tests Crash With libcublas.so Not Found

One sentence: Bazel tests crash because the venv gets GPU torch (from uv.lock), which shadows Docker’s CPU torch. GPU torch needs CUDA. No CUDA on the runner. Boom.

I wrote a 3-part series on this and still couldn’t explain the problem clearly. This post is the replacement. One page, complete picture.

CPU torch vs GPU torch

PyPI ships two flavors of torch for Linux x86_64:

The GPU wheel is the default on PyPI for Linux x86_64. If you pip install torch on a Linux machine, you get the GPU wheel. You need an explicit --index-url or --extra-index-url pointing to PyTorch’s CPU-only index to get the CPU wheel.

This distinction is the entire problem.

How pytest (non-Bazel) handles it

Simple. One shared torch installation. CI installs CPU torch via uv sync with the CPU index configured. All tests use it. Done.

No venv-per-test. No dependency resolution at test time. One torch, one place, CPU. Works.

How Bazel creates the problem

Each py_torch_test creates its own venv. Two sources of torch fight:

1. Docker image has CPU torch at /usr/local/.../site-packages/torch/ ✅
2. @pypi//torch from BUILD deps → resolved from uv.lock → GPU torch → installed into venv ❌

Python checks the venv first, system packages second. GPU torch shadows CPU torch. GPU torch tries to load CUDA libs. No CUDA on this machine.

import torch
  → Python finds torch in venv (GPU)
    → torch.__init__ loads libtorch_cuda.so
      → libtorch_cuda.so links libcublas.so
        → libcublas.so not found
          → ValueError 💥

Where does GPU torch come from?

Trace the chain:

BUILD file: deps = ["@pypi//torch"]
  → Bazel pypi hub
    → resolved from uv.lock
      → default PyPI index
        → platform = linux x86_64
          → default wheel = GPU torch (865 MB)

On macOS, PyPI only has CPU wheels for torch. So bazel test on your Mac works fine. The GPU wheel only exists for Linux, which is why this only blows up in CI.

The 5 environments

Why only bazel-torch-docker breaks

It’s the only environment where both conditions are true:

1. The venv has GPU torch (from uv.lock → PyPI default)
2. The machine has no CUDA

Every other environment dodges at least one condition. pytest CI never creates a venv with PyPI torch. Local Mac gets CPU torch because macOS. The bazel CI job filters torch out entirely. The GPU docker image actually has CUDA.

One environment. Two conditions. That’s the whole bug.

How it actually broke

Earlier versions of this post implied bazel-torch-docker was never green. That’s wrong. It was green for a month. Here’s what actually happened.

The timeline

You can verify this yourself on the nightly run dashboard. Compare last green on Apr 1 (125/125 pass) vs first red on Apr 2 (114/125 pass, 11 fail).

What happened on Apr 2

The automated gazelle drift bot (PR #28031, merged Mar 27) ran gazelle on mai_io/_node_cache. Gazelle saw the source files do import torch and import safetensors, so it added them to the BUILD:

# BEFORE (original migration, Mar 11)
deps = ["//lib/blobfile/src/blobfile"]

# AFTER (auto-fix by gazelle bot, Mar 27)
deps = [
    "//lib/blobfile/src/blobfile",
    "@pypi//safetensors",
    "@pypi//torch",        # ← the 865MB GPU wheel enters the dep graph
]

But it didn’t break on Mar 27. Mar 27 to Apr 1 was still green. Why?

The Bazel CI workflow caches by uv.lock hash. Between Mar 27 and Apr 1, the cache from the old uv.lock was still valid. The old cache had the previous dep resolution where @pypi//torch wasn’t in mai_io test venvs.

On Apr 2, multiple uv.lock changes merged (new deps from several PRs). The Bazel cache key changed. Fresh resolution. Now @pypi//torch (865MB GPU wheel) properly resolved into the mai_io test venvs. GPU torch shadowed Docker’s CPU torch. libcublas.so crash. 11 tests broke.

The irony

Gazelle did exactly what it’s supposed to do. The source code imports torch, so gazelle added @pypi//torch to deps. Correct behavior. The problem is that @pypi//torch on linux x86_64 means “865MB GPU wheel that needs CUDA.” Nobody expected a correct dep addition to break CI five days later via a cache miss.

The fix (and the 4 wrong ways to do it)

Root cause in one line: @pypi//torch in BUILD files → resolved from uv.lock → default PyPI → 865 MB GPU wheel on linux x86_64 → installed into Bazel venv → shadows Docker’s CPU torch → libcublas.so crash.

Five approaches. One shipped.

Solution 1: Pin torch to CPU index in pyproject.toml (IDEAL, NOT SHIPPED)

# pyproject.toml
[[tool.uv.index]]
name = "pytorch-cpu"
url = "https://download.pytorch.org/whl/cpu"
explicit = true

[tool.uv.sources]
torch = { index = "pytorch-cpu" }
torchvision = { index = "pytorch-cpu" }
torchaudio = { index = "pytorch-cpu" }

Then uv lock → uv.lock gets CPU wheel URLs → aspect_rules_py picks them up. One config change, fixes everything.

Why this works: aspect_rules_py’s uv extension is a dumb lockfile consumer. It reads uv.lock, takes the wheel URLs, creates http_file() repos. No index logic at the Bazel layer. If uv.lock has CPU wheel URLs, Bazel gets CPU torch. Period.

Why we didn’t ship this: Changing the torch index in pyproject.toml forces a full uv.lock regeneration. In a monorepo with ~400 transitive deps, that’s a massive diff that touches every team’s lockfile entries. The regeneration also surfaced unrelated resolution conflicts we didn’t want to untangle in the same PR. Correct fix, wrong blast radius.

Solution 2: Override @pypi//torch to empty target in MODULE.bazel ✅ SHIPPED

Use aspect_rules_py’s override mechanism to replace @pypi//torch with an empty py_library. Torch comes from system site-packages instead.

How uv.override_requirement works

Before the fix:

BUILD: deps = ["@pypi//torch"]
  → Bazel pypi hub (from uv.lock)
    → downloads 865MB GPU wheel
      → installed into venv
        → shadows system CPU torch
          → import torch → libcublas.so → 💥

After the fix:

MODULE.bazel: uv.override_requirement(requirement="torch", target="//bazel/torch")
  → @pypi//torch now = empty py_library (no wheel, zero bytes)
    → venv has no torch
      → include_system_site_packages=True
        → Python finds Docker's CPU torch at /usr/local/.../site-packages/
          → import torch → CPU torch → ✅

The config change (exact code)

# MODULE.bazel
uv.override_requirement(
    hub_name = "pypi",
    requirement = "torch",
    target = "//bazel/torch",
    venv_name = "default",
)

# bazel/torch/BUILD.bazel
py_library(
    name = "torch",
    srcs = [],
    visibility = ["//visibility:public"],
)

Two files. That’s the whole fix.

What uv.override_requirement does mechanically

It tells aspect_rules_py: “whenever you see @pypi//torch, don’t use the wheel from uv.lock. Use this Bazel target instead.”

Normal: @pypi//torch → http_file(url="...865MB.whl") After: @pypi//torch → //bazel/torch → py_library(srcs=[]) (empty)

The 865MB wheel is never downloaded. The venv gets an empty target. System CPU torch (from Docker) wins via include_system_site_packages=True.

The tradeoff

uv.override_requirement() runs at module extension time, before select() exists. Unconditional. Can’t do “empty on Linux, real wheel on macOS.”

macOS impact: local devs need pip install torch in their system Python. But py_torch_test already uses include_system_site_packages=True, so system torch is the intended source anyway. In practice, every dev who runs torch tests already has torch installed locally.

Score: 9/10. Minimal blast radius (2 files), no lockfile churn, works immediately on Linux CI. macOS “install torch yourself” is a non-issue for this team.

Solution 3: PYTHONPATH override in CI workflow YAML

--test_env=PYTHONPATH=/usr/local/lib/python3.12/site-packages

Force Python to find system CPU torch before venv GPU torch.

Score: 4/10. Duct tape. Fragile if Docker image path changes. Invisible to anyone reading BUILD files. Nobody will remember why this exists in 6 months.

Solution 4: Lazy import in mai_config/structured.py

Move import torch from module-level to inside functions. Tests that don’t use torch never trigger CUDA init.

Score: 7/10. Good hygiene, but only fixes 6 of 12 failures (the ones that import torch transitively through mai_config). Tests that directly import torch (mai_io/_node_cache) still fail.

Solution 5: Remove @pypi//torch from library BUILD deps

Remove @pypi//torch from py_library deps, rely on system site-packages only.

Score: 5/10. Breaks bazel build for anything that legitimately needs torch at analysis time. Also makes dependency graph a lie.

The ranking

* macOS: devs need torch in system Python. Already the case for anyone running torch tests.

The lesson for teammates

aspect_rules_py is a lockfile consumer, not an index resolver. The ideal fix is at the uv layer (Solution 2 in the table), but that required a full lockfile regen with massive blast radius.

What we shipped instead: uv.override_requirement to replace @pypi//torch with an empty target. Two files, zero lockfile churn, the 865MB GPU wheel never downloads. System CPU torch (already in the Docker image) takes over via include_system_site_packages=True.

If someone asks “why not just fix pyproject.toml?”, the answer is: we will, eventually. But that PR touches every team’s lockfile entries. The override is surgical, ships today, and buys time to do the lockfile migration properly.

一句话总结

只有 bazel-torch-docker 炸。因为只有它同时满足两个条件：venv 里有 GPU torch，机器上没有 CUDA。