Bazel Torch: How Bazel Resolves Python Deps and Why torch Breaks

TL;DR: Bazel’s Python dependency resolution is a clean three-step pipeline: uv.lock feeds @pypi// targets, gazelle wires them into BUILD files, and everything Just Works. Until one of your dependencies isn’t in uv.lock because it’s system-installed via Docker. Then the pipeline has a hole in it, and 62% of your packages fall through.

This is the third post in a series. Part 1 was observational — what happens when GPU shows up. Part 2 was the battle plan — the scoreboard and the strategy. This post goes deeper: how the resolution machinery actually works, where exactly it breaks, and the specific mechanics of each fix.

The Resolution Pipeline (When It Works)

Before we get to what’s broken, let’s trace how Python dependencies actually flow through a Bazel monorepo. It’s a three-stage pipeline, and when all three stages agree on reality, it’s genuinely elegant.

Stage 1: uv.lock — The Single Source of Truth

You declare dependencies in pyproject.toml, run uv lock, and get a lockfile. This part is standard Python tooling — nothing Bazel-specific yet. The lockfile pins every package to an exact version with hashes.

# pyproject.toml (workspace root)
[project]
dependencies = [
    "numpy>=1.26",
    "pydantic>=2.0",
]

uv.lock resolves these to exact versions, exact wheels, exact hashes. One lockfile for the entire monorepo.

Stage 2: @pypi// — Bazel Learns About pip Packages

Here’s where aspect_rules_py does its magic. It reads uv.lock and creates a Bazel repository called @pypi//. Every resolved package becomes a target:

@pypi//numpy      → numpy 1.26.4 (wheel, checksummed)
@pypi//pydantic   → pydantic 2.6.1 (wheel, checksummed)
@pypi//flask      → flask 3.0.2 (wheel, checksummed)

These targets are hermetic. Bazel downloads the exact wheel, verifies the hash, caches it. No network calls after the first fetch. No “works on my machine.” If the lockfile says numpy 1.26.4, every developer and every CI runner gets numpy 1.26.4.

This is the good part. This is what makes Bazel worth the complexity.

Stage 3: Gazelle — Auto-Wiring BUILD Files

The last stage is gazelle, which scans your Python source files, finds import statements, and generates BUILD files with the correct deps:

# my_package/analysis.py
import numpy as np
import pydantic

Gazelle sees those imports and generates:

# my_package/BUILD.bazel
py_library(
    name = "analysis",
    srcs = ["analysis.py"],
    deps = [
        "@pypi//numpy",
        "@pypi//pydantic",
    ],
)

Three stages. Lockfile resolves versions. @pypi// makes them Bazel targets. Gazelle wires targets into BUILD files. For pure Python packages — which is most of them — this pipeline is fast, correct, and requires almost zero manual intervention.

Now let’s break it.

Where the Pipeline Breaks

The pipeline has a simple contract: every import statement must resolve to either a @pypi// target or an in-repo target. If gazelle encounters an import it can’t map to either, the BUILD file can’t be generated. The package is stuck.

The Hole: Packages That Aren’t in uv.lock

Here’s our root pyproject.toml:

[project]
dependencies = [
    "torch; sys_platform == 'darwin'",
]

[tool.uv]
override-dependencies = [
    "flash-attn; sys_platform == 'never'",
    "triton; sys_platform == 'never'",
    "ucxx-cu12; sys_platform == 'never'",
]

Read that carefully.

torch only resolves on macOS. On Linux — where CI runs, where training happens, where it actually matters — torch isn’t in uv.lock. It’s system-installed inside the Docker image.

flash-attn is even wilder. sys_platform == 'never' is a platform that doesn’t exist. It’s a hack that tells uv “pretend this package isn’t real.” Because flash-attn can’t be pip-installed normally — it needs CUDA headers at compile time, specific GPU architectures, the whole native-build circus. So it lives in the Docker image and the lockfile pretends it doesn’t exist.

Now trace the pipeline:

1. uv.lock — no torch (on Linux), no flash-attn (anywhere)
2. @pypi// — no @pypi//torch, no @pypi//flash_attn
3. Gazelle encounters import torch — can’t resolve

The pipeline has a hole. And everything that imports torch — directly or transitively — falls through it.

The Current Workaround: gazelle:resolve Overrides

Today, we patch the hole with duct tape:

