Bazel Migration: Torch Packages, the Complete Decision Tree

Your package imports torch. You want to migrate it to Bazel. Here’s every case, every command, and every way it can go wrong.

Why Torch Is Special

In a normal Bazel Python build, dependencies come from @pypi// — Bazel downloads the wheel, hash-verifies it, and symlinks it into a hermetic sandbox. Your code sees exactly the deps declared in the BUILD file, nothing else. This is the whole point.

Torch breaks this model in three ways:

1. The wheel is ~2GB. Every CI run fetching a 2GB wheel is not a build system, it’s a CDN stress test.
2. Platform-specific CUDA bindings. The Linux wheel needs CUDA libraries that don’t exist in Bazel’s hermetic sandbox. The macOS wheel is CPU-only and fine.
3. Our uv.lock only has macOS wheels. The lockfile specifies sys_platform == 'darwin' for torch. On Linux, there’s no @pypi//torch to fetch — it simply doesn’t exist in the dependency graph.

On Linux, torch lives in Docker. System-installed, pip install --break-system-packages, done. But Bazel runs Python with the -I flag (isolated mode), which makes system packages invisible. So the torch that’s right there in /usr/local/lib/python3.12/site-packages/ might as well not exist.

The solution is a two-tier architecture: non-torch tests run hermetically like normal. Torch tests run inside Docker with a special Bazel configuration that punches a hole through the isolation layer.

The Three Components

1. The Docker Image

bazel/docker/Dockerfile.torch-test is minimal — Python 3.12, CPU torch, Bazelisk, uv. That’s it. No repo code baked in.

FROM python:3.12-slim AS base

RUN pip install --no-cache-dir --break-system-packages \
    torch --index-url https://download.pytorch.org/whl/cpu

RUN python3 -c "import torch; print(f'torch {torch.__version__}')"

RUN curl -fsSL https://github.com/bazelbuild/bazelisk/releases/latest/download/bazelisk-linux-amd64 \
    -o /usr/local/bin/bazel && chmod +x /usr/local/bin/bazel

RUN curl -LsSf https://astral.sh/uv/install.sh | sh

CPU-only torch. Because these are unit tests, not training runs. If your test needs a GPU, it doesn’t belong in Bazel yet (more on that later).

2. The Local Toolchain

MODULE.bazel registers Docker’s Python as a Bazel toolchain via local_runtime_repo:

local_runtime_repo(
    name = "local_python",
    interpreter_path = "/usr/local/bin/python3",
    on_failure = "skip",
)

on_failure = "skip" is the key. On macOS, /usr/local/bin/python3 isn’t the right interpreter — Bazel falls back to its normal hermetic toolchain. On Linux in Docker, it registers the system Python that has torch installed.

3. The py_test Macro

This is where the magic lives. Our custom py_test macro in bazel/rules/defs.bzl does two things:

For non-torch tests, it wraps aspect_rules_py’s py_test with pytest_main = True — standard stuff, just ensures tests run through pytest.

For torch tests, there’s a separate py_torch_test macro:

def py_torch_test(name, srcs, deps = [], tags = [], **kwargs):
    if "torch" not in tags:
        tags = tags + ["torch"]

    test_file_args = [native.package_name() + "/" + src for src in srcs]

    py_venv_test(
        name = name,
        srcs = srcs + [Label("//bazel/rules:pytest_runner.py")],
        main = Label("//bazel/rules:pytest_runner.py"),
        args = test_file_args + ["--ignore-glob=*pytest_main*", "-o", "asyncio_mode=auto"],
        deps = deps,
        tags = tags,
        include_system_site_packages = True,
        target_compatible_with = ["@platforms//os:linux"],
        **kwargs
    )

Three things happen here:

• include_system_site_packages = True — punches through Bazel’s isolation. System torch becomes visible.
• target_compatible_with = ["@platforms//os:linux"] — this test only exists on Linux. On macOS, Bazel skips it.
• tags = ["torch"] — CI uses this to split torch tests into the Docker job.

The pytest_runner.py is needed because pytest_main = True doesn’t work with py_venv_test. Instead, test file paths get passed as args to a shared runner script.

