From fb96171076715bedb33cd76c5106217302693c38 Mon Sep 17 00:00:00 2001 From: Mike Knepper Date: Tue, 30 Jun 2026 08:43:40 -0500 Subject: [PATCH] Ensure weights sum to non-zero value Signed-off-by: Mike Knepper --- .../data-designer-config/src/data_designer/config/models.py | 5 ++++- .../src/data_designer/config/sampler_params.py | 5 ++++- packages/data-designer-config/tests/config/test_models.py | 6 ++++++ .../tests/config/test_sampler_params.py | 6 ++++++ 4 files changed, 20 insertions(+), 2 deletions(-) diff --git a/packages/data-designer-config/src/data_designer/config/models.py b/packages/data-designer-config/src/data_designer/config/models.py index 509eef4c4..e6ead1123 100644 --- a/packages/data-designer-config/src/data_designer/config/models.py +++ b/packages/data-designer-config/src/data_designer/config/models.py @@ -335,7 +335,10 @@ class ManualDistributionParams(ConfigBase): @model_validator(mode="after") def _normalize_weights(self) -> Self: if self.weights is not None: - self.weights = [w / sum(self.weights) for w in self.weights] + total_weight = sum(self.weights) + if total_weight == 0: + raise ValueError("`weights` must sum to a non-zero value") + self.weights = [w / total_weight for w in self.weights] return self @model_validator(mode="after") diff --git a/packages/data-designer-config/src/data_designer/config/sampler_params.py b/packages/data-designer-config/src/data_designer/config/sampler_params.py index fa10892d5..23e282abf 100644 --- a/packages/data-designer-config/src/data_designer/config/sampler_params.py +++ b/packages/data-designer-config/src/data_designer/config/sampler_params.py @@ -74,7 +74,10 @@ class CategorySamplerParams(ConfigBase): @model_validator(mode="after") def _normalize_weights_if_needed(self) -> Self: if self.weights is not None: - self.weights = [w / sum(self.weights) for w in self.weights] + total_weight = sum(self.weights) + if total_weight == 0: + raise ValueError("'weights' must sum to a non-zero value") + self.weights = [w / total_weight for w in self.weights] return self @model_validator(mode="after") diff --git a/packages/data-designer-config/tests/config/test_models.py b/packages/data-designer-config/tests/config/test_models.py index 0cea2fa11..2083229b4 100644 --- a/packages/data-designer-config/tests/config/test_models.py +++ b/packages/data-designer-config/tests/config/test_models.py @@ -477,6 +477,12 @@ def test_manual_distribution_weight_normalization(): assert dist.params.weights == [0.5, 0.5] +@pytest.mark.parametrize("weights", ([0, 0], [1, -1])) +def test_manual_distribution_zero_sum_weights(weights): + with pytest.raises(ValidationError, match="`weights` must sum to a non-zero value"): + ManualDistributionParams(values=[0.2, 0.8], weights=weights) + + def test_invalid_manual_distribution(): # Empty list should fail with pytest.raises(ValidationError): diff --git a/packages/data-designer-config/tests/config/test_sampler_params.py b/packages/data-designer-config/tests/config/test_sampler_params.py index 03ac22f04..e80b1916d 100644 --- a/packages/data-designer-config/tests/config/test_sampler_params.py +++ b/packages/data-designer-config/tests/config/test_sampler_params.py @@ -29,6 +29,12 @@ def test_category_sampler_params(): CategorySamplerParams(values=["a", "b", "c"], weights=[10, 1]) +@pytest.mark.parametrize("weights", ([0, 0], [1, -1])) +def test_category_sampler_params_zero_sum_weights(weights): + with pytest.raises(ValidationError, match="'weights' must sum to a non-zero value"): + CategorySamplerParams(values=["a", "b"], weights=weights) + + def test_datetime_sampler_params(): params = DatetimeSamplerParams(start="2020-01-01", end="2025-01-01", unit="D") assert params.start == "2020-01-01"