NVIDIA-NeMo / NVIDIA-NeMo/Curator
Refactor benchmarking scripts
Open
@sarahyurick is already working on this.
Since Jan 16, 2026.
enhancement
- Dominant language
- Python
- Stars
- 1.8k
- Forks
- 328
- Avg merge
- 4d 5h
- Merged PRs (30d)
- 30
Description
https://github.com/NVIDIA-NeMo/Curator/pull/1382 helped clean up some repeated logic by creating a utils.py script with helper functions.
While working on that, I noticed other areas of improvement. For example, we pass the same arguments a lot of the time, so we could consider creating a helper function for it:
def add_common_benchmark_args(parser: argparse.ArgumentParser) -> None:
"""Add common benchmark arguments to an argument parser.
This adds the standard arguments used across most benchmarks:
- benchmark-results-path (required)
- input-path (required)
- output-path (optional, with default)
- executor (optional, with choices)
Args:
parser: ArgumentParser to add arguments to
"""
parser.add_argument("--benchmark-results-path", type=Path, required=True, help="Path to benchmark results")
parser.add_argument("--input-path", required=True, type=Path, help="Path to input data")
parser.add_argument("--output-path", type=Path, help="Output directory for results")
parser.add_argument("--executor", default="ray_data", choices=["xenna", "ray_data", "ray_actors"], help="Executor to use")
A lot of the success/failure logic is redundant between scripts too. Helper functions like:
def run_benchmark_main(
benchmark_func: Callable[[argparse.Namespace], dict[str, Any]],
args: argparse.Namespace,
benchmark_name: str,
) -> int:
"""Standard main function wrapper for benchmarks.
Handles the try-except-finally pattern and result writing that's common
across all benchmark scripts.
Args:
benchmark_func: Function that runs the benchmark and returns results dict
args: Parsed command-line arguments
benchmark_name: Name of the benchmark for logging
Returns:
Exit code (0 for success, 1 for failure)
"""
logger.info(f"=== {benchmark_name} Benchmark Starting ===")
logger.info(f"Arguments: {vars(args)}")
try:
results = benchmark_func(args)
except Exception as e: # noqa: BLE001
error_traceback = traceback.format_exc()
print(f"Benchmark failed: {e}")
logger.debug(f"Full traceback:\n{error_traceback}")
results = {
"params": vars(args),
"metrics": {
"is_success": False,
},
"tasks": [],
}
finally:
write_benchmark_results(results, args.benchmark_results_path)
# Return proper exit code based on success
return 0 if results["metrics"]["is_success"] else 1
could help us create a cleaner experience.
Contributor guide
First steps
- Read the whole issue, then the project's contributing guide.
- Comment on the issue to say you are picking it up — it saves two people doing the same work.
- Fork the repository and make your change on a branch.
- Open a pull request that references the issue number.
Assessment
This issue has not been assessed yet.