The Auto-Detection Upgrade (PR #23305)

The current approach requires engineers to manually write py_torch_test targets. PR #23305 merges this into the regular py_test macro with auto-detection:

_TORCH_DEP = "@pypi//torch"

def py_test(name, deps = [], **kwargs):
    if _TORCH_DEP in deps:
        _py_torch_test(name = name, deps = deps, **kwargs)
    else:
        _py_test(name = name, pytest_main = True, deps = deps, **kwargs)

If gazelle adds @pypi//torch to your test’s deps (because it saw import torch), the macro intercepts it, strips the @pypi//torch dep (torch comes from system), and switches to the py_venv_test path. You write a normal test, gazelle generates a normal BUILD file, the macro handles the rest.

The CI Split

Two jobs in .github/workflows/bazel_build_test.yml:

bazel job — bare runner, no Docker:

- name: Test all targets
  run: bazel test //... --test_tag_filters=-integration,-torch

bazel-torch-docker job — Docker container with system torch:

- name: Run torch tests in container
  run: |
    docker run --rm \
      -v "${{ github.workspace }}:/workspace" \
      -w /workspace \
      yolo-torch-test:${{ github.sha }} \
      bazel test //... --test_tag_filters=torch,-integration

The torch tag is the multiplexer. Non-torch job excludes it, Docker job includes it. Clean split.

The Decision Tree

Your package wants to migrate. Here’s how to figure out what kind of migration you’re doing.

Does your code import torch?
├── NO → Standard migration. Nothing special. Skip to "The Commands."
└── YES → Where does the import happen?
    ├── LIBRARY CODE (src/*.py)
    │   → You need gazelle:ignore on the import.
    │   → Tests that exercise torch paths need py_torch_test.
    │
    ├── TEST CODE ONLY (tests/*.py)
    │   → With auto-detection (#23305): Just run gazelle. Done.
    │   → Without auto-detection: Hand-write py_torch_test targets.
    │
    └── TRANSITIVE (your dep imports torch, you don't)
        → Auto-detection CANNOT help here. (See "Why" below.)
        → Manual annotation required: add @pypi//torch to test deps,
          or use py_torch_test.
        → But check first: does your test actually trigger the torch
          code path? If not, you're fine.

Case by Case

Case 1: Torch in Library Code

This is the tricky one. Your src/my_package/model.py has import torch at the top.

The problem: Gazelle sees the import and adds @pypi//torch to the py_library deps. But py_library doesn’t go through our custom macro — it’s the standard rule. On Linux, there’s no @pypi//torch wheel to fetch. Bazel blows up.

The fix: Add # gazelle:ignore torch to the source file:

import torch  # gazelle:ignore torch

This tells gazelle to skip that import when generating deps. The library target has no @pypi//torch dep. At runtime in Docker, torch is available via system site-packages anyway.

This is still needed even with auto-detection (PR #23305). The auto-detection only works for py_test targets, not py_library. If gazelle adds @pypi//torch to a library, there’s no macro to intercept it.

For tests in this package, use py_torch_test (current) or let auto-detection handle it (#23305).

Case 2: Torch in Test Code Only

Your library doesn’t import torch, but your tests do. tests/test_model.py has import torch.

With auto-detection (#23305): This is the happy path. Run gazelle. It generates:

py_test(
    name = "test_model",
    srcs = ["test_model.py"],
    deps = ["@pypi//torch", "//my_package"],
)

Our py_test macro sees @pypi//torch, auto-switches to the torch path. Done. No manual editing.

Without auto-detection (current): You need to:

1. Add # gazelle:exclude test_*.py to the test BUILD file
2. Hand-write py_torch_test targets:

load("//bazel/rules:defs.bzl", "py_torch_test")

py_torch_test(
    name = "test_model",
    srcs = ["test_model.py"],
    deps = ["//my_package"],
)

Note: no @pypi//torch in the deps. Torch comes from system. You only list your first-party deps and any other third-party deps.

Case 3: Transitive Torch Dependency

Your package depends on mai_config, which imports torch. Your code never touches torch directly.

This is the subtle one. And it exposes a fundamental limitation of the auto-detection approach.

The problem: The py_test macro auto-detects torch by checking if @pypi//torch is in the deps list. But it’s just string-matching on the literal list that gazelle wrote into the BUILD file. Your test depends on //lib/mai_config, not @pypi//torch. The macro sees deps = ["//lib/mai_config"] and takes the normal py_test path — hermetic mode, no system site-packages, no Docker.

Then at runtime, mai_config tries to import torch, and it’s not there.

test_my_thing.py → imports mai_config → mai_config imports torch → ModuleNotFoundError

Why the macro can’t see transitive deps: This isn’t a bug — it’s Bazel’s phase model. Macros run during the loading phase, which is basically string processing on BUILD files. Transitive dependency resolution happens during the analysis phase, after macros have already expanded. By the time Bazel knows that //lib/mai_config needs torch, the macro has already decided which rule to emit. There’s no callback, no second pass.

Can we fix this automatically? I went through every option. Short answer: no.

Every automatic solution either hits the phase-ordering wall (loading vs analysis) or requires reimplementing rules_py internals. The Bazel team designed this separation deliberately — macros are intentionally limited to loading-phase information for hermeticity and correctness. You’re not going to outsmart the build system here.

The fix is manual annotation. Two options depending on which infra you’re on:

With auto-detection (#23305) — add @pypi//torch as an explicit dep:

py_test(
    name = "test_my_thing",
    deps = [
        "//lib/mai_config",
        "@pypi//torch",  # keep — transitive torch dep, triggers venv path
    ],
)

The # keep comment prevents gazelle from removing it on the next run (since gazelle doesn’t see import torch in your test file).

Without auto-detection (current) — use py_torch_test:

py_torch_test(
    name = "test_my_thing",
    srcs = ["test_my_thing.py"],
    deps = ["//lib/mai_config"],
)

“But how do I know if my dep chain includes torch?” You mostly don’t — until the test fails in CI. This is the weakest part of the current system. A bazel query validator could catch it:

# Find py_test targets that transitively depend on a torch-using library
# but aren't tagged as torch tests
bazel query 'kind("py_test", rdeps(//..., @pypi//torch))' \
  | diff - <(bazel query 'attr(tags, torch, //...)')

This is a half-day of scripting to turn into a CI check. Not built yet, but it’s the right next step.

Is this even a big problem? Smaller than it looks. Wave 0 packages (20 packages) directly import torch — auto-detection handles them. Waves 1-3 (~10 packages) are the transitive cases. But many of those tests don’t actually exercise torch code paths — they import mai_config for its non-torch functionality. And if mai_config uses lazy imports (which Path A / PR #20799 introduced), import mai_config doesn’t trigger import torch at all. The actual number of tests that would silently fail is likely single digits.

Case 4: macOS Development

Torch tests are Linux-only via target_compatible_with = ["@platforms//os:linux"]. On macOS:

$ bazel test //my_package:test_model
INFO: Skipping target //my_package:test_model; build was requested, but it is incompatible

Not failed. Skipped. Silently. This surprises people but it’s by design — there’s no system torch on macOS to test against.

For local development: bazel test //my_package/... --test_tag_filters=-torch runs everything except torch tests. Torch tests only run in CI Docker.

Case 5: Linux CPU (CI)

This is the standard CI path. Docker container, CPU torch, bazel test //... --test_tag_filters=torch,-integration.

Most unit tests don’t need a GPU. If your test does torch.tensor([1, 2, 3]) and checks the output, CPU torch is fine. If your test needs torch.cuda.is_available() to return True, see the next case.

Case 6: GPU Tests

GPU tests are not in Bazel. Full stop. They run via run_gpu_tests.yml using uv run pytest on GPU runners. Don’t try to Bazel-ify them.

If your package has both CPU-testable and GPU-testable code, split them. CPU tests go through Bazel. GPU tests stay in the traditional CI pipeline.

The Commands

Step-by-step migration for a package at lib/my_package/:

Step 1: The Three-File Sync

Three files must be updated together. Miss one and CI catches it via check_bazel_coverage.py.

.bazelignore — remove your package:

-lib/my_package

Root BUILD.bazel — remove the gazelle exclude:

-# gazelle:exclude lib/my_package/**

bazel/BUILD.bazel — add to GAZELLE_DIRS:

 GAZELLE_DIRS = [
     "bazel",
     "lib/blobfile",
+    "lib/my_package",
 ]

Step 2: Handle Torch Imports in Library Code

If any .py file in src/ imports torch:

import torch  # gazelle:ignore torch
from torch import nn  # gazelle:ignore torch

Every line that imports from torch needs the ignore comment. This is tedious but necessary — there’s no file-level ignore directive.

Step 3: Run Gazelle

bazel run //bazel:gazelle

Review the generated BUILD files. Check that:

• Library targets (py_library) have no @pypi//torch in deps
• Test targets have the right deps
• No unexpected files got picked up

Step 4: Handle Torch Tests

With auto-detection (#23305): If gazelle generated py_test targets with @pypi//torch in deps, you’re done. The macro handles it.

Without auto-detection (current): For each test file that imports torch:

1. Add # gazelle:exclude test_*.py to the test directory’s BUILD file
2. Write py_torch_test targets manually
3. Make sure your non-torch test files still have gazelle-generated targets

Step 5: Add System Dependencies to Docker

If your package needs Python packages alongside torch that aren’t in the standard Docker image (e.g., omegaconf, antlr4-python3-runtime), add them to Dockerfile.torch-test:

RUN pip install --no-cache-dir --break-system-packages \
    torch --index-url https://download.pytorch.org/whl/cpu \
    omegaconf \
    antlr4-python3-runtime

Step 6: Validate

# Build everything
bazel build //lib/my_package/...

# Run non-torch tests locally
bazel test //lib/my_package/... --test_tag_filters=-torch

# Torch tests can't run locally (need Docker)
# CI validates them in bazel-torch-docker job

Step 7: Update Migration Status

python3 bazel/tools/find_bazel_ready_packages.py

This regenerates bazel/MIGRATION_STATUS.md and may cascade-unblock other packages that were waiting on yours.

The Platform Matrix

Gazelle Configuration

Two directives in the root BUILD.bazel make this work:

# gazelle:map_kind py_test py_test //bazel/rules:defs.bzl

This tells gazelle that when it writes py_test, it should load from our macro, not from @aspect_rules_py. Every gazelle-generated test goes through our custom rule.

# gazelle:resolve py torch @pypi//torch

This tells gazelle that import torch maps to @pypi//torch. Without this, gazelle might not know how to resolve the import (especially if torch isn’t in gazelle_python.yaml).

The Gotchas

1. gazelle:ignore torch is still needed on library files. Even with auto-detection. The py_test macro intercepts torch deps. The py_library rule does not. If gazelle adds @pypi//torch to a library target, Bazel tries to fetch a wheel that doesn’t exist on Linux.

2. The gazelle_python.yaml manifest is manually maintained. An upstream bug (aspect-build/rules_py#784) blocks auto-generation. If your package introduces a new third-party dep that gazelle can’t resolve, you may need to add a mapping manually.

3. Torch tests can’t run locally on macOS. target_compatible_with = ["@platforms//os:linux"] makes them invisible. Not failed — invisible. Use --test_tag_filters=-torch for local dev. This is by design but every new engineer asks about it.

4. New system deps need Docker image updates. If your torch test needs omegaconf or pydantic or whatever alongside torch, add it to Dockerfile.torch-test. Otherwise the test imports fine in your local venv but fails in CI Docker.

5. Optional torch imports confuse gazelle. try: import torch or if TYPE_CHECKING: import torch — gazelle still sees it. Two packages (audio_ops, resumable_ray_pipeline) are manually excluded for exactly this reason. If your package does conditional torch imports, you’ll need gazelle:ignore.

6. Permanent blockers exist. Packages depending on lib/bus (~28 packages via C++/Cython), lib/chz (metaprogramming breaks gazelle), mai_kernels, or simplertransformer (C++/CUDA) cannot migrate regardless of torch status. Check MIGRATION_STATUS.md before investing effort.

7. py_venv_test is a private API. It’s loaded from @aspect_rules_py//py/private/py_venv:defs.bzl. The private in the path means exactly what you think. It works today. It might break on a rules_py upgrade.

8. pytest_main = True doesn’t work with py_venv_test. That’s why torch tests use a shared pytest_runner.py script. Test file paths are passed as positional args. This is a workaround, not a feature.

9. Tags matter for CI routing. If your torch test doesn’t have the "torch" tag, it runs in the non-Docker job and fails because system torch isn’t available. The auto-detection macro adds the tag automatically. If you’re writing py_torch_test manually, double-check.

10. GPU tests stay outside Bazel. Don’t try to make them work. The Docker image has CPU torch only. GPU tests use uv run pytest in a separate CI workflow on GPU runners. When someone asks “can we Bazel-ify GPU tests?” the answer is “not yet, and maybe not ever.”

Why Auto-Detection Has a Ceiling

If you read Case 3 and thought “surely there’s a way to detect transitive torch deps automatically” — I went down that rabbit hole. Here’s why the ceiling exists and why it’s okay.

The Phase Wall

Bazel builds in three phases:

LOADING (macros run here)
  → reads BUILD files
  → expands macros to rule instantiations
  → string processing only, no dep graph

ANALYSIS (rules run here)
  → resolves dependency graph
  → selects toolchains
  → plans the build

EXECUTION
  → actually builds and tests things

Our py_test macro runs during loading. It sees the literal deps = [...] list from the BUILD file — a list of strings. It does not and cannot know what those strings transitively depend on. That information doesn’t exist yet.

By the time Bazel resolves //lib/mai_config → imports torch, the macro has already expanded into either py_test (hermetic) or py_venv_test (system site-packages). There’s no going back.

This isn’t a limitation of our implementation. It’s a fundamental property of Bazel’s architecture. Macros are intentionally restricted to loading-phase information — that’s what makes them fast and hermetic.

I Tried Everything

Every path either hits the loading-vs-analysis wall or requires reimplementing rules_py internals. The Bazel team separated these phases deliberately. You don’t get to peek at the analysis results from a macro. Full stop.

Why This Is Okay

Three reasons this doesn’t need a fancy solution:

1. The scope is small. Wave 0 (20 packages that directly import torch) is handled by auto-detection. Waves 1-3 (~10 packages with transitive torch) are the manual cases. Of those, most tests don’t even exercise the torch code path. Actual tests that would silently fail: likely single digits.

2. Lazy imports shrink the problem. mai_config used to do import torch at the top level — meaning every package that imports mai_config transitively needs torch. After Path A (PR #20799 — lazy imports), import mai_config no longer triggers import torch. The transitive torch chain is broken at the source. Other packages doing the same lazy import pattern further reduce the blast radius.

3. CI catches it. If you get it wrong, the test fails in CI with ModuleNotFoundError: No module named 'torch'. Noisy, but not silent. And a bazel query validator (half a day of scripting) can catch it before the test even runs:

# Tests that transitively need torch but aren't tagged
bazel query 'kind("py_test", rdeps(//..., @pypi//torch))' \
  except attr(tags, torch, //...)

The right answer is manual py_torch_test annotation for the ~10 transitive cases, plus a CI validator that catches mistakes. Not glamorous. But it’s honest engineering — solving the problem at the right layer instead of fighting the build system’s architecture.

What This Unlocked

Before the torch infrastructure, 22 packages were directly blocked by torch deps. Another ~30 were transitively blocked. That was 62% of the monorepo stuck outside Bazel.

After landing the system torch support and cascade-analyzing the dep graph:

From 24/101 packages on Bazel to potentially 54/101. That’s the power of removing a single blocker — the cascade effect through the dependency graph.

The remaining 35 packages are stuck behind C++/Cython blockers (lib/bus, lib/chz). That’s a different problem, for a different post.

The Shortest Version

1. Update three files: .bazelignore, root BUILD.bazel, bazel/BUILD.bazel
2. Add # gazelle:ignore torch to library source files that import torch
3. Run bazel run //bazel:gazelle
4. If pre-#23305: hand-write py_torch_test for torch test targets
5. If post-#23305: gazelle handles it, just review the diff
6. bazel build //your_package/... to verify
7. bazel test //your_package/... --test_tag_filters=-torch to run non-torch tests locally
8. Push. If Docker CI is green, you’re done.