Merge branch 'master' into srinathk10/format_batches

Author: srinathk10

ci/docker/forge.Dockerfile

@@ -61,16 +61,16 @@ wget -qO- https://astral.sh/uv/install.sh | sudo env UV_UNMANAGED_INSTALL="/usr/
 
 mkdir -p /usr/local/python
 # Install Python 3.9 using uv
-uv python install --install-dir /usr/local/python 3.9
-uv python pin 3.9
+UV_PYTHON_VERSION=3.9
+uv python install --install-dir /usr/local/python "$UV_PYTHON_VERSION"
 
 export UV_PYTHON_INSTALL_DIR=/usr/local/python
-# Make Python 3.9 from uv the default by creating symlinks
-PYTHON39_PATH=$(uv python find 3.9)
-echo $PYTHON39_PATH
-ln -s $PYTHON39_PATH /usr/local/bin/python3.9
-ln -s $PYTHON39_PATH /usr/local/bin/python3
-ln -s $PYTHON39_PATH /usr/local/bin/python
+# Make Python from uv the default by creating symlinks
+UV_PYTHON_BIN="$(uv python find --no-project "$UV_PYTHON_VERSION")"
+echo "uv python binary location: $UV_PYTHON_BIN"
+ln -s "$UV_PYTHON_BIN" "/usr/local/bin/python${UV_PYTHON_VERSION}"
+ln -s "$UV_PYTHON_BIN" /usr/local/bin/python3
+ln -s "$UV_PYTHON_BIN" /usr/local/bin/python
 
 # As a convention, we pin all python packages to a specific version. This
 # is to to make sure we can control version upgrades through code changes.

python/ray/_private/test_utils.py

