Skip to content

Add ray to each benchmarking entry and make sure no benchmarking/script has RayClient #1593

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
praateekmahajan merged 2 commits into from
Mar 11, 2026
Merged

Add ray to each benchmarking entry and make sure no benchmarking/script has RayClient #1593

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
b8f5bdd
Add ray to each entry and make sure no benchmarking/script has RayCli…
praateekmahajan Mar 10, 2026
74cf224
Other minor fixes in individual scripts
praateekmahajan Mar 11, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
16 changes: 14 additions & 2 deletions benchmarking/nightly-benchmark.yaml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -153,6 +153,10 @@ entries:
- number_of_domains_predicted
ping_on_failure:
- U022P8CDX40 # Sarah Yurick
ray:
num_cpus: 64
num_gpus: 4
enable_object_spilling: false
Comment on lines +156 to +159

@greptile-apps greptile-apps Bot Mar 11, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

num_gpus: 4 vs actual GPU usage — worth verifying

domain_classification_xenna and embedding_generation_xenna (line 220-223) now both reserve 4 GPUs on the Ray cluster. Neither entry passes a --gpus or similar argument to their scripts, so it isn't immediately obvious how many GPUs the workloads actually consume at runtime.

If these benchmarks only use 1 GPU internally (similar to the audio_fleurs case), 3 GPUs per run are reserved but idle for the entire benchmark duration, potentially blocking other concurrent jobs. It may be worth confirming the per-script GPU allocation before landing this so the reservation matches actual utilisation.

This same concern applies to the embedding_generation_xenna entry at line 220.

@praateekmahajan praateekmahajan Mar 11, 2026

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Currently it anyway uses 4 gpus we are just now being explicit it.

requirements:
# Observed throughput of 2900 docs/sec so we allow a 5% buffer to account for variability
- metric: throughput_docs_per_sec
Expand Down Expand Up @@ -213,6 +217,10 @@ entries:
ping_on_failure:
- U022P8CDX40 # Sarah Yurick
- U07JL5K0L10 # Praateek Mahajan
ray:
num_cpus: 64
num_gpus: 4
enable_object_spilling: false
requirements:
# Observed throughput of 8600 docs/sec so we allow a 5% buffer to account for variability
- metric: throughput_docs_per_sec
Expand Down Expand Up @@ -409,7 +417,7 @@ entries:
# ensure the total number of documents processed is correct
- metric: num_documents_processed
exact_value: 2119489
# account for stochastic filters
# account for stochastic filters
- metric: num_kept_documents
min_value: 2090470
max_value: 2090490
Expand Down Expand Up @@ -442,7 +450,7 @@ entries:
# ensure the total number of documents processed is correct
- metric: num_documents_processed
exact_value: 2119489
# account for stochastic filters
# account for stochastic filters
- metric: num_kept_documents
min_value: 2090470
max_value: 2090490
Expand Down Expand Up @@ -583,6 +591,10 @@ entries:
--split=dev
--wer-threshold=5.5
--gpus=1
ray:
num_cpus: 64
num_gpus: 4
enable_object_spilling: false
Comment on lines +594 to +597

@greptile-apps greptile-apps Bot Mar 10, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

GPU over-provisioning in audio_fleurs Ray config

The ray: block allocates num_gpus: 4 to the Ray cluster, but the script argument --gpus=1 (line 593) only uses 1 GPU. This reserves 3 additional GPUs on the cluster node for the duration of this benchmark run without utilizing them, potentially blocking other concurrent jobs from accessing those resources.

Every other GPU-using entry in this file (e.g. image_curation, video_embedding) consistently requests the same number of GPUs in the Ray config as are actually consumed by the workload. Consider aligning this entry:

Suggested change
ray:
num_cpus: 64
num_gpus: 4
enable_object_spilling: false
ray:
num_cpus: 64
num_gpus: 1
enable_object_spilling: false

