perf(ray-actor-pool): balance weighted RAFT tasks

[praateekmahajan]

Summary

• let file partitioning opt in to emitting a task weight; semantic KMeans uses input bytes
• weight fuzzy connected-components inputs by edge count
• balance weighted RAFT tasks deterministically with largest-first assignment
• preserve existing equal-count splitting when weights are absent

Exact deduplication and fuzzy LSH remain unchanged because they use the shuffle path rather than RAFT batching.

Motivation

Equal task counts can produce uneven per-GPU input sizes, leaving collective KMeans or connected-components workers waiting on the most heavily loaded actor.

Alternatives considered

• deterministic randomization: smaller, but probabilistic
• blocksize grouping: simple, but requires tuning and does not handle oversized files

Testing

• focused unit tests and pre-commit passed
• GPU-specific fuzzy unit tests were not run locally because cudf is unavailable

EOS benchmark

Run viewer and Comparision
All three relevant 8xH100 entries passed. Compared with the same-day main nightly:

Entry	Main	PR	Delta
fuzzy dedup	446.8s	449.1s	+0.5%
semdedup 20% dataset	214.9s	219.1s	+2.0%
semdedup 20% fit	1138.1s	1141.1s	+0.3%

There is no material speedup or slowdown. The full semdedup run reduced per-actor input-row standard deviation from 51,585 to 25, but the baseline imbalance was already only about 0.1% of the mean, so elapsed time remained unchanged.
This is directionally the intended behavior for skewed inputs while remaining neutral on already-balanced data.

[copy-pr-bot]

Auto-sync is disabled for draft pull requests in this repository. Workflows must be run manually.

Contributors can view more details about this message here.

[praateekmahajan]

/claude review
@greptileai review

[praateekmahajan]

/ok to test d6460ce

[greptile-apps]

Greptile Summary

The PR adds optional task-weight metadata and uses deterministic largest-first assignment to balance coordinated RAFT work while retaining equal-count splitting when weights are unavailable.

• KMeans partitions are weighted by input bytes.
• Fuzzy connected-components inputs are weighted by generated edge count.
• Weighted batches use accumulated weight with task count as a deterministic tie-breaker.
• Focused tests cover balancing, fallback behavior, zero-weight distribution, and metadata generation.

Confidence Score: 5/5

The PR appears safe to merge.

No blocking failure remains.

Important Files Changed

Filename	Overview
nemo_curator/backends/ray_actor_pool/executor.py	Adds validated weighted RAFT batching with deterministic load and task-count tie-breaking while preserving the existing unweighted path.
nemo_curator/stages/file_partitioning.py	Optionally records known partition byte totals as task metadata and omits the weight when size information is unavailable.
nemo_curator/stages/deduplication/semantic/kmeans.py	Enables file-size weighting for single-file KMeans input partitions.
nemo_curator/stages/deduplication/fuzzy/buckets_to_edges.py	Records generated edge count as scheduling weight for downstream coordinated processing.
tests/backends/ray_actor_pool/test_executor.py	Covers weighted assignment, legacy count splitting, and balanced distribution of zero-weight tasks.

Flowchart

%%{init: {'theme': 'neutral'}}%%
flowchart LR
    A[Partitioned input tasks] --> B{All tasks have weights?}
    B -- Yes --> C[Sort tasks largest-first]
    C --> D[Assign each task to least-loaded RAFT batch]
    D --> E[Submit balanced batches to RAFT actors]
    B -- No --> F[Split tasks equally by count]
    F --> E

Loading

_{Reviews (2): Last reviewed commit: "fix(ray-actor-pool): distribute zero-wei..." | Re-trigger Greptile}

[praateekmahajan]

/claude review
@greptileai review