# gazelle:resolve py torch
# gazelle:resolve py flash_attn
# gazelle:resolve py triton

These directives tell gazelle: “when you see import torch, don’t try to resolve it. Just… skip it.” The import still works at runtime because the Docker image has torch installed. But Bazel doesn’t know about it. The dependency exists in reality but not in the build graph.

This “works” in the sense that individual files can have BUILD targets despite importing torch. But it creates a deeper problem: any package that transitively depends on torch can’t be fully migrated, because somewhere in its dependency tree there’s an import gazelle can’t resolve, and the BUILD file generation either fails or produces an incomplete graph.

Result: 59 packages blocked. 62% of the monorepo.

The Resolution Dilemma

We’re stuck in a four-way bind:

None of these are good. But two of them are workable, if you squint.

Two Populations, Two Fixes

Here’s the insight that turns a wall into a door: not all 59 blocked packages actually need torch.

Fifty packages are blocked because of a transitive dependency through mai_config — a configuration package that happens to import torch at module load time. They never call torch.cuda.is_available(). They never create tensors. They read config values. And they’re blocked.

Nine packages actually use torch for GPU work. They need it for real.

Different problems. Different solutions.

Path A: Unblock the Innocent Bystanders

The Problem in Detail

Here’s what happens when you import mai_config today:

# mai_config/__init__.py (simplified)
import torch  # <-- THIS LINE. Right here. Module load time.

class DeviceConfig:
    def get_device(self):
        return torch.device("cuda" if torch.cuda.is_available() else "cpu")

class LoggingConfig:
    log_dir: str = "/var/log/mai"
    log_level: str = "INFO"
    # ... 200 more lines of config that have nothing to do with GPUs

The import torch at the top of the file means that every single module that imports mai_config — even if they only want LoggingConfig — transitively depends on torch. And because mai_config is the config package, everything imports it.

                    torch >= 2.7.0     numpy == 1.*
                         \              /
                          \            /
                     +------------------+
                     |   mai_config     |  <-- THE GATEWAY BLOCKER
                     | (torch + numpy)  |      blocks 22 direct, ~50 transitive
                     +--------+---------+
                              |
              +---------------+---------------+
              |               |               |
         mai_cluster    mai_preprocessors   mai_numerics
              |               |
         mai_distributed   mai_logger
              |               |
              +-------+-------+
                      |
          mai_layers -> mai_trainer -> mai_multimodal -> simplertransformer

One top-level import. Fifty blocked packages.

The Fix: Lazy Imports

The fix is conceptually trivial and practically fiddly. Move torch imports inside the functions that use them:

# BEFORE (blocks everything):
import torch

def get_device():
    return torch.device("cuda" if torch.cuda.is_available() else "cpu")

# AFTER (torch loads only when get_device() is called):
def get_device():
    import torch
    return torch.device("cuda" if torch.cuda.is_available() else "cpu")

After this change:

• import mai_config succeeds without torch installed
• mai_config.get_device() still works — it imports torch on demand
• Gazelle no longer sees a top-level import torch in mai_config
• ~50 packages that import mai_config for config purposes lose their transitive torch dependency

The ripple effect is enormous. One file change, ~50 packages unblocked.

What “Lazy” Looks Like in Practice

It’s not quite as simple as “move the import down.” You need to handle several patterns:

Pattern 1: Direct usage in functions — just move the import inside:

def detect_gpu():
    import torch
    return torch.cuda.is_available()

Pattern 2: Type annotations — use TYPE_CHECKING:

from __future__ import annotations
from typing import TYPE_CHECKING

if TYPE_CHECKING:
    import torch

def get_device() -> torch.device:
    import torch
    return torch.device("cuda")

Pattern 3: Module-level constants — extract to a separate submodule:

# mai_config/gpu_utils.py  (new file, explicitly depends on torch)
import torch
DEFAULT_DEVICE = torch.device("cpu")

# mai_config/__init__.py  (no longer imports torch)
# Users who need GPU utils import mai_config.gpu_utils explicitly

The goal is surgical: import mai_config must succeed without torch. Everything that actually needs torch can still get it — they just have to be explicit about it.

The Audit

Before touching any code, we need an audit of mai_config’s source:

1. Find every import torch and from torch import ...
2. Classify each one: top-level vs. inside function vs. inside TYPE_CHECKING
3. For top-level imports: what code path actually uses the torch object?
4. Can that code path be lazy without breaking the public API?

