diff --git a/doc/source/ray-core/scheduling/accelerators.rst b/doc/source/ray-core/scheduling/accelerators.rst index 487555b6928e..b5a99a58d6e2 100644 --- a/doc/source/ray-core/scheduling/accelerators.rst +++ b/doc/source/ray-core/scheduling/accelerators.rst @@ -79,10 +79,11 @@ If you need to, you can :ref:`override ` this. .. tip:: - You can set the ``ONEAPI_DEVICE_SELECTOR`` environment variable before starting a Ray node + You can set the ``ZE_AFFINITY_MASK`` environment variable before starting a Ray node to limit the Intel GPUs that are visible to Ray. - For example, ``ONEAPI_DEVICE_SELECTOR=1,3 ray start --head --num-gpus=2`` + For example, ``ZE_AFFINITY_MASK=1,3 ray start --head --num-gpus=2`` lets Ray only see devices 1 and 3. + ``ONEAPI_DEVICE_SELECTOR`` is still read as a fallback for backward compatibility. .. tab-item:: AWS Neuron Core :sync: AWS Neuron Core @@ -282,12 +283,12 @@ and assign accelerators to the task or actor by setting the corresponding enviro class GPUActor: def ping(self): print("GPU IDs: {}".format(ray.get_runtime_context().get_accelerator_ids()["GPU"])) - print("ONEAPI_DEVICE_SELECTOR: {}".format(os.environ["ONEAPI_DEVICE_SELECTOR"])) + print("ZE_AFFINITY_MASK: {}".format(os.environ["ZE_AFFINITY_MASK"])) @ray.remote(num_gpus=1) def gpu_task(): print("GPU IDs: {}".format(ray.get_runtime_context().get_accelerator_ids()["GPU"])) - print("ONEAPI_DEVICE_SELECTOR: {}".format(os.environ["ONEAPI_DEVICE_SELECTOR"])) + print("ZE_AFFINITY_MASK: {}".format(os.environ["ZE_AFFINITY_MASK"])) gpu_actor = GPUActor.remote() ray.get(gpu_actor.ping.remote()) @@ -298,9 +299,9 @@ and assign accelerators to the task or actor by setting the corresponding enviro :options: +MOCK (GPUActor pid=52420) GPU IDs: [0] - (GPUActor pid=52420) ONEAPI_DEVICE_SELECTOR: 0 + (GPUActor pid=52420) ZE_AFFINITY_MASK: 0 (gpu_task pid=51830) GPU IDs: [1] - (gpu_task pid=51830) ONEAPI_DEVICE_SELECTOR: 1 + (gpu_task pid=51830) ZE_AFFINITY_MASK: 1 .. tab-item:: AWS Neuron Core :sync: AWS Neuron Core diff --git a/doc/source/serve/llm/user-guides/sglang.md b/doc/source/serve/llm/user-guides/sglang.md index 8303dd08fdb7..a774aac22552 100644 --- a/doc/source/serve/llm/user-guides/sglang.md +++ b/doc/source/serve/llm/user-guides/sglang.md @@ -25,6 +25,7 @@ Set the following environment variable before running any example: - **CUDA:** `RAY_EXPERIMENTAL_NOSET_CUDA_VISIBLE_DEVICES=0` - **ROCm:** `RAY_EXPERIMENTAL_NOSET_HIP_VISIBLE_DEVICES=0` +- **Intel GPU:** `RAY_EXPERIMENTAL_NOSET_ZE_AFFINITY_MASK=0` ## Online serving (single node) diff --git a/java/dependencies.bzl b/java/dependencies.bzl index af6941d13266..e4b5efbc9c20 100644 --- a/java/dependencies.bzl +++ b/java/dependencies.bzl @@ -6,7 +6,7 @@ def gen_java_deps(): artifacts = [ "com.fasterxml.jackson.core:jackson-databind:2.18.8", "com.github.java-json-tools:json-schema-validator:2.2.14", - "com.google.code.gson:gson:2.9.1", + "com.google.code.gson:gson:2.11.0", "com.google.guava:guava:32.0.1-jre", "com.google.protobuf:protobuf-java:3.23.4", "com.google.protobuf:protobuf-java-util:3.23.4", diff --git a/python/ray/_private/accelerators/intel_gpu.py b/python/ray/_private/accelerators/intel_gpu.py index fe229990fd19..d40ae4f73b73 100644 --- a/python/ray/_private/accelerators/intel_gpu.py +++ b/python/ray/_private/accelerators/intel_gpu.py @@ -7,8 +7,10 @@ logger = logging.getLogger(__name__) +ZE_AFFINITY_MASK_ENV_VAR = "ZE_AFFINITY_MASK" +NOSET_ZE_AFFINITY_MASK_ENV_VAR = "RAY_EXPERIMENTAL_NOSET_ZE_AFFINITY_MASK" + ONEAPI_DEVICE_SELECTOR_ENV_VAR = "ONEAPI_DEVICE_SELECTOR" -NOSET_ONEAPI_DEVICE_SELECTOR_ENV_VAR = "RAY_EXPERIMENTAL_NOSET_ONEAPI_DEVICE_SELECTOR" ONEAPI_DEVICE_BACKEND_TYPE = "level_zero" ONEAPI_DEVICE_TYPE = "gpu" @@ -22,25 +24,29 @@ def get_resource_name() -> str: @staticmethod def get_visible_accelerator_ids_env_var() -> str: - return ONEAPI_DEVICE_SELECTOR_ENV_VAR + return ZE_AFFINITY_MASK_ENV_VAR @staticmethod def get_current_process_visible_accelerator_ids() -> Optional[List[str]]: - oneapi_visible_devices = os.environ.get( - IntelGPUAcceleratorManager.get_visible_accelerator_ids_env_var(), None - ) + # Primary: ZE_AFFINITY_MASK uses bare IDs ("0,1,2"), like CUDA_VISIBLE_DEVICES. + ze_mask = os.environ.get(ZE_AFFINITY_MASK_ENV_VAR, None) + if ze_mask is not None: + if ze_mask == "": + return [] + return list(ze_mask.split(",")) + + # Fallback: ONEAPI_DEVICE_SELECTOR for backward compatibility. + oneapi_visible_devices = os.environ.get(ONEAPI_DEVICE_SELECTOR_ENV_VAR, None) if oneapi_visible_devices is None: return None - - if oneapi_visible_devices == "": - return [] - - if oneapi_visible_devices == "NoDevFiles": + if oneapi_visible_devices == "" or oneapi_visible_devices == "NoDevFiles": return [] prefix = ONEAPI_DEVICE_BACKEND_TYPE + ":" - - return list(oneapi_visible_devices.split(prefix)[1].split(",")) + if prefix in oneapi_visible_devices: + return list(oneapi_visible_devices.split(prefix)[1].split(",")) + # bare IDs without prefix (e.g. "0,1") — accepted as-is + return list(oneapi_visible_devices.split(",")) @staticmethod def get_current_node_num_accelerators() -> int: @@ -95,10 +101,14 @@ def validate_resource_request_quantity( def set_current_process_visible_accelerator_ids( visible_xpu_devices: List[str], ) -> None: - if env_bool(NOSET_ONEAPI_DEVICE_SELECTOR_ENV_VAR, False): + if env_bool(NOSET_ZE_AFFINITY_MASK_ENV_VAR, False): return - prefix = ONEAPI_DEVICE_BACKEND_TYPE + ":" - os.environ[ - IntelGPUAcceleratorManager.get_visible_accelerator_ids_env_var() - ] = prefix + ",".join([str(i) for i in visible_xpu_devices]) + # ZE_AFFINITY_MASK masks devices at the Level Zero driver, below oneAPI/SYCL + # and torch-xpu, so it fully restricts visibility on its own. It uses bare + # IDs ("0,1,2") like CUDA_VISIBLE_DEVICES. Ray only sets this one env var: + # writing ONEAPI_DEVICE_SELECTOR too would collide with frameworks (e.g. + # SGLang) that manage it themselves, and would need a separate save/restore. + os.environ[ZE_AFFINITY_MASK_ENV_VAR] = ",".join( + [str(i) for i in visible_xpu_devices] + ) diff --git a/python/ray/dashboard/modules/metrics/dashboards/serve_llm_dashboard_panels.py b/python/ray/dashboard/modules/metrics/dashboards/serve_llm_dashboard_panels.py index ce4389ce0273..acc766be24bd 100644 --- a/python/ray/dashboard/modules/metrics/dashboards/serve_llm_dashboard_panels.py +++ b/python/ray/dashboard/modules/metrics/dashboards/serve_llm_dashboard_panels.py @@ -1,5 +1,7 @@ # ruff: noqa: E501 +from collections.abc import Sequence + from ray.dashboard.modules.metrics.dashboards.common import ( DashboardConfig, GridPos, @@ -109,6 +111,35 @@ def _ratio_with_join_and_guard( ) +def _summed_ratio_with_join_and_guard( + numerator_metrics: Sequence[str], + denominator_metric: str, +) -> str: + """Ratio of added rate metrics over one denominator, NaN guard + WorkerId join. + + The numerator rates are added before the division, so counters that each + contribute a share of the same denominator chart as a single ratio. + """ + return ( + "(\n" + " (\n" + " (\n" + + "\n +\n".join( + f" sum by(WorkerId) (rate({metric}{{{{{_VLLM_FILTER}}}}}[$interval]))" + for metric in numerator_metrics + ) + + "\n )\n" + " /\n" + f" sum by(WorkerId) (rate({denominator_metric}{{{{{_VLLM_FILTER}}}}}[$interval]))\n" + " )\n" + " and on(WorkerId)\n" + " (\n" + f" sum by(WorkerId) (rate({denominator_metric}{{{{{_VLLM_FILTER}}}}}[$interval])) > 0\n" + " )\n" + ")" + _WORKER_JOIN + ) + + # --------------------------------------------------------------------------- # Histogram helper: generates Mean / P50 / P90 panels for a given metric # --------------------------------------------------------------------------- @@ -557,7 +588,212 @@ def _histogram_panels( ] # =================================================================== -# Row 7: Token Distribution (collapsed) +# Row 7: Native CPU KV Offload (collapsed) +# =================================================================== +_kv_offload_panels = [ + Panel( + id=52, + title="KV Offload: Store Throughput", + description="GPU-to-CPU KV data/s.", + unit="GBs", + targets=[ + Target( + expr=( + "(" + + _rate_with_join("ray_vllm_kv_offload_store_bytes_total") + + ") / 1024 / 1024 / 1024" + ), + legend=_DEP_REPLICA, + ), + ], + fill=1, + linewidth=1, + stack=False, + grid_pos=GridPos(0, 103, 8, 8), + ), + Panel( + id=53, + title="KV Offload: Reload Throughput", + description="CPU-to-GPU KV data/s.", + unit="GBs", + targets=[ + Target( + expr=( + "(" + + _rate_with_join("ray_vllm_kv_offload_load_bytes_total") + + ") / 1024 / 1024 / 1024" + ), + legend=_DEP_REPLICA, + ), + ], + fill=1, + linewidth=1, + stack=False, + grid_pos=GridPos(8, 103, 8, 8), + ), + Panel( + id=54, + title="KV Offload: Store and Reload Operations/s", + description="KV store and reload operations/s.", + unit="ops", + targets=[ + Target( + expr=_rate_with_join("ray_vllm_kv_offload_store_size_count"), + legend="store — " + _DEP_REPLICA, + ), + Target( + expr=_rate_with_join("ray_vllm_kv_offload_load_size_count"), + legend="reload — " + _DEP_REPLICA, + ), + ], + fill=1, + linewidth=1, + stack=False, + grid_pos=GridPos(16, 103, 8, 8), + ), + Panel( + id=55, + title="KV Offload: Store Bandwidth", + description="GPU-to-CPU KV transfer bandwidth.", + unit="GBs", + targets=[ + Target( + expr=_ratio_with_join_and_guard( + "ray_vllm_kv_offload_store_bytes_total", + "ray_vllm_kv_offload_store_time_total", + scale="/ 1024 / 1024 / 1024", + ), + legend=_DEP_REPLICA, + ), + ], + fill=1, + linewidth=1, + stack=False, + grid_pos=GridPos(0, 111, 8, 8), + ), + Panel( + id=56, + title="KV Offload: Reload Bandwidth", + description="CPU-to-GPU KV transfer bandwidth.", + unit="GBs", + targets=[ + Target( + expr=_ratio_with_join_and_guard( + "ray_vllm_kv_offload_load_bytes_total", + "ray_vllm_kv_offload_load_time_total", + scale="/ 1024 / 1024 / 1024", + ), + legend=_DEP_REPLICA, + ), + ], + fill=1, + linewidth=1, + stack=False, + grid_pos=GridPos(8, 111, 8, 8), + ), + Panel( + id=57, + title="KV Offload: CPU Capacity Pinned by Transfers", + description="Share of the CPU KV pool pinned by in-flight transfers. Sustained values near 100% mean transfers may be dropped.", + unit="percentunit", + targets=[ + Target( + expr=_gauge_with_join("ray_vllm_kv_offload_cpu_cache_usage_perc"), + legend="total — " + _DEP_REPLICA, + ), + Target( + expr=_gauge_with_join("ray_vllm_kv_offload_cpu_cache_write_usage_perc"), + legend="stores — " + _DEP_REPLICA, + ), + Target( + expr=_gauge_with_join("ray_vllm_kv_offload_cpu_cache_read_usage_perc"), + legend="reloads — " + _DEP_REPLICA, + ), + ], + fill=1, + linewidth=1, + stack=False, + grid_pos=GridPos(16, 111, 8, 8), + ), + Panel( + id=58, + title="KV Offload: External Prefix Hit Rate", + description="Connector prefix-cache hit rate.", + unit="percent", + targets=[ + Target( + expr=( + "100 * " + + _ratio_with_join_and_guard( + "ray_vllm_external_prefix_cache_hits_total", + "ray_vllm_external_prefix_cache_queries_total", + ) + ), + legend=_DEP_REPLICA, + ), + ], + fill=1, + linewidth=1, + stack=False, + grid_pos=GridPos(0, 119, 8, 8), + ), + Panel( + id=60, + title="KV Offload: Overall Prefix Hit Rate", + description="Prompt tokens served from either cache tier, GPU or offloaded.", + unit="percent", + targets=[ + Target( + expr=( + "100 * " + + _summed_ratio_with_join_and_guard( + [ + "ray_vllm_prefix_cache_hits_total", + "ray_vllm_external_prefix_cache_hits_total", + ], + "ray_vllm_prefix_cache_queries_total", + ) + ), + legend="overall — " + _DEP_REPLICA, + ), + Target( + expr=( + "100 * " + + _ratio_with_join_and_guard( + "ray_vllm_prefix_cache_hits_total", + "ray_vllm_prefix_cache_queries_total", + ) + ), + legend="GPU only — " + _DEP_REPLICA, + ), + ], + fill=1, + linewidth=1, + stack=False, + grid_pos=GridPos(8, 119, 8, 8), + ), + Panel( + id=59, + title="KV Offload: Lookup Delay -- P90", + description="P90 offloaded-prefix lookup latency.", + unit="s", + targets=[ + Target( + expr=_percentile_with_join( + "ray_vllm_kv_offload_lookup_sync_delay_seconds", 0.9 + ), + legend=_DEP_REPLICA, + ), + ], + fill=1, + linewidth=1, + stack=False, + grid_pos=GridPos(16, 119, 8, 8), + ), +] + +# =================================================================== +# Row 8: Token Distribution (collapsed) # =================================================================== _WORKERID_FILTER = 'WorkerId=~"$workerid", {global_filters}' @@ -785,6 +1021,12 @@ def _histogram_panels( Row(title="Request Length", id=504, panels=_request_length_panels), Row(title="Scheduler", id=505, panels=_scheduler_panels), Row(title="NIXL", id=506, panels=_nixl_panels), + Row( + title="KV Cache Offload / Reload", + id=508, + panels=_kv_offload_panels, + collapsed=True, + ), Row( title="Token Distribution", id=507, diff --git a/python/ray/data/BUILD.bazel b/python/ray/data/BUILD.bazel index d650fa41479e..4255d9eddde5 100644 --- a/python/ray/data/BUILD.bazel +++ b/python/ray/data/BUILD.bazel @@ -1214,6 +1214,21 @@ py_test( ], ) +py_test( + name = "test_torch_inference", + size = "medium", + srcs = ["tests/test_torch_inference.py"], + tags = [ + "exclusive", + "gpu", + "team:data", + ], + deps = [ + ":conftest", + "//:ray_lib", + ], +) + py_test( name = "test_dynamic_block_split", size = "large", diff --git a/python/ray/data/_internal/arrow_block.py b/python/ray/data/_internal/arrow_block.py index 02ed2a774a76..2bd7c4e87116 100644 --- a/python/ray/data/_internal/arrow_block.py +++ b/python/ray/data/_internal/arrow_block.py @@ -290,11 +290,20 @@ def to_pandas(self) -> "pandas.DataFrame": # their own to_pandas_dtype() hooks. Note: native FixedShapeTensorType # subclasses BaseExtensionType but not ExtensionType, so we check the # broader BaseExtensionType. + # - Arrow's null type carries no type information, and pandas cannot box a + # non-null value into a null[pyarrow] column, so fillna and masked + # assignment raise ArrowInvalid (and can abort the worker from Arrow + # C++). Fall back to pandas' default conversion; PandasBlockAccessor + # .to_arrow() coerces all-null columns back to pa.null(), so the + # round-trip is unchanged. A column that is all-null in every block + # therefore stays null-typed rather than being promoted. def _types_mapper(t): if isinstance(t, pyarrow.BaseExtensionType) or pyarrow.types.is_dictionary( t ): return None + if pyarrow.types.is_null(t): + return None return pd.ArrowDtype(t) # Gated on enable_arrow_backed_pandas_conversion so callers can restore the diff --git a/python/ray/data/_internal/compute.py b/python/ray/data/_internal/compute.py index dafb209bc5a3..f0c4b1bb9a1f 100644 --- a/python/ray/data/_internal/compute.py +++ b/python/ray/data/_internal/compute.py @@ -114,7 +114,7 @@ def __init__( max_size: Optional[int] = None, initial_size: Optional[int] = None, max_tasks_in_flight_per_actor: Optional[int] = None, - enable_true_multi_threading: bool = False, + enable_true_multi_threading: Optional[bool] = None, ): """Construct ActorPoolStrategy for a Dataset transform. @@ -130,9 +130,10 @@ def __init__( opportunities for pipelining task dependency prefetching with computation and avoiding actor startup delays, but will also increase queueing delay. - enable_true_multi_threading: If enable_true_multi_threading=False, no more than 1 UDF - runs per actor. Otherwise, respects the `max_concurrency` argument. For more details, see - the `ActorPoolStrategy` class docstring. + enable_true_multi_threading: If enable_true_multi_threading=False, no more + than 1 UDF runs per actor. Otherwise, respects the `max_concurrency` argument. + By default, this flag is `None`, which gets translated to `False`. + For more details, see the `ActorPoolStrategy` class docstring. """ if size is not None: if size < 1: @@ -178,9 +179,23 @@ def __init__( self.max_tasks_in_flight_per_actor = max_tasks_in_flight_per_actor self.num_workers = 0 self.ready_to_total_workers_ratio = 0.8 - self.enable_true_multi_threading = enable_true_multi_threading + self._enable_true_multi_threading = enable_true_multi_threading + + @property + def enable_true_multi_threading(self) -> bool: + # backwards compatibility from serialization: instances pickled + # before this became a property carry the value under the public + # name instead. + return bool( + self.__dict__.get( + "_enable_true_multi_threading", + self.__dict__.get("enable_true_multi_threading"), + ) + ) def __eq__(self, other: Any) -> bool: + # intentionally compare resolved enable_true_multi_threading values + # because the two strategy classes are effectively the same. return isinstance(other, ActorPoolStrategy) and ( self.min_size == other.min_size and self.max_size == other.max_size diff --git a/python/ray/data/_internal/logical/operators/map_operator.py b/python/ray/data/_internal/logical/operators/map_operator.py index fbc1ccf7d0dc..3465ab3c553f 100644 --- a/python/ray/data/_internal/logical/operators/map_operator.py +++ b/python/ray/data/_internal/logical/operators/map_operator.py @@ -256,6 +256,42 @@ def __post_init__(self): "_name", self._get_operator_name(self.__class__.__name__, self.fn), ) + self._wrap_torch_inference() + + def _wrap_torch_inference(self) -> None: + """Detect a ``TorchInference`` UDF and wrap it in the managed + callable that drives collate/transfer/process/finalize. + """ + from ray.data._internal.utils.torch_inference import ( + is_torch_inference_class, + is_torch_inference_instance, + validate_torch_inference_op, + ) + + if is_torch_inference_instance(self.fn): + raise ValueError( + "Pass the `TorchInference` subclass to `map_batches`, " + "not an instance of it." + ) + if not is_torch_inference_class(self.fn): + return + + validate_torch_inference_op( + self.fn, + self.fn_args, + self.fn_kwargs, + self.compute, + self.ray_remote_args, + ) + self._set_torch_inference_udf() + + def _set_torch_inference_udf(self) -> None: + """Replace ``fn`` with the managed serial wrapper.""" + from ray.data._internal.utils.torch_inference import ( + make_torch_inference_callable, + ) + + object.__setattr__(self, "fn", make_torch_inference_callable(self.fn)) @dataclass(frozen=True, repr=False, eq=False) diff --git a/python/ray/data/_internal/utils/torch_inference.py b/python/ray/data/_internal/utils/torch_inference.py new file mode 100644 index 000000000000..e5747c0d06c1 --- /dev/null +++ b/python/ray/data/_internal/utils/torch_inference.py @@ -0,0 +1,261 @@ +"""Ray-managed execution of +:class:`~ray.data.util.torch_inference.TorchInference`. + +``Dataset.map_batches`` detects UDF classes that subclass +``TorchInference`` and wraps them in a managed callable that drives the +``collate`` -> host-to-device transfer -> ``process_on_device`` -> +device-to-host transfer -> ``finalize`` flow, so users only implement the +model-specific pieces. + +NOTE: This module must stay importable without ``torch``: detection runs for +every ``map_batches`` call, so anything needing torch is imported lazily. +""" + +import inspect +import logging +from collections.abc import Mapping +from typing import ( + TYPE_CHECKING, + Any, + Dict, + Iterable, + Optional, + Tuple, + Type, + Union, +) + +from ray.data.util.torch_inference import TorchInference + +if TYPE_CHECKING: + import torch + + from ray.data._internal.compute import ComputeStrategy + from ray.data.block import CallableClass, DataBatch + from ray.data.collate_fn import TensorBatchType + +logger = logging.getLogger(__name__) + +# The host side of the managed transfers: `collate` outputs must live here, +# and `process_on_device` outputs are moved back here before `finalize`. +_CPU_DEVICE = "cpu" + +# Methods a `TorchInference` subclass may override; validated (e.g. must +# not be async) before wrapping. +_TORCH_INFERENCE_METHODS = ( + "initialize", + "get_device", + "collate", + "process_on_device", + "finalize", +) + + +class _BaseTorchInferenceUDFWrapper: + """Marker base class for the Ray-managed ``TorchInference`` wrappers.""" + + +def is_torch_inference_class(fn: Any) -> bool: + """True iff ``fn`` is a ``TorchInference`` subclass (the class itself, + not an instance).""" + return isinstance(fn, type) and issubclass(fn, TorchInference) + + +def is_torch_inference_instance(fn: Any) -> bool: + """True iff ``fn`` is an *instance* of a ``TorchInference`` subclass + (users must pass the class, not an instance).""" + return isinstance(fn, TorchInference) + + +def validate_torch_inference_op( + cls: Type[TorchInference], + fn_args: Optional[Iterable[Any]], + fn_kwargs: Optional[Dict[str, Any]], + compute: "ComputeStrategy", + ray_remote_args: Dict[str, Any], +) -> None: + """Validate a ``map_batches`` call whose UDF is a ``TorchInference``. + + Raises for arguments the managed flow can't honor; warns for likely + misconfigurations. Does not modify any of its inputs. + (``fn_constructor_args``/``fn_constructor_kwargs`` are supported — they + are forwarded to ``initialize``.) + """ + from ray.data._internal.compute import ActorPoolStrategy + + assert isinstance(compute, ActorPoolStrategy) + + if fn_args or fn_kwargs: + raise ValueError( + "`fn_args` and `fn_kwargs` are not supported with a " + "`TorchInference`; its methods only take the batch." + ) + + for method_name in _TORCH_INFERENCE_METHODS: + method = getattr(cls, method_name, None) + if method is None: + continue + if inspect.iscoroutinefunction(method) or inspect.isasyncgenfunction(method): + raise TypeError( + f"`{cls.__name__}.{method_name}` must not be async; " + "`TorchInference` methods are called synchronously." + ) + + if "__call__" in cls.__dict__: + logger.warning( + f"`TorchInference` subclass `{cls.__name__}` defines " + "`__call__`, but it will not be used directly: batches flow " + "through the Ray-managed flow (`collate` -> `process_on_device` " + "-> `finalize`), driven by a Ray-provided `__call__`." + ) + + if not ray_remote_args.get("num_gpus"): + logger.warning( + f"`{cls.__name__}` is a `TorchInference` but `num_gpus` is " + "not set; the managed flow is GPU-only, so pass `num_gpus` to " + "`map_batches` so the actor is scheduled on a GPU." + ) + + +def split_batch_and_other( + ret: Union["TensorBatchType", Tuple["TensorBatchType", Any]], +) -> Tuple["TensorBatchType", Any]: + """Split a ``collate``/``process_on_device`` return into + ``(tensors, other)``. + + A 2-tuple always means ``(tensors, other)``; any other return is + ``(ret, None)``. (To return a batch that IS a pair of tensors, use a + list or also do: ``((tensors, tensors), None)``.) + """ + if isinstance(ret, tuple) and len(ret) == 2: + return ret[0], ret[1] + return ret, None + + +def _resolve_cuda_device(user: TorchInference) -> "torch.device": + """Resolve and validate ``user.get_device()`` into a concrete CUDA device.""" + import torch + + # `torch.device(...)` is idempotent, so a `get_device` that returns a + # device string still works. + device = torch.device(user.get_device()) + if device.type != "cuda": + raise ValueError( + f"`{type(user).__name__}.get_device()` must return a CUDA device; " + f"got `{device}`. The `TorchInference` flow is GPU-only." + ) + if torch.cuda.is_available() and device.index is None: + # Resolve "cuda" to a concrete index so different actor threads all + # deterministically resolve to the same device. + device = torch.device("cuda", torch.cuda.current_device()) + return device + + +def _validate_no_tensors_off_device( + batch: "TensorBatchType", device: "torch.device", requirement: str +) -> None: + """Raise if any tensor in ``batch`` is not on ``device``; ``requirement`` + is the caller's message prefix (what must hold, and why).""" + from ray.data._internal.utils.torch_utils import find_first_tensor_not_on_device + + off_device = find_first_tensor_not_on_device(batch, device) + if off_device is not None: + raise ValueError(f"{requirement}; found a tensor on `{off_device.device}`.") + + +def validate_collated_batch(collated: Any, user_cls: Type[TorchInference]) -> None: + """Validate the ``collate`` output: a ``TensorBatchType`` of CPU tensors.""" + import torch + + from ray.data.collate_fn import is_tensor_batch_type + + if not is_tensor_batch_type(collated): + raise ValueError( + f"`{user_cls.__name__}.collate` must return a `TensorBatchType` " + "(a `torch.Tensor`, sequence of tensors, or mapping of str to " + f"tensors), or a `(TensorBatchType, other)` tuple; got " + f"{type(collated)}." + ) + _validate_no_tensors_off_device( + collated, + torch.device(_CPU_DEVICE), + f"`{user_cls.__name__}.collate` must return CPU tensors (Ray " + "manages the transfer to the device)", + ) + + +def validate_processed_batch( + out: Any, device: "torch.device", user_cls: Type[TorchInference] +) -> None: + """Validate the ``process_on_device`` output: a ``TensorBatchType`` on + ``device``.""" + from ray.data.collate_fn import is_tensor_batch_type + + if not is_tensor_batch_type(out): + raise ValueError( + f"`{user_cls.__name__}.process_on_device` must return a " + "`TensorBatchType`, or a `(TensorBatchType, other)` tuple; got " + f"{type(out)}." + ) + _validate_no_tensors_off_device( + out, + device, + f"`{user_cls.__name__}.process_on_device` must return tensors on " + f"`{device}` (Ray manages the transfer back to the host)", + ) + + +def make_torch_inference_callable(user_cls: Type[TorchInference]) -> "CallableClass": + """Wrap a ``TorchInference`` subclass in a managed callable class. + + Per batch, the wrapper's ``__call__`` runs the flow serially: ``collate`` + (validated CPU tensors) -> synchronous host-to-device transfer -> + ``process_on_device`` (validated on-device tensors) -> synchronous + device-to-host transfer -> ``finalize``. Everything runs on the current + stream; there is no overlap between the stages. + """ + import torch + + from ray.data.util.torch_inference import TorchInference + from ray.data.util.torch_utils import move_tensors_to_device + + class _TorchInferenceUDFWrapper(_BaseTorchInferenceUDFWrapper): + def __init__(self, *args: Any, **kwargs: Any) -> None: + self._ti_user: TorchInference = user_cls(*args, **kwargs) + assert isinstance(self._ti_user, TorchInference) + + self._ti_device = _resolve_cuda_device(self._ti_user) + + def __repr__(self) -> str: + return repr(self._ti_user) + + @torch.no_grad() + def __call__(self, batch: "DataBatch") -> "DataBatch": + # Shallow copy: `collate` replacing keys in the batch mapping + # can't corrupt the `input_batch` handed to process_on_device/ + # finalize (the underlying arrays are shared, not copied). + input_batch = dict(batch) if isinstance(batch, Mapping) else batch + + collated, collated_other = split_batch_and_other( + self._ti_user.collate(batch) + ) + validate_collated_batch(collated, user_cls) + + moved = move_tensors_to_device( + collated, self._ti_device, non_blocking=False + ) + + out, output_other = split_batch_and_other( + self._ti_user.process_on_device(input_batch, moved, collated_other) + ) + validate_processed_batch(out, self._ti_device, user_cls) + + cpu_out = move_tensors_to_device(out, _CPU_DEVICE, non_blocking=False) + + return self._ti_user.finalize(input_batch, cpu_out, output_other) + + # Wrapping happens before the MapBatches logical op is built, so take the + # user's class name for operator naming (`_get_operator_name` uses + # `fn.__name__`). + _TorchInferenceUDFWrapper.__name__ = user_cls.__name__ + return _TorchInferenceUDFWrapper diff --git a/python/ray/data/_internal/utils/torch_utils.py b/python/ray/data/_internal/utils/torch_utils.py index aaa15f25424d..9af08ee8d031 100644 --- a/python/ray/data/_internal/utils/torch_utils.py +++ b/python/ray/data/_internal/utils/torch_utils.py @@ -1,5 +1,6 @@ from abc import ABC, abstractmethod -from typing import Any, Union +from collections.abc import Mapping +from typing import Any, Optional, Union import torch @@ -47,3 +48,31 @@ def __call__(self, batch: Union[TensorBatchType, CustomBatchType]) -> FinalizedD non_blocking=DEFAULT_TENSOR_NON_BLOCKING_TRANSFER, ) return FinalizedData(data=batch) + + +def _iter_tensors(batch: TensorBatchType): + """Yield every tensor in a TensorBatchType, recursively.""" + if isinstance(batch, torch.Tensor): + yield batch + elif isinstance(batch, Mapping): + for value in batch.values(): + yield from _iter_tensors(value) + elif isinstance(batch, (list, tuple)): + for value in batch: + yield from _iter_tensors(value) + + +def find_first_tensor_not_on_device( + batch: TensorBatchType, device: torch.device +) -> Optional[torch.Tensor]: + """Return the first tensor in ``batch`` not on ``device``, else None. + + A ``device`` without an index (e.g. plain ``cuda``) matches any index of + that device type. + """ + for tensor in _iter_tensors(batch): + if tensor.device.type != device.type or ( + device.index is not None and tensor.device.index != device.index + ): + return tensor + return None diff --git a/python/ray/data/tests/test_arrow_block.py b/python/ray/data/tests/test_arrow_block.py index 41c5ff4529b7..37e9381d5dba 100644 --- a/python/ray/data/tests/test_arrow_block.py +++ b/python/ray/data/tests/test_arrow_block.py @@ -347,6 +347,51 @@ def test_arrow_block_to_pandas_preserves_arrow_types_through_roundtrip( assert roundtripped.to_pydict() == {"x": expected_values} +def test_arrow_block_to_pandas_null_type_is_not_arrow_backed(): + # A block whose column is entirely null is typed pa.null(). Keeping that as + # null[pyarrow] makes the column unusable in pandas: fillna and masked + # assignment cannot box a non-null value into Arrow's null type. Fall back to + # pandas' default conversion instead, which still round-trips to pa.null(). + table = pa.table({"x": pa.array([None, None], type=pa.null())}) + + df = ArrowBlockAccessor(table).to_pandas() + assert not isinstance(df.dtypes["x"], pd.ArrowDtype) + assert df["x"].tolist() == [None, None] + + # Untouched, the column round-trips back to Arrow's null type. + roundtripped = BlockAccessor.for_block(df).to_arrow() + assert roundtripped.schema.field("x").type == pa.null() + assert roundtripped.to_pydict() == {"x": [None, None]} + + # Filling the nulls now works and yields the fill value's type. + filled = BlockAccessor.for_block(df.fillna({"x": 3.0})).to_arrow() + assert filled.to_pydict() == {"x": [3.0, 3.0]} + + +def test_pandas_udf_can_fill_per_block_null_columns(ray_start_regular_shared): + # One-row blocks type a column with no values as pa.null(), so a pandas UDF + # calling fillna used to fail with ArrowInvalid on whichever block happened + # to hold only nulls for that column. + ds = ray.data.from_items( + [ + {"a": 1.0, "b": 2.0}, + {"a": 3.0, "b": None}, + {"a": None, "b": 4.0}, + ], + override_num_blocks=3, + ) + + filled = ds.map_batches( + lambda df: df.fillna({"a": 0.0, "b": 0.0}), batch_format="pandas" + ) + + assert sorted(filled.take_all(), key=lambda row: row["a"]) == [ + {"a": 0.0, "b": 4.0}, + {"a": 1.0, "b": 2.0}, + {"a": 3.0, "b": 0.0}, + ] + + def test_arrow_block_to_pandas_opt_out_numpy_dtypes(restore_data_context): # https://github.com/ray-project/ray/issues/64765: opting out restores the # pre-2.56 numpy conversion, so standard Arrow types no longer become diff --git a/python/ray/data/tests/test_torch_inference.py b/python/ray/data/tests/test_torch_inference.py new file mode 100644 index 000000000000..2342a2266e67 --- /dev/null +++ b/python/ray/data/tests/test_torch_inference.py @@ -0,0 +1,538 @@ +import logging +from dataclasses import replace +from typing import Any + +import numpy as np +import pandas as pd +import pytest +import torch + +import ray +from ray.data._internal.logical.operators.map_operator import MapBatches +from ray.data._internal.utils.torch_inference import ( + is_torch_inference_class, + is_torch_inference_instance, + make_torch_inference_callable, +) +from ray.data._internal.utils.torch_utils import find_first_tensor_not_on_device +from ray.data.tests.conftest import * # noqa +from ray.data.util.torch_inference import TorchInference +from ray.tests.conftest import * # noqa + + +class MyInferenceActor(TorchInference): + # pyrefly: ignore[bad-override] # narrowing `*args/**kwargs` is the API. + def initialize(self): + self.scale = 2.0 + + def get_device(self): + return torch.device("cuda") + + def process_on_device(self, input_batch, collated_tensors, collated_other): + return {"y": collated_tensors["x"] * self.scale} + + +class PlainCallableActor: + def __call__(self, batch): + return batch + + +def _make_input_ds(n=32): + return ray.data.range(n).map_batches( + lambda b: {"x": b["id"].astype(np.float32)}, batch_size=None + ) + + +# ===== Base class defaults ===== + + +def test_init_calls_initialize(): + actor = MyInferenceActor() + assert actor.scale == 2.0 + + +def test_default_collate_converts_numpy_batch(): + actor = MyInferenceActor() + batch = { + "x": np.arange(6, dtype=np.float32).reshape(3, 2), + "id": np.arange(3, dtype=np.int64), + } + tensors = actor.collate(batch) + assert isinstance(tensors, dict) + assert set(tensors) == {"x", "id"} + x, ids = tensors["x"], tensors["id"] + assert isinstance(x, torch.Tensor) and isinstance(ids, torch.Tensor) + assert torch.equal(x, torch.as_tensor(batch["x"])) + assert x.dtype == torch.float32 + assert ids.dtype == torch.int64 + + +@pytest.mark.parametrize( + "bad_batch", + [ + pd.DataFrame({"x": [1.0, 2.0]}), + [1, 2, 3], + {"x": [1.0, 2.0]}, + {"x": torch.zeros(2)}, + ], +) +def test_default_collate_rejects_non_numpy_batch(bad_batch): + actor = MyInferenceActor() + with pytest.raises(TypeError, match="override `collate`"): + actor.collate(bad_batch) + + +def test_default_finalize_converts_tensor_mapping(): + actor = MyInferenceActor() + out = actor.finalize({}, {"y": torch.arange(3, dtype=torch.float32)}, None) + assert isinstance(out["y"], np.ndarray) + assert np.array_equal(out["y"], np.arange(3, dtype=np.float32)) + + +def test_default_finalize_converts_recursively(): + # The conversion preserves the dict/sequence structure and converts every + # tensor leaf. (Annotated Any: the nesting is deeper than the declared + # `TensorBatchType`, which the recursive default supports at runtime.) + actor = MyInferenceActor() + nested_output: Any = { + "flat": torch.zeros(2), + "chunks": [torch.zeros(2), torch.ones(2)], + "nested": {"inner": (torch.arange(3),)}, + } + out = actor.finalize({}, nested_output, None) + assert isinstance(out["flat"], np.ndarray) + assert isinstance(out["chunks"], list) + assert all(isinstance(chunk, np.ndarray) for chunk in out["chunks"]) + assert np.array_equal(out["chunks"][1], np.ones(2)) + assert isinstance(out["nested"]["inner"], tuple) + assert np.array_equal(out["nested"]["inner"][0], np.arange(3)) + + +@pytest.mark.parametrize( + "bad_output", + [ + "not tensors", + {"y": "not a tensor"}, + {"y": np.zeros(2)}, + {"y": [torch.zeros(2), None]}, + ], +) +def test_default_finalize_rejects_non_tensor_leaves(bad_output): + actor = MyInferenceActor() + with pytest.raises(TypeError, match="Override `finalize`"): + actor.finalize({}, bad_output, None) + + +def test_default_finalize_rejects_output_other(): + # The default can't know how to fold side data into the batch. + actor = MyInferenceActor() + with pytest.raises(TypeError, match="output_other"): + actor.finalize({}, {"y": torch.zeros(2)}, {"lengths": [1, 2]}) + + +def test_split_batch_and_other(): + from ray.data._internal.utils.torch_inference import split_batch_and_other + + tensors = {"x": torch.zeros(2)} + # Bare TensorBatchType -> no side data. + assert split_batch_and_other(tensors) == (tensors, None) + # A 2-tuple always means (tensors, other) — even when `other` is a + # tensor. A batch that IS a pair of tensors must be a list instead. + assert split_batch_and_other((tensors, {"dims": (2, 3)})) == ( + tensors, + {"dims": (2, 3)}, + ) + assert split_batch_and_other((tensors, None)) == (tensors, None) + t1, t2 = torch.zeros(2), torch.ones(2) + assert split_batch_and_other((t1, t2)) == (t1, t2) + # Non-2-tuples are never split. + triple = (t1, t2, torch.zeros(2)) + assert split_batch_and_other(triple) == (triple, None) + assert split_batch_and_other([t1, t2]) == ([t1, t2], None) + + +def test_default_get_device(): + class DefaultDevice(TorchInference): + def process_on_device(self, input_batch, collated_tensors, collated_other): + return collated_tensors + + actor = DefaultDevice() + if torch.cuda.is_available(): + assert actor.get_device() == torch.device("cuda") + else: + with pytest.raises(RuntimeError, match="CUDA is not available"): + actor.get_device() + + +def test_process_on_device_not_implemented(): + class NoProcess(TorchInference): + pass + + with pytest.raises(NotImplementedError, match="process_on_device"): + NoProcess().process_on_device({}, {}, None) + + +# ===== Detection ===== + + +def test_detection(): + assert is_torch_inference_class(MyInferenceActor) + assert is_torch_inference_class(TorchInference) + assert not is_torch_inference_class(PlainCallableActor) + assert not is_torch_inference_class(lambda b: b) + assert not is_torch_inference_class("not a class") + # Instances don't match the class check, but match the instance check. + actor = MyInferenceActor() + assert not is_torch_inference_class(actor) + assert is_torch_inference_instance(actor) + assert not is_torch_inference_instance(MyInferenceActor) + assert not is_torch_inference_instance(PlainCallableActor()) + + +def test_wrapper_not_detected_and_keeps_name(): + wrapper_cls = make_torch_inference_callable(MyInferenceActor) + # Composition: the wrapper is not a subclass, so plan rewrites re-running + # __post_init__ can't wrap it again. + assert not is_torch_inference_class(wrapper_cls) + assert wrapper_cls.__name__ == "MyInferenceActor" + + +def test_wrapper_rejects_non_cuda_device(): + # The device comes from the instance's `get_device()`, so the GPU-only + # check happens at actor init — before any CUDA state is touched, so it + # is testable without a GPU. + class CpuDevice(MyInferenceActor): + def get_device(self): + return torch.device("cpu") + + wrapper_cls = make_torch_inference_callable(CpuDevice) + with pytest.raises(ValueError, match="must return a CUDA device"): + wrapper_cls() + + +# ===== Tensor helpers ===== + + +def test_find_first_tensor_not_on_device(): + t_cpu = torch.zeros(2) + assert find_first_tensor_not_on_device({"a": t_cpu}, torch.device("cpu")) is None + assert find_first_tensor_not_on_device({"a": t_cpu}, torch.device("cuda")) is t_cpu + # Nested containers are searched; index-less specs match any index. + assert find_first_tensor_not_on_device([(t_cpu,)], torch.device("cuda:1")) is t_cpu + + +# ===== map_batches validation ===== + + +def test_instance_rejected(ray_start_regular_shared_2_cpus): + # `map_batches`'s generic UDF validation rejects the (non-callable) + # instance before the TorchInference-specific check; either way, + # passing an instance fails at map_batches time. + with pytest.raises(ValueError): + _make_input_ds().map_batches( + MyInferenceActor(), # pyrefly: ignore[bad-argument-type] + batch_size=8, + compute=ray.data.ActorPoolStrategy(size=1), + ) + + +@pytest.mark.parametrize( + "kwargs", + [ + {"fn_args": (1,)}, + {"fn_kwargs": {"a": 1}}, + ], +) +def test_fn_args_rejected(ray_start_regular_shared_2_cpus, kwargs): + with pytest.raises(ValueError, match="fn_args"): + _make_input_ds().map_batches( + MyInferenceActor, + batch_size=8, + compute=ray.data.ActorPoolStrategy(size=1), + **kwargs, + ) + + +def test_fn_constructor_args_forwarded_to_initialize(ray_start_regular_shared_2_cpus): + class Parameterized(TorchInference): + # pyrefly: ignore[bad-override] # narrowing `*args/**kwargs` is the API. + def initialize(self, scale, offset=0.0): + self.scale = scale + self.offset = offset + + def process_on_device(self, input_batch, collated_tensors, collated_other): + return collated_tensors + + # The base __init__ forwards constructor args to initialize. + actor = Parameterized(2.0, offset=1.0) + assert actor.scale == 2.0 and actor.offset == 1.0 + + # And map_batches accepts them for the wrapped actor. + ds = _make_input_ds().map_batches( + Parameterized, + batch_size=8, + compute=ray.data.ActorPoolStrategy(size=1), + fn_constructor_args=(2.0,), + fn_constructor_kwargs={"offset": 1.0}, + num_gpus=0.001, + ) + op = ds._logical_plan.dag + assert isinstance(op, MapBatches) + assert op.fn_constructor_args == (2.0,) + assert op.fn_constructor_kwargs == {"offset": 1.0} + + +def test_async_method_rejected(ray_start_regular_shared_2_cpus): + class AsyncCollate(MyInferenceActor): + # pyrefly: ignore[bad-override] # async is the defect under test. + async def collate(self, input_batch): + ... + + with pytest.raises(TypeError, match="async"): + _make_input_ds().map_batches( + AsyncCollate, + batch_size=8, + compute=ray.data.ActorPoolStrategy(size=1), + ) + + +def test_call_warns(ray_start_regular_shared_2_cpus, caplog, propagate_logs): + class WithCall(MyInferenceActor): + def __call__(self, batch): + return batch + + with caplog.at_level(logging.WARNING, logger="ray.data"): + _make_input_ds().map_batches( + WithCall, + batch_size=8, + compute=ray.data.ActorPoolStrategy(size=1), + num_gpus=0.001, + ) + assert "will not be used directly" in caplog.text + + +def test_missing_num_gpus_warns( + ray_start_regular_shared_2_cpus, caplog, propagate_logs +): + with caplog.at_level(logging.WARNING, logger="ray.data"): + _make_input_ds().map_batches( + MyInferenceActor, + batch_size=8, + compute=ray.data.ActorPoolStrategy(size=1), + ) + assert "num_gpus" in caplog.text + + +# ===== Wrapping ===== + + +def test_serial_wrap_applied(ray_start_regular_shared_2_cpus): + strategy = ray.data.ActorPoolStrategy(size=1) + ds = _make_input_ds().map_batches( + MyInferenceActor, + batch_size=8, + compute=strategy, + num_gpus=0.001, + ) + op = ds._logical_plan.dag + assert isinstance(op, MapBatches) + # The UDF is replaced by the managed wrapper, keeping the user's name. + assert op.fn is not MyInferenceActor + assert isinstance(op.fn, type) + assert op.fn.__name__ == "MyInferenceActor" + assert op.name == "MapBatches(MyInferenceActor)" + # The serial flow needs no concurrency normalization. + assert "max_concurrency" not in op.ray_remote_args + assert isinstance(op.compute, ray.data.ActorPoolStrategy) + assert op.compute.max_tasks_in_flight_per_actor is None + + +def test_replace_does_not_rewrap(ray_start_regular_shared_2_cpus): + # Logical-plan rewrites reconstruct operators with dataclasses.replace, + # which re-runs __post_init__ with the already-wrapped fn. + ds = _make_input_ds().map_batches( + MyInferenceActor, + batch_size=8, + compute=ray.data.ActorPoolStrategy(size=1), + num_gpus=0.001, + ) + op = ds._logical_plan.dag + assert isinstance(op, MapBatches) + op2 = replace(op, input_dependencies=op.input_dependencies) + assert op2.fn is op.fn + + +# ===== GPU end-to-end (runs in the GPU CI job; skipped without CUDA) ===== + +GPU_FEATURES = 64 +GPU_BATCH_SIZE = 256 +GPU_NUM_BATCHES = 8 +GPU_NUM_ROWS = GPU_NUM_BATCHES * GPU_BATCH_SIZE + + +def _make_gpu_source(): + """Rows where every element of row `id` is `1 + id`, so per-row sums are + unique, nonzero integers (exact in fp32) — a zeroed, torn, or cross-batch + read after the device transfer changes the checksum.""" + + def to_x(batch): + ids = np.asarray(batch["id"], dtype=np.int64) + vals = (1.0 + ids).astype(np.float32) + return {"id": ids, "data": np.repeat(vals[:, None], GPU_FEATURES, axis=1)} + + return ( + ray.data.range(GPU_NUM_ROWS, override_num_blocks=GPU_NUM_BATCHES) + .map_batches(to_x, batch_size=GPU_BATCH_SIZE) + .materialize() + ) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA") +def test_e2e_cuda(shutdown_only): + ray.init(num_cpus=2, num_gpus=1) + + # NOTE: Defined inside the test so cloudpickle serializes it by value + # (module-level test classes aren't importable from Ray workers). + class Predictor(TorchInference): + def collate(self, input_batch): + # Only the model input takes the tensor path; `id` stays in + # `input_batch`. The `(tensors, other)` form threads per-batch + # side data (here, the batch's min id) to process_on_device. + return ( + {"data": torch.from_numpy(input_batch["data"])}, + {"min_id": int(input_batch["id"].min())}, + ) + + def process_on_device(self, input_batch, collated_tensors, collated_other): + tensor = collated_tensors["data"] + # The managed flow must hand us device tensors, the untouched + # pre-collate batch, and collate's side data. + assert tensor.device.type == "cuda" + assert isinstance(input_batch["id"], np.ndarray) + assert collated_other == {"min_id": int(input_batch["id"].min())} + # Forward the side data on to finalize. + return ( + {"rowsum": tensor.sum(dim=1), "double": tensor[:, 0] * 2.0}, + collated_other, + ) + + def finalize(self, input_batch, output_tensors, output_other): + # process_on_device's side data arrives untouched. + assert output_other == {"min_id": int(input_batch["id"].min())} + return { + "id": input_batch["id"], + "rowsum": output_tensors["rowsum"].numpy(), + "double": output_tensors["double"].numpy(), + } + + ds = _make_gpu_source().map_batches( + Predictor, + batch_size=GPU_BATCH_SIZE, + batch_format="numpy", + zero_copy_batch=True, + compute=ray.data.ActorPoolStrategy(size=1), + num_gpus=1, + ) + + rows = sorted(ds.take_all(), key=lambda row: row["id"]) + ids = np.asarray([row["id"] for row in rows], dtype=np.int64) + rowsums = np.asarray([row["rowsum"] for row in rows]) + doubles = np.asarray([row["double"] for row in rows]) + + # Exactly-once id coverage (passthrough via `input_batch` intact). + assert np.array_equal(ids, np.arange(GPU_NUM_ROWS, dtype=np.int64)) + # Exact per-row checksums: the H2D delivered the right bytes. + assert np.array_equal(rowsums, (GPU_FEATURES * (1.0 + ids)).astype(np.float32)) + # Compute results survive the D2H. + assert np.array_equal(doubles, (2.0 * (1.0 + ids)).astype(np.float32)) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA") +def test_e2e_cuda_default_collate_and_finalize(shutdown_only): + # Only `process_on_device` implemented: the default `get_device` (cuda), + # `collate` (numpy -> tensors), and `finalize` (tensors -> numpy) carry + # the batch through the managed flow. `fn_constructor_args` reach + # `initialize`. + ray.init(num_cpus=2, num_gpus=1) + + class MinimalPredictor(TorchInference): + # pyrefly: ignore[bad-override] # narrowing `*args/**kwargs` is the API. + def initialize(self, scale): + self.scale = scale + + def process_on_device(self, input_batch, collated_tensors, collated_other): + return { + "id": collated_tensors["id"], + "rowsum": collated_tensors["data"].sum(dim=1) * self.scale, + } + + ds = _make_gpu_source().map_batches( + MinimalPredictor, + batch_size=GPU_BATCH_SIZE, + batch_format="numpy", + compute=ray.data.ActorPoolStrategy(size=1), + fn_constructor_args=(2.0,), + num_gpus=1, + ) + + rows = sorted(ds.take_all(), key=lambda row: row["id"]) + ids = np.asarray([row["id"] for row in rows], dtype=np.int64) + rowsums = np.asarray([row["rowsum"] for row in rows]) + assert np.array_equal(ids, np.arange(GPU_NUM_ROWS, dtype=np.int64)) + assert np.array_equal( + rowsums, (2.0 * GPU_FEATURES * (1.0 + ids)).astype(np.float32) + ) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA") +def test_collate_output_must_be_cpu(shutdown_only): + ray.init(num_cpus=2, num_gpus=1) + + class GpuCollate(TorchInference): + def collate(self, input_batch): + return {"data": torch.from_numpy(input_batch["data"]).cuda()} + + def process_on_device(self, input_batch, collated_tensors, collated_other): + return collated_tensors + + def finalize(self, input_batch, output_tensors, output_other): + return {"data": output_tensors["data"].numpy()} + + ds = _make_gpu_source().map_batches( + GpuCollate, + batch_size=GPU_BATCH_SIZE, + compute=ray.data.ActorPoolStrategy(size=1), + num_gpus=1, + ) + with pytest.raises(Exception, match="must return CPU tensors"): + ds.take_all() + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA") +def test_process_output_must_be_on_device(shutdown_only): + ray.init(num_cpus=2, num_gpus=1) + + class CpuProcess(TorchInference): + def collate(self, input_batch): + return {"data": torch.from_numpy(input_batch["data"])} + + def process_on_device(self, input_batch, collated_tensors, collated_other): + return {"data": collated_tensors["data"].cpu()} + + def finalize(self, input_batch, output_tensors, output_other): + return {"data": output_tensors["data"].numpy()} + + ds = _make_gpu_source().map_batches( + CpuProcess, + batch_size=GPU_BATCH_SIZE, + compute=ray.data.ActorPoolStrategy(size=1), + num_gpus=1, + ) + with pytest.raises(Exception, match="must return tensors on"): + ds.take_all() + + +if __name__ == "__main__": + import sys + + sys.exit(pytest.main(["-v", __file__])) diff --git a/python/ray/data/util/torch_inference.py b/python/ray/data/util/torch_inference.py new file mode 100644 index 000000000000..2c7762945831 --- /dev/null +++ b/python/ray/data/util/torch_inference.py @@ -0,0 +1,241 @@ +"""The user-facing base class for Torch batch inference with ``map_batches``.""" +from typing import TYPE_CHECKING, Any, Tuple, Union + +from ray.util.annotations import PublicAPI + +if TYPE_CHECKING: + import torch + + from ray.data.block import DataBatch + from ray.data.collate_fn import TensorBatchType + + +@PublicAPI(stability="alpha") +class TorchInference: + """Base class for PyTorch batch inference with + :meth:`~ray.data.Dataset.map_batches`. + + Subclass this and implement :meth:`process_on_device`. For every batch, + Ray Data runs: + + 1. :meth:`collate` — convert the batch into CPU torch tensors. + 2. Ray Data moves the tensors to the :meth:`get_device` device. + 3. :meth:`process_on_device` — compute on the device torch tensors. + 4. Ray Data moves the resulting torch tensors back to the CPU. + 5. :meth:`finalize` — convert them into the output batch. + + Override :meth:`collate`, :meth:`finalize`, or :meth:`get_device` only if + the defaults don't fit your data. + + **Passing non-tensor data between steps.** Only tensors go through the + managed device transfers, but :meth:`collate` and + :meth:`process_on_device` can each return a ``(tensors, other)`` tuple to + hand small side values (original shapes, padding lengths, strings, etc) + directly to the next step: ``collate``'s ``other`` arrives as + ``collated_other`` in :meth:`process_on_device`, and + ``process_on_device``'s ``other`` arrives as ``output_other`` in + :meth:`finalize`. Return a bare ``TensorBatchType`` and the next step + receives ``None`` instead. Ray Data passes ``other`` values through + untouched. + + Examples: + + Minimal — only :meth:`process_on_device`, everything else default: + + .. testcode:: + :skipif: True + + import numpy as np + import torch + import ray + from ray.data.util.torch_inference import TorchInference + + class MyInferenceActor(TorchInference): + + def initialize(self): + self.model = torch.nn.Identity().cuda().eval() + + def process_on_device( + self, input_batch, collated_tensors, collated_other + ): + out = self.model(collated_tensors["data"]) + return {"mean": out.float().mean(dim=1)} + + ds = ( + ray.data.from_numpy(np.ones((32, 100), dtype=np.float32)) + .map_batches( + MyInferenceActor, + batch_size=4, + compute=ray.data.ActorPoolStrategy(size=1), + num_gpus=1, + ) + ) + + .. note:: + Don't define ``__init__`` in your subclass — put setup code in + :meth:`initialize`, which Ray Data calls once when the actor starts. + + .. note:: + With the default :meth:`collate`, the ``batch_format`` given to + ``map_batches`` must be ``"default"`` or ``"numpy"``. + + The per-batch methods run under ``torch.no_grad()``; gradients are not + recorded. + """ + + def __init__(self, *args: Any, **kwargs: Any): + """Forward the constructor arguments to :meth:`initialize`. + + Args: + *args: Forwarded to :meth:`initialize`. + **kwargs: Forwarded to :meth:`initialize`. + """ + self.initialize(*args, **kwargs) + + def initialize(self, *args: Any, **kwargs: Any) -> None: + """Initialize actor state, such as the model. + + Called once when the actor starts. Override this instead of defining + ``__init__``. + + Args: + *args: The ``fn_constructor_args`` given to + :meth:`~ray.data.Dataset.map_batches`. + **kwargs: The ``fn_constructor_kwargs`` given to + :meth:`~ray.data.Dataset.map_batches`. + """ + pass + + def get_device(self) -> "torch.device": + """Return the device batches are processed on. + + Called once when the actor starts. The default returns + ``torch.device("cuda")``. + + Returns: + The device that :meth:`collate` outputs are moved to and that + :meth:`process_on_device` runs on. Must be a CUDA device. + """ + import torch + + if not torch.cuda.is_available(): + raise RuntimeError( + "CUDA is not available on this system. The default TorchInference " + "flow is GPU-only and requires a CUDA-capable device." + ) + return torch.device("cuda") + + def collate( + self, input_batch: "DataBatch" + ) -> Union["TensorBatchType", Tuple["TensorBatchType", Any]]: + """Convert an input batch into CPU tensors. + + Nested tensor sequences (e.g. ``Dict[str, List[Tensor]]``) are treated + as chunks of one logical tensor and concatenated along the batch + dimension during the device transfer. Use flat structures to preserve + tensor shapes as-is. + + Args: + input_batch: The batch to convert, in the ``batch_format`` given + to :meth:`~ray.data.Dataset.map_batches`. The default + implementation only supports NumPy batches + (``Dict[str, np.ndarray]``) and raises ``TypeError`` for + other batch formats. + + Returns: + The batch as Torch tensors, optionally with side data: + + - ``tensors``: the tensors must be on the CPU; Ray Data moves + them to the :meth:`get_device` device before calling + :meth:`process_on_device`. + - ``(tensors, other)``: additionally hand ``other`` — any + non-tensor side value, passed through untouched — to + :meth:`process_on_device` as ``collated_other``. When only + ``tensors`` is returned, ``collated_other`` is ``None``. + """ + import numpy as np + + from ray.data.util.torch_utils import ( + _get_type_str, + convert_ndarray_batch_to_torch_tensor_batch, + ) + + if not isinstance(input_batch, dict) or not all( + isinstance(column, np.ndarray) for column in input_batch.values() + ): + raise TypeError( + "The default `collate` only supports NumPy batches " + f"(`Dict[str, np.ndarray]`); got {_get_type_str(input_batch)}. " + 'Use `batch_format="numpy"` in `map_batches`, or override ' + "`collate` to convert the batch yourself." + ) + return convert_ndarray_batch_to_torch_tensor_batch(input_batch) + + def process_on_device( + self, + input_batch: "DataBatch", + collated_tensors: "TensorBatchType", + collated_other: Any, + ) -> Union["TensorBatchType", Tuple["TensorBatchType", Any]]: + """Process a batch of device tensors and return device tensors. + + This is the only method a subclass must implement. + + Args: + input_batch: The untouched input batch (pre-:meth:`collate`), for + reading fields that don't go through the tensor path. + collated_tensors: The tensors returned by :meth:`collate`, moved + to the :meth:`get_device` device. + collated_other: The side data returned by :meth:`collate`, or + ``None`` if :meth:`collate` returned only tensors. + + Returns: + The resulting Torch tensors, optionally with side data: + + - ``tensors``: the tensors must be on the device; Ray Data moves + them back to the CPU before calling :meth:`finalize`. + - ``(tensors, other)``: additionally hand ``other`` — any + non-tensor side value, passed through untouched — to + :meth:`finalize` as ``output_other``. When only ``tensors`` is + returned, ``output_other`` is ``None``. + """ + raise NotImplementedError( + f"`{type(self).__name__}` must implement `process_on_device`." + ) + + def finalize( + self, + input_batch: "DataBatch", + output_tensors: "TensorBatchType", + output_other: Any, + ) -> "DataBatch": + """Convert the output tensors into the batch ``map_batches`` returns. + + Args: + input_batch: The untouched input batch (pre-:meth:`collate`), for + passing fields through to the output. + output_tensors: The tensors returned by :meth:`process_on_device`, + already moved back to the CPU. + output_other: The side data returned by :meth:`process_on_device`, + or ``None`` if it returned only tensors. + + Returns: + The output batch, in a format + :meth:`~ray.data.Dataset.map_batches` accepts. The default + implementation recursively converts every Torch tensor into a + NumPy array, preserving the surrounding dict/sequence structure + (e.g. ``Dict[str, torch.Tensor]`` becomes + ``Dict[str, np.ndarray]``), and raises ``TypeError`` for + non-tensor values — or for a non-``None`` ``output_other``, which + it doesn't know how to fold into the batch; override ``finalize`` + to consume it. + """ + from ray.data.util.torch_utils import convert_tensors_to_numpy + + if output_other is not None: + raise TypeError( + "The default `finalize` doesn't know how to fold " + "`output_other` into the output batch. Override `finalize` " + "to consume the side data returned by `process_on_device`." + ) + return convert_tensors_to_numpy(output_tensors) diff --git a/python/ray/data/util/torch_utils.py b/python/ray/data/util/torch_utils.py index 7a14a657ff31..5c8e6867a670 100644 --- a/python/ray/data/util/torch_utils.py +++ b/python/ray/data/util/torch_utils.py @@ -519,3 +519,23 @@ def move_tensors_to_device( "Dict[str, torch.Tensor], " "Mapping[str, List/Tuple[torch.Tensor]]" ) + + +def convert_tensors_to_numpy(batch: TensorBatchType) -> Any: + """Recursively convert every Torch tensor in ``batch`` to a NumPy array, + preserving the surrounding dict/sequence structure. + + Raises ``TypeError`` on non-tensor leaves (the default + ``TorchInference.finalize`` error path). + """ + if _is_tensor(batch): + return batch.detach().cpu().numpy() + elif isinstance(batch, Mapping): + return {k: convert_tensors_to_numpy(v) for k, v in batch.items()} + elif isinstance(batch, (list, tuple)): + return type(batch)(convert_tensors_to_numpy(v) for v in batch) + raise TypeError( + "The default `finalize` only supports Torch tensors nested in " + f"dicts/sequences; got {_get_type_str(batch)}. Override `finalize` " + "to convert the output yourself." + ) diff --git a/python/ray/tests/accelerators/test_intel_gpu.py b/python/ray/tests/accelerators/test_intel_gpu.py index b74dd5296265..9e099d532c84 100644 --- a/python/ray/tests/accelerators/test_intel_gpu.py +++ b/python/ray/tests/accelerators/test_intel_gpu.py @@ -12,7 +12,24 @@ from ray.util.accelerators import INTEL_MAX_1100, INTEL_MAX_1550 -def test_visible_intel_gpu_ids(shutdown_only): +@pytest.fixture(autouse=False) +def clean_accelerator_env(): + """Restore all Intel GPU env vars after each test, even if the test fails.""" + keys = ( + "ZE_AFFINITY_MASK", + "ONEAPI_DEVICE_SELECTOR", + "RAY_EXPERIMENTAL_NOSET_ZE_AFFINITY_MASK", + ) + saved = {k: os.environ.get(k) for k in keys} + yield + for k, v in saved.items(): + if v is None: + os.environ.pop(k, None) + else: + os.environ[k] = v + + +def test_visible_intel_gpu_ids(shutdown_only, clean_accelerator_env): with patch.object(Accelerator, "get_current_node_num_accelerators", return_value=4): os.environ["ONEAPI_DEVICE_SELECTOR"] = "level_zero:0,1,2" # Delete the cache so it can be re-populated the next time @@ -25,7 +42,7 @@ def test_visible_intel_gpu_ids(shutdown_only): assert ray.available_resources()["GPU"] == 3 -def test_visible_intel_gpu_type(shutdown_only): +def test_visible_intel_gpu_type(shutdown_only, clean_accelerator_env): with patch.object( Accelerator, "get_current_node_num_accelerators", return_value=4 ), patch.object( @@ -80,11 +97,26 @@ def test_get_current_node_accelerator_type(): def test_intel_gpu_accelerator_manager_api(): assert Accelerator.get_resource_name() == "GPU" - assert Accelerator.get_visible_accelerator_ids_env_var() == "ONEAPI_DEVICE_SELECTOR" + # Primary env var is now ZE_AFFINITY_MASK (bare IDs, subprocess-safe). + assert Accelerator.get_visible_accelerator_ids_env_var() == "ZE_AFFINITY_MASK" assert Accelerator.validate_resource_request_quantity(0.1) == (True, None) -def test_get_current_process_visible_accelerator_ids(): +def test_get_current_process_visible_accelerator_ids(clean_accelerator_env): + # ZE_AFFINITY_MASK is the primary — read it first. + os.environ["ZE_AFFINITY_MASK"] = "0,1,2" + assert Accelerator.get_current_process_visible_accelerator_ids() == ["0", "1", "2"] + + # No vars set at all — must return None (not empty list). + del os.environ["ZE_AFFINITY_MASK"] + assert Accelerator.get_current_process_visible_accelerator_ids() is None + + # Empty string means "no devices allowed". + os.environ["ZE_AFFINITY_MASK"] = "" + assert Accelerator.get_current_process_visible_accelerator_ids() == [] + del os.environ["ZE_AFFINITY_MASK"] + + # Backward compat: falls back to ONEAPI_DEVICE_SELECTOR when ZE_AFFINITY_MASK absent. os.environ["ONEAPI_DEVICE_SELECTOR"] = "level_zero:0,1,2" assert Accelerator.get_current_process_visible_accelerator_ids() == ["0", "1", "2"] @@ -97,20 +129,76 @@ def test_get_current_process_visible_accelerator_ids(): os.environ["ONEAPI_DEVICE_SELECTOR"] = "NoDevFiles" assert Accelerator.get_current_process_visible_accelerator_ids() == [] - del os.environ["ONEAPI_DEVICE_SELECTOR"] - - -def test_set_current_process_visible_accelerator_ids(): - Accelerator.set_current_process_visible_accelerator_ids(["0"]) - assert os.environ["ONEAPI_DEVICE_SELECTOR"] == "level_zero:0" +@pytest.mark.parametrize( + "physical_ids, expected_ze", + [ + # No devices visible — ZE_AFFINITY_MASK is set to "" (parsed back as []). + ([], ""), + # GPU0 only. + (["0"], "0"), + # GPU0 + GPU1. + (["0", "1"], "0,1"), + # GPU1 only — physical id is carried as-is (bare, like CUDA_VISIBLE_DEVICES). + (["1"], "1"), + # Non-contiguous physical ids. + (["1", "3"], "1,3"), + ], +) +def test_set_current_process_visible_accelerator_ids( + clean_accelerator_env, physical_ids, expected_ze +): + # Ray only sets ZE_AFFINITY_MASK (bare physical ids). It does not write + # ONEAPI_DEVICE_SELECTOR, leaving that var free for frameworks like SGLang. + Accelerator.set_current_process_visible_accelerator_ids(physical_ids) + assert os.environ["ZE_AFFINITY_MASK"] == expected_ze + assert "ONEAPI_DEVICE_SELECTOR" not in os.environ + + +def test_set_current_process_visible_accelerator_ids_roundtrip(clean_accelerator_env): + # What the setter writes must read back as the original physical ids, so a + # reused worker sees exactly the devices Ray assigned. + os.environ.pop("ONEAPI_DEVICE_SELECTOR", None) + Accelerator.set_current_process_visible_accelerator_ids(["1", "3"]) + assert Accelerator.get_current_process_visible_accelerator_ids() == ["1", "3"] + + +def test_set_visible_accelerator_ids_noset_ze(clean_accelerator_env): + # RAY_EXPERIMENTAL_NOSET_ZE_AFFINITY_MASK suppresses the only var Ray sets, so + # neither ZE_AFFINITY_MASK nor ONEAPI_DEVICE_SELECTOR is touched. + os.environ.pop("ONEAPI_DEVICE_SELECTOR", None) + os.environ["RAY_EXPERIMENTAL_NOSET_ZE_AFFINITY_MASK"] = "1" + Accelerator.set_current_process_visible_accelerator_ids(["1", "3"]) + assert "ZE_AFFINITY_MASK" not in os.environ + assert "ONEAPI_DEVICE_SELECTOR" not in os.environ + + +def test_set_visible_accelerator_ids_leaves_oneapi_untouched(clean_accelerator_env): + # A pre-existing ONEAPI_DEVICE_SELECTOR (e.g. set by the user or a framework) + # must survive: Ray sets only ZE_AFFINITY_MASK and never overwrites ONEAPI. + os.environ["ONEAPI_DEVICE_SELECTOR"] = "level_zero:2" Accelerator.set_current_process_visible_accelerator_ids(["0", "1"]) - assert os.environ["ONEAPI_DEVICE_SELECTOR"] == "level_zero:0,1" - - Accelerator.set_current_process_visible_accelerator_ids(["0", "1", "2"]) - assert os.environ["ONEAPI_DEVICE_SELECTOR"] == "level_zero:0,1,2" - - del os.environ["ONEAPI_DEVICE_SELECTOR"] + assert os.environ["ZE_AFFINITY_MASK"] == "0,1" + assert os.environ["ONEAPI_DEVICE_SELECTOR"] == "level_zero:2" + + +def test_visible_accelerator_env_var_restored(clean_accelerator_env): + # The only var Ray sets is the primary one from get_visible_accelerator_ids_env_var, + # so Ray's existing save/restore path fully cleans it up on reused workers. + assert Accelerator.get_visible_accelerator_ids_env_var() == "ZE_AFFINITY_MASK" + os.environ.pop("ZE_AFFINITY_MASK", None) + + env_var = Accelerator.get_visible_accelerator_ids_env_var() + saved = os.environ.get(env_var) # None + Accelerator.set_current_process_visible_accelerator_ids(["1", "3"]) + assert os.environ["ZE_AFFINITY_MASK"] == "1,3" + + # Restore, mirroring reset_visible_accelerator_env_vars. + if saved is None: + os.environ.pop(env_var, None) + else: + os.environ[env_var] = saved + assert "ZE_AFFINITY_MASK" not in os.environ if __name__ == "__main__": diff --git a/python/ray/tests/accelerators/test_intel_gpu_e2e.py b/python/ray/tests/accelerators/test_intel_gpu_e2e.py index 1b1426595b5a..22da881c473c 100644 --- a/python/ray/tests/accelerators/test_intel_gpu_e2e.py +++ b/python/ray/tests/accelerators/test_intel_gpu_e2e.py @@ -94,7 +94,8 @@ def gpu_task() -> Dict[str, Any]: return { "gpu_ids": gpu_ids, "pid": os.getpid(), - "oneapi_selector": os.environ.get("ONEAPI_DEVICE_SELECTOR"), + "ze_affinity_mask": os.environ.get("ZE_AFFINITY_MASK"), + "selector": os.environ.get("ONEAPI_DEVICE_SELECTOR"), } @@ -106,6 +107,7 @@ def cluster_probe_task() -> Dict[str, Any]: "node_ip": ray.util.get_node_ip_address(), "worker_id": context.get_worker_id(), "gpu_ids": context.get_accelerator_ids().get("GPU", []), + "ze_affinity_mask": os.environ.get("ZE_AFFINITY_MASK"), "selector": os.environ.get("ONEAPI_DEVICE_SELECTOR"), } @@ -118,7 +120,7 @@ def assert_valid_gpu_binding(result: Dict[str, Any], label: str) -> None: def _validate_gpu_binding_common( - result: Dict[str, Any], label: str, selector_key: str = "oneapi_selector" + result: Dict[str, Any], label: str, ze_key: str = "ze_affinity_mask" ) -> int: """Validate basic GPU binding properties shared by single- and multi-GPU tests.""" @@ -127,17 +129,30 @@ def _validate_gpu_binding_common( primary_gpu_id = int(gpu_ids[0]) - selector = result.get(selector_key) - assert selector, f"ONEAPI_DEVICE_SELECTOR not set in environment for {label}." - selector_lower = selector.lower() - assert ( - "level_zero:" in selector_lower - ), f"ONEAPI_DEVICE_SELECTOR should target GPU devices for {label}, got: {selector}." - - selector_gpu_ids = {int(match) for match in re.findall(r"\b\d+\b", selector_lower)} + # ZE_AFFINITY_MASK carries the physical device IDs (bare, e.g. "0" or "1,3"). + # This is the primary env var and must match what Ray assigned. + ze_mask = result.get(ze_key) + assert ze_mask is not None, f"ZE_AFFINITY_MASK not set in environment for {label}." + ze_gpu_ids = {int(x) for x in ze_mask.split(",")} assert ( - primary_gpu_id in selector_gpu_ids - ), f"ONEAPI_DEVICE_SELECTOR does not reference bound GPU id for {label}: {selector}." + primary_gpu_id in ze_gpu_ids + ), f"ZE_AFFINITY_MASK does not reference bound GPU id for {label}: {ze_mask}." + + # ONEAPI_DEVICE_SELECTOR carries re-indexed sequential IDs (e.g. "level_zero:0") + # because ZE_AFFINITY_MASK has already re-indexed devices before it applies. + # Check it is present and has the correct format — but do not compare its IDs + # to the physical gpu_ids since they are intentionally different. + selector = result.get("selector") + if selector is not None: + assert ( + "level_zero:" in selector.lower() + ), f"ONEAPI_DEVICE_SELECTOR should target GPU devices for {label}, got: {selector}." + expected_count = len(gpu_ids) + selector_ids = re.findall(r"\b\d+\b", selector) + assert len(selector_ids) == expected_count, ( + f"ONEAPI_DEVICE_SELECTOR should have {expected_count} re-indexed id(s) " + f"for {label}, got: {selector}." + ) return primary_gpu_id @@ -215,7 +230,7 @@ def test_scale_out_task_distribution(ray_gpu_session, num_nodes) -> None: } for result in probe_results: - _validate_gpu_binding_common(result, "scale-out probe task", "selector") + _validate_gpu_binding_common(result, "scale-out probe task", "ze_affinity_mask") assert len(node_ids) == num_nodes or len(node_ips) == num_nodes, ( f"Expected probe tasks to execute on {num_nodes} distinct nodes, " diff --git a/python/ray/tests/accelerators/test_tpu.py b/python/ray/tests/accelerators/test_tpu.py index 242504e21db3..dfa2a1bbb903 100644 --- a/python/ray/tests/accelerators/test_tpu.py +++ b/python/ray/tests/accelerators/test_tpu.py @@ -104,9 +104,15 @@ def test_autodetect_num_tpus_vfio_mixed_groups(mock_list, mock_glob): "/sys/kernel/iommu_groups/10/devices": ["0000:01:00.0"], "/sys/kernel/iommu_groups/96/devices": ["0016:03:00.2"], } + # Build keys with os.path.join so they match production path construction + # on both POSIX and Windows (where join inserts backslashes). vendor_results = { - "/sys/kernel/iommu_groups/10/devices/0000:01:00.0/vendor": "0x1ae0\n", - "/sys/kernel/iommu_groups/96/devices/0016:03:00.2/vendor": "0x15b3\n", + os.path.join( + "/sys/kernel/iommu_groups/10/devices", "0000:01:00.0", "vendor" + ): "0x1ae0\n", + os.path.join( + "/sys/kernel/iommu_groups/96/devices", "0016:03:00.2", "vendor" + ): "0x15b3\n", } def fake_listdir(path): diff --git a/rllib/BUILD.bazel b/rllib/BUILD.bazel index 05fcb65ecb69..3e753ef3fdc0 100644 --- a/rllib/BUILD.bazel +++ b/rllib/BUILD.bazel @@ -2328,6 +2328,17 @@ py_test( deps = [":conftest"], ) +py_test( + name = "test_torch_squashed_gaussian", + size = "small", + srcs = ["core/distribution/torch/tests/test_torch_squashed_gaussian.py"], + tags = [ + "core", + "team:rllib", + ], + deps = [":conftest"], +) + # Default Models py_test( name = "test_base_models", diff --git a/rllib/algorithms/cql/torch/default_cql_torch_rl_module.py b/rllib/algorithms/cql/torch/default_cql_torch_rl_module.py index 1c2e7a7a2301..8c3a7cae691c 100644 --- a/rllib/algorithms/cql/torch/default_cql_torch_rl_module.py +++ b/rllib/algorithms/cql/torch/default_cql_torch_rl_module.py @@ -165,13 +165,11 @@ def _repeat_actions( action_logits = self.pi(pi_encoder_outs[ENCODER_OUT]) # Generate the squashed Gaussian from the model's logits. action_dist = self.get_train_action_dist_cls().from_logits(action_logits) - # Sample the actions. Note, we want to make a backward pass through - # these actions. - output[Columns.ACTIONS] = action_dist.rsample() - # Compute the action log-probabilities. - output[Columns.ACTION_LOGP] = action_dist.logp( - output[Columns.ACTIONS] - ).view(batch_size, num_actions, 1) + # Sample the actions (reparameterized, for backprop) together with + # their log-probs, so logp is exact even if the policy saturates. + actions, action_logp = action_dist.rsample_and_logp() + output[Columns.ACTIONS] = actions + output[Columns.ACTION_LOGP] = action_logp.view(batch_size, num_actions, 1) else: output[Columns.ACTIONS] = actions diff --git a/rllib/algorithms/sac/torch/default_sac_torch_rl_module.py b/rllib/algorithms/sac/torch/default_sac_torch_rl_module.py index 2c198bf7eaa1..ca65cf7359dc 100644 --- a/rllib/algorithms/sac/torch/default_sac_torch_rl_module.py +++ b/rllib/algorithms/sac/torch/default_sac_torch_rl_module.py @@ -153,16 +153,15 @@ def _forward_train_continuous(self, batch: Dict[str, Any]) -> Dict[str, Any]: # Sample actions for the current state. Note that we need to apply the # reparameterization trick (`rsample()` instead of `sample()`) to avoid the # expectation over actions. - actions_resampled = action_dist_curr.rsample() - # Compute the log probabilities for the current state (for the critic loss). - output["logp_resampled"] = action_dist_curr.logp(actions_resampled) - - # Sample actions for the next state. - actions_next_resampled = action_dist_next.sample().detach() - # Compute the log probabilities for the next state. - output["logp_next_resampled"] = ( - action_dist_next.logp(actions_next_resampled) - ).detach() + ( + actions_resampled, + output["logp_resampled"], + ) = action_dist_curr.rsample_and_logp() + + # Sample actions for the next state + actions_next_resampled, logp_next_resampled = action_dist_next.sample_and_logp() + actions_next_resampled = actions_next_resampled.detach() + output["logp_next_resampled"] = logp_next_resampled.detach() # Compute Q-values for the current policy in the current state with # the sampled actions. diff --git a/rllib/algorithms/tqc/torch/default_tqc_torch_rl_module.py b/rllib/algorithms/tqc/torch/default_tqc_torch_rl_module.py index 8cfb1f5b9dbb..ab5b0f44687b 100644 --- a/rllib/algorithms/tqc/torch/default_tqc_torch_rl_module.py +++ b/rllib/algorithms/tqc/torch/default_tqc_torch_rl_module.py @@ -96,8 +96,7 @@ def _forward_train(self, batch: Dict[str, Any]) -> Dict[str, Any]: # Sample actions from current policy for current observations action_dist_class = self.catalog.get_action_dist_cls(framework=self.framework) action_dist_curr = action_dist_class.from_logits(pi_out) - actions_curr = action_dist_curr.rsample() - logp_curr = action_dist_curr.logp(actions_curr) + actions_curr, logp_curr = action_dist_curr.rsample_and_logp() output["actions_curr"] = actions_curr output["logp_curr"] = logp_curr @@ -128,8 +127,7 @@ def _forward_train(self, batch: Dict[str, Any]) -> Dict[str, Any]: # Sample actions for next state action_dist_next = action_dist_class.from_logits(pi_out_next) - actions_next = action_dist_next.rsample() - logp_next = action_dist_next.logp(actions_next) + actions_next, logp_next = action_dist_next.rsample_and_logp() output["actions_next"] = actions_next output["logp_next"] = logp_next diff --git a/rllib/core/distribution/torch/tests/__init__.py b/rllib/core/distribution/torch/tests/__init__.py new file mode 100644 index 000000000000..e69de29bb2d1 diff --git a/rllib/core/distribution/torch/tests/test_torch_squashed_gaussian.py b/rllib/core/distribution/torch/tests/test_torch_squashed_gaussian.py new file mode 100644 index 000000000000..88800c1a22bc --- /dev/null +++ b/rllib/core/distribution/torch/tests/test_torch_squashed_gaussian.py @@ -0,0 +1,30 @@ +import unittest + +import torch + +from ray.rllib.core.distribution.torch.torch_distribution import TorchSquashedGaussian + + +class TestTorchSquashedGaussian(unittest.TestCase): + def test_logp_is_positive_for_a_saturated_policy(self): + # Regression test for #65036 + # A confident policy (large mean, small std) puts almost all of its mass + # near the action boundary, so its own samples have a large *positive* + # log-density. If we regress, we return a large + # negative value instead. The fix must return a finite, positive logp. + loc = torch.full((1024, 6), 10.0) # mean well past tanh's linear range + log_std = torch.full((1024, 6), -2.0) # small std -> confident policy + dist = TorchSquashedGaussian.from_logits(torch.cat([loc, log_std], dim=-1)) + + _, logp = dist.rsample_and_logp() + + self.assertTrue(torch.isfinite(logp).all()) + self.assertGreater(logp.mean().item(), 0.0) + + +if __name__ == "__main__": + import sys + + import pytest + + sys.exit(pytest.main(["-v", __file__])) diff --git a/rllib/core/distribution/torch/torch_distribution.py b/rllib/core/distribution/torch/torch_distribution.py index d0b94828a9c0..f1e09c422573 100644 --- a/rllib/core/distribution/torch/torch_distribution.py +++ b/rllib/core/distribution/torch/torch_distribution.py @@ -4,6 +4,7 @@ already be familiar with. """ import abc +import math from typing import Dict, Iterable, List, Optional import gymnasium as gym @@ -67,6 +68,14 @@ def rsample( ) return rsample + def sample_and_logp(self, *, sample_shape=None): + sample = self.sample(sample_shape=sample_shape) + return sample, self.logp(sample) + + def rsample_and_logp(self, *, sample_shape=None): + rsample = self.rsample(sample_shape=sample_shape) + return rsample, self.logp(rsample) + @classmethod @override(Distribution) def from_logits(cls, logits: TensorType, **kwargs) -> "TorchDistribution": @@ -263,35 +272,56 @@ def _get_torch_distribution(self, loc, scale) -> "torch.distributions.Distributi def sample( self, *, sample_shape=None ) -> Union[TensorType, Tuple[TensorType, TensorType]]: - # Sample from the Normal distribution. sample = super().sample( sample_shape=sample_shape if sample_shape is not None else torch.Size() ) - # Return the squashed sample. return self._squash(sample) @override(TorchDistribution) def rsample( self, *, sample_shape=None ) -> Union[TensorType, Tuple[TensorType, TensorType]]: - # Sample from the Normal distribution. sample = super().rsample( sample_shape=sample_shape if sample_shape is not None else torch.Size() ) - # Return the squashed sample. return self._squash(sample) @override(TorchDistribution) def logp(self, value: TensorType, **kwargs) -> TensorType: - # Unsquash value. - value = self._unsquash(value) - # Get log-probabilities from Normal distribution. - logp = super().logp(value, **kwargs) - # Clip the log probabilities as a safeguard and sum. - logp = torch.clamp(logp, -100, 100).sum(-1) - # Return the log probabilities for squashed Normal. - value = torch.tanh(value) - return logp - torch.log(1 - value**2 + SMALL_NUMBER).sum(-1) + """Returns the log probability of `value` under the squashed Gaussian. + + Exact for interior actions. + `rsample_and_logp()` / `sample_and_logp()` are always exact. + """ + return self._squashed_logp(self._unsquash(value)) + + def _squashed_logp(self, unsquashed: TensorType) -> TensorType: + # Normal log-prob of the pre-squash sample minus the tanh Jacobian + # log(1 - tanh(x)^2), via the numerically stable identity + # log(1 - tanh(x)^2) = 2 * (log(2) - x - softplus(-2x)). + # This is a known SAC implementation trick. + # See Haarnoja et al. 2018 (arXiv:1801.01290), Appendix C + # Eq. 21 for the tanh change of variables, and OpenAI Spinning Up's SAC + # (spinup/algos/pytorch/sac/core.py) for this numerically stable form. + logp = super().logp(unsquashed).sum(-1) + jacobian = 2.0 * ( + math.log(2.0) - unsquashed - nn.functional.softplus(-2.0 * unsquashed) + ) + return logp - jacobian.sum(-1) + + @override(TorchDistribution) + def rsample_and_logp(self, *, sample_shape=None): + unsquashed = super().rsample( + sample_shape=sample_shape if sample_shape is not None else torch.Size() + ) + return self._squash(unsquashed), self._squashed_logp(unsquashed) + + @override(TorchDistribution) + def sample_and_logp(self, *, sample_shape=None): + unsquashed = super().sample( + sample_shape=sample_shape if sample_shape is not None else torch.Size() + ) + return self._squash(unsquashed), self._squashed_logp(unsquashed) @override(TorchDistribution) def entropy(self) -> TensorType: diff --git a/src/ray/common/status_or.h b/src/ray/common/status_or.h index a048e9f4652c..626c84c60b02 100644 --- a/src/ray/common/status_or.h +++ b/src/ray/common/status_or.h @@ -94,8 +94,14 @@ class StatusOr { return *this; } if (rhs.ok()) { - status_ = Status::OK(); - AssignValue(rhs.value()); + // Reuse the existing value if we already hold one; otherwise construct it + // and set status_ only after the construction succeeds. + if (ok()) { + get() = rhs.value(); + } else { + MakeValue(rhs.value()); + status_ = Status::OK(); + } return *this; } AssignStatus(rhs.status()); @@ -115,8 +121,13 @@ class StatusOr { return *this; } if (rhs.ok()) { - status_ = Status::OK(); - AssignValue(std::move(rhs).value()); + // See the copy-assignment operator. + if (ok()) { + get() = std::move(rhs).value(); + } else { + MakeValue(std::move(rhs).value()); + status_ = Status::OK(); + } return *this; } AssignStatus(rhs.status()); @@ -260,15 +271,6 @@ class StatusOr { new (&data_) T(std::forward(arg)...); } - // Assign value to current status or. - template - void AssignValue(U &&value) { - if (ok()) { - ClearValue(); - } - MakeValue(std::forward(value)); - } - // Assign status to current status or. void AssignStatus(Status s) { if (ok()) { @@ -348,9 +350,30 @@ T &&StatusOr::value() && { template void StatusOr::swap(StatusOr &rhs) { - using std::swap; - swap(status_, rhs.status_); - swap(data_, rhs.data_); + // data_ is only a live T when ok(), so handle each state combination: swap the + // values (or statuses) when both sides match, and move the lone value across + // when only one side holds it, constructing into the other's raw storage. + if (ok()) { + if (rhs.ok()) { + using std::swap; + swap(get(), rhs.get()); + } else { + new (&rhs.data_) T(std::move(get())); + get().~T(); + status_ = std::move(rhs.status_); + rhs.status_ = Status::OK(); + } + } else { + if (rhs.ok()) { + new (&data_) T(std::move(rhs.get())); + rhs.get().~T(); + rhs.status_ = std::move(status_); + status_ = Status::OK(); + } else { + using std::swap; + swap(status_, rhs.status_); + } + } } template diff --git a/src/ray/common/tests/status_or_test.cc b/src/ray/common/tests/status_or_test.cc index 5c20ab4e387a..76f4bd128f75 100644 --- a/src/ray/common/tests/status_or_test.cc +++ b/src/ray/common/tests/status_or_test.cc @@ -17,6 +17,7 @@ #include #include +#include #include #include "ray/common/tests/testing.h" @@ -43,6 +44,35 @@ class Derived : public Base { StatusOr GetErrorStatus() { return Status::Invalid("Invalid error status."); } StatusOr GetValue() { return 1; } +// Tracks live instances by identity so a test can assert, deterministically and +// without a sanitizer, both that every constructed value is destroyed exactly +// once and that no operation runs a destructor on the wrong storage. `live` +// holds the address of each live instance: the destructor flags tearing down an +// address that was never recorded (destroying unconstructed or already-dead +// storage), and each scope asserts `live` is empty on exit (catching a value +// whose destructor was skipped). A single net counter cannot see the swap bug — +// the skipped destructor of one operand and the spurious ops on the other's +// unconstructed storage cancel out in the count. +struct Counted { + static std::multiset live; + int value; + explicit Counted(int v = 0) : value(v) { live.insert(this); } + Counted(const Counted &o) : value(o.value) { live.insert(this); } + Counted(Counted &&o) noexcept : value(o.value) { live.insert(this); } + // Assignment reuses the existing object, so liveness is unchanged. + Counted &operator=(const Counted &o) = default; + Counted &operator=(Counted &&o) noexcept = default; + ~Counted() { + auto it = live.find(this); + EXPECT_TRUE(it != live.end()) << "destroyed an instance that was not live"; + if (it != live.end()) { + live.erase(it); + } + } +}; + +std::multiset Counted::live; + } // namespace TEST(StatusOrTest, AssignTest) { @@ -270,4 +300,109 @@ TEST(StatusOrTest, MoveAssignment) { } } +// Cover the error/value combinations for copy and move assignment. Each inner +// scope must end with Counted::live empty (see Counted above). +TEST(StatusOrTest, AssignmentAcrossStates) { + // error <- value: destination has no value yet, so it must not destroy its + // unconstructed storage before taking the new value. + { + Counted::live.clear(); + { + StatusOr dst = Status::InvalidArgument("error"); + StatusOr src{Counted{1}}; + dst = src; + ASSERT_TRUE(dst.ok()); + EXPECT_EQ(dst.value().value, 1); + } + EXPECT_TRUE(Counted::live.empty()); + } + // value <- error: destination holds a value that must be destroyed. + { + Counted::live.clear(); + { + StatusOr dst{Counted{2}}; + StatusOr src = Status::InvalidArgument("error"); + dst = src; + EXPECT_FALSE(dst.ok()); + } + EXPECT_TRUE(Counted::live.empty()); + } + // error <- value via move: the move counterpart of the first case, and the one + // that regresses if move assignment sets status_ before constructing the value. + { + Counted::live.clear(); + { + StatusOr dst = Status::InvalidArgument("error"); + StatusOr src{Counted{3}}; + dst = std::move(src); + ASSERT_TRUE(dst.ok()); + EXPECT_EQ(dst.value().value, 3); + } + EXPECT_TRUE(Counted::live.empty()); + } + // value <- value via move. + { + Counted::live.clear(); + { + StatusOr dst{Counted{4}}; + StatusOr src{Counted{5}}; + dst = std::move(src); + ASSERT_TRUE(dst.ok()); + EXPECT_EQ(dst.value().value, 5); + } + EXPECT_TRUE(Counted::live.empty()); + } +} + +TEST(StatusOrTest, Swap) { + // value <-> error: exercises both a live and an unconstructed operand. + { + Counted::live.clear(); + { + StatusOr a{Counted{1}}; + StatusOr b = Status::InvalidArgument("error"); + a.swap(b); + EXPECT_FALSE(a.ok()); + ASSERT_TRUE(b.ok()); + EXPECT_EQ(b.value().value, 1); + } + EXPECT_TRUE(Counted::live.empty()); + } + // error <-> value: the mirror of the case above, covering the other + // move-across branch (this side unconstructed, rhs holds the value). + { + Counted::live.clear(); + { + StatusOr a = Status::InvalidArgument("error"); + StatusOr b{Counted{1}}; + a.swap(b); + ASSERT_TRUE(a.ok()); + EXPECT_EQ(a.value().value, 1); + EXPECT_FALSE(b.ok()); + } + EXPECT_TRUE(Counted::live.empty()); + } + // value <-> value (via the free swap function). + { + Counted::live.clear(); + { + StatusOr a{Counted{1}}; + StatusOr b{Counted{2}}; + swap(a, b); + ASSERT_TRUE(a.ok() && b.ok()); + EXPECT_EQ(a.value().value, 2); + EXPECT_EQ(b.value().value, 1); + } + EXPECT_TRUE(Counted::live.empty()); + } + // error <-> error. + { + StatusOr a = Status::InvalidArgument("a"); + StatusOr b = Status::NotFound("b"); + a.swap(b); + EXPECT_EQ(a.code(), StatusCode::NotFound); + EXPECT_EQ(b.code(), StatusCode::InvalidArgument); + } +} + } // namespace ray