Bazel Torch: The py_torch_test Pattern

TL;DR: PR #22703 establishes a two-tier test architecture for the yolo monorepo’s Bazel migration: hermetic py_test for everything that doesn’t need torch, and a new py_torch_test macro (backed by py_venv_test) for tests that do. Torch-tagged tests run inside a Docker container where torch is system-installed. This post explains every file changed, every concept involved, and every design decision – the kind of deep-dive you’d want before fielding questions from teammates.

This is a study guide, not a war story. If you’re looking for the narrative version, check Part 3 which covers the dependency resolution problem. This post is the “how it actually works, file by file” reference.

The core problem in 30 seconds

The monorepo is migrating to Bazel. Bazel’s Python rules (aspect_rules_py) run Python in isolated mode (-I flag), which means Python only sees packages that Bazel explicitly manages. That’s great for hermeticity. It’s terrible for torch.

Torch can’t be a normal Bazel pip dependency because:

• It’s ~2GB – downloading it per-build is a non-starter
• It’s platform-specific – CPU vs CUDA wheels, Linux vs macOS, CUDA 11 vs 12
• GPU tests need CUDA drivers – which aren’t pip-installable
• It has sdist-only transitive deps – aspect_rules_py can’t build sdists

So torch lives in Docker images as a system package. But isolated mode means Bazel’s Python can’t see system packages. That’s the fundamental tension this PR resolves.

The two-tier architecture

The solution splits tests into two tiers. Every test falls into exactly one.

Tier 1: Hermetic tests (no torch)

This is the default. If your test doesn’t need torch, it goes here. Standard hermeticity – reproducible, cacheable, fast.

Tier 2: System-visible tests (torch)

This is the new tier. Tests tagged torch run inside a Docker container where torch is a system package, and the -I flag is stripped so Python can see it.

The tag-based filtering is the glue. CI runs two jobs against the same //... target pattern – one excludes torch, the other includes only torch. Every test ends up in exactly one job.

Concept: hermetic vs non-hermetic Python

This distinction is the whole backbone of the PR, so let’s get precise.

What -I (isolated mode) does

When Python starts with the -I flag, it:

1. Ignores PYTHONPATH environment variable
2. Ignores user site-packages (~/.local/lib/python3.x/site-packages)
3. Ignores system site-packages (/usr/lib/python3/dist-packages, etc.)
4. Only sees the paths that the caller explicitly passes

In Bazel terms: py_test from aspect_rules_py always passes -I. The only packages visible are the ones declared in deps and resolved through @pypi//. This is hermetic – nothing leaks in from the host system.

What py_venv_test with include_system_site_packages=True does

Two things:

1. Strips the -I flag from the Python invocation
2. Sets include-system-site-packages = true in the venv’s pyvenv.cfg

This means Python creates a virtual environment that inherits system packages. If torch is installed at /usr/local/lib/python3.12/site-packages/torch/, the venv can see it. Combined with Bazel-managed packages in the venv itself, the test gets both worlds: hermetic Bazel deps + system torch.

Why this is a controlled compromise, not a hack

The non-hermetic tier is constrained:

• Only runs on Linux – target_compatible_with ensures macOS never tries it
• Only in Docker – the Docker image is version-controlled, so “system” is actually deterministic
• User packages excluded – include_user_site_packages=False blocks ~/.local/ leakage
• Tag-filtered – you can’t accidentally mix the tiers

It’s not fully hermetic, but it’s reproducible – which is what actually matters.

Concept: py_test vs py_venv_test

This confused me initially, so let’s be explicit.

py_test is the public API from aspect_rules_py. It always passes -I. It supports pytest_main = True for automatic pytest integration. It’s what you use for normal tests.

py_venv_test is the private rule that py_test is built on. It’s loaded from:

load("@aspect_rules_py//py/private/py_venv:defs.bzl", "py_venv_test")

