Add math modality benchmarking support

[ronjer30]

Benchmark changes:

• Add math_pipeline_benchmark.py to run benchmark
• Add prepare_math_benchmark_data.py to enrich FineMath-4+ with Common Crawl binary content for offline benchmarking
• Add math_preprocess, math_preprocess_classifier, and math_preprocess_llm_cleanup entries to nightly-benchmark.yaml
• Install lynx in Dockerfiles for HTML-to-text extraction

Stage fixes:

• Fix deepcopy/pickle crash in MathContentExtractor by adding getstate/setstate to strip unpickleable threading.Lock and magic.Magic before serialization
• Add S3 transport to CommonCrawlWARCReader via boto3 as an alternative to HTTPS, activated by use_s3=True or CC_USE_S3 env var, this addresses latency and throttling issues revealed during benchmarking.
• Add boto3 to math_cpu dependency

Description

Usage

# Add snippet demonstrating usage

Checklist

• I am familiar with the Contributing Guide.
• New or Existing tests cover these changes.
• The documentation is up to date with these changes.

[greptile-apps]

Four separate full dataset scans — consolidate into one pass

compute_extraction_metrics calls .count() four times on filtered Ray datasets. Each .count() triggers a full scan of all output JSONL files (read → filter → count). For any non-trivial benchmark output this quadruples the I/O cost of the post-processing step and inflates the wall-clock time visible in the results.

A single aggregation pass would be both faster and more accurate:

def compute_extraction_metrics(output_dir: str) -> dict:
    metrics = {}
    try:
        jsonl_files = get_all_file_paths_under(output_dir, keep_extensions=[".jsonl"])
        if not jsonl_files:
            return metrics

        ds = ray.data.read_json(jsonl_files)

        def _count_types(batch: dict) -> dict:
            types = batch.get("type", [])
            texts = batch.get("text", [])
            return {
                "html": [sum(1 for t in types if t == "html")],
                "text": [sum(1 for t in types if t == "text")],
                "notebook": [sum(1 for t in types if t == "notebook")],
                "html_empty": [
                    sum(1 for t, tx in zip(types, texts) if t == "html" and not (tx or "").strip())
                ],
            }

        agg = ds.map_batches(_count_types, batch_format="numpy").sum(
            ["html", "text", "notebook", "html_empty"]
        )
        metrics["type_html_count"] = int(agg["html"])
        metrics["type_text_count"] = int(agg["text"])
        metrics["type_notebook_count"] = int(agg["notebook"])
        metrics["html_empty_text_count"] = int(agg["html_empty"])

    except Exception as e:
        logger.warning(f"Could not compute extraction metrics: {e}")

    return metrics

(The exact Ray aggregation API may need adjustment to match the version in use, but the principle of a single pass applies.)

[greptile-apps]

boto3 version pin missing in math_cpu — inconsistent with inference-server

The inference-server extra pins boto3>=1.35, but math_cpu (and transitively math_cuda12) adds boto3 with no version constraint. When only the math_cpu extra is installed, pip/uv may resolve an older boto3 that lacks features or fixes present in 1.35+.

