Bazel Migration: Cracking the Biggest Blocker, 4 Cython Files

2026/03/10

BUILDbazelcc-toolchainmigrationmonorepopython⭐

lib/bus blocks 28 packages from migrating to Bazel. It’s listed in our migration tracker under “Packages Blocked by Non-Python Code” with the label “Contains: C++.” It’s the single biggest remaining blocker after we solved torch.

I went to look at the actual C++ source files. There aren’t any.

The Problem

Our monorepo has 104 workspace packages. 27 are on Bazel. Another ~21 are unblocked now that we solved torch. mai_kernels (actual CUDA/C++ – see that post) is getting migrated via the Docker bridge approach.

But lib/bus is different. It sits at the root of a dependency chain that looks like this:

lib/bus
  └── lib/oai_api         (depends on lib/bus + conversation)
        └── common_utils   (depends on lib/oai_api + conversation + mai_preprocessors + prompt_registry)
              └── ~28 packages (caas, judges, research_utils, rocket, safety, swe_world, ...)

The chain goes: lib/bus -> lib/oai_api -> common_utils -> everything. If lib/bus can’t migrate, lib/oai_api can’t migrate. If lib/oai_api can’t migrate, common_utils can’t migrate. And common_utils is the shared utility package that half the monorepo depends on.

From the migration tracker:

Blocker Packages transitively blocked
torch (solved) ~21
lib/bus ~28
mai_kernels (in progress) ~10
simplertransformer overlaps with lib/bus chain

lib/bus is the biggest single blocker left. Solve it, and the blocked count drops from 47 to roughly 19.

The Investigation

The migration tracker says lib/bus “Contains: C++”. That label comes from find_bazel_ready_packages.py, which scans for file extensions: .cpp, .cc, .cu, .h, .pyx. If it finds any, it marks the package as blocked by non-Python code.

So I did what I probably should have done months ago. I counted the files.

$ find lib/bus -name "*.cpp" -o -name "*.cc" -o -name "*.cxx" -o -name "*.c" -o -name "*.cu"
# (nothing)

$ find lib/bus -name "*.h" -o -name "*.hpp"
# (nothing)

$ find lib/bus -name "*.py" | wc -l
420

$ find lib/bus -name "*.pyx"
lib/bus/src/oai/multimodal_token/mmbatch/batch.pyx
lib/bus/src/oai/multimodal_token/token_schema/token_schema.pyx
lib/bus/src/oai/multimodal_token/token_schema/token_ops.pyx
lib/bus/src/oai/multimodal_token/toklist/util_token.pyx

420 Python files. 4 Cython files. Zero C++ source files. Zero CUDA files. Zero headers.

The “C++” label? It comes from one line in setup.py:

extensions = [
    Extension(
        name="oai.multimodal_token.token_schema.*",
        sources=["src/oai/multimodal_token/token_schema/*.pyx"],
        include_dirs=[np.get_include()],
        language="c++",
    ),
    Extension(
        name="oai.multimodal_token.mmbatch.*",
        sources=["src/oai/multimodal_token/mmbatch/*.pyx"],
        include_dirs=[np.get_include()],
        language="c++",
    ),
    Extension(
        name="oai.multimodal_token.toklist.*",
        sources=["src/oai/multimodal_token/toklist/*.pyx"],
        include_dirs=[np.get_include()],
        language="c++",
    ),
]

language="c++" tells Cython to generate .cpp code instead of .c code from the .pyx source. That’s it. It’s a codegen flag for the generated output, not a declaration that there’s C++ source code. The migration scanner saw .pyx files and a language="c++" flag and filed it under “C++.”

The scariest blocker in our monorepo was four files cosplaying as a C++ library.

The Plot Twist: What Do Those 4 Files Actually Do?

Before deciding how to deal with them, I needed to understand what they’re doing. Maybe there’s a reason they’re Cython. Maybe they’re doing FFI into some native library. Maybe they’re wrapping CUDA kernels.