sink_data:
- name: slack
ping_on_failure:
Expand Down
2 changes: 1 addition & 1 deletion benchmarking/scripts/arxiv_e2e_pipeline_benchmark.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -480,7 +480,7 @@ def main() -> int:
},
"tasks": [],
}
success_code = 0
success_code = 1
try:
results = run_benchmark(args)
success_code = 0 if results["metrics"]["is_success"] else 1
Expand Down
9 changes: 5 additions & 4 deletions benchmarking/scripts/audio_fleurs_benchmark.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -23,9 +23,8 @@
from typing import Any

from loguru import logger
from utils import write_benchmark_results
from utils import setup_executor, write_benchmark_results

from nemo_curator.backends.xenna import XennaExecutor
from nemo_curator.pipeline import Pipeline
from nemo_curator.stages.audio.common import GetAudioDurationStage, PreserveByValueStage
from nemo_curator.stages.audio.datasets.fleurs.create_initial_manifest import CreateInitialManifestFleursStage
Expand All @@ -44,6 +43,7 @@ def run_audio_fleurs_benchmark( # noqa: PLR0913
split: str,
wer_threshold: float,
gpus: int,

@ayushdg ayushdg Mar 11, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

unrelated to the PR but this arg is confusing IMO. The Cli claims its the number of GPUs to use but actually its the resource requirement for the inference stage

executor: str = "xenna",
**kwargs, # noqa: ARG001
) -> dict[str, Any]:
"""Run the audio fleurs benchmark and collect comprehensive metrics."""
Expand All @@ -65,7 +65,7 @@ def run_audio_fleurs_benchmark( # noqa: PLR0913
logger.info(f"WER threshold: {wer_threshold}")
logger.info(f"GPUs: {gpus}")

executor = XennaExecutor()
executor_obj = setup_executor(executor)
pipeline = Pipeline(name="audio_inference", description="Inference audio and filter by WER threshold.")

# Add stages
Expand Down Expand Up @@ -106,7 +106,7 @@ def run_audio_fleurs_benchmark( # noqa: PLR0913
)
)

results = pipeline.run(executor)
results = pipeline.run(executor_obj)

logger.success("Benchmark completed successfully")

Expand All @@ -126,6 +126,7 @@ def main() -> int:
parser.add_argument("--lang", default="hy_am", help="Language code")
parser.add_argument("--split", default="dev", help="Dataset split to use")
parser.add_argument("--wer-threshold", type=float, default=5.5, help="WER threshold for filtering")
parser.add_argument("--executor", default="xenna", choices=["xenna", "ray_data"], help="Executor to use")
parser.add_argument("--gpus", type=int, default=1, help="Number of GPUs to use")

args = parser.parse_args()
Expand Down
8 changes: 5 additions & 3 deletions benchmarking/scripts/dedup_removal_benchmark.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -124,9 +124,11 @@ def run_removal_benchmark( # noqa: PLR0913
for k, v in TaskPerfUtils.aggregate_task_metrics(workflow_run_result).items()
if k.endswith("_process_time_mean")
}
io_percentage = round(
(task_metrics["jsonl_reader"] + task_metrics["parquet_writer"]) * 100 / sum(task_metrics.values()), 2
)
reader_key = f"{input_filetype}_reader"

@ayushdg ayushdg Mar 11, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thanks!

writer_key = f"{output_filetype}_writer"
io_time = task_metrics.get(reader_key, 0) + task_metrics.get(writer_key, 0)
total_time = sum(task_metrics.values())
io_percentage = round(io_time * 100 / total_time, 2) if total_time > 0 else 0

