NVIDIA-NeMo / NVIDIA-NeMo/Curator

Refactor benchmarking scripts

Open
#1,384 0 comments 0 reactions 1 assignee View on GitHub

@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

Open the contributing guide

First steps

  1. Read the whole issue, then the project's contributing guide.
  2. Comment on the issue to say you are picking it up — it saves two people doing the same work.
  3. Fork the repository and make your change on a branch.
  4. Open a pull request that references the issue number.

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.