Suggested change

	# Math Curation Dependencies
	math_cpu = [
	"nemo_curator[text_cpu]", # Math examples use text processing utilities
	"boto3>=1.35", # S3 range requests for Common Crawl WARC fetching

[greptile-apps]

Redundant full-dataset scan in compute_llm_cleanup_metrics

ds.count() on line 220 triggers a complete scan of all output JSONL files. Then iter_rows() on line 223 immediately triggers a second full scan of the same dataset. The total variable will equal len(text_lengths) after the list comprehension completes, so the first scan is entirely wasteful.

Suggested change

	ds = ray.data.read_json(jsonl_files)
	total = ds.count()
	metrics["num_chunks_processed"] = total
	text_lengths = [len(row.get("cleaned_text") or "") for row in ds.select_columns(["cleaned_text"]).iter_rows()]
	if text_lengths:
	metrics["avg_output_text_length"] = sum(text_lengths) / len(text_lengths)
	ds = ray.data.read_json(jsonl_files)
	text_lengths = [len(row.get("cleaned_text") or "") for row in ds.select_columns(["cleaned_text"]).iter_rows()]
	total = len(text_lengths)
	metrics["num_chunks_processed"] = total

[greptile-apps]

All classifier scores loaded into memory at once

The list comprehension on line 194 materializes the entire finemath_int_scores column into a Python list before any aggregation is done. For a large benchmark output the scores list could consume significant memory. The subsequent sum(1 for s in scores if s == i) passes also re-iterate the same in-memory list six times.

Using iter_batches with incremental counters avoids holding all scores at once:

score_counts = [0] * 6
score_sum = 0
total = 0
for batch in ds.select_columns(["finemath_int_scores"]).iter_batches(batch_format="numpy"):
    for s in batch["finemath_int_scores"]:
        score_counts[min(s, 5)] += 1
        score_sum += s
        total += 1

if total > 0:
    metrics["mean_finemath_score"] = score_sum / total
    for i in range(6):
        metrics[f"score_distribution_{i}"] = score_counts[i]
    metrics["docs_score_ge_3"] = sum(score_counts[MIN_HIGH_QUALITY_SCORE:])

[greptile-apps]

Deepcopy test doesn't verify lock independence

The test confirms the cloned lock is not None, but the key property of __getstate__/__setstate__ is that the cloned object gets a fresh, independent lock — not the same object shared with the original. Without asserting cloned._lock is not extractor._lock, a naive implementation that simply copies the reference (rather than creating a new lock) would still pass this test.

Suggested change

	def test_deepcopy_extractor_with_lock(self) -> None:
	extractor = MathContentExtractor()
	cloned = copy.deepcopy(extractor)
	assert cloned is not extractor
	assert cloned._lock is not None
	def test_deepcopy_extractor_with_lock(self) -> None:
	extractor = MathContentExtractor()
	cloned = copy.deepcopy(extractor)
	assert cloned is not extractor
	assert cloned._lock is not None
	assert cloned._lock is not extractor._lock

[greptile-apps]

extraction_success_rate silently absent on complete extraction failure

The metric is only added when num_output_documents > 0. If all documents fail extraction (a real failure scenario worth benchmarking), num_output_documents is 0, so the condition and num_output_documents > 0 is false and extraction_success_rate is never written to metrics. The YAML benchmark config lists extraction_success_rate as an expected additional_metrics entry, so a missing key may cause silent mis-reporting or downstream errors in the benchmark framework.

The correct rate when output is zero is 0.0 (100% failure rate), not an absent metric. Only the division-by-zero case (num_input_documents == 0) warrants omitting the metric:

Suggested change

	}
	if num_input_documents > 0 and num_output_documents > 0:
	metrics["extraction_success_rate"] = num_output_documents / num_input_documents
	if num_input_documents > 0:
	metrics["extraction_success_rate"] = num_output_documents / num_input_documents

[greptile-apps]

iter_rows() materializes one Python dict per row

ds.iter_rows() yields one Python dict per document before the list comprehension extracts only the integer length. For a large LLM-cleanup run with many chunks this is noticeably slower and more memory-intensive than iter_batches. Since the previous double-scan issue (ds.count() + iter_rows()) has been fixed, consider converting the remaining iter_rows() call to an iter_batches loop to stay consistent with how compute_classifier_metrics was already refactored:

total = 0
sum_len = 0
no_content = 0
for batch in ds.iter_batches(batch_format="numpy"):
    for text in batch.get("cleaned_text", []):
        length = len(text or "")
        sum_len += length
        if length == 0:
            no_content += 1
        total += 1

metrics["num_chunks_processed"] = total
if total:
    metrics["avg_output_text_length"] = sum_len / total
metrics["no_useful_content_count"] = no_content

