fix: Salt failure is a hard error

src/inference_endpoint/commands/benchmark/execute.py

@@ -367,6 +367,12 @@ def _load_datasets(
         except Exception as e:
             raise SetupError(f"Failed to load dataset: {e}") from e
 
+        # Fail fast on a warmup dataset that salt cannot bust — at load time,
+        # before any worker/aggregator subprocess is spawned.
+        warmup = config.settings.warmup
+        if warmup.enabled and warmup.salt:
+            dataloader.validate_saltable()
+
         if perf_cfg.accuracy_config is not None:
             accuracy_config = perf_cfg.accuracy_config
             if accuracy_config.num_repeats != 1:

src/inference_endpoint/dataset_manager/dataset.py

@@ -30,6 +30,7 @@
 from datasets import load_dataset, load_from_disk
 
 from ..config.schema import APIType, ModelParams
+from ..exceptions import DatasetValidationError
 from .transforms import (
     ColumnFilter,
     Transform,
@@ -257,6 +258,23 @@ def load_from_huggingface(
     return ds[split].to_pandas()
 
 
+def _salt_violation(sample: Any) -> str | None:
+    """Return a human-readable reason a sample cannot be salted, or None if it can.
+
+    Salt requires a dict sample with a str 'prompt' and no 'input_tokens' (which
+    adapters send verbatim, so a salted 'prompt' would not reach the server).
+    """
+    if not isinstance(sample, dict):
+        return f"is a {type(sample).__name__}, not a dict"
+    if "input_tokens" in sample:
+        return "has 'input_tokens' (salt cannot bust a pre-tokenized cache)"
+    if "prompt" not in sample:
+        return "has no 'prompt' field"
+    if not isinstance(sample["prompt"], str):
+        return f"has a 'prompt' of type {type(sample['prompt']).__name__}, not str"
+    return None
+
+
 class Dataset:
     """Class for loading and managing benchmark datasets.
 
@@ -437,55 +455,57 @@ def load_sample(self, index: int) -> Any:
             data = self._apply_salt(data)
         return data
 
+    def validate_saltable(self) -> None:
+        """Raise if any loaded sample cannot be salted.
+
+        salt requires a dict sample with a text ('str') 'prompt' and no
+        'input_tokens' (adapters send those verbatim, so a salted 'prompt' would
+        never reach the server). A sample salt cannot bust would silently defeat
+        cache-busting, so it is rejected rather than skipped. Called before any
+        load is issued — at benchmark setup and again from with_salt().
+
+        Raises:
+            DatasetValidationError: naming the first offending sample.
+        """
+        if self.data is None:
+            return
+        for i, sample in enumerate(self.data):
+            reason = _salt_violation(sample)
+            if reason is not None:
+                raise DatasetValidationError(
+                    f"salt=True requires every sample to be a dict with a text "
+                    f"'prompt' and no 'input_tokens', but sample {i} {reason}. "
+                    f"Disable salt (--warmup-salt / warmup.salt: false) or use a "
+                    f"text-prompt dataset."
+                )
+
     def with_salt(self, rng: random.Random) -> "Dataset":
         """Return a shallow copy of this dataset that salts each load_sample() call.
 
         The returned dataset shares the same loaded data — no re-loading needed.
         Each load_sample() call on the returned dataset prepends a unique hex salt
-        derived from rng to the prompt field, preventing KV-cache reuse.
+        derived from rng to the 'prompt' field, preventing KV-cache reuse.
+
+        Validates every sample first (see validate_saltable), so a dataset salt
+        cannot bust fails here rather than silently issuing unsalted prompts.
+
+        Raises:
+            DatasetValidationError: if any sample cannot be salted.
         """
+        self.validate_saltable()
         clone = copy.copy(self)
         clone._salt_rng = rng
         return clone
 
-    def _apply_salt(self, data: Any) -> Any:
-        """Prepend a unique salt to the prompt field of a sample dict."""
+    def _apply_salt(self, data: dict[str, Any]) -> dict[str, Any]:
+        """Prepend a unique salt to the 'prompt' field.
+
+        with_salt() has validated every sample, so ``data`` is guaranteed to be a
+        dict with a str 'prompt' and no 'input_tokens'.
+        """
         assert self._salt_rng is not None
-        if not isinstance(data, dict):
-            return data
-        if "input_tokens" in data and "prompt" not in data:
-            self.logger.warning(
-                "salt=True: sample has 'input_tokens' but no 'prompt' — "
-                "salt cannot be applied to pre-tokenized input; KV-cache reuse may not be prevented"
-            )
-            return data
-        if "input_tokens" in data and "prompt" in data:
-            self.logger.warning(
-                "salt=True: sample has both 'input_tokens' and 'prompt' — "
-                "salt applied to 'prompt' only; adapters that use 'input_tokens' "
-                "directly will still reuse the KV cache"
-            )
-        if "prompt" not in data:
-            return data
-        prompt = data["prompt"]
         salt = self._salt_rng.randbytes(8).hex()
-        if isinstance(prompt, str):
-            return {**data, "prompt": f"[{salt}] {prompt}"}
-        if isinstance(prompt, list) and prompt:
-            # Find the first text part at any index (image-first prompts place text at index 1+)
-            for i, part in enumerate(prompt):
-                if isinstance(part, dict) and part.get("type") == "text":
-                    salted_parts = [
-                        *prompt[:i],
-                        {**part, "text": f"[{salt}] {part['text']}"},
-                        *prompt[i + 1 :],
-                    ]
-                    return {**data, "prompt": salted_parts}
-            self.logger.warning(
-                "salt=True: multimodal prompt has no text part — "
-                "salt cannot be applied; KV-cache reuse may not be prevented"
-            )
-        return data  # unsupported prompt type — skip salting
+        return {**data, "prompt": f"[{salt}] {data['prompt']}"}
 
     def num_samples(self) -> int:
         assert self.data is not None, "Dataset not loaded. Call load() first."

tests/unit/commands/test_benchmark.py

@@ -77,7 +77,10 @@
 from inference_endpoint.dataset_manager.dataset import Dataset
 from inference_endpoint.endpoint_client.config import HTTPClientConfig
 from inference_endpoint.evaluation.scoring import Scorer
-from inference_endpoint.exceptions import InputValidationError, SetupError
+from inference_endpoint.exceptions import (
+    InputValidationError,
+    SetupError,
+)
 from inference_endpoint.load_generator.sample_order import create_sample_order
 from inference_endpoint.load_generator.session import PhaseType
 from inference_endpoint.metrics.metric import Throughput
@@ -221,6 +224,46 @@ def test_dataset_string_coercion(
             assert ds.accuracy_config.eval_method == acc_eval_method
 
 
+@pytest.mark.unit
+class TestLoadDatasetsSaltValidation:
+    """_load_datasets validates salt-compatibility at dataset-load time — before
+    any worker/aggregator subprocess is spawned — when warmup salt is enabled.
+    """
+
+    def _config(self, tmp_path: Path, warmup: WarmupConfig) -> OfflineConfig:
+        ds = tmp_path / "perf.jsonl"
+        ds.write_text('{"prompt": "hello world"}\n{"prompt": "second prompt"}\n')
+        return OfflineConfig(
+            endpoint_config={"endpoints": ["http://test:8000"]},
+            model_params={"name": "test-model"},
+            datasets=[{"path": str(ds)}],
+            settings=OfflineSettings(
+                client=HTTPClientConfig(
+                    num_workers=1, warmup_connections=0, max_connections=10
+                ),
+                warmup=warmup,
+            ),
+        )
+
+    @patch.object(Dataset, "validate_saltable")
+    def test_validates_when_warmup_salt_enabled(self, mock_validate, tmp_path):
+        config = self._config(tmp_path, WarmupConfig(enabled=True, salt=True))
+        _load_datasets(config, tmp_path, TestMode.PERF)
+        mock_validate.assert_called_once()
+
+    @patch.object(Dataset, "validate_saltable")
+    def test_skips_validation_when_warmup_disabled(self, mock_validate, tmp_path):
+        config = self._config(tmp_path, WarmupConfig(enabled=False, salt=True))
+        _load_datasets(config, tmp_path, TestMode.PERF)
+        mock_validate.assert_not_called()
+
+    @patch.object(Dataset, "validate_saltable")
+    def test_skips_validation_when_salt_off(self, mock_validate, tmp_path):
+        config = self._config(tmp_path, WarmupConfig(enabled=True, salt=False))
+        _load_datasets(config, tmp_path, TestMode.PERF)
+        mock_validate.assert_not_called()
+
+
 class TestCommandHandlers:
     """Test offline/online/from_config handlers (mock run_benchmark)."""
 

tests/unit/dataset_manager/test_salted_dataset.py

@@ -17,11 +17,11 @@
 
 import random
 import re
-from unittest.mock import MagicMock
 
 import pandas as pd
 import pytest
 from inference_endpoint.dataset_manager.dataset import Dataset
+from inference_endpoint.exceptions import DatasetValidationError
 
 
 def _make_loaded_dataset(rows: list[dict]) -> Dataset:
@@ -31,7 +31,6 @@ def _make_loaded_dataset(rows: list[dict]) -> Dataset:
     ds.transforms = None
     ds.repeats = 1
     ds.data = list(rows)
-    ds.logger = MagicMock()
     ds._salt_rng = None
     return ds
 
@@ -143,77 +142,84 @@ def test_seeded_rng_is_reproducible(self):
 
 
 @pytest.mark.unit
-class TestSaltPassthrough:
-    """Samples without a 'prompt' key, or non-dict samples, are passed through unchanged."""
+class TestSaltValidation:
+    """with_salt() hard-errors up front unless every sample has a text 'prompt'.
 
-    def test_dict_without_prompt_key_is_unchanged(self):
+    salt=True guarantees a KV-cache-busting prefix; a sample it cannot salt is a
+    configuration error, not something to skip silently. Validation runs in
+    with_salt() (before any load is issued), so the error names the offending
+    sample and no partial warmup runs against an unsalted dataset.
+    """
+
+    def test_dict_without_prompt_key_raises(self):
         inner = _make_loaded_dataset([{"question": "what is 2+2?", "answer": "4"}])
-        sd = inner.with_salt(random.Random())
-        assert sd.load_sample(0) == {"question": "what is 2+2?", "answer": "4"}
+        with pytest.raises(DatasetValidationError, match="prompt"):
+            inner.with_salt(random.Random())
 
-    def test_empty_dict_is_unchanged(self):
+    def test_empty_dict_raises(self):
         inner = _make_loaded_dataset([{}])
-        sd = inner.with_salt(random.Random())
-        assert sd.load_sample(0) == {}
+        with pytest.raises(DatasetValidationError, match="prompt"):
+            inner.with_salt(random.Random())
 
-    def test_non_dict_sample_is_returned_as_is(self):
+    def test_non_dict_sample_raises(self):
         inner = _make_loaded_dataset([{"prompt": "x"}])
         inner.data = ["raw string sample"]
-        sd = inner.with_salt(random.Random())
-        assert sd.load_sample(0) == "raw string sample"
+        with pytest.raises(DatasetValidationError, match="dict"):
+            inner.with_salt(random.Random())
 
-    def test_multimodal_list_prompt_first_text_part_is_salted(self):
+    def test_multimodal_list_prompt_raises(self):
         content_parts = [
             {"type": "text", "text": "describe this image"},
             {"type": "image_url", "image_url": {"url": "data:image/png;base64,abc"}},
         ]
         inner = _make_loaded_dataset([{"prompt": content_parts}])
-        sd = inner.with_salt(random.Random())
-        parts = sd.load_sample(0)["prompt"]
-        assert isinstance(parts, list)
-        assert len(parts) == 2
-        assert re.match(r"^\[([0-9a-f]{16})\] describe this image$", parts[0]["text"])
-        assert parts[1] == content_parts[1]
-
-    def test_multimodal_image_first_text_at_index_1_is_salted(self):
-        content_parts = [
-            {"type": "image_url", "image_url": {"url": "data:image/png;base64,abc"}},
-            {"type": "text", "text": "what do you see?"},
-        ]
-        inner = _make_loaded_dataset([{"prompt": content_parts}])
-        sd = inner.with_salt(random.Random())
-        parts = sd.load_sample(0)["prompt"]
-        assert parts[0] == content_parts[0]
-        assert re.match(r"^\[([0-9a-f]{16})\] what do you see\?$", parts[1]["text"])
+        with pytest.raises(DatasetValidationError, match="str"):
+            inner.with_salt(random.Random())
 
-    def test_multimodal_list_prompt_original_not_mutated(self):
-        content_parts = [{"type": "text", "text": "original text"}]
-        inner = _make_loaded_dataset([{"prompt": content_parts}])
-        sd = inner.with_salt(random.Random())
-        sd.load_sample(0)
-        assert inner.data[0]["prompt"][0]["text"] == "original text"
-
-    def test_unknown_prompt_type_is_not_salted(self):
+    def test_non_str_prompt_raises(self):
         inner = _make_loaded_dataset([{"prompt": 42}])
-        sd = inner.with_salt(random.Random())
-        assert sd.load_sample(0) == {"prompt": 42}
+        with pytest.raises(DatasetValidationError, match="str"):
+            inner.with_salt(random.Random())
 
-    def test_input_tokens_only_warns_and_passes_through(self):
+    def test_input_tokens_only_raises(self):
         inner = _make_loaded_dataset([{"input_tokens": [1, 2, 3]}])
-        sd = inner.with_salt(random.Random())
-        result = sd.load_sample(0)
-        assert result == {"input_tokens": [1, 2, 3]}
-        sd.logger.warning.assert_called_once()
-        assert "input_tokens" in sd.logger.warning.call_args[0][0]
+        with pytest.raises(DatasetValidationError, match="input_tokens"):
+            inner.with_salt(random.Random())
 
-    def test_input_tokens_and_prompt_warns_and_salts_prompt(self):
+    def test_input_tokens_and_prompt_raises(self):
         inner = _make_loaded_dataset([{"input_tokens": [1, 2, 3], "prompt": "hello"}])
+        with pytest.raises(DatasetValidationError, match="input_tokens"):
+            inner.with_salt(random.Random())
+
+    def test_error_names_offending_sample_index(self):
+        inner = _make_loaded_dataset([{"prompt": "ok"}, {"prompt": 42}])
+        with pytest.raises(DatasetValidationError, match=r"\b1\b"):
+            inner.with_salt(random.Random())
+
+    def test_valid_str_prompt_dataset_does_not_raise(self):
+        inner = _make_loaded_dataset([{"prompt": "a"}, {"prompt": "b"}])
         sd = inner.with_salt(random.Random())
-        result = sd.load_sample(0)
-        assert result["input_tokens"] == [1, 2, 3]
-        assert result["prompt"].startswith("[")
-        sd.logger.warning.assert_called_once()
-        assert "input_tokens" in sd.logger.warning.call_args[0][0]
+        assert sd.load_sample(0)["prompt"].startswith("[")
+
+    def test_data_none_does_not_raise(self):
+        inner = _make_loaded_dataset([{"prompt": "x"}])
+        inner.data = None
+        # No samples to salt (e.g. EmptyDataset) — nothing to validate.
+        assert inner.with_salt(random.Random())._salt_rng is not None
+
+    def test_data_empty_list_does_not_raise(self):
+        inner = _make_loaded_dataset([])
+        # Zero samples — no violation, so no error.
+        assert inner.with_salt(random.Random())._salt_rng is not None
+
+    def test_validate_saltable_noop_on_valid(self):
+        inner = _make_loaded_dataset([{"prompt": "a"}, {"prompt": "b"}])
+        assert inner.validate_saltable() is None
+
+    def test_validate_saltable_raises_on_bad_sample(self):
+        inner = _make_loaded_dataset([{"prompt": "ok"}, {"input_tokens": [1, 2]}])
+        with pytest.raises(DatasetValidationError, match="input_tokens"):
+            inner.validate_saltable()
 
 
 @pytest.mark.unit