Developer Guide

September 16, 2026 · View on GitHub

This guide covers everything you need to know to contribute to Datarax development.

Development Environment Setup

Datarax uses uv as its package manager for all installation, development, and deployment tasks.

Quick Start

# Install uv if not already installed
pip install uv

# Run the automatic setup script
./setup.sh

# Activate the environment
source activate.sh

Setup Script Options

The setup.sh script provides several options:

Option	Description
`--backend {auto,cpu,cuda12,metal}`	Choose the backend policy (default: `auto`)
`--with-benchmarks`	Also install the competitor-framework benchmark extra
`--python <version>`	Create the environment with a specific Python version
`--recreate`	Remove the existing `.venv` before syncing
`--force-clean`	Remove `.venv`, the generated `.datarax.env`, and repo-local test artifacts
`--dry-run`	Print the resolved backend and uv commands without changing files
`--help, -h`	Show help message

Example usage:

./setup.sh                        # Standard setup with auto backend detection
./setup.sh --backend cpu          # Force a CPU-only environment
./setup.sh --backend cuda12       # Force the CUDA 12 backend
./setup.sh --with-benchmarks      # Add the competitor-framework benchmark extra
./setup.sh --recreate             # Rebuild .venv from scratch

Linux CUDA development uses JAX's uv-managed CUDA runtime via the cuda12 extra; the setup does not rely on a system CUDA toolkit or custom LD_LIBRARY_PATH injection.

Files Created by Setup

File	Purpose
`.venv/`	Virtual environment directory
`.datarax.env`	Generated backend configuration (user-owned `.env` is never modified)
`uv.lock`	Dependency lock file

activate.sh is checked into the repository (not generated by setup) and loads .datarax.env when sourced.

Package Management

Installing Dependencies

Datarax defines dependencies in pyproject.toml using optional dependency groups:

# Set up the environment the repository's tooling expects (detects the backend)
./setup.sh

# Or sync extras yourself. uv sync installs exactly the extras you name and removes
# the others, so list every extra you need in one command:
uv sync --extra dev --extra test --extra data --extra docs
uv sync --extra all          # Linux with CUDA 12: dev, test, data, docs and cuda12
uv sync --extra all-cpu      # every extra except a GPU backend

Adding New Dependencies

# Add a runtime dependency (edit pyproject.toml manually)
# Then sync with the extras you use (a bare `uv sync` removes them):
./setup.sh

# Or use uv add for development:
uv add package_name

Installing Multiple Extras

Important: uv sync installs exactly the extras you pass and removes the rest.

# ✅ Correct: multiple --extra flags for uv sync
uv sync --extra dev --extra test --extra data

# ✅ Recommended: use compound extras defined in pyproject.toml
uv sync --extra all      # includes dev, test, data, docs, cuda12
uv sync --extra all-cpu  # includes dev, test, data, docs (no GPU backend)

# ❌ Wrong: comma-separated values with --extra flag
# uv sync --extra dev,test,data  # This will ERROR!

Dependency Groups

Group	Contents
`dev`	Build tools, linters, type checkers, pytest plugins
`test`	Testing dependencies (pytest, coverage, etc.)
`docs`	Documentation tools (MkDocs, mkdocstrings)
`data`	Data loading libraries (datasets, tensorflow-datasets)
`cuda12`	CUDA 12 support for JAX
`all`	All of the above

Type Checking

Datarax uses Pyright for static type checking. Configuration is in pyproject.toml:

[tool.pyright]
exclude = ["example_data", ".deprecated", "**/__pycache__", "**/.venv"]
include = ["src", "tests", "examples", "scripts", "benchmarks"]
pythonVersion = "3.12"
typeCheckingMode = "basic"

Code in src/, tests/, examples/, scripts/, and benchmarks/ is type-checked. Certain rules are relaxed to accommodate JAX's dynamic typing patterns.

Running Type Checks

# Run Pyright manually
uv run pyright

# Through pre-commit
uv run pre-commit run pyright --all-files

Type Annotation Guidelines

When writing new code:

Add type annotations to all function signatures (parameters and return types)
Use proper generics for container types (e.g., list[int] instead of list)
Avoid Any whenever possible; use specific types or TypeVar for generic code
Handle None explicitly with Optional[T] or T | None syntax
Use jax.Array for JAX array types

Common Type Checking Issues

Optional Types: Always check if a value can be None before accessing attributes
JAX Arrays: Use jax.Array for JAX array types
Type Narrowing: Use appropriate guards (isinstance(), etc.) to narrow types properly
Union Types: Ensure all operations are valid for all possible types in a union

Code Style

Datarax follows standard Python code style practices enforced by Ruff:

Setting	Value
Line length	100 characters
Quote style	Double quotes
Docstring convention	Google style
Import sorting	isort-compatible
Target Python	3.12+

Running Linters

# Check for issues
uv run ruff check .

# Auto-fix issues
uv run ruff check --fix .

# Format code
uv run ruff format .

Ruff Configuration