That private in the path is significant – it means this isn’t a stable public API. The # buildifier: disable=bzl-visibility comment in our code explicitly acknowledges this.

Why use the private API? Because the public py_test has no way to disable -I. There’s no hermetic = False flag, no include_system_site_packages parameter. The only rule that supports it is py_venv_test.

Tradeoff: If aspect_rules_py refactors their internal structure, our macro might break. That’s the cost of using a private API. The benefit is that we don’t need to fork or patch the rules.

What py_venv_test doesn’t support: pytest_main = True. This is why we need a custom pytest runner (more on that below).

The toolchain registration (MODULE.bazel)

For Bazel to use the Docker container’s Python (which has torch in its site-packages), we need to tell Bazel about it. That’s done with a local runtime toolchain – three steps in MODULE.bazel.

Step 1: Point Bazel at the system Python

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

This says: “There’s a Python interpreter at /usr/local/bin/python3. Use it.” The on_failure = "skip" is crucial – on a macOS dev machine, this path might not exist or might point to a Python without torch. Instead of failing the build, Bazel silently skips this toolchain and falls back to the hermetic one.

Step 2: Create toolchain definitions

local_runtime_toolchains_repo(
    name = "local_toolchains",
    runtimes = ["system_python3"],
)

This generates the actual Bazel toolchain() targets from the runtime repo.

Step 3: Register them

register_toolchains("@local_toolchains//:all")

This goes before the hermetic toolchain registrations. In Bazel, toolchain resolution is first-match-wins. So:

• In Docker: /usr/local/bin/python3 exists -> local toolchain matches -> uses system Python with torch
• On macOS: path doesn’t exist or fails -> on_failure = "skip" -> falls through to hermetic toolchain

All three steps use dev_dependency = True. This means they only apply when running Bazel in the yolo workspace – not when some other module depends on yolo.

Why this matters

Without the local toolchain, py_venv_test would create a venv from the hermetic Python (downloaded by Bazel). That Python’s site-packages wouldn’t have torch. The venv would set include-system-site-packages = true but “system” would be the hermetic Python’s empty site-packages, not Docker’s.

The local toolchain makes “system Python” mean “Docker’s Python with torch.” That’s the connection.

The py_torch_test macro (bazel/rules/defs.bzl)

Here’s the actual macro, with every line explained:

def py_torch_test(name, srcs, deps = [], data = [], args = [], tags = [], **kwargs):
    # Auto-add pytest as a dep (callers shouldn't need to remember this)
    if "@pypi//pytest" not in deps:
        deps = deps + ["@pypi//pytest"]

    # Auto-tag with "torch" for CI filtering
    if "torch" not in tags:
        tags = tags + ["torch"]

    # Build full runfiles paths for each test source file
    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",
        ] + args,
        deps = deps,
        data = data,
        tags = tags,
        include_system_site_packages = True,
        include_user_site_packages = False,
        target_compatible_with = ["@platforms//os:linux"],
        **kwargs
    )

Label() vs string paths

srcs = srcs + [Label("//bazel/rules:pytest_runner.py")],
main = Label("//bazel/rules:pytest_runner.py"),

In Bazel macros, string paths resolve relative to the caller’s BUILD file. So if a test in dataprocessing/item_response_theory/tests/BUILD.bazel calls py_torch_test, a string path like "pytest_runner.py" would look for the file in that directory.

Label("//bazel/rules:pytest_runner.py") resolves relative to the .bzl file’s own package – bazel/rules/. This is a common Bazel macro pitfall: always use Label() for paths that point to the macro’s own package.

native.package_name() for test file paths

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

The shared pytest_runner.py needs to know where to find the test files in Bazel’s runfiles tree. native.package_name() returns the caller’s Bazel package (e.g., dataprocessing/item_response_theory/tests), which is the runfiles path prefix. So test_models.py becomes dataprocessing/item_response_theory/tests/test_models.py.