This audit is the single highest-leverage hour of work in the entire migration. If the answer is “yes, we can make torch lazy in mai_config,” we cut the blocked package count from 59 to ~9 in a single PR.

Path B: Make torch Actually Work in Bazel

Nine packages genuinely import torch for GPU computation. They create tensors, run CUDA kernels, do gradient calculations. They need torch. There’s no lazy-importing around that.

For these packages, we need to fix the pipeline — make @pypi//torch a real, resolvable target.

The Platform Split

This is where it gets fun. “Make torch work in Bazel” is actually two completely different engineering problems, one per platform.

macOS: Probably Fine?

On Mac, torch comes from PyPI as a normal wheel. A big wheel (68MB, arm64), but a normal one. The uv.lock does resolve torch on macOS (remember the sys_platform == 'darwin' marker). So @pypi//torch should just… exist.

We haven’t tested it yet. But the mechanism is already there:

# Hypothetical BUILD target (auto-generated from uv.lock)
# @pypi//torch -> torch-2.7.0-cp312-cp312-macosx_14_0_arm64.whl
py_test(
    name = "smoke_test",
    srcs = ["smoke_test.py"],
    deps = ["@pypi//torch"],  # Should resolve on macOS
)

This is a 15-minute spike. Write the test, run it, see what happens. If it works, macOS development is unblocked for torch-using packages today. If it doesn’t work, we learn exactly where the gap is — maybe wheel extraction fails, maybe the binary extensions can’t find the right dylibs, maybe something else entirely.

Either way: data is better than speculation.

Linux: The Hard One

On Linux, torch isn’t in uv.lock. Period. The Docker image provides it. So @pypi//torch doesn’t exist, and we need a different mechanism.

The core question: how do you tell Bazel “this Python package is provided by the environment, don’t try to resolve it”?

This is the system-provided package override problem, and it doesn’t have a clean solution today. But there are approaches:

Approach 1: aspect_rules_py uv extension overrides

aspect_rules_py (the Bazel ruleset that reads uv.lock) supports override configuration in MODULE.bazel. The idea: tell the extension “for torch, don’t look in uv.lock — use this target instead.”

We’d create a wrapper target that points at the system-installed torch:

# BUILD.bazel (somewhere in the repo)
py_library(
    name = "system_torch",
    imports = ["/usr/local/lib/python3.12/site-packages"],
    # Or however system packages are exposed
)

Then in MODULE.bazel, override the torch resolution to point at this target instead of @pypi//torch.

This is unexplored territory. The override mechanism exists, but nobody’s documented using it for system-provided GPU packages. Needs a spike.

Approach 2: The Symlink Hack (Already Exists)

We already do something like this for C++ builds:

# Real production code:
torch_result = repository_ctx.execute(
    [python, "-c", "import torch; print(torch.__file__)"]
)
repository_ctx.symlink(torch_dir, "torch")

This runs Python inside a repository rule to discover where torch is installed, then symlinks it into Bazel’s sandbox. It works for C++ linking. For Python py_library deps, we’d need a Python-specific version of this trick — create a py_library target that wraps the symlinked path.

It’s ugly. The comment in the codebase calls it “a horrible hack.” But it proves the concept: you can bridge system-installed packages into Bazel’s dependency graph. We just need a cleaner version of it.

Approach 3: Dual lockfiles

Maintain a separate uv.lock that resolves torch on Linux (using PyPI’s CUDA-enabled wheels). Use this lockfile for Bazel resolution on Linux CI. The regular lockfile stays as-is for the Docker-based workflow.

Downsides: two lockfiles to maintain, potential version skew between Docker-torch and lockfile-torch, and CUDA wheel versions may not match what’s in the Docker image.

This is the least appealing option, but it’s worth noting as a fallback.

What About flash-attn, triton, and Friends?

Same problem, harder version. These packages use sys_platform == 'never' to exclude themselves from resolution entirely. They can’t even be pip-installed on most machines — they need CUDA headers, specific GPU architectures, and multi-hour compilation times.

Whatever mechanism we build for system-torch on Linux, it needs to generalize to these packages too. If we solve torch, we solve the pattern. flash-attn, triton, and ucxx-cu12 all slot into the same override mechanism.

This is why the system-override spike matters even though it only directly affects ~9 packages today. It establishes the pattern for every system-provided GPU package.