Key Ruff settings in pyproject.toml:

[tool.ruff]
line-length = 100
target-version = "py312"

[tool.ruff.format]
quote-style = "double"
indent-style = "space"

[tool.ruff.lint.pydocstyle]
convention = "google"

Pre-commit Hooks

Pre-commit hooks run automatically on every commit to ensure code quality.

Setup

# Install pre-commit hooks (done automatically by setup.sh)
uv run pre-commit install

# Run all hooks manually
uv run pre-commit run --all-files

Configured Hooks

The pipeline runs a number of hooks, including those below. See .pre-commit-config.yaml for the authoritative, complete list.

Hook	Purpose
`sort_pyproject`	Keep `pyproject.toml` sorted
`validate-pyproject`	Validate `pyproject.toml` against the packaging schema
`trailing-whitespace`	Remove trailing whitespace
`end-of-file-fixer`	Ensure files end with newline
`check-yaml` / `check-toml` / `check-json`	Validate config file syntax
`check-added-large-files`	Prevent large files
`check-merge-conflict`	Catch unresolved merge markers
`debug-statements`	Catch leftover debugger calls
`mixed-line-ending`	Normalize line endings
`ruff` / `ruff-format`	Linting with auto-fix and formatting
`ruff-docstring-whitespace`	Docstring whitespace checks
`ruff-best-practices-check` / `ruff-critical-check`	Extra Ruff rule sets
`ruff-github-actions-check`	Lint GitHub Actions workflows
`check-file-length`	Enforce a maximum file length
`lint-imports`	Enforce import-layer boundaries
`interrogate`	Docstring coverage
`pydoclint`	Docstring sections match signatures and raise statements
`xenon`	Complexity thresholds
`pylint-duplicate-code`	Detect duplicated code
`vulture`	Detect dead code
`pyright`	Type checking
`bandit`	Security scanning
`nbqa-ruff`	Notebook linting
`shellcheck`	Shell script linting

Skipping Hooks

If you need to skip hooks temporarily (not recommended):

git commit --no-verify -m "message"

Testing

Running Tests

# Run all tests (CPU-only, most stable)
JAX_PLATFORMS=cpu uv run pytest

# Run specific test module
JAX_PLATFORMS=cpu uv run pytest tests/sources/test_memory_source_module.py

# Run with verbose output
uv run pytest -v

# Run with coverage
uv run pytest --cov=src/datarax --cov-report=html

Test Categories

Tests use pytest markers for categorization:

Marker	Description
`@pytest.mark.unit`	Unit tests
`@pytest.mark.integration`	Integration tests
`@pytest.mark.e2e`	End-to-end tests
`@pytest.mark.accelerator(kind="gpu")`	Tests that need a GPU backend (substrax plugin)
`@pytest.mark.devices(count)`	Tests that need at least `count` devices (substrax plugin)
`@pytest.mark.slow`	Slow-running tests
`@pytest.mark.benchmark`	Performance benchmarks
`@pytest.mark.tfds`	TensorFlow Datasets tests
`@pytest.mark.hf`	HuggingFace Datasets tests

Running Specific Test Types

# Run on a GPU; tests marked accelerator skip on a CPU run
DATARAX_TEST_JAX_PLATFORMS=cuda uv run pytest

# Run only integration tests
uv run pytest -m integration

# Run only unit tests (fast)
uv run pytest -m unit

# Run benchmarks
uv run pytest -m benchmark --benchmark-autosave

Test Directory Structure

Tests mirror the source structure:

tests/
├── augment/         # Augmentation tests
├── batching/        # Batch processing tests
├── benchmarks/      # Performance benchmark tests
├── checkpoint/      # Checkpoint tests
├── cli/             # CLI tests
├── config/          # Configuration tests
├── control/         # Control flow tests
├── core/            # Core functionality tests
├── data/            # Test data and fixtures
├── distributed/     # Distributed training tests
├── examples/        # Example validation tests
├── fixtures/        # Shared pytest fixtures
├── integration/     # End-to-end tests
├── memory/          # Memory management tests
├── monitoring/      # Monitoring tests
├── operators/       # Pipeline operator tests
├── performance/     # Performance tests
├── pipeline/        # Pipeline / DAG execution tests
├── samplers/        # Sampling tests
├── scripts/         # Test helper scripts
├── sharding/        # Sharding tests
├── sources/         # Data source tests
├── test_common/     # Common testing utilities
├── transforms/      # Transform tests (neural network ops)
├── utils/           # Utility function tests
└── conftest.py      # Pytest configuration

Writing New Tests

Place tests in the directory matching the module being tested
Name files test_<component>.py
Name test functions test_<behavior>()
Use appropriate markers for hardware requirements
Create standalone tests that don't depend on other test files

Example:

import numpy as np
import pytest
from datarax.sources import MemorySource, MemorySourceConfig

@pytest.mark.unit
def test_memory_source_initialization():
    """Test that MemorySource initializes correctly."""
    config = MemorySourceConfig()
    data = {"x": np.array([1, 2, 3])}
    source = MemorySource(config, data=data)
    assert source is not None
    assert len(source) == 3

