[Data] Fix SF 100 release test

[owenowenisme]

Description

Fix aggregator actors OOM for SF‑100 map_groups on an autoscaling cluster.

Main problems:

• Even with “scheduling_strategy”: “SPREAD”, autoscaling causes greedy placement on the same nodes (node fills → autoscale → schedule on new node → autoscale), so actors cluster on the first few nodes.
• Available memory calculations didn’t exclude the Object Store, a known issue ([Core] Actor/Task Memory resource allocation should be separate from Object Store allocation  #49344).
• Aggregator resource miscalculation: we observed ≈1.7 GiB spikes, but only allocated 1.13 GiB per actor.Previously, aggregator memory was computed as dataset_size/num_aggregators + dataset_size/num_partitions (input + output).
  This is incorrect because partition size is determined by num_partitions, not num_aggregators, and some aggregators receive multiple partitions via partition_id % num_aggregators, leading to under‑allocation.For SF‑100, dataset size ≈82.1 GiB.
  The old formula gives (82.1/128) + (82.1/200) ≈ 1.13 GiB. For an aggregator receiving 2 partitions, the max memory should be (82.1/200 per‑partition) × 2 (input + output) × 2 (ceil(num_partitions/num_aggregators)) ≈ 1.76 GiB, matching the observed ≈1.7 GiB peak.

This fails only on autoscaling clusters because the resource budget is tighter than fixed‑size clusters (≈18 nodes vs 32 nodes), and aggregators aren’t well spread.

Job URL: https://console.anyscale-staging.com/o/anyscale-internal/jobs/prodjob_z2pll84knil1iexjhwfcmislan

Fixes

• Correct the calculation of aggregators from dataset_size/num_aggregators + dataset_size/num_partitions to (82.1/ 200)(per partition size) * 2( Input and output) * 2 (ceil (num_partition/num_aggregators))
  -Compute CPU from real available memory (exclude Object Store proportion). (May not be the best way but this is what we can do right now.)

• Manually exclude Object Store Memory proportion when calculating target_num_cpus for aggregators

        worker_heap_memory_proportion = 1 - DEFAULT_OBJECT_STORE_MEMORY_PROPORTION
        target_num_cpus = min(
            cap,
            estimated_aggregator_memory_required
            / (4 * GiB * worker_heap_memory_proportion),
        )

Related issues

Closes #58734

[alexeykudinkin]

Let's keep these line-items separate as they help estimate individual components

[owenowenisme]

Not sure what this mean, but I restore the info to show shuffle & output memory info

[alexeykudinkin]

Let's preserve the structure of estimating individual components (shuffle, output, working memory)

[alexeykudinkin]

What i meant is following

• We still want to estimate total memory required based on corresponding components (shuffle, output, etc)
• Core change you're making is moving from dataset_size / num_aggregators to partition_byte_size * max_partitions, which is just a rounded version

Let's keep the overall structure and just make the actual change you want to make by updating shuffle component estimate.

BTW, we can add some skew factor (like say 30% extra as padding against skews)

[alexeykudinkin]

Previous comment got lost -- make this a class var so that we can patch it if needed

[alexeykudinkin]

Suggested change

	SKEW_FACTOR = 1.3
	SHUFFLE_AGGREGATOR_MEMORY_ESTIMATE_SKEW_FACTOR = 1.3

[cursor]

Bug: Test expected values not updated for new memory formula

The test_hash_shuffle_operator_remote_args test cases have expected memory values that match the old memory calculation formula, but HashShuffleOperator._estimate_aggregator_memory_allocation was changed to use a new formula. The new formula multiplies partition_byte_size * max_partitions_for_aggregator for both input and output (plus a 1.3 skew factor), resulting in significantly larger memory values. For example, Case 1 expects 671088640 but the new formula produces approximately 1,395,864,372. These tests will fail until the expected values are updated.

python/ray/data/tests/test_hash_shuffle.py#L392-L393

ray/python/ray/data/tests/test_hash_shuffle.py

Lines 392 to 393 in 3f4a317

	"num_cpus": 0.16,
	"memory": 671088640,

python/ray/data/tests/test_hash_shuffle.py#L440-L441

ray/python/ray/data/tests/test_hash_shuffle.py

Lines 440 to 441 in 3f4a317

	"num_cpus": 0.16, # ~0.6Gb / 4Gb = ~0.16
	"memory": 687865856,

 

---