[Data] Reject non-finite weights in Dataset.mix

python/ray/data/_internal/execution/operators/mix_operator.py

@@ -14,6 +14,7 @@
 from ray.data._internal.logical.operators.n_ary_operator import (
     MixStoppingCondition,
     estimate_num_mix_outputs,
+    validate_mix_weights,
 )
 from ray.data._internal.stats import StatsDict
 from ray.data.context import DataContext
@@ -39,8 +40,7 @@ def __init__(
     ):
         assert len(input_ops) >= 1
         assert len(weights) == len(input_ops)
-        if any(w <= 0 for w in weights):
-            raise ValueError("Weights must be positive.")
+        validate_mix_weights(weights)
 
         total_weight = sum(weights)
         self._weights = [w / total_weight for w in weights]

python/ray/data/_internal/logical/operators/n_ary_operator.py

@@ -1,4 +1,5 @@
 import enum
+import math
 from dataclasses import dataclass, field
 from typing import TYPE_CHECKING, List, Optional
 
@@ -37,6 +38,14 @@ class MixStoppingCondition(enum.Enum):
     STOP_ON_LONGEST_DROP = "stop_on_longest_drop"
 
 
+def validate_mix_weights(weights: List[float]) -> None:
+    # NaN and +inf slip through `weight <= 0`, then poison the normalized weights.
+    if any(not math.isfinite(weight) for weight in weights):
+        raise ValueError(f"Weights must be finite. Got weights: {weights}")
+    if any(weight <= 0 for weight in weights):
+        raise ValueError(f"Weights must be positive. Got weights: {weights}")
+
+
 def estimate_num_mix_outputs(
     per_input_counts: List[Optional[int]],
     weights: List[float],
@@ -149,8 +158,7 @@ def __init__(
                 f"Number of input operators ({len(input_dependencies)}) must match "
                 f"number of weights ({len(weights)})."
             )
-        if any(weight <= 0 for weight in weights):
-            raise ValueError(f"Weights must be positive. Got weights: {weights}")
+        validate_mix_weights(weights)
 
         for input_op in input_dependencies:
             assert isinstance(input_op, LogicalOperator), input_op

python/ray/data/dataset.py

@@ -3531,7 +3531,8 @@ def mix(
 
         Raises:
             ValueError: If the length of ``weights`` doesn't match the
-                number of datasets.
+                number of datasets, or if any weight is not positive and
+                finite.
         """
         datasets = [self] + list(other)
 

python/ray/data/tests/test_mix.py

@@ -17,6 +17,49 @@ def _make_ds(source_id, num_rows, rows_per_block):
     )
 
 
+@pytest.mark.parametrize(
+    "bad_weight,expected_match",
+    [
+        (float("nan"), "finite"),
+        (float("inf"), "finite"),
+        (float("-inf"), "finite"),
+        (0.0, "positive"),
+        (-1.0, "positive"),
+    ],
+)
+def test_mix_rejects_invalid_weights(
+    ray_start_10_cpus_shared, bad_weight, expected_match
+):
+    ds1 = ray.data.range(10)
+    ds2 = ray.data.range(10)
+    with pytest.raises(ValueError, match=expected_match):
+        ds1.mix(ds2, weights=[bad_weight, 1.0])
+
+
+@pytest.mark.parametrize(
+    "bad_weight,expected_match",
+    [
+        (float("nan"), "finite"),
+        (float("inf"), "finite"),
+        (float("-inf"), "finite"),
+        (0.0, "positive"),
+        (-1.0, "positive"),
+    ],
+)
+def test_mix_operator_rejects_invalid_weights(bad_weight, expected_match):
+    from unittest.mock import MagicMock
+
+    from ray.data._internal.execution.interfaces import PhysicalOperator
+
+    input_ops = [MagicMock(spec=PhysicalOperator) for _ in range(2)]
+    with pytest.raises(ValueError, match=expected_match):
+        MixOperator(
+            ray.data.DataContext.get_current(),
+            *input_ops,
+            weights=[bad_weight, 1.0],
+        )
+
+
 @pytest.mark.parametrize("weights", [[1, 1], None])
 def test_mix_equal_weights(ray_start_10_cpus_shared, weights):
     """Equal weights with uniform blocks should produce 50/50 every 2 blocks.