Bazel Migration: mai_config, the First Torch Package

2026/03/06

BUILDbazelcimonorepopythontorch

mai_config is a Python package that imports torch. Until this week, that single fact made it unmigrateable. Torch is excluded from Bazel’s @pypi// index — it’s system-installed in Docker, not pip-managed — so any package that touches torch was stuck in the “not ready” pile. 62% of the monorepo, blocked by one dependency.

Then PR #21169 landed py_torch_test — a macro that exposes Docker’s system torch to Bazel tests via include_system_site_packages=True. The wall came down. Torch packages can migrate now.

mai_config is the first one through the door. Here’s what that looked like — including the four CI iterations it took to get green.

The Migration Steps

The recipe itself isn’t complicated. That’s the point — most of the work was done by the py_torch_test infrastructure. The actual migration is:

  1. Remove mai_config from .bazelignore and gazelle:exclude
  2. Add to GAZELLE_DIRS in bazel/BUILD.bazel
  3. Add python_root directive at mai_config/src/BUILD.bazel
  4. Run gazelle to generate BUILD files
  5. Add # gazelle:ignore torch to suppress gazelle trying to add @pypi//torch as a dep
  6. Convert torch-importing tests from py_test → py_torch_test
  7. Leave non-torch tests (like test_language_utils) as regular py_test

Step 5 is worth calling out. Gazelle reads your imports and tries to resolve them against the pip index. It sees import torch in structured.py and goes “ah, @pypi//torch.” But torch isn’t in @pypi// — there’s no wheel to download, it’s system-installed. So you tell gazelle to ignore it:

import torch  # gazelle:ignore torch

That’s a Python comment. Zero runtime impact. Gazelle sees it and backs off. The actual torch dependency comes from py_torch_test’s include_system_site_packages=True at test time.

Four Iterations to Green

Here’s where the fun starts. “Fun.”

Iteration 1: The Unused Import

Buildifier lint caught an unused load() statement in mai_config/tests/BUILD.bazel. The file loads py_test at the top — because that’s what gazelle generated — but after I converted all the torch tests to py_torch_test, nothing in the file actually uses py_test anymore. Dead import. One-line delete.

This is the kind of failure that takes 70 minutes of CI to discover and 3 seconds to fix. Classic.

Iteration 2: Gazelle Drift From Another Package

The gazelle idempotency check failed. But not because of anything I did.

This check runs gazelle across the entire repo and diffs the output. If any BUILD file would change, it fails. The idea is: BUILD files should always match what gazelle would generate. If they don’t, someone committed without re-running gazelle.

The drift was in lib/mai_nano. Someone on main had added splits.py and test_magic_mock_packages.py without re-running gazelle. Pre-existing issue. But the idempotency check is global — if any package has drift, it fails for everyone.

Had to fix mai_nano’s BUILD files in our PR. Ten lines of adding targets that gazelle would’ve generated anyway. Not ideal to mix into an unrelated PR, but the alternative was waiting for someone else to fix it.

Lesson: gazelle idempotency is a repo-wide invariant. Any PR that touches gazelle-managed files will surface everyone else’s drift.

Iteration 3: The __init__.py Trap and a Trailing Space

Two failures. One juicy, one annoying.

The annoying one: MIGRATION_STATUS.md had a trailing space. The generation script (find_bazel_ready_packages.py) outputs a space after caas_server because its reason string is empty. Pre-commit catches it. Trimmed the space, moved on.

The juicy one: Adding __init__.py to mai_config/tests/ broke tests.

Here’s what happened. I added __init__.py so gazelle could discover the test files as a Python package. Standard practice, right? Except mai_config’s tests assert on fully-qualified class names. Think Hydra configs where _target_ is the full dotted path to a class:

# Before __init__.py:
assert config._target_ == "test_config_tooling.MyConfig"

# After __init__.py:
# Python now sees tests/ as a package, so:
assert config._target_ == "tests.test_config_tooling.MyConfig"  # BOOM

Adding __init__.py to a directory changes its Python module namespace. Files that were previously top-level modules become submodules of the new package. Anything that depends on __qualname__, __module__, or Hydra’s _target_ will break.

The fix: delete __init__.py. Our BUILD file already has # gazelle:exclude test_*.py and manually defines all test targets, so gazelle doesn’t need the __init__.py for discovery. We were adding a file we didn’t need, and it was breaking things we didn’t expect.

Iteration 4: Green

45 checks passing. 0 failures. Merged.

The Namespace Trap, Generalized

The __init__.py thing deserves its own section because it will bite other migrations.

Here’s the mental model. In Python, a directory without __init__.py is a namespace package (or just a directory). Files in it have short module names — test_foo is just test_foo. The moment you add __init__.py, Python treats the directory as a regular package. Now test_foo is tests.test_foo (or whatever the directory is called).

Most code doesn’t care. But anything that inspects class identity — type(obj).__qualname__, type(obj).__module__, Hydra’s _target_ resolution, pickle, some serialization frameworks — will see different strings before and after.

Rule of thumb for Bazel migration: if your BUILD file manually manages test targets with gazelle:exclude, you don’t need __init__.py. Don’t add it just because it feels right. Add it only if gazelle actually needs it for auto-discovery, and if you do, check whether any tests care about fully-qualified names.

What It Unblocked

After regenerating MIGRATION_STATUS.md:

One package migrated. The dominoes start falling.

The Review Cheat Sheet

If you’re reviewing this PR — or doing your own torch package migration — here are the questions that’ll come up:

“Why py_torch_test instead of just py_test?” Torch isn’t in Bazel’s pip index. On Linux CI, it’s system-installed in Docker. py_torch_test wraps py_venv_test with include_system_site_packages=True so tests can see Docker’s torch. On macOS, these tests are skipped via target_compatible_with = Linux.

“Why gazelle:ignore instead of adding torch to the pip index?” Torch is 2GB+ and platform-specific. The Bazel uv extension doesn’t have Linux wheels in uv.lock (we use system torch). Adding it to @pypi// would mean maintaining platform-specific wheel config for something that still wouldn’t work — there’s no wheel to download.

“Does this change any runtime behavior?” No. The only change to Python source is adding # gazelle:ignore torch comments to two import lines. That’s a Python comment.

“What about the tests/__init__.py thing — will it happen again?” Yes. Any package with tests that assert on fully-qualified class names will break if you add __init__.py to a previously non-package test directory. Check before you add. If your BUILD manually manages targets, you probably don’t need it.

“The mai_nano fix — shouldn’t that be a separate PR?” Ideally, but gazelle idempotency is a global check. If any package has drift, every PR fails. We had to fix it to get green. It’s ten lines of targets gazelle would have generated anyway.

What’s Next

mai_cluster is ready. Then the Wave 0-3 cascade — mai_layers, mai_trainer, mai_distributed, and friends. Each one follows the same recipe. The __init__.py trap and gazelle drift are the two things to watch for. The actual migration steps are mechanical.

The hard part was building the infrastructure. The system torch support, the py_torch_test macro, the platform split. All of that is done. Now it’s just… migrating packages. One at a time. Watching the number go up.

PR: #23401