@@ -1065,20 +1065,6 @@ def fetch_prometheus_metric_timeseries(
     return samples_by_name
 
 
-def raw_metrics(info: RayContext) -> Dict[str, List[Any]]:
-    """Return prometheus metrics from a RayContext
-
-    Args:
-        info: Ray context returned from ray.init()
-
-    Returns:
-        Dict from metric name to a list of samples for the metrics
-    """
-    metrics_page = "localhost:{}".format(info.address_info["metrics_export_port"])
-    print("Fetch metrics from", metrics_page)
-    return fetch_prometheus_metrics([metrics_page])
-
-
 def raw_metric_timeseries(
     info: RayContext, result: PrometheusTimeseries
 ) -> Dict[str, List[Any]]:

python/ray/data/grouped_data.py

@@ -1,3 +1,4 @@
+from collections.abc import Iterator as IteratorABC
 from functools import partial
 from typing import Any, Callable, Dict, Iterable, Iterator, List, Optional, Tuple, Union
 
@@ -538,7 +539,10 @@ def std(
 
 
 def _apply_udf_to_groups(
-    udf: Callable[[DataBatch, ...], DataBatch],
+    udf: Union[
+        Callable[[DataBatch, ...], DataBatch],
+        Callable[[DataBatch, ...], Iterator[DataBatch]],
+    ],
     block: Block,
     keys: List[str],
     batch_format: Optional[str],
@@ -548,7 +552,8 @@ def _apply_udf_to_groups(
     """Apply UDF to groups of rows having the same set of values of the specified
     columns (keys).
 
-    NOTE: This function is defined at module level to avoid capturing closures and make it serializable."""
+    NOTE: This function is defined at module level to avoid capturing closures and make it serializable.
+    """
     block_accessor = BlockAccessor.for_block(block)
 
     boundaries = block_accessor._get_group_boundaries_sorted(keys)
@@ -560,7 +565,17 @@ def _apply_udf_to_groups(
         # Convert corresponding block of each group to batch format here,
         # because the block format here can be different from batch format
         # (e.g. block is Arrow format, and batch is NumPy format).
-        yield udf(group_block_accessor.to_batch_format(batch_format), *args, **kwargs)
+        result = udf(
+            group_block_accessor.to_batch_format(batch_format), *args, **kwargs
+        )
+
+        # Check if the UDF returned an iterator/generator.
+        if isinstance(result, IteratorABC):
+            # If so, yield each item from the iterator.
+            yield from result
+        else:
+            # Otherwise, yield the single result.
+            yield result
 
 
 # Backwards compatibility alias.

python/ray/data/tests/test_groupby_e2e.py

@@ -1,7 +1,7 @@
 import itertools
 import random
 import time
-from typing import Optional
+from typing import Iterator, Optional
 
 import numpy as np
 import pandas as pd
@@ -1142,6 +1142,40 @@ def func(x, y):
     assert "MapBatches(func)" in ds.__repr__()
 
 
+def test_map_groups_generator_udf(ray_start_regular_shared_2_cpus):
+    """
+    Tests that map_groups supports UDFs that return generators (iterators).
+    """
+    ds = ray.data.from_items(
+        [
+            {"group": 1, "data": 10},
+            {"group": 1, "data": 20},
+            {"group": 2, "data": 30},
+        ]
+    )
+
+    def generator_udf(df: pd.DataFrame) -> Iterator[pd.DataFrame]:
+        # For each group, yield two DataFrames.
+        # 1. A DataFrame where 'data' is multiplied by 2.
+        yield df.assign(data=df["data"] * 2)
+        # 2. A DataFrame where 'data' is multiplied by 3.
+        yield df.assign(data=df["data"] * 3)
+
+    # Apply the generator UDF to the grouped data.
+    result_ds = ds.groupby("group").map_groups(generator_udf)
+
+    # The final dataset should contain all results from all yields.
+    # Group 1 -> data: [20, 40] and [30, 60]
+    # Group 2 -> data: [60] and [90]
+    expected_data = sorted([20, 40, 30, 60, 60, 90])
+
+    # Collect and sort the actual data to ensure correctness regardless of order.
+    actual_data = sorted([row["data"] for row in result_ds.take_all()])
+
+    assert actual_data == expected_data
+    assert result_ds.count() == 6
+
+
 if __name__ == "__main__":
     import sys
 

python/ray/data/tests/test_join.py

@@ -8,7 +8,7 @@
 import ray
 from ray._private.arrow_utils import get_pyarrow_version
 from ray.data._internal.logical.operators.join_operator import JoinType
-from ray.data._internal.util import MiB
+from ray.data._internal.util import MiB, rows_same
 from ray.data.context import DataContext
 from ray.data.dataset import Dataset
 from ray.exceptions import RayTaskError
@@ -263,11 +263,7 @@ def test_simple_self_join(ray_start_regular_shared_2_cpus, left_suffix, right_su
 
         assert 'Field "double" exists 2 times' in str(exc_info.value.cause)
     else:
-
-        joined_pd = pd.DataFrame(joined.take_all())
-
-        # Sort resulting frame and reset index (to be able to compare with expected one)
-        joined_pd_sorted = joined_pd.sort_values(by=["id"]).reset_index(drop=True)
+        joined_pd = joined.to_pandas()
 
         # Join using Pandas (to assert against)
         expected_pd = doubles_pd.join(
@@ -278,7 +274,7 @@ def test_simple_self_join(ray_start_regular_shared_2_cpus, left_suffix, right_su
             rsuffix=right_suffix,
         ).reset_index(drop=True)
 
-        pd.testing.assert_frame_equal(expected_pd, joined_pd_sorted)
+        assert rows_same(expected_pd, joined_pd), "Expected contents to be same"
 
 
 def test_invalid_join_config(ray_start_regular_shared_2_cpus):

python/ray/tests/test_memory_pressure.py

@@ -13,7 +13,7 @@
 )
 from ray._private.grpc_utils import init_grpc_channel
 from ray._private.state_api_test_utils import verify_failed_task
-from ray._private.test_utils import raw_metrics
+from ray._private.test_utils import PrometheusTimeseries, raw_metric_timeseries
 from ray._private.utils import get_used_memory
 from ray.util.state.state_manager import StateDataSourceClient
 
@@ -118,8 +118,10 @@ def get_additional_bytes_to_reach_memory_usage_pct(pct: float) -> int:
     return bytes_needed
 
 
-def has_metric_tagged_with_value(addr, tag, value) -> bool:
-    metrics = raw_metrics(addr)
+def has_metric_tagged_with_value(
+    addr, tag, value, timeseries: PrometheusTimeseries
+) -> bool:
+    metrics = raw_metric_timeseries(addr, timeseries)
     for name, samples in metrics.items():
         for sample in samples:
             if tag in set(sample.labels.values()) and sample.value == value:
@@ -145,13 +147,15 @@ def test_restartable_actor_throws_oom_error(ray_with_memory_monitor, restartable
     with pytest.raises(ray.exceptions.OutOfMemoryError):
         ray.get(leaker.allocate.remote(bytes_to_alloc, memory_monitor_refresh_ms * 3))
 
+    timeseries = PrometheusTimeseries()
     wait_for_condition(
         has_metric_tagged_with_value,
         timeout=10,
         retry_interval_ms=100,
         addr=addr,
         tag="MemoryManager.ActorEviction.Total",
         value=2.0 if restartable else 1.0,
+        timeseries=timeseries,
     )
 
     wait_for_condition(
@@ -161,6 +165,7 @@ def test_restartable_actor_throws_oom_error(ray_with_memory_monitor, restartable
         addr=addr,
         tag="Leaker.__init__",
         value=2.0 if restartable else 1.0,
+        timeseries=timeseries,
     )
 
 
@@ -180,13 +185,15 @@ def test_restartable_actor_oom_retry_off_throws_oom_error(
     with pytest.raises(ray.exceptions.OutOfMemoryError) as _:
         ray.get(leaker.allocate.remote(bytes_to_alloc, memory_monitor_refresh_ms * 3))
 
+    timeseries = PrometheusTimeseries()
     wait_for_condition(
         has_metric_tagged_with_value,
         timeout=10,
         retry_interval_ms=100,
         addr=addr,
         tag="MemoryManager.ActorEviction.Total",
         value=2.0,
+        timeseries=timeseries,
     )
     wait_for_condition(
         has_metric_tagged_with_value,
@@ -195,6 +202,7 @@ def test_restartable_actor_oom_retry_off_throws_oom_error(
         addr=addr,
         tag="Leaker.__init__",
         value=2.0,
+        timeseries=timeseries,
     )
 
 
@@ -210,13 +218,15 @@ def test_non_retryable_task_killed_by_memory_monitor_with_oom_error(
     with pytest.raises(ray.exceptions.OutOfMemoryError) as _:
         ray.get(allocate_memory.options(max_retries=0).remote(bytes_to_alloc))
 
+    timeseries = PrometheusTimeseries()
     wait_for_condition(
         has_metric_tagged_with_value,
         timeout=10,
         retry_interval_ms=100,
         addr=addr,
         tag="MemoryManager.TaskEviction.Total",
         value=1.0,
+        timeseries=timeseries,
     )
     wait_for_condition(
         has_metric_tagged_with_value,
@@ -225,6 +235,7 @@ def test_non_retryable_task_killed_by_memory_monitor_with_oom_error(
         addr=addr,
         tag="allocate_memory",
         value=1.0,
+        timeseries=timeseries,
     )
 
 
@@ -372,13 +383,15 @@ def test_task_oom_no_oom_retry_fails_immediately(
             )
         )
 
+    timeseries = PrometheusTimeseries()
     wait_for_condition(
         has_metric_tagged_with_value,
         timeout=10,
         retry_interval_ms=100,
         addr=addr,
         tag="MemoryManager.TaskEviction.Total",
         value=1.0,
+        timeseries=timeseries,
     )
     wait_for_condition(
         has_metric_tagged_with_value,
@@ -387,6 +400,7 @@ def test_task_oom_no_oom_retry_fails_immediately(
         addr=addr,
         tag="allocate_memory",
         value=1.0,
+        timeseries=timeseries,
     )
 
 
@@ -411,13 +425,15 @@ def test_task_oom_only_uses_oom_retry(
             )
         )
 
+    timeseries = PrometheusTimeseries()
     wait_for_condition(
         has_metric_tagged_with_value,
         timeout=10,
         retry_interval_ms=100,
         addr=addr,
         tag="MemoryManager.TaskEviction.Total",
         value=task_oom_retries + 1,
+        timeseries=timeseries,
     )
     wait_for_condition(
         has_metric_tagged_with_value,
@@ -426,6 +442,7 @@ def test_task_oom_only_uses_oom_retry(
         addr=addr,
         tag="allocate_memory",
         value=task_oom_retries + 1,
+        timeseries=timeseries,
     )
 
 
@@ -502,6 +519,7 @@ def infinite_retry_task():
             time.sleep(5)
 
     with ray.init() as addr:
+        timeseries = PrometheusTimeseries()
         with pytest.raises(ray.exceptions.OutOfMemoryError) as _:
             ray.get(infinite_retry_task.remote())
 
@@ -512,6 +530,7 @@ def infinite_retry_task():
             addr=addr,
             tag="MemoryManager.TaskEviction.Total",
             value=1.0,
+            timeseries=timeseries,
         )