[greptile-apps]

Module-level spec_from_file_location can return None

importlib.util.spec_from_file_location returns None when the target file is not found (e.g. when _TUTORIAL_PATH does not exist in the deployed benchmarking environment). The very next line passes None directly into module_from_spec, which raises:

AttributeError: 'NoneType' object has no attribute 'submodule_search_locations'

Because this code runs at module level, the entire script fails to import — not just at runtime when the benchmark is executed — producing a deeply confusing traceback that does not mention the missing file. A guard with an early, descriptive error would make the failure actionable:

_spec = importlib.util.spec_from_file_location("cc_index_lookup", _TUTORIAL_PATH)
if _spec is None:
    raise FileNotFoundError(
        f"Could not load tutorial module from {_TUTORIAL_PATH}. "
        "Ensure the tutorials directory is present relative to the benchmarking scripts."
    )
_mod = importlib.util.module_from_spec(_spec)
_spec.loader.exec_module(_mod)

[greptile-apps]

numpy scalar in metrics dict may cause JSON serialization failure

iter_batches(batch_format="numpy") yields numpy.int64 scalars for each s. The accumulator score_sum starts as a Python int but becomes numpy.int64 after the first score_sum += s. Consequently score_sum / total produces a numpy.float64, and metrics["mean_finemath_score"] is stored as a numpy.float64.

If write_benchmark_results serializes metrics with the standard json.dumps, this fails:

TypeError: Object of type float64 is not JSON serializable

The simplest fix is to cast at the assignment site:

Suggested change

	for batch in ds.iter_batches(batch_format="numpy"):
	for s in batch["finemath_int_scores"]:
	score_counts[min(s, 5)] += 1
	score_sum += s
	total += 1
	if total > 0:
	metrics["mean_finemath_score"] = score_sum / total
	for batch in ds.iter_batches(batch_format="numpy"):
	for s in batch["finemath_int_scores"]:
	score_counts[min(int(s), 5)] += 1
	score_sum += int(s)
	total += 1
	if total > 0:
	metrics["mean_finemath_score"] = score_sum / total

Similarly, totals[k] += row[k] in compute_extraction_metrics and compute_llm_cleanup_metrics accumulates numpy scalars from iter_rows() into the totals dict, which then flow directly into metrics. Casting each row[k] to int before the += would prevent the same issue in those functions.

[greptile-apps]

RuntimeError from missing boto3 silently swallowed

The outer except Exception block catches all exceptions including the RuntimeError raised by _get_s3_client() when boto3 is not installed. Because self._s3_client is never set when the import fails, every call to _read_warc_record_s3 will:

1. Enter _get_s3_client(), acquire the lock, attempt import boto3, fail, and raise RuntimeError.
2. Have that RuntimeError caught here and downgraded to a logger.warning.
3. Return None.

This means the entire output dataset will be silently empty — with thousands of "S3 fetch failed for …: boto3 is not installed" warnings — rather than producing a clear, early failure. The test_use_s3_without_boto3_raises_clear_error test verifies that _get_s3_client() raises correctly, but nothing tests the end-to-end behaviour through _read_warc_record_s3.

The fix is to let non-transient configuration errors propagate out of _read_warc_record_s3:

except RuntimeError:
    raise  # Re-raise configuration errors (e.g. missing boto3) immediately
except Exception as e:  # noqa: BLE001
    logger.warning(f"S3 fetch failed for {filename}: {e}")
    return None

[sarahyurick]

It looks like there is no option for "xenna" vs "ray_data" for this one?

[ronjer30]

Thanks, explicitly added "xenna" as the default, not tested on ray_data.

[sarahyurick]

Normally we don't want individual benchmarking scripts to initialize the Ray client on their own (#1593). I am wondering if there is a better way to do this.

[ronjer30]

Agreed, removed this,ray.init is needed to set the env variables correctly before ray.put is called

[sarahyurick]

It can be a top-level import?