How Our GPU Tests Actually Work

I work on CI infrastructure for a 100+ package Python monorepo. Our GPU tests run on Slurm clusters via SSH bridges, get sharded across nodes with bin-packing algorithms, and discover test metadata by parsing Python source code into abstract syntax trees.

I didn’t understand half of those words when I started this job.

Recently I reviewed a PR that replaced a grep-based test discovery pipeline with Python’s ast module. It was a clean change. The before/after was obvious. But as I tried to understand why it was better, I realized I couldn’t explain the foundational concepts the whole system is built on. conftest.py? pytest markers? __pycache__? AST parsing? I used these things daily. I couldn’t have told you how any of them worked.

This post is the explanation I wish I’d had. We’ll trace the full path from “PR gets pushed” to “GPU tests run on Slurm,” stopping to explain each concept as it becomes relevant. It’s a companion to The 500 Lines That Run Every GPU Test (which covers the Slurm execution bridge) and Stop Grepping Your Python Tests (which deep-dives the AST vs grep tradeoff).

The problem: which tests need GPUs?

Not all tests are equal. Some tests check that a config parser handles malformed YAML. Some tests train a model across 8 GPUs using NCCL collective operations. These cannot run on the same hardware.

Our monorepo has 100+ Python packages. GitHub Actions runners – where our CI runs – don’t have GPUs. CPU tests run there. But GPU tests need to be shipped to a Slurm cluster via SSH, run inside Docker containers with CUDA and 8x H100s, and the results piped back.

Before any of that can happen, though, the CI needs to answer a deceptively simple question: which test files actually need GPUs?

That question is harder than it sounds. And the answer involves every concept in this post.

__pycache__ – the thing you keep deleting

Let’s start with the simplest concept because it’ll come up immediately when we start walking the filesystem.

When Python runs a .py file, it compiles it to bytecode and caches the result in a __pycache__/ directory as .pyc files. Next time you import that module, Python checks the .pyc timestamp and skips recompilation if the source hasn’t changed. That’s it. It’s a performance cache. The .o files of Python.

Why does CI care? Because any tool that walks the filesystem looking for test files will stumble into __pycache__/ directories and find .pyc files that match test_*.py patterns. Our test config verifier has an explicit exclusion for this:

# .github/ci/python/verify_test_configs.py
def get_all_test_files(root: Path):
    ignore_dirs = {
        ".git",
        ".venv",
        "__pycache__",
        "node_modules",
        "bazel-bin",
        "bazel-out",
    }

    for dirpath, dirnames, filenames in os.walk(root):
        # Modify dirnames in-place to skip ignored directories
        dirnames[:] = [d for d in dirnames if d not in ignore_dirs]

The .gitignore excludes it twice – line 1 and line 72 – because apparently one exclusion felt insufficient:

__pycache__
# ... 70 lines later ...
__pycache__*

This is one of those things nobody thinks about until they write a filesystem walker and get garbage results. You discover __pycache__ the way you discover gravity: by hitting the ground.

Pytest markers – sticky notes on your tests

Now we need to talk about how tests declare what they need. The mechanism is pytest markers, and there are two ways to apply them.

Way 1: Decorate a single function.

# mai_stats/tests/test_stats_logger.py
class TestStatsLogger:
    @pytest.mark.require_cuda_1gpu
    def test_add_stat_basic(self, stats_logger: StatsLogger) -> None:
        num = torch.tensor(10.0, device="cuda")
        denum = torch.tensor(2.0, device="cuda")

This says “this specific test needs one CUDA GPU.” Straightforward.

Way 2: Set pytestmark at the file level.

# mai_kernels/tests/unit/test_ctests.py
pytestmark = pytest.mark.falcon_gpu_nightly

This stamps every test in the entire file. No decorator needed on individual functions. It’s like a rubber stamp for the whole folder – one line, and everything in the file inherits it.

Here’s the critical insight that took me a while to internalize: markers by themselves do absolutely nothing. They’re metadata. Labels. They don’t skip tests, they don’t route tests to GPU nodes, they don’t do anything at all.

@pytest.mark.require_cuda_1gpu is not a built-in pytest feature. It’s a string someone made up. pytest will happily let you write @pytest.mark.pineapple_on_pizza and nothing will happen. It’s just a tag.