They’re not.

File 1: token_schema.pyx (117 lines)

A Token class and a HashToken subclass. Eight fields:

@cython.no_gc
cdef class Token:
    def __init__(self, int modality, object data, int offset=0, int total=1,
                 object target=None, float target_weight=1.0, bint mask=True,
                 bint skip_setting_target=False) -> None:
        self.modality = modality
        self.data = data
        self.offset = offset
        self.total = total
        self.target = target
        self.target_weight = target_weight
        self.mask = mask
        self.skip_setting_target = skip_setting_target

Plus __eq__, __repr__, __getstate__, __setstate__. A dataclass. With pickle support. In Cython.

The corresponding .pxd file declares the fields:

cpdef enum TokenType:
    TEXT = 0
    IMAGE = 1
    AUDIO = 2
    VIDEO = 3
    # ... 25 modality types total

@cython.no_gc
cdef class Token:
    cdef public int modality
    cdef public object data
    cdef public int offset, total
    cdef public object target
    cdef public float target_weight
    cdef public bool mask, skip_setting_target

@cython.no_gc disables garbage collection tracking for this class – a performance hint. The cdef public declarations are what you’d get from __slots__ in Python. The cpdef enum is a StrEnum (or IntEnum, depending on your vibe).

File 2: token_ops.pyx (215 lines)

Helper functions for Token manipulation:

The pattern is: Python-facing wrapper function calls a cdef (C-level) inner function. The inner function uses typed declarations for loop variables:

cdef list[int] _token_list_to_list_int(
    list[Token] tokens,
    bool raise_on_non_text,
    bool skip_non_text,
):
    cdef int i, l = len(tokens), n = 0
    cdef Token token
    for i in range(l):
        token = tokens[i]
        if token.modality == TokenType.TEXT:
            n += 1

Tight loop with typed variables. Cython compiles this to C that avoids Python object protocol overhead. It’s fast. But it’s doing list iteration, not GPU computation.

File 3: batch.pyx (196 lines)

ExpandedDataBatch – a container class with 12 numpy array fields. Plus process_and_expand_data_batch(), which processes a batch of token chunks into arrays for model input. The heaviest function in the set: nested loops over batches and tokens, type-checking each token, filling numpy arrays.

cdef class ExpandedDataBatch:
    cdef public cnp.ndarray tokens
    cdef public cnp.ndarray targets
    cdef public cnp.ndarray modalities
    cdef public cnp.ndarray totals
    cdef public cnp.ndarray mask
    cdef public cnp.ndarray targets_mask
    cdef public cnp.ndarray offsets
    cdef public cnp.ndarray modality_data_map
    cdef public list data
    cdef public list metadata
    cdef public list data_total
    cdef public list data_indices

Again: typed field declarations (equivalent to __slots__), numpy memoryviews for fast array access. No FFI. No C libraries. No CUDA.

File 4: util_token.pyx (199 lines)

toklist_to_list_tokens() – converts a TokList (the newer token format) to list[Token] (the older format). Dispatches on span type: IntTokSpan, NumpyArraySpan, ProcessedImagePatchSpan, etc. Complex but purely Python-level logic compiled for speed.

The summary

File Lines What it does Why Cython?
token_schema.pyx 117 Token dataclass + enum Fast attribute access, no GC overhead
token_ops.pyx 215 Token list operations Typed loops, avoid Python dispatch
batch.pyx 196 Batch expansion for model input Tight nested loops, numpy memoryviews
util_token.pyx 199 TokList -> Token conversion Span dispatch, typed iteration
Total 727 Data structures + batch ops Performance optimization

727 lines of Cython. Zero FFI. Zero C library calls. Zero CUDA. These files implement what @dataclass(slots=True) + IntEnum + tight loops do, compiled via Cython for speed.

The question is: how much speed?

The Performance Question

