FazBrowse GitHub Viewer | Trending |
URL:
| Home
Tools: [Download Repo ZIP]   [Original HTTPS Page]

fix: Use a valid ShuffleStrategy for the Ray DatasetContext by YHC66 · Pull Request #6942 · feast-dev/feast · GitHub

  1. feast-dev
  2. feast

Repository navigation

fix: Use a valid ShuffleStrategy for the Ray DatasetContext #6942

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
ntkathole merged 1 commit into from
Oct 7, 2026
+45 −7
Conversation 1 Commits 1 Checks 23 Files changed 3
Merged

fix: Use a valid ShuffleStrategy for the Ray DatasetContext #6942

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension .py  (3) All 1 file type selected
Viewed files
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Unified
Split
Hide whitespace
Diff view
Unified
Split
Hide whitespace
4 changes: 2 additions & 2 deletions sdk/python/feast/infra/offline_stores/contrib/ray_offline_store/ray.py
Show comments
View file
Open in desktop
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters. Learn more about bidirectional Unicode characters
Show hidden characters
Original file line number Diff line number Diff line change
Expand Up @@ -12,7 +12,7 @@
import pyarrow as pa
import ray
from ray.data import Dataset
from ray.data.context import DatasetContext
from ray.data.context import DatasetContext, ShuffleStrategy

from feast.data_source import DataSource
from feast.dataframe import DataFrameEngine, FeastDataFrame
Expand Down Expand Up @@ -460,7 +460,7 @@ def configure_ray_context(self) -> None:
else 2
)
ctx.max_parallelism = self.available_cpus * multiplier
ctx.shuffle_strategy = "sort" # type: ignore
ctx.shuffle_strategy = ShuffleStrategy.SORT_SHUFFLE_PULL_BASED
ctx.enable_tensor_extension_casting = False

if not getattr(self.config, "enable_ray_logging", False):
Expand Down
8 changes: 4 additions & 4 deletions sdk/python/feast/infra/ray_initializer.py
Show comments
View file
Open in desktop
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters. Learn more about bidirectional Unicode characters
Show hidden characters
Original file line number Diff line number Diff line change
Expand Up @@ -25,7 +25,7 @@
from typing import Any, Dict, List, Optional, Union

import ray
from ray.data.context import DatasetContext
from ray.data.context import DatasetContext, ShuffleStrategy

logger = logging.getLogger(__name__)

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

# Configure DatasetContext
ctx = DatasetContext.get_current()
ctx.shuffle_strategy = "sort" # type: ignore
ctx.shuffle_strategy = ShuffleStrategy.SORT_SHUFFLE_PULL_BASED
ctx.enable_tensor_extension_casting = False
if hasattr(ctx, "enable_arrow_backed_pandas_conversion"):
ctx.enable_arrow_backed_pandas_conversion = False
Expand Down Expand Up @@ -760,7 +760,7 @@ def _initialize_remote_ray(config: Any, enable_logging: bool = False) -> None:

# Configure DatasetContext
ctx = DatasetContext.get_current()
ctx.shuffle_strategy = "sort" # type: ignore
ctx.shuffle_strategy = ShuffleStrategy.SORT_SHUFFLE_PULL_BASED
ctx.enable_tensor_extension_casting = False
if hasattr(ctx, "enable_arrow_backed_pandas_conversion"):
ctx.enable_arrow_backed_pandas_conversion = False
Expand Down Expand Up @@ -855,7 +855,7 @@ def ensure_ray_initialized(
logger.info("Ray is already initialized externally, using existing cluster")
# Configure DatasetContext even if Ray is already initialized
ctx = DatasetContext.get_current()
ctx.shuffle_strategy = "sort" # type: ignore
ctx.shuffle_strategy = ShuffleStrategy.SORT_SHUFFLE_PULL_BASED
ctx.enable_tensor_extension_casting = False
if hasattr(ctx, "enable_arrow_backed_pandas_conversion"):
ctx.enable_arrow_backed_pandas_conversion = False
Expand Down
40 changes: 39 additions & 1 deletion sdk/python/tests/component/ray/test_resource_scheduling.py
Show comments
View file
Open in desktop
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters. Learn more about bidirectional Unicode characters
Show hidden characters
Original file line number Diff line number Diff line change
Expand Up @@ -20,6 +20,7 @@
import pyarrow as pa
import pytest
import ray
from ray.data.context import DatasetContext, ShuffleStrategy

from feast.infra.compute_engines.dag.model import DAGFormat
from feast.infra.compute_engines.dag.node import DAGNode
Expand All @@ -31,7 +32,7 @@
RayOfflineStoreConfig,
RayResourceManager,
)
from feast.infra.ray_initializer import CodeFlareRayWrapper
from feast.infra.ray_initializer import CodeFlareRayWrapper, ensure_ray_initialized
from feast.infra.ray_shared_utils import RemoteDatasetProxy

# ---------------------------------------------------------------------------
Expand Down Expand Up @@ -492,6 +493,43 @@ def test_available_gpus_zero_when_ray_not_initialized(self):
assert mgr.available_gpus == 0


# ---------------------------------------------------------------------------
# DatasetContext shuffle strategy — must be a valid ShuffleStrategy member
# (Ray >= 2.59 rejects unknown values such as "sort")
# ---------------------------------------------------------------------------


@pytest.fixture
def restore_shuffle_strategy():
ctx = DatasetContext.get_current()
original = ctx.shuffle_strategy
yield
ctx.shuffle_strategy = original


class TestShuffleStrategy:
def test_configure_ray_context_sets_sort_shuffle(self, restore_shuffle_strategy):
with patch("ray.is_initialized", return_value=False):
mgr = RayResourceManager()
mgr.configure_ray_context()
assert (
DatasetContext.get_current().shuffle_strategy
== ShuffleStrategy.SORT_SHUFFLE_PULL_BASED
)

def test_ensure_ray_initialized_sets_sort_shuffle(self, restore_shuffle_strategy):
with (
patch("feast.infra.ray_initializer._ray_initialized", False),
patch("feast.infra.ray_initializer.ray") as mock_ray,
):
mock_ray.is_initialized.return_value = True
ensure_ray_initialized(RayComputeEngineConfig())
assert (
DatasetContext.get_current().shuffle_strategy
== ShuffleStrategy.SORT_SHUFFLE_PULL_BASED
)


# ---------------------------------------------------------------------------
# safe_batch_processor / _is_empty_batch — format-aware empty detection
# Regression tests for AttributeError when gpu_batch_format is "numpy"/"pyarrow"
Expand Down
Loading
Back | FazBrowse Home | New Git URL