Something else has to read those tags and act on them. That something is conftest.py.

conftest.py – pytest’s hidden inheritance system

conftest.py is a special filename that pytest auto-discovers by walking up the directory tree from each test file. If you have:

mai_layers/
  tests/
    conftest.py        <- applies to all tests under tests/
    cpu_tests/
      conftest.py      <- applies to cpu_tests/ and below
    distributed/
      conftest.py      <- applies to distributed/ and below
    test_something.py

pytest loads all of them, starting from the root. Child conftest files can override or extend parent behavior. It works like CSS specificity, or like how .bashrc inherits from .profile. Each directory can have its own layer of configuration.

Three things go in a conftest: fixtures, hooks, and marker definitions. Each one does something different, and our codebase uses all three heavily.

Fixtures: shared test setup

Fixtures are pytest’s dependency injection system. Instead of each test building its own test environment, you declare what you need and conftest provides it.

Here’s a real one from our mai_job package – it spins up a Ray cluster with a head node and two workers:

# mai_job/tests/conftest.py
@pytest.fixture(scope="function")
def ray_cluster(temp_dir_factory):
    """Starts a Ray cluster with a head node and 2 worker nodes."""
    env_vars = {
        "PYTHONPATH": os.getcwd(),
        "RAY_agent_register_timeout_ms": "120000",
    }
    cluster = _start_ray_cluster_head(
        temp_dir_factory,
        env_vars,
        head_resources={
            f"{RAY_METAL_NODE_RESOURCE_NAME}:head": 1,
            "node:head": 1,
        },
        head_node_name="head",
        include_dashboard=True,
    )

    cluster.add_node(num_cpus=2, num_gpus=0, ...)
    cluster.add_node(num_cpus=2, num_gpus=0, ...)
    cluster.wait_for_nodes()
    wait_for_port(cluster.ray_client_port)

    os.environ["RAY_ADDRESS"] = str(cluster.address)
    yield cluster

    ray.shutdown()
    cluster.shutdown()

That’s 265 lines of conftest just for mai_job. A test that needs a Ray cluster just writes def test_something(ray_cluster): and the whole setup/teardown happens automatically. The fixture name in the function signature is the request.

Hooks: modifying pytest’s behavior

This is where markers get their power. A hook is a function with a specific name that pytest calls at specific points in its lifecycle. The hook in mai_layers/tests/conftest.py is what makes @pytest.mark.require_cuda_1gpu actually do something:

# mai_layers/tests/conftest.py
def pytest_runtest_setup(item):
    if "require_cuda_1gpu" in item.keywords and not torch.cuda.is_available():
        pytest.skip("CUDA is not available")
    elif "require_cuda_8gpu" in item.keywords and not torch.cuda.device_count() >= 8:
        pytest.skip("8 CUDA devices are not available")

Before each test runs, pytest calls pytest_runtest_setup. This hook checks: does this test have the require_cuda_1gpu marker? If so, is CUDA available? If not, skip it.

That’s the whole mechanism. The marker is a sticky note. The hook is the person who reads the sticky note and acts on it. Without this conftest hook, require_cuda_1gpu would be meaningless – just metadata floating in the void.

Another hook in the same conftest assigns each pytest-xdist worker its own GPU:

def pytest_configure(config):
    worker_id = os.environ.get("PYTEST_XDIST_WORKER")
    if worker_id is not None:
        tmp = [int(x) for x in re.findall(r"\d+", worker_id)]
        num_gpus = torch.cuda.device_count()
        gpu_id = tmp[0]
        torch.cuda.set_device(tmp[0])

When tests run in parallel (via pytest-xdist), each worker process gets an ID like gw0, gw1, etc. This hook maps gw0 -> GPU 0, gw1 -> GPU 1. Eight workers, eight GPUs, no conflicts.

Marker inheritance and programmatic markers

Conftest can also apply markers to tests that didn’t ask for them. Here’s mai_cluster using file-level pytestmark in a conftest to mark an entire directory as sequential:

# mai_cluster/tests/cpu_tests/ray_utils/conftest.py
pytestmark = pytest.mark.sequential

