tests: data_prep tests fail to collect due to missing cosmos_xenna dependency

[andrewwhitecdw]

Problem

Running the test suite in a fresh environment fails to collect most tests under tests/data_prep/ because src/nemotron/data_prep/utils/hf_env.py imports from cosmos_xenna:

ModuleNotFoundError: No module named 'cosmos_xenna'

Affected test files (from pytest tests/ -m "not integration"):

• tests/data_prep/test_bin_assignment.py
• tests/data_prep/test_binidx_stage_resume.py
• tests/data_prep/test_jsonl_pipeline.py
• tests/data_prep/test_packing_algorithms.py
• tests/data_prep/test_packing_builder.py
• tests/data_prep/test_rl_omni_ensure_raw_dir.py
• tests/data_prep/test_row_splitting.py
• tests/data_prep/test_size.py
• tests/data_prep/test_stage_keys.py
• tests/data_prep/test_transforms.py
• tests/data_prep/test_wandb_hook.py

Context

cosmos-xenna is listed in pyproject.toml under the xenna and byob extras, and also appears in the main dependency list, but it does not seem to be installed by uv pip install -e . in the default configuration. uv.lock does contain entries for it.

Question

How would you like this resolved? A few options:

1. Make cosmos-xenna a mandatory dependency in the default install.
2. Make the import in hf_env.py lazy/conditional so tests that don't need it can collect without it installed.
3. Document that contributors need to install with a specific extra (e.g., uv pip install -e "[byob]") to run the data_prep tests.
4. Something else?

Happy to open a PR once the preferred direction is clear.