Implementation:NVIDIA NeMo Aligner Conftest Fixtures
| Implementation Details | |
|---|---|
| Name | Conftest_Fixtures |
| Type | Pattern Doc |
| Module | conftest (project root) |
| Repository | NeMo Aligner |
| Last Updated | 2026-02-08 00:00 GMT |
Overview
Pytest configuration and shared fixtures for the NeMo Aligner test suite, providing model initialization helpers, device selection, MPI test filtering, and session lifecycle management.
Description
The root-level conftest.py defines the test infrastructure for the entire NeMo Aligner project. It provides several categories of functionality:
Custom command-line options: The pytest_addoption hook registers --cpu (to run tests on CPU instead of GPU) and --mpi (to run only MPI-marked tests).
Device selection fixtures: The device fixture returns "CPU" or "GPU" based on the --cpu flag, and the run_only_on_device_fixture (autouse) automatically skips tests marked with @pytest.mark.run_only_on("CPU") or @pytest.mark.run_only_on("GPU") when the current device does not match.
Model parallel initialization: The init_model_parallel fixture yields an initializer function wrapping Utils.initialize_model_parallel and automatically tears down model parallelism via Utils.destroy_model_parallel after each test.
Pre-built model fixtures: The dummy_gpt_model fixture creates a minimal single-GPU MegatronGPTModel with a 1-layer, 128-hidden-size architecture for unit testing. The dummy_actor_gpt_model_with_pp fixture creates a 2-device pipeline-parallel MegatronGPTActorModel with PPO configuration for integration testing.
MPI test collection: The pytest_collection_modifyitems hook ensures that when --mpi is passed, only MPI-marked tests run, and when it is not passed, MPI tests are skipped.
Session lifecycle: The pytest_sessionstart hook removes a SUCCESS_FILE marker at the start of a run (on local rank 0 only), and pytest_sessionfinish writes the marker on success and cleans up NCCL process groups to suppress warnings from PyTorch >= 2.4.
Usage
This conftest.py is automatically loaded by pytest for all tests in the repository. Individual fixtures can be requested by name in test function signatures. The --cpu and --mpi flags are used on the command line to control test selection.
Code Reference
Source Location
- Repository: NeMo Aligner
- File:
conftest.py(L1-420)
Signature
# Pytest hooks
def pytest_addoption(parser):
"""Registers --cpu and --mpi command-line options."""
...
def pytest_configure(config):
"""Registers custom markers: run_only_on, mpi."""
...
def pytest_collection_modifyitems(config, items):
"""Filters tests based on --mpi flag."""
...
def pytest_sessionstart(session):
"""Removes SUCCESS_FILE marker at session start (rank 0 only)."""
...
def pytest_sessionfinish(session, exitstatus):
"""Destroys NCCL process group and writes SUCCESS_FILE on success."""
...
# Fixtures
@pytest.fixture
def device(request):
"""Returns 'CPU' or 'GPU' based on --cpu flag."""
...
@pytest.fixture(autouse=True)
def run_only_on_device_fixture(request, device):
"""Skips tests whose run_only_on marker does not match current device."""
...
@pytest.fixture
def init_model_parallel():
"""Yields an initializer function; tears down model parallelism on cleanup."""
...
@pytest.fixture
def llama3_tokenizer():
"""Returns an AutoTokenizer for meta-llama/Meta-Llama-3-8b."""
...
@pytest.fixture
def dummy_gpt_model(init_model_parallel):
"""Creates a minimal 1-layer MegatronGPTModel for unit testing."""
...
@pytest.fixture
def dummy_actor_gpt_model_with_pp():
"""Creates a 2-device pipeline-parallel MegatronGPTActorModel with PPO config."""
...
Import
# conftest.py is auto-discovered by pytest; no explicit import needed.
# Fixtures are injected by name into test function signatures:
def test_example(init_model_parallel, dummy_gpt_model):
...
I/O Contract
Inputs
| Fixture / Hook | Inputs | Description |
|---|---|---|
pytest_addoption |
parser | Registers --cpu (bool) and --mpi (bool) CLI options
|
device |
request (pytest) | Reads --cpu option from the pytest config
|
init_model_parallel |
(none) | Accepts *args, **kwargs forwarded to Utils.initialize_model_parallel (e.g., tensor_model_parallel_size, pipeline_model_parallel_size)
|
dummy_gpt_model |
init_model_parallel | Depends on init_model_parallel fixture; initializes with TP=1, PP=1
|
dummy_actor_gpt_model_with_pp |
(none) | Self-contained; initializes distributed training internally with 2 devices, PP=2 |
Outputs
| Fixture / Hook | Type | Description |
|---|---|---|
device |
str | "CPU" or "GPU"
|
init_model_parallel |
Callable | A function that calls Utils.initialize_model_parallel(*args, **kwargs)
|
llama3_tokenizer |
AutoTokenizer | HuggingFace tokenizer for Meta-Llama-3-8B |
dummy_gpt_model |
MegatronGPTModel | Minimal GPT model (1 layer, hidden=128, heads=2, seq_len=512) |
dummy_actor_gpt_model_with_pp |
MegatronGPTActorModel | Pipeline-parallel actor model (2 layers, hidden=128, heads=4, PP=2) with PPO config |
| SUCCESS_FILE | file side-effect | Written to PYTEST_SUCCESS in the project root on successful test session completion
|
Usage Examples
# Running tests with specific options
# pytest --cpu # Run on CPU
# pytest --mpi # Run only MPI-marked tests
# pytest # Default: GPU, skip MPI tests
# Using the init_model_parallel fixture
def test_my_feature(init_model_parallel):
init_model_parallel(1, 1) # tensor_parallel=1, pipeline_parallel=1
# ... test code ...
# Teardown happens automatically
# Using the dummy_gpt_model fixture
def test_forward_pass(dummy_gpt_model):
model = dummy_gpt_model
# model is a fully initialized MegatronGPTModel ready for testing
assert model.cfg.num_layers == 1
assert model.cfg.hidden_size == 128
# Using the llama3_tokenizer fixture
def test_tokenization(llama3_tokenizer):
tokens = llama3_tokenizer.text_to_ids("Hello world")
assert len(tokens) > 0
# Marking a test for device-specific execution
@pytest.mark.run_only_on("GPU")
def test_gpu_only_feature(device):
assert device == "GPU"
# ... GPU-specific test ...
# Marking a test as MPI
@pytest.mark.mpi
def test_mpi_communication():
# Only runs when --mpi is passed
...