[data] Include total memory in aggregator memory calculation (#58002)

## Description
Currently, the aggregators don't consider total available memory usage
in scheduling, and will default to 1GiB if the `estimated_dataset_bytes`
is None. This can be troublesome for smaller machines. See
#57979 for details. This PR
uses total available memory to calculate memory per aggregator.

## Related issues
Fixes #57979

## Additional information
> Optional: Add implementation details, API changes, usage examples,
screenshots, etc.

---------

Signed-off-by: iamjustinhsu <[email protected]>

Author: iamjustinhsu

python/ray/data/_internal/execution/operators/hash_shuffle.py

@@ -1090,11 +1090,19 @@ def _get_default_aggregator_ray_remote_args(
             )
         else:
             # NOTE: In cases when we're unable to estimate dataset size,
-            #       we simply fallback to request
-            #       ``DEFAULT_HASH_SHUFFLE_AGGREGATOR_MEMORY_ALLOCATION`` worth of
-            #       memory for every Aggregator
-            estimated_aggregator_memory_required = (
-                DEFAULT_HASH_SHUFFLE_AGGREGATOR_MEMORY_ALLOCATION
+            #       we simply fallback to request the minimum of:
+            #       - conservative 50% of total available memory for a join operation.
+            #       - ``DEFAULT_HASH_SHUFFLE_AGGREGATOR_MEMORY_ALLOCATION`` worth of
+            #       memory for every Aggregator.
+
+            max_memory_per_aggregator = (
+                total_available_cluster_resources.memory / num_aggregators
+            )
+            modest_memory_per_aggregator = max_memory_per_aggregator / 2
+
+            estimated_aggregator_memory_required = min(
+                modest_memory_per_aggregator,
+                DEFAULT_HASH_SHUFFLE_AGGREGATOR_MEMORY_ALLOCATION,
             )
 
         remote_args = {

python/ray/data/_internal/execution/util.py

@@ -36,7 +36,7 @@ def make_ref_bundles(simple_data: List[List[Any]]) -> List["RefBundle"]:
     return output
 
 
-memory_units = ["B", "KB", "MB", "GB", "TB", "PB"]
+memory_units = ["B", "KiB", "MiB", "GiB", "TiB", "PiB"]
 
 
 def memory_string(num_bytes: float) -> str:

python/ray/data/tests/test_executor_resource_management.py

@@ -53,11 +53,11 @@ def test_execution_resources(ray_start_10_cpus_shared):
     )
     assert (
         repr(r4)
-        == "ExecutionResources(cpu=1, gpu=1, object_store_memory=100.0MB, memory=0.0B)"
+        == "ExecutionResources(cpu=1, gpu=1, object_store_memory=100.0MiB, memory=0.0B)"
     )
     assert (
         repr(r5)
-        == "ExecutionResources(cpu=1, gpu=1, object_store_memory=1.0GB, memory=64.0MB)"
+        == "ExecutionResources(cpu=1, gpu=1, object_store_memory=1.0GiB, memory=64.0MiB)"
     )
     assert (
         repr(unlimited)
@@ -66,8 +66,8 @@ def test_execution_resources(ray_start_10_cpus_shared):
 
     # Test object_store_memory_str.
     assert r3.object_store_memory_str() == "0.0B"
-    assert r4.object_store_memory_str() == "100.0MB"
-    assert r5.object_store_memory_str() == "1.0GB"
+    assert r4.object_store_memory_str() == "100.0MiB"
+    assert r5.object_store_memory_str() == "1.0GiB"
     assert unlimited.object_store_memory_str() == "inf"
 
     # Test add.