Cython exists for a reason. Compiling Python-like code to C gives you:

  1. No Python dispatch overhead – cdef function calls skip the Python method resolution protocol
  2. Typed local variables – cdef int i is a C int, not a Python object
  3. @cython.no_gc – objects aren’t tracked by the garbage collector
  4. cdef class fields – C struct member access instead of __dict__ lookup (same benefit as __slots__)
  5. Numpy memoryviews – cdef cnp.int64_t[:, :] gives C-level array indexing without Python overhead

For token_schema.pyx, the gains are mostly from (3) and (4) – creating millions of Token objects without GC tracking, with fast field access. A @dataclass(slots=True) gets you (4) but not (3).

For batch.pyx, the gains are from (2) and (5) – the inner loops of process_and_expand_data_batch process batch_size * context_size tokens, each touching 8+ arrays. With context sizes of 8K-128K and batch sizes of 8-32, that’s millions of iterations. Pure Python would add significant overhead per iteration.

But here’s the counterargument: these operations happen during data preprocessing, not during the actual training loop. They run on CPU, ahead of GPU computation. The training loop is GPU-bound, not data-prep-bound. If the data pipeline can keep up with the GPU, the absolute speed of token processing doesn’t matter.

Whether it can keep up depends on your context size, batch size, and how many dataloader workers you have. For small context sizes and moderate batches, pure Python is almost certainly fast enough. For 128K context with complex multimodal batches, the Cython version might matter. The honest answer is: nobody’s benchmarked it, because nobody’s tried the pure Python version.

The Options

Three paths forward. Each with different tradeoffs.

Option 1: Pure Python Rewrite

Replace the 4 .pyx files with equivalent Python. The native code problem gets deleted, not solved.

What it looks like:

# token_schema.pyx → token_schema.py
from enum import IntEnum
from dataclasses import dataclass

class TokenType(IntEnum):
    TEXT = 0
    IMAGE = 1
    AUDIO = 2
    VIDEO = 3
    # ... 25 modality types

@dataclass(slots=True)
class Token:
    modality: int
    data: object
    offset: int = 0
    total: int = 1
    target: object = None
    target_weight: float = 1.0
    mask: bool = True
    skip_setting_target: bool = False

The batch.pyx rewrite is more involved – you lose numpy memoryviews, so you’d replace cdef cnp.int64_t[:, :] with standard numpy indexing. Slightly slower but functionally identical.

What it costs:

What it gains:

Option 2: Docker Bridge

Same approach as mai_kernels: Docker compiles Cython during image build, Bazel manages Python, include_system_site_packages=True bridges the gap.

What it looks like:

What it costs:

What it gains:

Option 3: Cython Rules in Bazel

Teach Bazel to build .pyx files natively. Custom genrule + cc_library pipeline, or a community Cython ruleset.

What it looks like:

# Hypothetical BUILD.bazel
load("@rules_cython//:cython.bzl", "cython_library")

cython_library(
    name = "token_schema",
    srcs = ["token_schema.pyx"],
    pxd_srcs = ["token_schema.pxd"],
    deps = ["@pypi//numpy"],
    language = "c++",
)

What it costs:

What it gains:

The Decision Framework

Here’s how I’m thinking about this:

Factor Pure Python Docker Bridge Cython Rules
Unblocks 28 packages Immediately After CI work After rule work
Engineering effort 1-2 days ~1 day 1-2 weeks
Ongoing maintenance Lowest Medium Highest
Performance risk Unknown Zero Zero
Bazel hermeticity Full Partial Full
New Bazel infrastructure None None CC toolchain + Cython rules
Precedent New Proven (mai_kernels) New and fragile

The Docker bridge is the safe choice. Proven pattern, zero performance risk, moderate effort. But it adds complexity to every test that touches lib/bus – and because lib/bus is upstream of ~28 packages, that’s a lot of tests.

The pure Python rewrite is the bold choice. It eliminates the problem entirely. No Cython in the build system, no Docker bridge needed, standard Bazel migration. The risk is performance regression in token preprocessing. The mitigation is: benchmark first, rewrite second.