include_system_site_packages = True

The whole point. Makes Docker’s torch visible.

include_user_site_packages = False

Blocks ~/.local/lib/python3.x/site-packages/. We don’t want random user-level pip installs leaking into tests.

target_compatible_with = ["@platforms//os:linux"]

These tests only make sense on Linux (in Docker). On macOS, bazel test //... silently skips them – they appear as “SKIPPED” in the output, not “FAILED.”

Auto-tagging torch

if "torch" not in tags:
    tags = tags + ["torch"]

CI uses --test_tag_filters to route tests. Adding torch automatically means callers don’t need to remember the tagging convention.

**kwargs passthrough

Allows callers to pass additional Bazel attributes (like size, timeout, env) without the macro needing to enumerate them.

The pytest runner (bazel/rules/pytest_runner.py)

Since py_venv_test doesn’t support pytest_main = True, we need a manual pytest entry point. Here’s what it does:

import sys
import pytest

args = [
    "--verbose",
    "-rN",                          # no summary of passed tests
    "--ignore=external",            # ignore Bazel's external/ directory
    "--ignore-glob=*pytest_main*",  # ignore any pytest_main.py files
]

# JUnit XML output -- Bazel sets this env var
xml_output = os.environ.get("XML_OUTPUT_FILE")
if xml_output:
    target_name = os.environ.get("BAZEL_TARGET", "unknown")
    args += [f"--junitxml={xml_output}", f"-o junit_suite_name={target_name}"]

# Test filtering -- Bazel's --test_filter maps to this env var
test_filter = os.environ.get("TESTBRIDGE_TEST_ONLY")
if test_filter:
    args += ["-k", test_filter]

# Append test file paths (passed as command-line args by the macro)
args += sys.argv[1:]

sys.exit(pytest.main(args))

Key details:

• XML_OUTPUT_FILE – Bazel expects test runners to produce JUnit XML at this path. Without it, Bazel can’t display per-test results.
• BAZEL_TARGET – Used for the JUnit suite name so test results are traceable to the Bazel target.
• TESTBRIDGE_TEST_ONLY – When you run bazel test --test_filter=test_foo, Bazel sets this env var. The runner translates it to pytest’s -k flag.
• sys.argv[1:] – These are the test file paths from test_file_args in the macro, plus any extra args the caller specified.

The # gazelle:ignore torch annotation

In models.py:

import torch  # gazelle:ignore torch

This is a build-time only annotation. Here’s what it does and doesn’t do:

Why? Because torch comes from the Docker image’s system Python, not from Bazel’s @pypi// mechanism. If Gazelle added @pypi//torch, Bazel would try to download torch as a pip wheel – which is the exact thing we’re avoiding.

Think of it this way: the annotation is a note to the build system, not to Python. It says “I know this import exists, and I know you can’t find it in @pypi//. Trust me, it’ll be there at runtime.”

Related: # gazelle:resolve in BUILD files

# gazelle:resolve py item_response_theory //dataprocessing/item_response_theory/src/item_response_theory

This is the opposite problem. Here, Gazelle can resolve the import but to the wrong target (or can’t find it at all because of the non-standard python_root layout). The gazelle:resolve directive tells it exactly which Bazel target maps to the item_response_theory module.

Related: # gazelle:exclude test_*.py

# gazelle:exclude test_*.py

Gazelle would normally auto-generate py_test targets for files matching test_*.py. But we’re using py_torch_test – a custom macro Gazelle doesn’t know about. So we exclude the test files from auto-generation and write the targets manually.

The antlr4 sdist workaround (MODULE.bazel)

This is a separate but related issue. omegaconf (a torch-ecosystem config library) transitively depends on antlr4-python3-runtime. That package is distributed as an sdist only – no prebuilt wheels on PyPI.

aspect_rules_py can only install wheels, not build sdists. (Building sdists would require a C compiler toolchain, which is a whole separate problem – see the sdist toolchain post.)