return {
"metrics": {
Expand Down
17 changes: 7 additions & 10 deletions benchmarking/scripts/exact_dedup_identification_benchmark.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -173,6 +173,13 @@ def main() -> int:
logger.info("=== Exact Duplicate Identification Benchmark Starting ===")
logger.info(f"Arguments: {vars(args)}")

results = {
"params": vars(args),
"metrics": {
"is_success": False,
},
"tasks": [],
}
try:
results = run_exact_duplicate_identification_benchmark(
input_path=args.input_path,
Expand All @@ -187,16 +194,6 @@ def main() -> int:
rmm_pool_size=args.rmm_pool_size,
spill_memory_limit=args.spill_memory_limit,
)

except Exception as e:
print(f"Benchmark failed: {e}")
results = {
"params": vars(args),
"metrics": {
"is_success": False,
},
"tasks": [],
}
finally:
write_benchmark_results(results, args.benchmark_results_path)
Comment on lines 183 to 198

@greptile-apps greptile-apps Bot Mar 11, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Unhandled NameError can propagate past finally

run_exact_duplicate_identification_benchmark() has an internal except Exception block that sets success = False, but it does not define workflow_result when the exception occurs before line 77 (workflow_result = workflow.run(...)). The return at line 118 then references workflow_result, causing a NameError that escapes the inner function.

Previously, the outer except Exception in main() would have caught this and written clean fallback results. With the new try/finally-only pattern, the NameError propagates past finally (fallback results are still written correctly), but main() never returns — it exits via an unhandled Python exception traceback instead of a clean exit code.

If an unambiguous exit code is desired even in this edge case, consider adding a narrow guard in the inner function:

return {
    ...
    "tasks": workflow_result if success else [],
}


Expand Down
11 changes: 5 additions & 6 deletions benchmarking/scripts/fasttext_filter_benchmark.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -165,6 +165,11 @@ def main() -> int:
logger.info("=== FastText Filter Benchmark Starting ===")
logger.info(f"Arguments: {vars(args)}")

results = {
"params": vars(args),
"metrics": {"is_success": False},
"tasks": [],
}
try:
results = run_fasttext_filter_benchmark(
input_path=args.input_path,
Expand All @@ -176,12 +181,6 @@ def main() -> int:
fasttext_quality_model_path=args.fasttext_quality_model_path,
overrides=args.overrides,
)
except Exception:
results = {
"params": vars(args),
"metrics": {"is_success": False},
"tasks": [],
}
finally:
write_benchmark_results(results, args.benchmark_results_path)

Expand Down
24 changes: 7 additions & 17 deletions benchmarking/scripts/image_pipeline_benchmark.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -27,7 +27,6 @@
from loguru import logger
from utils import setup_executor, write_benchmark_results

from nemo_curator.core.client import RayClient
from nemo_curator.pipeline import Pipeline
from nemo_curator.stages.file_partitioning import FilePartitioningStage
from nemo_curator.stages.image.embedders.clip_embedder import ImageEmbeddingStage
Expand Down Expand Up @@ -172,9 +171,6 @@ def run_image_pipeline_benchmark(args: argparse.Namespace) -> dict[str, Any]:

def main() -> int:
"""Main entry point for image pipeline benchmark."""
ray_client = RayClient()
ray_client.start()

@ayushdg ayushdg Mar 10, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Isn't this a noop if cluster is started already?

@praateekmahajan praateekmahajan Mar 11, 2026

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Having RayClient inside benchmarking/scripts/*.py should be a no-op however for sake of consistency I would recommend users to use benchmarking/run.py to do the heavy lifting since it'll log details related to env / ray cluster / object store as well, which otherwise would be missed and wouldn't ensure if a run is comparable to another


parser = argparse.ArgumentParser(
description="Image curation pipeline benchmark with embedding generation and quality scoring"
)
Expand Down Expand Up @@ -255,23 +251,17 @@ def main() -> int:
logger.info("=== Image Pipeline Benchmark Starting ===")
logger.info(f"Arguments: {vars(args)}")

results = {
"params": vars(args),
"metrics": {
"is_success": False,
},
"tasks": [],
}
try:
results = run_image_pipeline_benchmark(args)

except Exception as e:
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)
ray_client.stop()

# Return proper exit code based on success
return 0 if results["metrics"]["is_success"] else 1
Expand Down
19 changes: 7 additions & 12 deletions benchmarking/scripts/modifier_benchmark.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -133,25 +133,20 @@ def main() -> int:
logger.info("=== Modify Benchmark Starting ===")
logger.info(f"Arguments: {vars(args)}")

results = {
"params": vars(args),
"metrics": {
"is_success": False,
},
"tasks": [],
}
try:
results = run_modify_benchmark(
input_path=args.input_path,
output_path=args.output_path,
executor_name=args.executor,
benchmark_results_path=args.benchmark_results_path,
)

except Exception as e:
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)

Expand Down
4 changes: 0 additions & 4 deletions benchmarking/scripts/multimodal_mint1t_benchmark.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -23,7 +23,6 @@
from loguru import logger
from utils import collect_parquet_output_metrics, setup_executor, validate_parquet_ordering, write_benchmark_results

from nemo_curator.core.client import RayClient
from nemo_curator.pipeline import Pipeline
from nemo_curator.stages.interleaved.io import InterleavedParquetWriterStage, WebdatasetReader
from nemo_curator.stages.interleaved.stages import InterleavedAspectRatioFilterStage
Expand Down Expand Up @@ -156,8 +155,6 @@ def main() -> int:
parser.set_defaults(materialize_on_write=False, materialize_on_read=False)
args = parser.parse_args()

ray_client = RayClient()
ray_client.start()
try:
results = run_benchmark(args)
except Exception as e:
Expand All @@ -170,7 +167,6 @@ def main() -> int:
}
finally:
write_benchmark_results(results, args.benchmark_results_path)
ray_client.stop()

return 0 if results["metrics"]["is_success"] else 1

Expand Down
19 changes: 7 additions & 12 deletions benchmarking/scripts/score_filter_benchmark.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -156,6 +156,13 @@ def main() -> int:
logger.info("=== ScoreFilter Benchmark Starting ===")
logger.info(f"Arguments: {vars(args)}")

results = {
"params": vars(args),
"metrics": {
"is_success": False,
},
"tasks": [],
}
try:
results = run_score_filter_benchmark(
input_path=args.input_path,
Expand All @@ -165,18 +172,6 @@ def main() -> int:
yaml_config=args.yaml_config,
overrides=args.overrides,
)

except Exception as e:
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)

Expand Down
11 changes: 8 additions & 3 deletions benchmarking/scripts/semdedup_identification_benchmark.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -106,6 +106,10 @@ def run_semdedup_identification_benchmark( # noqa: PLR0913
num_duplicates = workflow_run_result.metadata.get("num_duplicates")

# Calculate percentage times
kmeans_read_percent_time = None
kmeans_write_percent_time = None
kmeans_fit_predict_percent_time = None
kmeans_percent_time = None
pairwise_percent_time = None
if workflow_total_time:
# we get read / fit / write time from task_metrics
Expand All @@ -116,9 +120,10 @@ def run_semdedup_identification_benchmark( # noqa: PLR0913
# while this is just sum of mean time taken across actors across the three steps
_kmeans_time_taken = kmeans_read_time + kmeans_write_time + kmeans_fit_predict_time

kmeans_read_percent_time = round((kmeans_read_time / _kmeans_time_taken) * 100, 2)
kmeans_write_percent_time = round((kmeans_write_time / _kmeans_time_taken) * 100, 2)
kmeans_fit_predict_percent_time = round((kmeans_fit_predict_time / _kmeans_time_taken) * 100, 2)
if _kmeans_time_taken > 0:
kmeans_read_percent_time = round((kmeans_read_time / _kmeans_time_taken) * 100, 2)
kmeans_write_percent_time = round((kmeans_write_time / _kmeans_time_taken) * 100, 2)
kmeans_fit_predict_percent_time = round((kmeans_fit_predict_time / _kmeans_time_taken) * 100, 2)

kmeans_percent_time = round((kmeans_time / workflow_total_time) * 100, 2)
pairwise_percent_time = round((pairwise_time / workflow_total_time) * 100, 2)
Expand Down
19 changes: 7 additions & 12 deletions benchmarking/scripts/video_pipeline_benchmark.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -155,20 +155,15 @@ def main() -> int:
logger.info("=== Video Pipeline Benchmark Starting ===")
logger.info(f"Arguments: {vars(args)}")

results = {
"params": vars(args),
"metrics": {
"is_success": False,
},
"tasks": [],
}
try:
results = run_video_pipeline_benchmark(args)

except Exception as e:
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)

Expand Down
Loading