The Cython rules option is the “correct” option that nobody should actually pick. The maintenance burden of custom Cython build rules in Bazel – for four files that implement dataclasses – is absurd. You’d spend more time maintaining the build rules than the Cython code itself.

What I’d recommend

Start with the pure Python rewrite. Here’s why:

  1. The 4 files are 727 lines total. This is not a massive rewrite. The Cython code is Python-adjacent – the .pyx syntax is close enough to Python that much of it can be mechanically converted.

  2. The Cython features used are limited. cdef class -> @dataclass(slots=True). cpdef enum -> IntEnum. cdef typed locals -> just… variables. Numpy memoryviews -> numpy array indexing. No raw C pointers. No libc imports. No inline C.

  3. The performance question has a clear test. Write the pure Python version. Benchmark process_and_expand_data_batch() with realistic inputs (128K context, batch size 32, multimodal tokens). If pure Python is within 2x of Cython: ship it. If not: fall back to Docker bridge.

  4. The maintenance win is permanent. Once the Cython is gone, lib/bus is just Python. New engineers don’t need to learn Cython. setup.py goes away. The build step goes from setuptools + cythonize + gcc to… nothing.

  5. 28 packages unblocked. That’s the real number. Not unblocked “in theory after we set up Docker bridge CI.” Unblocked immediately, with standard Gazelle-generated BUILD files.

If the benchmark shows the pure Python version is too slow for large context sizes, the Docker bridge is still there as a fallback. But I’d bet against it – these are CPU-side preprocessing operations, and the GPU is the bottleneck.

The Inventory

For the record, here’s the complete lib/bus dependency situation:

lib/bus itself:

The poison chain:

lib/bus                              → 420 .py + 4 .pyx
  └── lib/oai_api                    → depends on lib/bus + conversation
        └── common_utils             → depends on lib/oai_api + conversation + 2 more
              ├── caas
              ├── judges
              ├── research_utils
              ├── rocket (also blocked by mai_kernels, simplertransformer)
              ├── safety
              ├── swe_world
              ├── launcher
              ├── dispatcher
              ├── llm_generator
              └── ... (~18 more packages)

From the migration tracker (current state):

Status Package count
Already on Bazel 27
Blocked by dependencies 47
Blocked by non-Python code 4
Blocked by excluded external deps 21
Manually excluded 4
Total 104

Solving lib/bus doesn’t directly unblock all 28 packages – some have additional blockers (like conversation depending on mai_config which depends on torch). But it removes the single biggest root cause blocker. Once lib/bus migrates, lib/oai_api can migrate, common_utils can migrate, and the cascade begins.

What’s Next

  1. Benchmark the pure Python path. Write a conversion of batch.pyx (the most performance-sensitive file) and measure against realistic workloads. If it’s within 2x: proceed with full rewrite. If not: Docker bridge.

  2. If rewriting: do it in dependency order. token_schema.pxd/pyx first (it’s the base class everything else depends on), then token_ops.pyx, then batch.pyx, then util_token.pyx.

  3. If Docker bridge: follow the mai_kernels pattern. py_library for Python, include_system_site_packages for compiled modules, manual BUILD files with gazelle:ignore.

  4. Either way: unblock the chain. lib/bus -> lib/oai_api -> common_utils -> 28 packages. Each step is its own PR, its own migration, its own CI run.

The Lesson

The package that blocks 28 others from migrating to Bazel contains zero C++ source files. The “C++” label in the migration tracker came from a Cython codegen flag. The actual native code is 727 lines of Cython implementing token dataclasses and batch operations – the kind of thing @dataclass(slots=True) was designed for.

Before you scope out a multi-week project to add native compilation support to your Bazel setup – go read the source code. Count the actual .cpp files. Count the actual .pyx files. Check what those .pyx files are doing. Sometimes the monster under the bed is a pillow with a scary pillowcase.

Four files. 727 lines. 28 packages blocked. And not a single line of C++.