Every test in that directory inherits sequential without declaring it. The conftest acts as a directory-wide stamp.

Even more powerful: pytest_collection_modifyitems lets you add markers programmatically based on file path. Here’s mai_evaluator’s GPU tests conftest:

# mai_evaluator/tests/gpu_tests/conftest.py
def pytest_collection_modifyitems(config, items):
    for item in items:
        if "gpu_tests" in str(item.fspath):
            item.add_marker(pytest.mark.require_cuda_8gpu)

After pytest collects all test items, this hook runs through them and adds require_cuda_8gpu to anything living under a gpu_tests/ directory. The tests themselves don’t need any marker at all – the conftest handles it based on where the file lives.

Directory enforcement

sgl_server takes this further. Its conftest doesn’t just add markers – it kills pytest entirely if a test file is in the wrong directory:

# sgl_server/tests/conftest.py
def pytest_collect_file(file_path: Path, parent) -> None:
    tests_root = Path(__file__).parent
    allowed_dirs = {"cpu_only", "gpu_only", "redis_and_gpu"}

    if file_path.suffix == ".py" and file_path.name.startswith("test_"):
        relative = file_path.relative_to(tests_root)
        if len(relative.parts) > 0:
            top_dir = relative.parts[0]
            if top_dir not in allowed_dirs:
                error_msg = [
                    "ERROR: Test file in invalid location!",
                    "All sgl_server tests must be in one of:",
                    f"  - {tests_root}/cpu_only/",
                    f"  - {tests_root}/gpu_only/",
                    f"  - {tests_root}/redis_and_gpu/",
                    f"Found: {file_path}",
                ]
                pytest.exit("\n".join(error_msg), returncode=1)

Hard error. Not a skip, not a warning – pytest.exit with returncode 1. If you drop a test file in sgl_server/tests/somewhere_else/, CI dies. This is the conftest saying “I don’t trust humans to put files in the right place, so I’ll enforce it mechanically.”

The autouse fixture: global state police

One more conftest pattern worth highlighting. Both mai_layers and mai_evaluator have this:

# mai_layers/tests/conftest.py
_TORCH_GLOBAL_DEFAULTS = {
    "allow_tf32": torch.backends.cuda.matmul.allow_tf32,
    "cudnn_benchmark": torch.backends.cudnn.benchmark,
    "grad_enabled": torch.is_grad_enabled(),
    "default_dtype": torch.get_default_dtype(),
    # ... more globals
}

@pytest.fixture(autouse=True)
def check_torch_globals(request: pytest.FixtureRequest):
    yield  # run the test

    errors = []
    if torch.backends.cuda.matmul.allow_tf32 != _TORCH_GLOBAL_DEFAULTS["allow_tf32"]:
        errors.append("torch.backends.cuda.matmul.allow_tf32 was modified")
    # ... check each global

    if errors:
        raise RuntimeError(
            f"Test {request.node.nodeid} modified global state:\n  "
            + "\n  ".join(errors)
        )

autouse=True means this fixture runs for every test without being requested. It snapshots PyTorch’s global state before the test, yields (letting the test run), then checks that nothing changed. If your test modifies torch.backends.cudnn.deterministic and doesn’t restore it, this fixture will blow up.

This matters because GPU tests run in parallel. If test A sets allow_tf32 = True and doesn’t clean up, test B might get silently different numerical results. The conftest acts as a state machine guardian – verifying each test is a good citizen.

The config layer: guest lists and seating charts

Conftest handles the how. But the what – which packages run on GPU at all – is controlled by two config files.

GPU_TEST_PKGS: the guest list

In run_gpu_tests.yml, there’s a hardcoded list:

GPU_TEST_PKGS: >-
  mai_cluster
  mai_distributed
  mai_evaluator/tests/gpu_tests
  mai_io/tests/gpu_tests
  mai_job/tests/gpu_tests/
  mai_kernels/tests/functional
  mai_kernels/tests/st_functional
  mai_kernels/tests/unit
  mai_layers
  mai_multimodal
  mai_trainer
  mai_stats
  sgl_server/tests/gpu_only

If your package isn’t on this list, it doesn’t get GPU CI. Period. No marker, no conftest, no clever trick will save you. This is the bouncer at the door.

