diff --git a/python/ray/data/_internal/execution/operators/mix_operator.py b/python/ray/data/_internal/execution/operators/mix_operator.py index 90387fc35cb2..cc3951242819 100644 --- a/python/ray/data/_internal/execution/operators/mix_operator.py +++ b/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] diff --git a/python/ray/data/_internal/logical/operators/n_ary_operator.py b/python/ray/data/_internal/logical/operators/n_ary_operator.py index eab23a32530e..97182396e67e 100644 --- a/python/ray/data/_internal/logical/operators/n_ary_operator.py +++ b/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 diff --git a/python/ray/data/dataset.py b/python/ray/data/dataset.py index 6f598dcf4a71..4bf559de1911 100644 --- a/python/ray/data/dataset.py +++ b/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) diff --git a/python/ray/data/tests/test_mix.py b/python/ray/data/tests/test_mix.py index 92415dfe473d..b22881d25f9e 100644 --- a/python/ray/data/tests/test_mix.py +++ b/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.