Datarax GPU Testing Guide

This document describes how to run Datarax tests on GPU hardware.

Prerequisites

1. NVIDIA GPU with a driver that supports CUDA 12
2. Python 3.12 or 3.13
3. The environment from ./setup.sh --backend cuda12 (JAX's CUDA runtime comes from the cuda12 extra, so no system CUDA toolkit is needed)

Setting Up the Environment

The recommended way to set up the environment is to use the main setup script:

# Set up the development environment with automatic backend detection
./setup.sh

# Or force the CUDA 12 backend explicitly
./setup.sh --backend cuda12

# Activate the virtual environment (loads the generated .datarax.env)
source activate.sh

This approach:

• Detects NVIDIA GPUs automatically, or accepts an explicit --backend cuda12
• Uses JAX's uv-managed CUDA runtime via the cuda12 extra — no system CUDA toolkit or custom LD_LIBRARY_PATH injection is required
• Writes backend configuration to the generated .datarax.env (your user-owned .env is never modified)
• Relies on the committed activate.sh (checked into the repository, not generated by setup) to load that configuration when sourced

Running GPU Tests

We provide a dedicated script for running tests on GPU:

# Run all GPU-specific tests
bash scripts/run_gpu_tests.sh

This script will:

1. Activate the project environment and check for GPU availability
2. Ask the test run for CUDA with DATARAX_TEST_JAX_PLATFORMS=cuda
3. Run the GPU-marked tests on the GPU

Manual GPU Testing

If you want more control over which tests to run on GPU, you can:

# Test runs use the CPU with eight emulated devices unless they ask for an
# accelerator; a JAX_PLATFORMS inherited from your shell does not change that.
export DATARAX_TEST_JAX_PLATFORMS="cuda"

# Run every test on the GPU
uv run pytest

# Run a specific test directory on the GPU
uv run pytest tests/operators/

# Run only the tests that need a GPU backend
uv run pytest -m accelerator

Troubleshooting

If you encounter issues with GPU tests:

1. Verify GPU is detected:

   uv run python scripts/check_gpu.py
   

2. Check CUDA installation:

   nvidia-smi
   

3. Memory issues: Adjust memory fraction if tests fail due to OOM errors:

   export XLA_CLIENT_MEM_FRACTION=0.5
   

4. GPU acceleration not used: Ensure JAX is using the GPU:

   JAX_PLATFORMS=cuda python -c "import jax; print(jax.devices())"
   

How GPU Testing Works

The GPU testing infrastructure consists of:

1. Test backend (tests/jax_test_environment.py): tests run on emulated CPU devices unless DATARAX_TEST_JAX_PLATFORMS asks for an accelerator; an inherited JAX_PLATFORMS does not move them onto a GPU.

2. Shell Script (scripts/run_gpu_tests.sh):

   • Verifies GPU availability using scripts/check_gpu.py
   • Sets DATARAX_TEST_JAX_PLATFORMS=cuda
   • Runs pytest over tests/

3. Test Markers: the substrax pytest plugin's @pytest.mark.accelerator(kind="gpu") skips a test unless the run uses a GPU backend, and @pytest.mark.devices(count) skips it below count visible devices. Most tests run on any device.

Adding New GPU Tests

Most Datarax tests are device-agnostic and run on whatever JAX backend the run selected. Declare a requirement with a marker from the substrax pytest plugin when a test needs it:

• Needs a GPU backend: @pytest.mark.accelerator(kind="gpu")
• Needs several devices: @pytest.mark.devices(count), or @pytest.mark.devices(count, kind="gpu") for several GPUs

Example

import jax
import pytest


@pytest.mark.devices(2, kind="gpu")
def test_multi_gpu_sharding():
    """Runs only when the run selected a GPU backend with at least two GPUs."""
    devices = jax.devices("gpu")
    ...

Test File Location

Place GPU-intensive tests in appropriate directories:

• tests/distributed/ - Multi-device and sharding tests
• tests/sharding/ - Data sharding tests
• tests/benchmarks/ - Performance benchmarks

Testing Status

The GPU testing infrastructure supports:

• Automatic device detection: Tests adapt to available hardware
• Selective test execution: -m accelerator runs only the tests that need an accelerator
• Memory management: Environment variables control GPU memory allocation

Running Full GPU Test Suite

# Run all tests on GPU
DATARAX_TEST_JAX_PLATFORMS=cuda uv run pytest tests/

# Run with memory limits (useful for shared GPUs)
XLA_CLIENT_MEM_FRACTION=0.5 DATARAX_TEST_JAX_PLATFORMS=cuda uv run pytest tests/

For more testing information, see the Testing Guide.