Building and Packaging

Building the Package

# Build source distribution and wheel
uv run python -m build

# Build outputs go to dist/ (version is derived dynamically from git tags)
ls dist/
# datarax-<version>.tar.gz
# datarax-<version>-py3-none-any.whl

Package Configuration

Build settings in pyproject.toml:

[build-system]
build-backend = "hatchling.build"
requires = ["hatchling>=1.18"]

[tool.hatch.build.targets.wheel]
packages = ["src/datarax"]

GPU/CUDA Support

Automatic Detection

The setup script automatically detects NVIDIA GPUs and configures CUDA support.

Manual GPU Setup

# Select the CUDA 12 backend explicitly, then load the environment
./setup.sh --backend cuda12
source activate.sh

# Rebuild the environment from scratch if it is broken
./setup.sh --backend cuda12 --recreate

Environment Variables for GPU

The generated .datarax.env file configures JAX for GPU:

# GPU configuration
export JAX_PLATFORMS="cuda,cpu"
export XLA_PYTHON_CLIENT_PREALLOCATE="false"
export XLA_CLIENT_MEM_FRACTION="0.8"

Testing GPU Support

# Check GPU availability
python -c "import jax; print(jax.devices())"

# Run the tests that need a GPU backend
DATARAX_TEST_JAX_PLATFORMS=cuda uv run pytest -m accelerator

Datarax provides Docker images for development, testing, and benchmarking across CPU/GPU/TPU platforms. See the Docker guide for build instructions, GPU passthrough, and cloud deployment (Vertex AI, SkyPilot).

Utility Scripts

Located in scripts/:

Script	Purpose
`run_tests.sh`	Run tests with auto GPU detection
`run_gpu_tests.sh`	Run GPU-specific tests with CUDA config
`run_full_benchmark.sh`	Run comparative benchmarks via the `benchmarks.cli` module
`run_all_examples_on_gpu.sh`	Run all examples on GPU
`run_typecheck.sh`	Run pyright type checking
`check_gpu.py`	Check GPU availability
`check_sync.py`	Check py/ipynb notebook sync
`validate_examples.py`	Validate example file structure
`jupytext_converter.py`	Convert between .py and .ipynb formats
`generate_docs.py`	Generate documentation from source
`generate_baselines.py`	Generate benchmark baseline data
`verify_docs.py`	Verify code blocks in markdown docs
`distributed_test_runner.py`	Distributed test runner for Vertex AI
`submit_vertex_job.py`	Submit jobs to Vertex AI

Running Scripts

# Run tests (auto-detects GPU)
./run_tests.sh

# Run tests with specific device
./run_tests.sh --device=cpu

# Check GPU
uv run python scripts/check_gpu.py

# Validate examples
uv run python scripts/validate_examples.py --verbose

# Check notebook sync
uv run python scripts/check_sync.py --verbose

Environment Variables

Key environment variables for development:

Variable	Purpose	Default
`JAX_PLATFORMS`	JAX device platforms	`cpu` or `cuda,cpu`
`JAX_ENABLE_X64`	Enable 64-bit floats	`0`
`XLA_PYTHON_CLIENT_PREALLOCATE`	GPU memory preallocation	`false`
`XLA_CLIENT_MEM_FRACTION`	GPU memory fraction	`0.8`
`TF_CPP_MIN_LOG_LEVEL`	TensorFlow logging level	`1`

Documentation

Building Documentation

# Serve documentation locally
uv run mkdocs serve

# Build static documentation
uv run mkdocs build

Documentation Structure

docs/
├── index.md                 # Home page
├── getting_started/         # Installation and quick start
├── user_guide/              # User documentation
│   ├── data_sources.md
│   ├── dag_construction.md
│   ├── distributed_training.md
│   └── ...
├── examples/                # Example documentation
├── core/, operators/, ...   # API reference pages
├── api_reference/           # Consolidated API reference
└── contributing/            # Contribution guidelines
    ├── contributing_guide.md
    ├── dev_guide.md                      # This guide
    ├── testing_guide.md
    ├── test_structure.md
    ├── gpu_testing.md
    ├── type_issues_guide.md
    ├── example_documentation_design.md
    └── performance_optimization_guide.md

Troubleshooting

Common Issues

Import errors after installation:

# Rebuild the environment
./setup.sh --recreate

GPU not detected:

# Check NVIDIA drivers
nvidia-smi

# Rebuild the environment with the CUDA 12 backend
./setup.sh --backend cuda12 --recreate

Pre-commit hook failures:

# Update hooks
uv run pre-commit autoupdate

# Run specific hook
uv run pre-commit run <hook-id> --all-files

Type checking errors:

# Run with verbose output
uv run pyright --verbose

# Check specific file
uv run pyright src/datarax/module.py

Getting Help

Check existing GitHub Issues
Read the API documentation
Review test files for usage examples