The fix is a dependency override in MODULE.bazel:

uv.override_requirement(
    hub_name = "pypi",
    requirement = "antlr4-python3-runtime",
    target = "//bazel:empty_py_lib",
)

And the empty target it points to in bazel/BUILD.bazel:

py_library(
    name = "empty_py_lib",
    srcs = [],
)

This tells Bazel: “whenever something depends on @pypi//antlr4_python3_runtime, give it this empty library instead.”

Why is this safe? antlr4-python3-runtime is used by omegaconf at its own build time – it’s part of omegaconf’s grammar parser, not something your code calls directly. And omegaconf itself is installed in the Docker image as a system package. So at test runtime, omegaconf (with antlr4) is already available via system site-packages. The empty override only affects Bazel’s dependency resolution, not runtime behavior.

The Docker image (Dockerfile.torch-test)

FROM python:3.12-slim AS base

# System dependencies
RUN apt-get update && apt-get install -y curl git build-essential

# torch CPU-only (from PyTorch's dedicated index)
RUN pip install --break-system-packages torch \
    --index-url https://download.pytorch.org/whl/cpu

# omegaconf + antlr4 (needed by torch-ecosystem packages)
RUN pip install --break-system-packages omegaconf antlr4-python3-runtime

# Bazelisk (auto-downloads correct Bazel version)
RUN curl -fsSL https://github.com/bazelbuild/bazelisk/releases/download/v1.25.0/bazelisk-linux-amd64 \
    -o /usr/local/bin/bazel && chmod +x /usr/local/bin/bazel

# uv (for lockfile resolution)
RUN curl -LsSf https://astral.sh/uv/install.sh | sh

Why --break-system-packages?

python:3.12-slim uses PEP 668’s “externally managed” Python. By default, pip install refuses to modify it – the intent is to protect system Python from accidental breakage. We want to install into system site-packages (that’s the whole design), so --break-system-packages overrides the protection.

Why CPU-only torch?

The CI tests that run in this image are unit tests, not training jobs. They don’t need GPUs. CPU-only torch is ~200MB vs ~2GB for the CUDA version. Faster image builds, smaller image size.

Why the repo isn’t baked in

The Dockerfile doesn’t COPY the repo. Instead, CI mounts it at runtime:

docker run --rm -v "${{ github.workspace }}:/workspace" -w /workspace \
    yolo-torch-test:${{ github.sha }} \
    bazel test //...

This keeps the Docker image stable and cacheable. The image changes only when system dependencies change (torch version, omegaconf version, etc.), not on every commit. The repo is the volatile part, so it’s mounted.

The CI workflow (bazel_build_test.yml)

Two parallel jobs against the same test universe:

Job 1: bazel (hermetic)

bazel test //... --test_tag_filters=-integration,-torch

Runs on a regular CI runner. No Docker. Excludes torch-tagged tests (they’d fail – no torch) and integration tests (they need infrastructure).

Job 2: bazel-torch-docker

# Build the Docker image
docker build -f bazel/docker/Dockerfile.torch-test -t yolo-torch-test:${{ github.sha }} .

# Run torch tests inside it
docker run --rm \
    -v "${{ github.workspace }}:/workspace" \
    -w /workspace \
    yolo-torch-test:${{ github.sha }} \
    bazel test //... --test_tag_filters=torch,-integration --test_output=all

Builds the image, mounts the repo, runs only torch-tagged tests. The -integration exclusion applies here too – integration tests are a separate concern.

The tag math

Every test ends up in exactly one job, or intentionally in neither (integration tests run separately).

The BUILD files for item_response_theory

This is the first package migrated using the torch pattern, so its BUILD files serve as the template for future packages.

src/BUILD.bazel – Python root marker

# gazelle:python_root

