Skip to content

Commit 5580ab4

Browse files
ntkatholentkathole
committed
feat: Allow to set gpu for ray
Signed-off-by: ntkathole <[email protected]>
1 parent 608b105 commit 5580ab4

9 files changed

Lines changed: 986 additions & 20 deletions

File tree

‎docs/reference/compute-engine/ray.md‎

Lines changed: 88 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -28,6 +28,7 @@ The Ray compute engine provides:
2828
- **Lazy Evaluation**: Deferred execution for optimal performance
2929
- **Resource Management**: Automatic scaling and resource optimization
3030
- **Point-in-Time Joins**: Efficient temporal joins for historical feature retrieval
31+
- **GPU Support**: Schedule transformation workers on GPU nodes via `num_gpus` config (all modes including KubeRay)
3132

3233
## Architecture
3334

@@ -87,6 +88,9 @@ batch_engine:
8788
| `enable_distributed_joins` | boolean | true | Enable distributed joins for large datasets |
8889
| `staging_location` | string | None | Remote path for batch materialization jobs |
8990
| `ray_conf` | dict | None | Ray configuration parameters (memory, CPU limits) |
91+
| `num_gpus` | float | None | Number of GPUs to request per worker task. Requires GPU nodes in the Ray cluster. Fractional values (e.g. `0.5`) are supported. Supported in all modes including KubeRay. |
92+
| `gpu_batch_format` | string | `"pandas"` | Batch format for `map_batches` when `num_gpus` is set. Use `"numpy"` or `"pyarrow"` for GPU-native libraries (e.g. cuDF, PyTorch). |
93+
| `worker_task_options` | dict | None | Arbitrary Ray `.options()` kwargs applied to every worker task. See [Worker Resource Scheduling](#worker-resource-scheduling) for the full reference. |
9094

9195
### Mode Detection Precedence
9296

@@ -344,13 +348,97 @@ import ray
344348
# Check cluster resources
345349
resources = ray.cluster_resources()
346350
print(f"Available CPUs: {resources.get('CPU', 0)}")
351+
print(f"Available GPUs: {resources.get('GPU', 0)}")
347352
print(f"Available memory: {resources.get('memory', 0) / 1e9:.2f} GB")
348353
349354
# Monitor job progress
350355
job = store.get_historical_features(...)
351356
# Ray compute engine provides built-in progress tracking
352357
```
353358

359+
## Worker Resource Scheduling
360+
361+
`worker_task_options` is a passthrough dict of [Ray `.options()` kwargs](https://docs.ray.io/en/latest/ray-core/api/doc/ray.remote_function.RemoteFunction.options.html) applied to every worker task Feast dispatches. It pairs with `ray_conf` (cluster-level `ray.init` options) — `worker_task_options` targets individual worker tasks. Options are forwarded at two levels so Ray schedules correctly:
362+
363+
1. On the `@ray.remote` orchestration task via `.options(**worker_task_options)` — controls node selection.
364+
2. Inside `map_batches` for the scheduling-relevant subset (`num_gpus`, `num_cpus`, `accelerator_type`, `resources`) — controls which nodes run the data workers.
365+
366+
This is supported across **all execution modes**: local, remote, and KubeRay.
367+
368+
### Common `worker_task_options` keys
369+
370+
| Key | Type | Description |
371+
|-----|------|-------------|
372+
| `num_cpus` | float | CPUs per task (default: 1). Fractional values supported. |
373+
| `memory` | int | Heap memory in **bytes** (e.g. `8589934592` for 8 GB). |
374+
| `accelerator_type` | string | Pin tasks to a specific GPU model — `"A100"`, `"T4"`, `"V100"`, etc. Useful on KubeRay clusters with mixed GPU node pools. |
375+
| `resources` | dict | Custom/Kubernetes extended resource labels, e.g. `{"intel.com/gpu": 1}`. |
376+
| `runtime_env` | dict | Per-task [Ray runtime environment](https://docs.ray.io/en/latest/ray-core/handling-dependencies.html) — `pip`, `conda`, `env_vars`, `working_dir`, etc. For KubeRay, use this to install packages on worker pods without rebuilding images. |
377+
| `max_retries` | int | Task retry count on worker failure (default: 3). |
378+
| `scheduling_strategy` | string | `"DEFAULT"`, `"SPREAD"`, or a placement group strategy. |
379+
380+
> For the full list of supported keys see the [Ray RemoteFunction.options() API docs](https://docs.ray.io/en/latest/ray-core/api/doc/ray.remote_function.RemoteFunction.options.html).
381+
382+
### GPU support
383+
384+
`num_gpus` is the only first-class GPU field because it also drives `gpu_batch_format` selection inside Feast. Set it directly rather than inside `worker_task_options`:
385+
386+
```yaml
387+
batch_engine:
388+
type: ray.engine
389+
num_gpus: 1 # GPUs per task (fractional values like 0.5 supported)
390+
gpu_batch_format: numpy # numpy/pyarrow for GPU-native libs (cuDF, PyTorch)
391+
```
392+
393+
When `num_gpus` is set your transformation UDF runs on a GPU worker:
394+
395+
```python
396+
import cudf # RAPIDS cuDF – GPU-accelerated DataFrame library
397+
398+
def gpu_transform(batch):
399+
gpu_df = cudf.from_pandas(batch)
400+
gpu_df["score"] = gpu_df["raw_value"] * 2.0
401+
return gpu_df.to_pandas()
402+
```
403+
404+
### Full example — KubeRay with GPU + all common options
405+
406+
```yaml
407+
batch_engine:
408+
type: ray.engine
409+
use_kuberay: true
410+
kuberay_conf:
411+
cluster_name: "feast-gpu-cluster"
412+
namespace: "feast-system"
413+
auth_token: "${RAY_AUTH_TOKEN}"
414+
auth_server: "https://api.openshift.com:6443"
415+
num_gpus: 1
416+
gpu_batch_format: numpy
417+
worker_task_options:
418+
num_cpus: 4
419+
memory: 8589934592 # 8 GB
420+
accelerator_type: "A100" # pin to A100 nodes on mixed GPU pool
421+
max_retries: 5
422+
runtime_env:
423+
pip:
424+
- cudf-cu12==24.10.0
425+
- torch==2.4.0
426+
env_vars:
427+
CUDA_VISIBLE_DEVICES: "0"
428+
```
429+
430+
### Checking cluster resources
431+
432+
```python
433+
import ray
434+
435+
ray.init(address="auto")
436+
resources = ray.cluster_resources()
437+
print(f"Available CPUs: {resources.get('CPU', 0)}")
438+
print(f"Available GPUs: {resources.get('GPU', 0)}")
439+
print(f"Available memory: {resources.get('memory', 0) / 1e9:.2f} GB")
440+
```
441+
354442
## Integration Examples
355443

356444
### With Spark Offline Store

‎docs/reference/offline-stores/ray.md‎

Lines changed: 10 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -31,6 +31,7 @@ The Ray offline store provides:
3131
- Efficient data filtering and column selection
3232
- Timestamp-based data processing with timezone awareness
3333
- Enterprise-ready KubeRay cluster support via CodeFlare SDK
34+
- **GPU support**: schedule worker tasks on GPU nodes via `num_gpus` config (all modes including KubeRay)
3435

3536

3637
## Functionality Matrix
@@ -246,6 +247,9 @@ batch_engine:
246247
| `max_parallelism_multiplier` | int | 2 | Parallelism as multiple of CPU cores |
247248
| `target_partition_size_mb` | int | 64 | Target partition size (MB) |
248249
| `window_size_for_joins` | string | "1H" | Time window for distributed joins |
250+
| `num_gpus` | float | None | GPUs per worker task. Supported in all modes. See [Worker Resource Scheduling](../compute-engine/ray.md#worker-resource-scheduling). |
251+
| `gpu_batch_format` | string | `"pandas"` | Batch format for `map_batches` when `num_gpus` is set (`"numpy"` or `"pyarrow"` for GPU-native libs). |
252+
| `worker_task_options` | dict | None | Arbitrary Ray `.options()` kwargs (num_cpus, memory, accelerator_type, resources, runtime_env, …). See [Worker Resource Scheduling](../compute-engine/ray.md#worker-resource-scheduling) for the full reference. |
249253

250254
#### Mode Detection Precedence
251255

@@ -542,6 +546,12 @@ python your_feast_script.py
542546
- Secure communication between client and Ray cluster
543547
- Automatic cluster discovery
544548

549+
### GPU Support
550+
551+
The Ray offline store supports GPU scheduling via the `num_gpus` and `gpu_batch_format` config options. This works across all execution modes (local, remote, and KubeRay).
552+
553+
For full configuration details, examples, and KubeRay GPU setup, see the [Ray Compute Engine GPU Support](../compute-engine/ray.md#gpu-support) section.
554+
545555
### Data Source Validation
546556

547557
The Ray offline store validates data sources to ensure compatibility:

‎sdk/python/feast/infra/compute_engines/ray/config.py‎

Lines changed: 45 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -46,6 +46,51 @@ class RayComputeEngineConfig(FeastConfigBaseModel):
4646
enable_optimization: bool = True
4747
"""Enable automatic performance optimizations."""
4848

49+
# Worker task resource configuration
50+
num_gpus: Optional[float] = None
51+
"""Number of GPUs to request per worker task. Requires GPU nodes in the
52+
Ray cluster. Fractional values (e.g. 0.5) are supported by Ray for GPU
53+
sharing. Supported in all modes: local, remote, and KubeRay."""
54+
55+
gpu_batch_format: str = "pandas"
56+
"""Batch format for map_batches when num_gpus is set. Use 'numpy' or
57+
'pyarrow' for GPU-native libraries (e.g. cuDF, PyTorch). Defaults to
58+
'pandas'."""
59+
60+
worker_task_options: Optional[Dict[str, Any]] = None
61+
"""Arbitrary Ray task options passed verbatim to @ray.remote .options()
62+
and map_batches for every worker task Feast dispatches. This is the
63+
escape hatch for any Ray or CodeFlare SDK scheduling parameter not
64+
covered by the dedicated fields above.
65+
66+
Pairs with ray_conf (which configures ray.init) — worker_task_options
67+
targets the individual worker tasks rather than the cluster connection.
68+
69+
Common keys (see https://docs.ray.io/en/latest/ray-core/api/doc/ray.remote_function.RemoteFunction.options.html):
70+
num_cpus (float) – CPUs per task (default: 1)
71+
memory (int) – Heap memory in bytes (e.g. 8 * 1024**3 for 8 GB)
72+
accelerator_type (str) – Specific GPU model, e.g. 'A100', 'T4', 'V100'.
73+
Pins tasks to nodes advertising that type. Useful
74+
on KubeRay clusters with mixed GPU pools.
75+
resources (dict) – Custom/extended resource labels, e.g.
76+
{'intel.com/gpu': 1} for Kubernetes extended resources.
77+
runtime_env (dict) – Per-task runtime environment (pip, conda, env_vars,
78+
working_dir, …). For KubeRay use this to install
79+
extra packages on workers without rebuilding images.
80+
max_retries (int) – Task retry count on worker failure (default: 3).
81+
scheduling_strategy (str) – 'DEFAULT', 'SPREAD', or a placement group strategy.
82+
83+
Example:
84+
worker_task_options:
85+
num_cpus: 4
86+
memory: 8589934592 # 8 GB
87+
accelerator_type: "A100"
88+
max_retries: 5
89+
runtime_env:
90+
pip: ["cudf-cu12==24.10.0"]
91+
env_vars: {CUDA_VISIBLE_DEVICES: "0"}
92+
"""
93+
4994
@property
5095
def window_size_timedelta(self) -> timedelta:
5196
"""Convert window size string to timedelta."""

‎sdk/python/feast/infra/compute_engines/ray/nodes.py‎

Lines changed: 29 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,7 @@
11
import json
22
import logging
33
from datetime import datetime, timedelta, timezone
4-
from typing import Dict, List, Optional, Set, Union
4+
from typing import Any, Dict, List, Optional, Set, Union
55

66
import dill
77
import pandas as pd
@@ -688,10 +688,36 @@ def apply_transformation_with_serialized_udf(
688688

689689
return transformed_batch
690690

691+
num_gpus = getattr(self.config, "num_gpus", None) or None
692+
task_options: Dict[str, Any] = dict(
693+
getattr(self.config, "worker_task_options", None) or {}
694+
)
695+
if num_gpus:
696+
task_options["num_gpus"] = num_gpus
697+
batch_format = (
698+
getattr(self.config, "gpu_batch_format", "pandas")
699+
if num_gpus
700+
else "pandas"
701+
)
702+
# Only the scheduling-relevant subset flows into map_batches
703+
_MAP_BATCHES_RESOURCE_KEYS = {
704+
"num_gpus",
705+
"num_cpus",
706+
"accelerator_type",
707+
"resources",
708+
}
709+
map_kwargs: Dict[str, Any] = {
710+
"batch_format": batch_format,
711+
"concurrency": self.config.max_workers or 12,
712+
**{
713+
k: v
714+
for k, v in task_options.items()
715+
if k in _MAP_BATCHES_RESOURCE_KEYS
716+
},
717+
}
691718
transformed_dataset = dataset.map_batches(
692719
apply_transformation_with_serialized_udf,
693-
batch_format="pandas",
694-
concurrency=self.config.max_workers or 12,
720+
**map_kwargs,
695721
)
696722

697723
return DAGValue(

‎sdk/python/feast/infra/compute_engines/ray/utils.py‎

Lines changed: 41 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -5,6 +5,7 @@
55
import logging
66
from typing import Callable, Dict, Union
77

8+
import numpy as np
89
import pandas as pd
910
import pyarrow as pa
1011

@@ -79,22 +80,54 @@ def write_to_online_store(
7980
logger.error(f"Failed to write to online store for {feature_view.name}: {e}")
8081

8182

83+
# Ray Data batch type: pandas DataFrame, numpy dict, or pyarrow Table
84+
BatchType = Union[pd.DataFrame, Dict[str, np.ndarray], pa.Table]
85+
86+
87+
def _is_empty_batch(batch: BatchType) -> bool:
88+
"""Return True if the batch contains no rows, regardless of Ray Data batch format.
89+
90+
Ray Data delivers batches in three formats depending on the batch_format
91+
argument passed to map_batches:
92+
- "pandas" → pd.DataFrame (.empty attribute)
93+
- "numpy" → Dict[str, np.ndarray] (check length of first array)
94+
- "pyarrow" → pa.Table (.num_rows attribute)
95+
"""
96+
if isinstance(batch, pd.DataFrame):
97+
return batch.empty
98+
if isinstance(batch, dict):
99+
if not batch:
100+
return True
101+
first = next(iter(batch.values()))
102+
return len(first) == 0
103+
if isinstance(batch, pa.Table):
104+
return batch.num_rows == 0
105+
return False
106+
107+
82108
def safe_batch_processor(
83-
func: Callable[[pd.DataFrame], pd.DataFrame],
84-
) -> Callable[[pd.DataFrame], pd.DataFrame]:
109+
func: Callable[[BatchType], BatchType],
110+
) -> Callable[[BatchType], BatchType]:
85111
"""
86-
Decorator for batch processing functions that handles empty batches and errors gracefully.
112+
Decorator for batch processing functions that handles empty batches and
113+
exceptions gracefully across all Ray Data batch formats.
114+
115+
Ray Data can deliver batches as a pandas DataFrame (batch_format="pandas"),
116+
a Dict[str, np.ndarray] (batch_format="numpy"), or a pa.Table
117+
(batch_format="pyarrow"). The decorator handles all three so that callers
118+
using gpu_batch_format="numpy" or "pyarrow" do not crash on the empty-batch
119+
check.
87120
88121
Args:
89-
func: Function that processes a pandas DataFrame batch
122+
func: Batch processing function. Receives and returns the same batch
123+
type that Ray Data passes (pandas, numpy dict, or pyarrow Table).
90124
91125
Returns:
92-
Wrapped function that handles empty batches and exceptions
126+
Wrapped function that skips empty batches and swallows exceptions.
93127
"""
94128

95-
def wrapper(batch: pd.DataFrame) -> pd.DataFrame:
96-
# Handle empty batches
97-
if batch.empty:
129+
def wrapper(batch: BatchType) -> BatchType:
130+
if _is_empty_batch(batch):
98131
return batch
99132

100133
try:

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

Lines changed: 48 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -365,6 +365,51 @@ class RayOfflineStoreConfig(FeastConfigBaseModel):
365365
kuberay_conf: Optional[Dict[str, Any]] = None
366366
"""KubeRay/CodeFlare configuration parameters (passed to CodeFlare SDK)"""
367367

368+
# Worker task resource configuration
369+
num_gpus: Optional[float] = None
370+
"""Number of GPUs to request per worker task. Requires GPU nodes in the
371+
Ray cluster. Fractional values (e.g. 0.5) are supported by Ray for GPU
372+
sharing. Supported in all modes: local, remote, and KubeRay."""
373+
374+
gpu_batch_format: str = "pandas"
375+
"""Batch format for map_batches when num_gpus is set. Use 'numpy' or
376+
'pyarrow' for GPU-native libraries (e.g. cuDF, PyTorch). Defaults to
377+
'pandas'."""
378+
379+
worker_task_options: Optional[Dict[str, Any]] = None
380+
"""Arbitrary Ray task options passed verbatim to @ray.remote .options()
381+
and map_batches for every worker task Feast dispatches. This is the
382+
escape hatch for any Ray or CodeFlare SDK scheduling parameter not
383+
covered by the dedicated fields above.
384+
385+
Pairs with ray_conf (which configures ray.init) — worker_task_options
386+
targets the individual worker tasks rather than the cluster connection.
387+
388+
Common keys (see https://docs.ray.io/en/latest/ray-core/api/doc/ray.remote_function.RemoteFunction.options.html):
389+
num_cpus (float) – CPUs per task (default: 1)
390+
memory (int) – Heap memory in bytes (e.g. 8 * 1024**3 for 8 GB)
391+
accelerator_type (str) – Specific GPU model, e.g. 'A100', 'T4', 'V100'.
392+
Pins tasks to nodes advertising that type. Useful
393+
on KubeRay clusters with mixed GPU pools.
394+
resources (dict) – Custom/extended resource labels, e.g.
395+
{'intel.com/gpu': 1} for Kubernetes extended resources.
396+
runtime_env (dict) – Per-task runtime environment (pip, conda, env_vars,
397+
working_dir, …). For KubeRay use this to install
398+
extra packages on workers without rebuilding images.
399+
max_retries (int) – Task retry count on worker failure (default: 3).
400+
scheduling_strategy (str) – 'DEFAULT', 'SPREAD', or a placement group strategy.
401+
402+
Example:
403+
worker_task_options:
404+
num_cpus: 4
405+
memory: 8589934592 # 8 GB
406+
accelerator_type: "A100"
407+
max_retries: 5
408+
runtime_env:
409+
pip: ["cudf-cu12==24.10.0"]
410+
env_vars: {CUDA_VISIBLE_DEVICES: "0"}
411+
"""
412+
368413

369414
class RayResourceManager:
370415
"""
@@ -382,12 +427,14 @@ def __init__(self, config: Optional[RayOfflineStoreConfig] = None) -> None:
382427
self.cluster_resources = {"CPU": 4, "memory": 8 * 1024**3}
383428
self.available_memory = 8 * 1024**3
384429
self.available_cpus = 4
430+
self.available_gpus = 0
385431
self.num_nodes = 1
386432
return
387433

388434
self.cluster_resources = ray.cluster_resources()
389435
self.available_memory = self.cluster_resources.get("memory", 8 * 1024**3)
390436
self.available_cpus = int(self.cluster_resources.get("CPU", 4))
437+
self.available_gpus = int(self.cluster_resources.get("GPU", 0))
391438
self.num_nodes = len(ray.nodes())
392439

393440
def configure_ray_context(self) -> None:
@@ -421,6 +468,7 @@ def configure_ray_context(self) -> None:
421468
if getattr(self.config, "enable_ray_logging", False):
422469
logger.info(
423470
f"Configured Ray context: {self.available_cpus} CPUs, "
471+
f"{self.available_gpus} GPUs, "
424472
f"{self.available_memory // 1024**3}GB memory, {self.num_nodes} nodes"
425473
)
426474

0 commit comments

Comments
 (0)