From 6c794d1083c0e1df974981f99edd3353a39f29eb Mon Sep 17 00:00:00 2001 From: Artur Niederfahrenhorst Date: Mon, 3 Aug 2026 21:24:16 +0200 Subject: [PATCH 1/8] [RLlib] Fix squashed-Gaussian log-prob corruption for saturated policies (#65036) TorchSquashedGaussian.logp() recovered the pre-squash value via atanh(clamp(action)), which caps the recovered value at ~7.25 and corrupts logp (large negative) once the policy mean saturates tanh. This breaks SAC/TQC alpha tuning and the critic entropy target on envs where the policy saturates (e.g. Humanoid). This is also part of the picture why TQC can't solve humanoid. The solution involves some magic math ( see https://github.com/openai/spinningup/blob/master/spinup/algos/pytorch/sac/core.py#L59-L60 ) Adds a regression unit test. --------- Signed-off-by: Artur Niederfahrenhorst Co-authored-by: Claude Opus 4.8 (1M context) --- rllib/BUILD.bazel | 11 ++++ .../cql/torch/default_cql_torch_rl_module.py | 12 ++-- .../sac/torch/default_sac_torch_rl_module.py | 19 +++---- .../tqc/torch/default_tqc_torch_rl_module.py | 6 +- .../core/distribution/torch/tests/__init__.py | 0 .../tests/test_torch_squashed_gaussian.py | 30 ++++++++++ .../distribution/torch/torch_distribution.py | 56 ++++++++++++++----- 7 files changed, 100 insertions(+), 34 deletions(-) create mode 100644 rllib/core/distribution/torch/tests/__init__.py create mode 100644 rllib/core/distribution/torch/tests/test_torch_squashed_gaussian.py 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: From 2ac83e966a1fc917133b6d09ae747fe431f58d94 Mon Sep 17 00:00:00 2001 From: YangJie Date: Tue, 4 Aug 2026 04:00:25 +0800 Subject: [PATCH 2/8] [Core] Fix UB in StatusOr swap and assignment on error-state operands (#64799) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ## Description `ray::StatusOr` keeps its value in a union member `data_` that is only a live object when `ok()`; in the error state `data_` is left unconstructed. Two paths ignored that and touched the unconstructed storage: - `swap()` swapped `data_` unconditionally, move-constructing from / assigning into unconstructed storage when either operand was an error. - `operator=(const&)` and `operator=(&&)` set `status_` to OK before calling `AssignValue`, which destroys the old value only when `ok()`. With the status flipped first, it ran `data_.~T()` on unconstructed storage whenever the destination was previously an error. Both assignment operators now branch on the current state: when this already holds a value, assign through `T`'s own `operator=` to reuse its resources and stay exception-safe (a throwing assignment leaves the object intact instead of status-OK-with-no-value); otherwise placement-new the value and only then set `status_` to OK. `swap` is specialized for the four `ok()`/error combinations: swap the underlying values (or statuses) directly when both sides match, and move the lone value across when only one side holds it — O(1) for swappable `T`, with no touch of unconstructed storage. Added tests covering assignment and swap across the error/value combinations using a type with a live-instance counter, so a stray or missing ctor/dtor is caught in a normal test run without relying on a sanitizer. Each branch was verified to fail against the pre-fix code. ## Related issues Fixes #64798 ## Additional information `bazel test //src/ray/common/tests:status_or_test` passes, including under `--config=asan`. Verified the new tests fail when the fix is reverted. --------- Signed-off-by: yangjie01 Co-authored-by: Chi-Sheng Liu --- src/ray/common/status_or.h | 55 +++++++--- src/ray/common/tests/status_or_test.cc | 135 +++++++++++++++++++++++++ 2 files changed, 174 insertions(+), 16 deletions(-) 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 From 02704400c9c123f8b77199cee3ec0363b9c1697b Mon Sep 17 00:00:00 2001 From: kahlun Date: Mon, 3 Aug 2026 15:27:59 -0700 Subject: [PATCH 3/8] Intel gpu ze affinity mask (#64440) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ## Root cause Verl SGLang crashes on Intel GPU because Ray sets ONEAPI_DEVICE_SELECTOR=level_zero:N in workers, but SGLang suppresses Ray's env var via RAY_EXPERIMENTAL_NOSET_ONEAPI_DEVICE_SELECTOR and sets its own bare IDs (e.g. ONEAPI_DEVICE_SELECTOR=0,1). That format is invalid — the level_zero: prefix is required — causing an immediate torch import crash in SGLang subprocesses. ## Changes **python/ray/_private/accelerators/intel_gpu.py** - ZE_AFFINITY_MASK is now the primary env var (bare IDs, consistent with CUDA_VISIBLE_DEVICES) - get_current_process_visible_accelerator_ids() reads ZE_AFFINITY_MASK first, falls back to ONEAPI_DEVICE_SELECTOR for backward compat (handles both level_zero:0,1 and bare 0,1 formats) - set_current_process_visible_accelerator_ids() writes both ZE_AFFINITY_MASK (physical bare IDs) and ONEAPI_DEVICE_SELECTOR (level_zero: prefixed re-indexed sequential IDs) - Added RAY_EXPERIMENTAL_NOSET_ZE_AFFINITY_MASK flag; existing RAY_EXPERIMENTAL_NOSET_ONEAPI_DEVICE_SELECTOR still works independently **Bug found and fixed:** initial implementation wrote ONEAPI_DEVICE_SELECTOR=level_zero:1 for physical GPU 1, but ZE_AFFINITY_MASK re-indexes devices before ONEAPI_DEVICE_SELECTOR applies — intersection was empty, worker saw 0 GPUs. Fixed: ONEAPI_DEVICE_SELECTOR always uses sequential re-indexed IDs (level_zero:0, level_zero:0,1) regardless of physical IDs. Hardware-confirmed on 2x Intel GPU: GPU 1 before fix: ZE=1 + ONEAPI=level_zero:1 → 0 devices visible GPU 1 after fix: ZE=1 + ONEAPI=level_zero:0 → 1 device visible ✓ **doc/source/ray-core/scheduling/accelerators.rst** - Updated tip and code examples to use ZE_AFFINITY_MASK - Added backward compat note for ONEAPI_DEVICE_SELECTOR users **doc/source/serve/llm/user-guides/sglang.md** - Added Intel GPU prerequisite entry: RAY_EXPERIMENTAL_NOSET_ZE_AFFINITY_MASK=1 and RAY_EXPERIMENTAL_NOSET_ONEAPI_DEVICE_SELECTOR=1 (both needed since Ray now writes both vars) **python/ray/tests/accelerators/test_intel_gpu.py** - Added parametrize cases for non-zero physical GPU IDs (["1"], ["1","3"]) — these are the cases that catch the re-indexing bug - Replaced manual del os.environ cleanup with clean_accelerator_env fixture that restores env vars even on test failure - Split noset flag tests into two independent functions **python/ray/tests/accelerators/test_intel_gpu_e2e.py** - Updated gpu_task and cluster_probe_task to return ze_affinity_mask - Updated _validate_gpu_binding_common to check ZE_AFFINITY_MASK for physical GPU IDs (ONEAPI_DEVICE_SELECTOR now carries re-indexed IDs so physical ID comparison against it would be wrong) ## Test plan - `python -m pytest python/ray/tests/accelerators/test_intel_gpu.py -v` No GPU needed, passes on any machine - `RAY_PYTEST_USE_GPU=1 python -m pytest python/ray/tests/accelerators/test_intel_gpu_e2e.py -v` Requires Intel GPU + dpctl. Single-node 2x GPU tests pass. Multi-node test (test_scale_out_task_distribution) is implemented but not yet validated — no multi-node Intel GPU cluster available. --------- Signed-off-by: Kah Lun teoh Co-authored-by: Edward Oakes --- .../ray-core/scheduling/accelerators.rst | 13 +- doc/source/serve/llm/user-guides/sglang.md | 1 + python/ray/_private/accelerators/intel_gpu.py | 44 ++++--- .../ray/tests/accelerators/test_intel_gpu.py | 120 +++++++++++++++--- .../tests/accelerators/test_intel_gpu_e2e.py | 41 ++++-- 5 files changed, 167 insertions(+), 52 deletions(-) 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/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/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, " From e79c7acf1c310ddc76b55bba6314044a7657f537 Mon Sep 17 00:00:00 2001 From: Rueian Date: Mon, 3 Aug 2026 15:53:06 -0700 Subject: [PATCH 4/8] [core][TPU] Use POSIX paths for VFIO sysfs vendor checks on windows (#65182) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ## Description Fix the test_tpu.py on windows ``` [2026-08-03T18:43:13Z] ================================== FAILURES =================================== --   | [2026-08-03T18:43:13Z] _________________ test_autodetect_num_tpus_vfio_mixed_groups __________________   | [2026-08-03T18:43:13Z]   | [2026-08-03T18:43:13Z] path = '/sys/kernel/iommu_groups/10/devices\\0000:01:00.0\\vendor', args = ()   | [2026-08-03T18:43:13Z] kwargs = {'encoding': 'ascii'}   | [2026-08-03T18:43:13Z]   | [2026-08-03T18:43:13Z] def fake_open(path, *args, **kwargs):   | [2026-08-03T18:43:13Z] try:   | [2026-08-03T18:43:13Z] > return mock.mock_open(read_data=vendor_results[path])()   | [2026-08-03T18:43:13Z] E KeyError: '/sys/kernel/iommu_groups/10/devices\\0000:01:00.0\\vendor'   | [2026-08-03T18:43:13Z]   | [2026-08-03T18:43:13Z] python\ray\tests\accelerators\test_tpu.py:120: KeyError   | [2026-08-03T18:43:13Z]   | [2026-08-03T18:43:13Z] During handling of the above exception, another exception occurred:   | [2026-08-03T18:43:13Z]   | [2026-08-03T18:43:13Z] mock_list =   | [2026-08-03T18:43:13Z] mock_glob =   | [2026-08-03T18:43:13Z]   | [2026-08-03T18:43:13Z] @patch("glob.glob")   | [2026-08-03T18:43:13Z] @patch("os.listdir")   | [2026-08-03T18:43:13Z] def test_autodetect_num_tpus_vfio_mixed_groups(mock_list, mock_glob):   | [2026-08-03T18:43:13Z] # Two VFIO groups: one is a Google TPU (vendor 0x1ae0), the other is the   | [2026-08-03T18:43:13Z] # BlueField-3 SoC (vendor 0x15b3). Only the TPU-backed group is counted.   | [2026-08-03T18:43:13Z] mock_glob.return_value = []   | [2026-08-03T18:43:13Z] listdir_results = {   | [2026-08-03T18:43:13Z] "/dev/vfio": ["vfio", "10", "96"],   | [2026-08-03T18:43:13Z] "/sys/kernel/iommu_groups/10/devices": ["0000:01:00.0"],   | [2026-08-03T18:43:13Z] "/sys/kernel/iommu_groups/96/devices": ["0016:03:00.2"],   | [2026-08-03T18:43:13Z] }   | [2026-08-03T18:43:13Z] vendor_results = {   | [2026-08-03T18:43:13Z] "/sys/kernel/iommu_groups/10/devices/0000:01:00.0/vendor": "0x1ae0\n",   | [2026-08-03T18:43:13Z] "/sys/kernel/iommu_groups/96/devices/0016:03:00.2/vendor": "0x15b3\n",   | [2026-08-03T18:43:13Z] }   | [2026-08-03T18:43:13Z]   | [2026-08-03T18:43:13Z] def fake_listdir(path):   | [2026-08-03T18:43:13Z] try:   | [2026-08-03T18:43:13Z] return listdir_results[path]   | [2026-08-03T18:43:13Z] except KeyError:   | [2026-08-03T18:43:13Z] raise AssertionError(f"unexpected listdir: {path}")   | [2026-08-03T18:43:13Z]   | [2026-08-03T18:43:13Z] def fake_open(path, *args, **kwargs):   | [2026-08-03T18:43:13Z] try:   | [2026-08-03T18:43:13Z] return mock.mock_open(read_data=vendor_results[path])()   | [2026-08-03T18:43:13Z] except KeyError:   | [2026-08-03T18:43:13Z] raise AssertionError(f"unexpected open: {path}")   | [2026-08-03T18:43:13Z]   | [2026-08-03T18:43:13Z] mock_list.side_effect = fake_listdir   | [2026-08-03T18:43:13Z] with patch("builtins.open", side_effect=fake_open):   | [2026-08-03T18:43:13Z] TPUAcceleratorManager.get_current_node_num_accelerators.cache_clear()   | [2026-08-03T18:43:13Z] > assert TPUAcceleratorManager.get_current_node_num_accelerators() == 1   | [2026-08-03T18:43:13Z]   | [2026-08-03T18:43:13Z] python\ray\tests\accelerators\test_tpu.py:127:   | [2026-08-03T18:43:13Z] _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _   | [2026-08-03T18:43:13Z] C:\rayci\python\ray\_private\accelerators\tpu.py:678: in get_current_node_num_accelerators   | [2026-08-03T18:43:13Z] if _is_vfio_group_a_tpu(group):   | [2026-08-03T18:43:13Z] C:\rayci\python\ray\_private\accelerators\tpu.py:601: in _is_vfio_group_a_tpu   | [2026-08-03T18:43:13Z] with open(vendor_path, encoding="ascii") as f:   | [2026-08-03T18:43:13Z] C:\Miniconda3\lib\unittest\mock.py:1114: in __call__   | [2026-08-03T18:43:13Z] return self._mock_call(*args, **kwargs)   | [2026-08-03T18:43:13Z] C:\Miniconda3\lib\unittest\mock.py:1118: in _mock_call   | [2026-08-03T18:43:13Z] return self._execute_mock_call(*args, **kwargs)   | [2026-08-03T18:43:13Z] C:\Miniconda3\lib\unittest\mock.py:1179: in _execute_mock_call   | [2026-08-03T18:43:13Z] result = effect(*args, **kwargs)   | [2026-08-03T18:43:13Z] _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _   | [2026-08-03T18:43:13Z]   | [2026-08-03T18:43:13Z] path = '/sys/kernel/iommu_groups/10/devices\\0000:01:00.0\\vendor', args = ()   | [2026-08-03T18:43:13Z] kwargs = {'encoding': 'ascii'}   | [2026-08-03T18:43:13Z]   | [2026-08-03T18:43:13Z] def fake_open(path, *args, **kwargs):   | [2026-08-03T18:43:13Z] try:   | [2026-08-03T18:43:13Z] return mock.mock_open(read_data=vendor_results[path])()   | [2026-08-03T18:43:13Z] except KeyError:   | [2026-08-03T18:43:13Z] > raise AssertionError(f"unexpected open: {path}")   | [2026-08-03T18:43:13Z] E AssertionError: unexpected open: /sys/kernel/iommu_groups/10/devices\0000:01:00.0\vendor   | [2026-08-03T18:43:13Z]   | [2026-08-03T18:43:13Z] python\ray\tests\accelerators\test_tpu.py:122: AssertionError ``` Signed-off-by: Rueian Huang --- python/ray/tests/accelerators/test_tpu.py | 10 ++++++++-- 1 file changed, 8 insertions(+), 2 deletions(-) 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): From cc289a4b516c399a9b62a1b274f2b46e18b14be5 Mon Sep 17 00:00:00 2001 From: Jeffrey Wang Date: Mon, 3 Aug 2026 16:22:43 -0700 Subject: [PATCH 5/8] [llm][kv][14/N] Add KV cache offload/reload dashboard (#65122) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ## Description Adds a collapsed Serve LLM grafana row for native vLLM KV offload/reload observability: GPU ↔ CPU throughput, operation rate, transfer bandwidth, CPU transfer pressure, external prefix-cache hit rate, and P90 lookup delay. Panels join vLLM worker metrics with serve deployment/replica labels. Screenshot 2026-07-29 at 5 29 57 PM ## Related issues > Link related issues: "Fixes #1234", "Closes #1234", or "Related to #1234". ## Additional information > Optional: Add implementation details, API changes, usage examples, screenshots, etc. --------- Signed-off-by: Jeffrey Wang --- .../dashboards/serve_llm_dashboard_panels.py | 244 +++++++++++++++++- 1 file changed, 243 insertions(+), 1 deletion(-) 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, From 26bedee6aa9fa0658f6833dadba1ec75eab8518b Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?T=C3=B4=20Qu=E1=BB=91c=20Kh=C3=A1nh?= <164204415+tqKhanh1712@users.noreply.github.com> Date: Tue, 4 Aug 2026 06:54:54 +0700 Subject: [PATCH 6/8] [java] Upgrade com.google.code.gson from 2.9.1 to 2.11.0 (#64273) (#65131) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ## Why are these changes needed? Fixes #64273 The Java dependency `com.google.code.gson:gson` is currently pinned to version `2.9.1`, which is outdated. This PR upgrades it to `2.11.0`, which includes bug fixes, performance improvements, and security enhancements over the previous version. Key changes in gson 2.10.0–2.11.0: - Improved `TypeToken` API safety - Fixed edge cases in JSON parsing - Performance improvements in serialization/deserialization ## Changes - Updated `java/dependencies.bzl`: bumped `com.google.code.gson:gson` from `2.9.1` to `2.11.0` - No lock file update required (Ray uses dynamic Maven resolution via `rules_jvm_external`, no `maven_install.json` pins gson) ## Checks - [ ] I've run `scripts/format.sh` to lint the changes in this PR. - [ ] I've made sure the tests are passing. ## Testing Dependency version bump only. No API-level changes to Ray source code. Gson 2.11.0 is backward compatible with 2.9.1 for all usages present in this codebase. Signed-off-by: tqKhanh1712 --- java/dependencies.bzl | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) 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", From 850f0be149c57d841f1f6ba6d73b29f2789b00a3 Mon Sep 17 00:00:00 2001 From: Goutam Date: Mon, 3 Aug 2026 16:56:18 -0700 Subject: [PATCH 7/8] [Data] Don't convert Arrow null columns to null[pyarrow] in to_pandas (#65187) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ## Why are these changes needed? Follow-up to #64765 / #63017. `ArrowBlockAccessor.to_pandas()` maps Arrow types to `pd.ArrowDtype` via a `_types_mapper`, with carve-outs for extension and dictionary types. Arrow's `null` type needs the same carve-out. `null[pyarrow]` is unusable from pandas: the type carries no type information, so pandas cannot box a non-null value into such a column. `fillna` and masked assignment raise `ArrowInvalid`, and in some pyarrow/pandas combinations the failure comes out of Arrow C++ and aborts the worker process rather than raising. This is easy to hit in practice because the type is assigned per block, not per dataset. A column that has values overall can still be entirely null within one block — which is common with small blocks — so a pandas UDF like `map_batches(lambda df: df.fillna(...), batch_format="pandas")` fails on whichever block happens to hold only nulls: ```python 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, ) ds.map_batches(lambda df: df.fillna({"a": 0.0, "b": 0.0}), batch_format="pandas").take_all() # ArrowInvalid on the block where "a" (or "b") is all-null ``` The fix returns `None` from `_types_mapper` for null-typed columns, falling back to pandas' default conversion. The Arrow round-trip is unchanged: `PandasBlockAccessor.to_arrow()` already coerces all-null columns back to `pa.null()`. The one behavioral consequence, noted in a comment: a column that is all-null in *every* block stays null-typed instead of being promoted — same as before this change. ## Related issue number Follow-up to #64765. Two tests added: - `test_arrow_block_to_pandas_null_type_is_not_arrow_backed` — a `pa.null()` column is not `pd.ArrowDtype`, still round-trips back to `pa.null()` untouched, and `fillna` now works and adopts the fill value's type. - `test_pandas_udf_can_fill_per_block_null_columns` — end-to-end regression test for the `map_batches` + `fillna` failure above. Signed-off-by: Goutam Co-authored-by: Claude Opus 5 (1M context) --- python/ray/data/_internal/arrow_block.py | 9 +++++ python/ray/data/tests/test_arrow_block.py | 45 +++++++++++++++++++++++ 2 files changed, 54 insertions(+) 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/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 From 77abc0cc956e85db0f25f2c0ce811ccefd346868 Mon Sep 17 00:00:00 2001 From: Daniel Shin <88547237+kyuds@users.noreply.github.com> Date: Mon, 3 Aug 2026 17:36:34 -0700 Subject: [PATCH 8/8] [Data] Managed Torch Inference (#65157) ## Motivation Batch inference via `map_batches` with a GPU actor has a lot of the same boilerplate code: move data to and from the device, convert data into tensors that can be processed by the GPU, convert the resultant output tensors back to a format that Ray Data can process. A lot of this is duplicated code, and similar to how we have sensible defaults in other parts of our API (such as `DefaultCollateFn` for `iter_torch_batches`), we should provide a short-hand framework for batch inference on `map_batches`. ## API The user has to subclass this class: ``` @PublicAPI(stability="alpha") class TorchInference: def initialize(self, *args: Any, **kwargs: Any) -> None: pass def get_device(self) -> "torch.device": ... def collate( self, input_batch: "DataBatch" ) -> Union["TensorBatchType", Tuple["TensorBatchType", Any]]: ... def process_on_device( self, input_batch: "DataBatch", collated_tensors: "TensorBatchType", collated_other: Any, ) -> Union["TensorBatchType", Tuple["TensorBatchType", Any]]: raise NotImplementedError() def finalize( self, input_batch: "DataBatch", output_tensors: "TensorBatchType", output_other: Any, ) -> "DataBatch": ... ``` and mandatorily implement `process_on_device` but can also optionally implement the other 4 methods. On the Ray Data side, we manage moving data to and from device and also providing sensible defaults for `get_device`, `collate` and `finalize`. ### The Defaults - `get_device`: `torch.device("cuda")` - `collate`: take a numpy batch, convert and concatenate tensors as necessary. - `fianlize`: recursively take every Tensor in the batch and convert them to numpy via `.numpy()`. ## What Happens Underneath? After initialization, for each batch, the following occurs: 1. The `DataBatch` is fed into `collate` 2. The tensor output of `collate` is moved to device. 3. `process_on_device` is called on GPU device. 4. The tensor output of `process_on_device` is moved back to CPU. 5. `finalize` is called on CPU tensor. Each stage may need to reference 1) additional information created in the steps that are not tensors and 2) some other columns in the original batch. Therefore, for each function invocation, we provide the original batch as `input_batch` and any other miscellaneous output as `output_other`. --------- Signed-off-by: Daniel Shin Signed-off-by: Daniel Shin <88547237+kyuds@users.noreply.github.com> Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com> --- python/ray/data/BUILD.bazel | 15 + python/ray/data/_internal/compute.py | 25 +- .../logical/operators/map_operator.py | 36 ++ .../data/_internal/utils/torch_inference.py | 261 +++++++++ .../ray/data/_internal/utils/torch_utils.py | 31 +- python/ray/data/tests/test_torch_inference.py | 538 ++++++++++++++++++ python/ray/data/util/torch_inference.py | 241 ++++++++ python/ray/data/util/torch_utils.py | 20 + 8 files changed, 1161 insertions(+), 6 deletions(-) create mode 100644 python/ray/data/_internal/utils/torch_inference.py create mode 100644 python/ray/data/tests/test_torch_inference.py create mode 100644 python/ray/data/util/torch_inference.py 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/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_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." + )