That’s the whole file. This tells Gazelle that src/ is a Python source root. Without it, Gazelle would think the import path is dataprocessing.item_response_theory.src.item_response_theory.models instead of just item_response_theory.models.

src/item_response_theory/BUILD.bazel – Library target

py_library(
    name = "item_response_theory",
    srcs = ["__init__.py", "models.py"],
    imports = [".."],
    visibility = ["//dataprocessing/item_response_theory:__subpackages__"],
    deps = [
        "//lib/mai_nano/src/mai_nano",
        "@pypi//numpy",
        "@pypi//scipy",
    ],
)

Key observations:

• No @pypi//torch in deps – that’s the gazelle:ignore working
• imports = [".."] – adds the parent directory (src/) to Python’s sys.path. This is what makes import item_response_theory work – without it, Python would need import src.item_response_theory
• visibility is scoped to subpackages of item_response_theory – other packages can’t depend on it without an explicit visibility grant

tests/BUILD.bazel – Test targets

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

# gazelle:resolve py item_response_theory //dataprocessing/item_response_theory/src/item_response_theory
# gazelle:exclude test_*.py

py_torch_test(
    name = "test_models",
    srcs = ["test_models.py"],
    deps = [
        "//dataprocessing/item_response_theory/src/item_response_theory",
        "@pypi//numpy",
        "@pypi//scipy",
    ],
)

Three Gazelle annotations working together:

1. gazelle:resolve – maps import item_response_theory to the correct Bazel target (needed because of the non-standard python_root layout)
2. gazelle:exclude test_*.py – stops Gazelle from auto-generating py_test targets that would conflict with our manual py_torch_test
3. The load statement brings in our custom macro

The coverage checker fix (check_bazel_coverage.py)

# Before:
f"labels(srcs, kind('py_library|py_test|py_binary', {target_pattern}))"

# After:
f"labels(srcs, kind('py_library|py_test|py_binary|py_venv_test', {target_pattern}))"

The coverage checker uses bazel query to find all Python source files tracked by Bazel. The kind() filter matches targets by rule type. Without py_venv_test in the filter, torch test targets are invisible – the coverage checker would report their source files as “untracked” and flag them as missing coverage.

One line change, but important: without it, the coverage checker and the test infrastructure disagree about what exists.

The GAZELLE_DIRS entry (bazel/BUILD.bazel)

GAZELLE_DIRS = [
    ...
    "dataprocessing/item_response_theory/",
    ...
]

This is the whitelist of directories where the migration tooling runs Gazelle. Adding a directory here signals: “this package is being Bazel-managed.” It’s the administrative step that enables BUILD file generation.

All 13 files changed – complete reference

Tradeoffs and alternatives

Private API usage (py_venv_test)

Tradeoff: Using @aspect_rules_py//py/private/py_venv:defs.bzl means we depend on an internal implementation detail. If aspect_rules_py refactors their internal package structure, the import path breaks.

Alternative considered: Fork aspect_rules_py and add a public flag. Too much maintenance burden for one flag.

Alternative considered: Patch py_test to support include_system_site_packages. Would require an upstream contribution and waiting for a release.

Mitigation: The # buildifier: disable=bzl-visibility comment makes the choice explicit. If the API breaks, the fix is a single import path change in one file.

CPU-only torch in Docker

Tradeoff: Tests can’t exercise GPU code paths. torch.cuda.is_available() returns False.

Alternative: Use a CUDA-capable Docker image. Significantly larger (~8GB), needs GPU runners, more expensive CI. Not worth it for unit tests.

Mitigation: GPU-specific tests should use the integration tag and run on GPU runners separately.

Empty antlr4 override

Tradeoff: If any code actually imports antlr4 at runtime (not just omegaconf’s internals), it would fail in the hermetic tier.

Mitigation: antlr4 is only used by omegaconf’s parser generator, which runs at omegaconf’s build time. No application code imports it directly.

Docker image rebuild on every CI run