Notice the granularity: some entries are full packages (mai_layers, mai_trainer), but others are specific subdirectories (mai_evaluator/tests/gpu_tests, sgl_server/tests/gpu_only). This is because some packages have both CPU and GPU tests. You only want the GPU subdirectory on the guest list.

test_exec_mode.yml: the seating chart

Once a test is on the guest list, how does it run? This is controlled by test_exec_mode.yml:

serial:
  - mai_cluster/tests/gpu_tests
  - mai_distributed/tests/distributed
  - mai_trainer/tests/flow
  - mai_trainer/tests/hooks/test_dummy_step_determinism.py

isolated:
  - mai_layers/tests/distributed
  - mai_trainer/tests/checkpointing
  - sgl_server/tests/gpu_only/loading/test_model_load.py
  - mai_layers/tests/numerics_check

SKIP:
  - mai_kernels/tests/performance
  - mai_cluster/tests/cpu_tests/
  - mai_layers/tests/configs/
  - mai_trainer/tests/configs
  - mai_trainer/tests/hooks

Four modes:

• serial: One test at a time, using ALL GPUs. For distributed training tests that need the entire node.
• isolated: Separate process for each test. For tests that leak global state, hold ports, or spawn subprocesses that conflict.
• normal: (Default, not listed.) Parallel execution, one GPU per test via pytest-xdist.
• SKIP: Don’t run in GPU CI. Period.

The SKIP section is interesting because it catches CPU test directories that live inside GPU packages. mai_trainer/tests/configs is a CPU test that lives under mai_trainer, which is in GPU_TEST_PKGS. Without the SKIP entry, those CPU tests would wastefully run on GPU nodes.

And there’s a specificity rule: longest matching prefix wins. mai_trainer/tests/hooks is SKIP (CPU tests). But mai_trainer/tests/hooks/test_dummy_step_determinism.py is serial (it’s a GPU test that checks training determinism). The more specific path override takes precedence.

Python’s ast module: reading code as data

Now the question: how does CI read these markers from test files without running them?

The naive approach is grep. And it works – for a while. But grep reads bytes, not semantics. It can’t tell the difference between:

pytestmark = pytest.mark.slow

and:

config = "pytestmark = pytest.mark.slow"

Identical bytes. Completely different meanings. One sets a marker. The other is a string literal that happens to contain marker-like text. Grep matches both.

Python’s ast module (Abstract Syntax Tree) reads source code as a data structure. ast.parse() turns Python into a tree of nodes:

import ast

code = 'pytestmark = pytest.mark.slow'
tree = ast.parse(code)
# Module -> Assign -> target: Name('pytestmark'), value: Attribute('pytest.mark.slow')

vs:

code = 'config = "pytestmark = pytest.mark.slow"'
tree = ast.parse(code)
# Module -> Assign -> target: Name('config'), value: Constant(str)

The AST sees these as completely different structures. The first is an assignment to pytestmark with an attribute chain value. The second is an assignment to config with a string constant value. False positives are eliminated by construction, not by regex tuning.

ast is in the standard library. No pip install. ast.parse() gives you the tree, ast.walk() traverses it.

How our manifest parser uses it

The real code in ci/src/ci/plan/manifest.py parses test files to extract marker information. Here’s the function that reads file-level pytestmark assignments:

def _get_file_level_decorators(node: ast.Assign) -> Decorators | None:
    if len(node.targets) != 1:
        return None

    target = node.targets[0]
    if not isinstance(target, ast.Name) or target.id != "pytestmark":
        return None
    if not isinstance(node.value, (ast.List, ast.Tuple)):
        return None

    return _get_marks_from_decorators(node.value.elts)

Step by step: Is it an assignment? Does it assign to the name pytestmark? Is the value a list or tuple? If all three pass, parse the markers from the elements. A string literal containing pytestmark fails at the second check – it would be assigning to config or whatever the variable is named.

And here’s how it extracts decorator names from the dotted chain pytest.mark.runs_on:

