Skip to content

Commit 6259f5c

Browse files
YHC66ntkathole
YHC66
authored and
ntkathole
committed
fix: Use a valid ShuffleStrategy for the Ray DatasetContext
Feast set `DatasetContext.shuffle_strategy = "sort"`, which is not a member of `ray.data.context.ShuffleStrategy`. Up to Ray 2.58 the setter stored the raw string and Ray fell through to the pull-based sort shuffle. Ray 2.59 coerces the value with `ShuffleStrategy(value)`, so `ensure_ray_initialized()` and `RayResourceManager.configure_ray_context()` now raise `ValueError: 'sort' is not a valid ShuffleStrategy`. Use `ShuffleStrategy.SORT_SHUFFLE_PULL_BASED`, which keeps the behavior on older Ray versions and is accepted by Ray 2.59. Signed-off-by: Yihang Chen <[email protected]>
1 parent b9c65f1 commit 6259f5c

3 files changed

Lines changed: 45 additions & 7 deletions

File tree

‎sdk/python/feast/infra/offline_stores/contrib/ray_offline_store/ray.py‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -12,7 +12,7 @@
1212
import pyarrow as pa
1313
import ray
1414
from ray.data import Dataset
15-
from ray.data.context import DatasetContext
15+
from ray.data.context import DatasetContext, ShuffleStrategy
1616

1717
from feast.data_source import DataSource
1818
from feast.dataframe import DataFrameEngine, FeastDataFrame
@@ -460,7 +460,7 @@ def configure_ray_context(self) -> None:
460460
else 2
461461
)
462462
ctx.max_parallelism = self.available_cpus * multiplier
463-
ctx.shuffle_strategy = "sort" # type: ignore
463+
ctx.shuffle_strategy = ShuffleStrategy.SORT_SHUFFLE_PULL_BASED
464464
ctx.enable_tensor_extension_casting = False
465465

466466
if not getattr(self.config, "enable_ray_logging", False):

‎sdk/python/feast/infra/ray_initializer.py‎

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -25,7 +25,7 @@
2525
from typing import Any, Dict, List, Optional, Union
2626

2727
import ray
28-
from ray.data.context import DatasetContext
28+
from ray.data.context import DatasetContext, ShuffleStrategy
2929

3030
logger = logging.getLogger(__name__)
3131

@@ -698,7 +698,7 @@ def _initialize_local_ray(config: Any, enable_logging: bool = False) -> None:
698698

699699
# Configure DatasetContext
700700
ctx = DatasetContext.get_current()
701-
ctx.shuffle_strategy = "sort" # type: ignore
701+
ctx.shuffle_strategy = ShuffleStrategy.SORT_SHUFFLE_PULL_BASED
702702
ctx.enable_tensor_extension_casting = False
703703
if hasattr(ctx, "enable_arrow_backed_pandas_conversion"):
704704
ctx.enable_arrow_backed_pandas_conversion = False
@@ -760,7 +760,7 @@ def _initialize_remote_ray(config: Any, enable_logging: bool = False) -> None:
760760

761761
# Configure DatasetContext
762762
ctx = DatasetContext.get_current()
763-
ctx.shuffle_strategy = "sort" # type: ignore
763+
ctx.shuffle_strategy = ShuffleStrategy.SORT_SHUFFLE_PULL_BASED
764764
ctx.enable_tensor_extension_casting = False
765765
if hasattr(ctx, "enable_arrow_backed_pandas_conversion"):
766766
ctx.enable_arrow_backed_pandas_conversion = False
@@ -855,7 +855,7 @@ def ensure_ray_initialized(
855855
logger.info("Ray is already initialized externally, using existing cluster")
856856
# Configure DatasetContext even if Ray is already initialized
857857
ctx = DatasetContext.get_current()
858-
ctx.shuffle_strategy = "sort" # type: ignore
858+
ctx.shuffle_strategy = ShuffleStrategy.SORT_SHUFFLE_PULL_BASED
859859
ctx.enable_tensor_extension_casting = False
860860
if hasattr(ctx, "enable_arrow_backed_pandas_conversion"):
861861
ctx.enable_arrow_backed_pandas_conversion = False

‎sdk/python/tests/component/ray/test_resource_scheduling.py‎

Lines changed: 39 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -20,6 +20,7 @@
2020
import pyarrow as pa
2121
import pytest
2222
import ray
23+
from ray.data.context import DatasetContext, ShuffleStrategy
2324

2425
from feast.infra.compute_engines.dag.model import DAGFormat
2526
from feast.infra.compute_engines.dag.node import DAGNode
@@ -31,7 +32,7 @@
3132
RayOfflineStoreConfig,
3233
RayResourceManager,
3334
)
34-
from feast.infra.ray_initializer import CodeFlareRayWrapper
35+
from feast.infra.ray_initializer import CodeFlareRayWrapper, ensure_ray_initialized
3536
from feast.infra.ray_shared_utils import RemoteDatasetProxy
3637

3738
# ---------------------------------------------------------------------------
@@ -492,6 +493,43 @@ def test_available_gpus_zero_when_ray_not_initialized(self):
492493
assert mgr.available_gpus == 0
493494

494495

496+
# ---------------------------------------------------------------------------
497+
# DatasetContext shuffle strategy — must be a valid ShuffleStrategy member
498+
# (Ray >= 2.59 rejects unknown values such as "sort")
499+
# ---------------------------------------------------------------------------
500+
501+
502+
@pytest.fixture
503+
def restore_shuffle_strategy():
504+
ctx = DatasetContext.get_current()
505+
original = ctx.shuffle_strategy
506+
yield
507+
ctx.shuffle_strategy = original
508+
509+
510+
class TestShuffleStrategy:
511+
def test_configure_ray_context_sets_sort_shuffle(self, restore_shuffle_strategy):
512+
with patch("ray.is_initialized", return_value=False):
513+
mgr = RayResourceManager()
514+
mgr.configure_ray_context()
515+
assert (
516+
DatasetContext.get_current().shuffle_strategy
517+
== ShuffleStrategy.SORT_SHUFFLE_PULL_BASED
518+
)
519+
520+
def test_ensure_ray_initialized_sets_sort_shuffle(self, restore_shuffle_strategy):
521+
with (
522+
patch("feast.infra.ray_initializer._ray_initialized", False),
523+
patch("feast.infra.ray_initializer.ray") as mock_ray,
524+
):
525+
mock_ray.is_initialized.return_value = True
526+
ensure_ray_initialized(RayComputeEngineConfig())
527+
assert (
528+
DatasetContext.get_current().shuffle_strategy
529+
== ShuffleStrategy.SORT_SHUFFLE_PULL_BASED
530+
)
531+
532+
495533
# ---------------------------------------------------------------------------
496534
# safe_batch_processor / _is_empty_batch — format-aware empty detection
497535
# Regression tests for AttributeError when gpu_batch_format is "numpy"/"pyarrow"

0 commit comments

Comments
 (0)