Tradeoff: docker build runs on every CI trigger. Even with layer caching, this adds ~1-2 minutes.

Alternative: Pre-build and push to a registry. More infrastructure to manage, but faster CI. This is a likely future optimization.

How to add the next torch package

The whole point of the pattern is that it’s repeatable. Here’s the checklist:

1. Add # gazelle:ignore torch to any source file that imports torch
2. Create BUILD files using py_torch_test for test targets (copy the item_response_theory template)
3. Add to GAZELLE_DIRS in bazel/BUILD.bazel
4. If new system deps are needed (e.g., transformers, flash-attn), add them to Dockerfile.torch-test
5. Run bazel test //your/package/... --test_tag_filters=torch in the Docker image to verify

That’s it. The macro, runner, toolchain, and CI pipeline are already in place.

Q&A for teammate conversations

“Why can’t we just add torch as a normal pip dependency?”

torch is ~2GB, platform-specific, and often needs CUDA. Bazel’s pip integration would download it on every build. System-installing in Docker gives us control over the exact version and CUDA pairing.

“What does gazelle:ignore torch actually do?”

It tells Gazelle not to add @pypi//torch to deps. Torch still works at runtime – it comes from Docker’s system Python. The annotation only affects BUILD file generation.

“Why py_venv_test instead of py_test?”

py_test always passes -I (isolated mode) which blocks system packages. py_venv_test supports include_system_site_packages=True which strips -I and makes system torch visible.

“Isn’t py_venv_test a private API? Is that risky?”

Yes, it’s private. Conscious tradeoff – there’s no public API for this. If it breaks, the fix is one import path change in defs.bzl.

“What happens if I run these tests on macOS?”

They’re silently skipped. target_compatible_with = ["@platforms//os:linux"] means Bazel marks them as incompatible. And the local toolchain has on_failure = "skip" so it falls back gracefully.

“Why Label() instead of a plain string for pytest_runner.py?”

In macros, strings resolve relative to the caller’s BUILD file. Label() resolves relative to the macro’s own package. Without it, tests would look for the runner in the wrong directory.

“Why a custom pytest runner instead of pytest_main=True?”

py_venv_test doesn’t support pytest_main. It’s a lower-level rule. The custom runner handles JUnit XML output, test filtering, and standard pytest flags.

“Why override antlr4 with an empty target?”

It’s sdist-only (no wheels). aspect_rules_py can’t build sdists. It’s a transitive dep of omegaconf, which is system-installed in Docker anyway. The override is safe because no application code imports antlr4 directly.

“How do I add my torch-dependent package?”

Add gazelle:ignore torch to imports, write BUILD files with py_torch_test, add to GAZELLE_DIRS, and add any new system deps to the Dockerfile. The infrastructure is already in place.

“What’s the difference between the two CI jobs?”

bazel runs hermetic tests on a regular runner (--test_tag_filters=-torch). bazel-torch-docker runs torch tests in Docker (--test_tag_filters=torch). Tag filtering routes each test to exactly one job.

Mental model: the whole system in one diagram

                    CI Trigger
                       │
              ┌────────┴────────┐
              ▼                 ▼
         bazel job         bazel-torch-docker job
              │                 │
              │            ┌────┴─────┐
              │            │  Docker   │
              │            │  - torch  │
              │            │  - omegaconf
              │            │  - system │
              │            │    Python │
              │            └────┬─────┘
              │                 │
     --test_tag_filters    --test_tag_filters
       =-torch,-integ        =torch,-integ
              │                 │
         ┌────┴────┐      ┌────┴────┐
         │ py_test  │      │py_venv_ │
         │ (-I flag)│      │test     │
         │ hermetic │      │(no -I)  │
         │ @pypi//  │      │@pypi// +│
         │ only     │      │system   │
         └─────────┘      └─────────┘

Two worlds. Same codebase. Tag-based routing. Clean separation.