def _get_decorator_name(decorator: ast.expr) -> str | None:
    if isinstance(decorator, ast.Name):
        return decorator.id                    # @slow
    elif isinstance(decorator, ast.Attribute):
        parts = []
        current = decorator
        while isinstance(current, ast.Attribute):
            parts.append(current.attr)         # runs_on, mark, ...
            current = current.value
        if isinstance(current, ast.Name):
            parts.append(current.id)           # pytest
            return ".".join(reversed(parts))   # pytest.mark.runs_on
    elif isinstance(decorator, ast.Call):
        return _get_decorator_name(decorator.func)
    return None

It walks the attribute chain backwards: runs_on -> mark -> pytest, then reverses to get the dotted name. This handles @pytest.mark.runs_on, @slow, and @pytest.mark.runs_on(node="condor") (the Call case recursing into the function).

The output is a list of TestTarget objects:

@dataclasses.dataclass(frozen=True, slots=True)
class TestTarget:
    test_path: str
    skip: bool
    runs_on: str
    run_mode: str
    test_class: str | None = None
    test_function: str | None = None
    test_method: str | None = None

Each target knows its file path, its runner requirement, and its execution mode. This is what the sharding algorithm consumes to build test plans.

The full pipeline

All these pieces connect into a pipeline that runs on every PR push and nightly trigger. Here’s the flow:

PR push / nightly trigger
        |
   find-changed action (ci.cli find_changed)
   "what packages changed?"
        |
   +----+--------------------+
   |                         |
CPU workflow              GPU workflow
run_python_tests          run_gpu_tests
   |                         |
cpu_tests.yml            GPU_TEST_PKGS (directory list)
(package -> path         test_exec_mode.yml (serial/isolated/SKIP)
 mapping)                    |
   |                    AST manifest parser (ci.cli manifest)
pytest on               parses markers -> TestTarget objects
GitHub Actions               |
runners                 Weighted sharding (bin-pack by duration)
                             |
                        run_gpu_tests.py -> Slurm nodes
                        (pytest per shard, per exec mode)

Step 1: find_changed. The CI figures out which packages were modified in this PR. If only mai_config changed, no point running mai_kernels GPU tests.

Step 2: filter to GPU_TEST_PKGS. The changed files are filtered against the guest list. uv run python -m ci.cli tests --filter $GPU_TEST_PKGS produces a JSON array of test file paths.

Step 3: apply test_exec_mode. Each test file gets its execution mode by longest-prefix match. SKIP’d tests are dropped. The rest are tagged serial, isolated, or normal.

Step 4: shard. The test files are bin-packed into shards based on historical duration data from gpu_test_durations.json. The goal is equal-weight shards so that no single shard becomes a bottleneck. Max 12 shards, minimum 40 tests per shard.

Step 5: execute. Each shard runs on a GPU node inside a Docker container. Serial tests get the full node. Normal tests get pytest-xdist with one GPU per worker. Isolated tests get their own process. Results flow back to GitHub via the SSH bridge described in The 500 Lines That Run Every GPU Test.

The CPU side is simpler: cpu_tests.yml maps each package to its test paths and flags, and pytest runs directly on GitHub Actions runners. No sharding, no Slurm, no Docker. Just uv run python -um pytest on a VM.

Where this intersects with Bazel and torch

One more thing. A teammate pointed out that the GPU_TEST_PKGS list is useful beyond just running tests – it’s also a signal for which packages need torch.

If a package is in GPU_TEST_PKGS, it almost certainly needs CUDA torch for its tests. If it’s in the SKIP section, it’s an exception (CPU tests living inside a GPU package). If it’s not in the list at all, it probably needs either CPU-only torch or no torch.

His _check_needs_torch function takes a different approach – it reads uv.lock to check if a package has a direct torch dependency. Package-level, not test-level. But the signal correlates: packages that declare torch as a dependency are usually the same ones whose tests live in GPU_TEST_PKGS.

This matters for the Bazel migration, for pyright CI, for any system that needs to know “does this package touch torch?” The config files we built for GPU CI turn out to be a useful source of truth for completely unrelated problems.

Cheat sheet

The GPU CI system is conftest hooks giving markers meaning, YAML configs routing tests to nodes, and AST parsing reading source code as data. Each layer does one thing. Together they answer the question we started with: which tests need GPUs, and how should they run?