The Mechanics, Visualized

Here’s the full pipeline with both paths marked:

    pyproject.toml
         |
         v
    uv lock (resolves deps)
         |
         v
    uv.lock
    ├── numpy==1.26.4          ✓ resolved
    ├── pydantic==2.6.1        ✓ resolved
    ├── torch==2.7.0 (macOS)   ✓ resolved (darwin only!)
    ├── torch (Linux)          ✗ NOT RESOLVED (Docker provides it)
    ├── flash-attn             ✗ NOT RESOLVED (sys_platform=='never')
    └── triton                 ✗ NOT RESOLVED (sys_platform=='never')
         |
         v
    aspect_rules_py → @pypi// repository
    ├── @pypi//numpy           ✓ target exists
    ├── @pypi//pydantic        ✓ target exists
    ├── @pypi//torch (macOS)   ✓ target exists
    ├── @pypi//torch (Linux)   ✗ NO TARGET  ←── Path B fixes this
    └── @pypi//flash_attn      ✗ NO TARGET
         |
         v
    gazelle (scans imports → generates BUILD)
    ├── import numpy           → deps = ["@pypi//numpy"]     ✓
    ├── import torch           → deps = ["@pypi//torch"]     ✗ (Linux)
    └── import mai_config      → transitive torch dep        ✗
                                                              ^
                                                              |
                                               Path A fixes this
                                           (remove transitive dep)

Path A eliminates the need for @pypi//torch in 50 packages by severing the transitive chain. Path B creates a valid @pypi//torch (or equivalent) target for the 9 packages that actually need it. Together, they close the hole in the pipeline.

The Execution Order

These paths aren’t independent. Path A is the prerequisite — both logically and strategically.

Why Path A first:

1. Higher ROI — one PR unblocks ~50 packages vs. a complex infrastructure spike for ~9
2. No infrastructure changes — it’s a code change in one package, not a Bazel extension investigation
3. De-risks Path B — with 50 packages already unblocked, the pressure on the torch override mechanism drops. We can take our time getting it right instead of rushing a “horrible hack v2”
4. Validates the approach — if mai_config can’t be made torch-lazy (unlikely but possible), we need to know before investing in Path B

The sequence:

Week 1:  Audit mai_config torch usage
         Smoke-test @pypi//torch on macOS (parallel)

Week 2:  PR: lazy torch imports in mai_config
         Run gazelle on newly-unblocked packages

Week 3+: Spike: system-torch override mechanism for Linux
         Migrate ~9 genuine torch-using packages

Path A is a code change. Path B is an infrastructure change. Code changes ship faster, fail cheaper, and don’t require understanding Bazel extension internals.

The Pattern Under the Pattern

Step back and look at what we’re really dealing with.

The Python AI/ML ecosystem was built on an assumption: the environment provides GPU dependencies. You build a Docker image with CUDA and torch pre-installed, you pip install your application code on top, and everything finds everything else because it’s all in the same site-packages.

Bazel was built on the opposite assumption: the build system provides all dependencies. Nothing from the environment leaks in. Everything is declared, fetched, checksummed, sandboxed.

These two assumptions are fundamentally incompatible. Our three-stage pipeline (uv.lock → @pypi// → gazelle) works beautifully when all dependencies live in the “build system provides” world. It breaks the moment a dependency lives in the “environment provides” world.

Path A doesn’t resolve this tension — it sidesteps it. “Most packages don’t actually need torch, so stop pretending they do.” Elegant, but only works for the innocent bystanders.

Path B confronts it head-on. “Teach Bazel that some packages are environment-provided.” This is the harder, more important problem. And it’s the one the broader Bazel+Python+GPU community hasn’t cleanly solved yet.

We’re not the only ones hitting this. Every team running Bazel on GPU workloads discovers the same gap. The tooling will catch up eventually — aspect_rules_py is actively developed, rules_python is evolving, and the “system-provided package” concept keeps coming up in Bazel GitHub issues.

In the meantime, we’ll build the bridge ourselves. One lazy import and one system override at a time.

This is part 3 in a series on Bazel + GPU in a Python monorepo. Part 1: When GPU Showed Up covers the observational side. Part 2: 62% of Our Packages Are Blocked lays out the battle plan. All code samples and dependency graphs are from a real 95-package production monorepo.