fix: Salt failure is a hard error

src/inference_endpoint/commands/benchmark/cli.py

@@ -67,9 +67,14 @@ def _run(
                 f"{'.'.join(str(p) for p in err['loc'])}: {err['msg']}"
                 for err in e.errors()
             )
-            raise DatasetValidationError(f"Invalid --dataset: {msgs}") from e
+            # --dataset parse failures aren't yet mapped to a specific Reason.
+            raise DatasetValidationError(
+                DatasetValidationError.Reason.UNSPECIFIED, f"Invalid --dataset: {msgs}"
+            ) from e
         except ValueError as e:
-            raise DatasetValidationError(f"Invalid --dataset: {e}") from e
+            raise DatasetValidationError(
+                DatasetValidationError.Reason.UNSPECIFIED, f"Invalid --dataset: {e}"
+            ) from e
     if config.audit is None:
         run_benchmark(config, mode)
         return

src/inference_endpoint/commands/benchmark/execute.py

@@ -401,6 +401,14 @@ 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. with_salt() runs
+        # the same check later; this earlier call is deliberate (not redundant),
+        # so an invalid dataset aborts before the subprocess fan-out.
+        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
             scorer_cls, extractor_cls = _resolve_accuracy_components(

src/inference_endpoint/config/schema.py

@@ -759,10 +759,19 @@ class WarmupConfig(BaseModel):
         bool,
         cyclopts.Parameter(
             alias="--warmup-salt",
-            help="Prepend a unique random hex salt to each warmup prompt",
+            help=(
+                "Prepend a unique random hex salt to each warmup prompt. Requires "
+                "text-'prompt' samples; enabling it on a pre-tokenized "
+                "('input_tokens') dataset is a hard error."
+            ),
         ),
     ] = Field(
-        True, description="Prepend a unique random hex salt to each warmup prompt"
+        False,
+        description=(
+            "Prepend a unique random hex salt to each warmup prompt. Requires "
+            "text-'prompt' samples; enabling it on a pre-tokenized "
+            "('input_tokens') dataset is a hard error."
+        ),
     )
     drain: Annotated[
         bool,

src/inference_endpoint/config/templates/concurrency_template_full.yaml

@@ -93,7 +93,7 @@ settings:
   warmup:
     enabled: false  # Enable warmup phase before performance run
     n_requests: null  # Warmup request count (None = full dataset once)
-    salt: true  # Prepend a unique random hex salt to each warmup prompt
+    salt: false  # Prepend a unique random hex salt to each warmup prompt. Requires text-'prompt' samples; enabling it on a pre-tokenized ('input_tokens') dataset is a hard error.
     drain: false  # Drain in-flight warmup requests before starting the performance phase
     warmup_random_seed: 42  # RNG seed for warmup scheduling and sample ordering
   profiling:

src/inference_endpoint/config/templates/offline_template_full.yaml

@@ -93,7 +93,7 @@ settings:
   warmup:
     enabled: false  # Enable warmup phase before performance run
     n_requests: null  # Warmup request count (None = full dataset once)
-    salt: true  # Prepend a unique random hex salt to each warmup prompt
+    salt: false  # Prepend a unique random hex salt to each warmup prompt. Requires text-'prompt' samples; enabling it on a pre-tokenized ('input_tokens') dataset is a hard error.
     drain: false  # Drain in-flight warmup requests before starting the performance phase
     warmup_random_seed: 42  # RNG seed for warmup scheduling and sample ordering
   profiling:

src/inference_endpoint/config/templates/online_template_full.yaml

@@ -94,7 +94,7 @@ settings:
   warmup:
     enabled: false  # Enable warmup phase before performance run
     n_requests: null  # Warmup request count (None = full dataset once)
-    salt: true  # Prepend a unique random hex salt to each warmup prompt
+    salt: false  # Prepend a unique random hex salt to each warmup prompt. Requires text-'prompt' samples; enabling it on a pre-tokenized ('input_tokens') dataset is a hard error.
     drain: false  # Drain in-flight warmup requests before starting the performance phase
     warmup_random_seed: 42  # RNG seed for warmup scheduling and sample ordering
   profiling:

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,24 @@ def load_from_huggingface(
     return ds[split].to_pandas()
 
 
+def _can_salt(sample: Any) -> DatasetValidationError.Reason | None:
+    """Return the 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).
+    """
+    Reason = DatasetValidationError.Reason
+    if not isinstance(sample, dict):
+        return Reason.TYPE_MISMATCH
+    if "input_tokens" in sample:
+        return Reason.INPUT_TOKENS_SHADOWING
+    if "prompt" not in sample:
+        return Reason.PROMPT_MISSING
+    if not isinstance(sample["prompt"], str):
+        return Reason.PROMPT_TYPE_MISMATCH
+    return None
+
+
 class Dataset:
     """Class for loading and managing benchmark datasets.
 
@@ -440,55 +459,62 @@ 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 non-saltable sample is an error, not a silent
+        skip: skipping would leave the KV cache un-busted. Every sample is
+        checked — a single invalid item fails the run, because the seeded warmup
+        subset can draw any index and salt correctness is all-or-nothing. Called
+        before any load is issued — at benchmark setup and again from with_salt().
+
+        Raises:
+            DatasetValidationError: naming the first offending sample. The index
+                is into the loaded, post-transform sample order, not the source
+                file line.
+        """
+        assert self.data is not None, "Dataset not loaded. Call load() first."
+        for i, sample in enumerate(self.data):
+            reason = _can_salt(sample)
+            if reason is not None:
+                raise DatasetValidationError(
+                    reason,
+                    detail=(
+                        f"sample {i} (index into the loaded, post-transform "
+                        f"order); disable salt (--warmup-salt / warmup.salt: "
+                        f"false) or use a 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): a non-saltable
+        dataset raises here, before any load is issued.
+
+        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."

src/inference_endpoint/exceptions.py

@@ -15,6 +15,8 @@
 
 """Custom exceptions for CLI error handling."""
 
+from enum import Enum
+
 
 class CLIError(Exception):
     """Base exception for CLI errors.
@@ -37,9 +39,32 @@ class InputValidationError(CLIError):
 
 
 class DatasetValidationError(InputValidationError):
-    """Invalid --dataset string or dataset configuration."""
+    """Invalid --dataset string or dataset configuration.
 
-    pass
+    The failure category is a ``Reason``; ``detail`` carries the specifics
+    (offending sample index, remediation hint, parser error text).
+    """
+
+    class Reason(Enum):
+        """Why a dataset failed validation."""
+
+        TYPE_MISMATCH = "sample is not a dict"
+        INPUT_TOKENS_SHADOWING = (
+            "sample has 'input_tokens'; salt cannot bust a pre-tokenized cache"
+        )
+        PROMPT_MISSING = "sample has no 'prompt' field"
+        PROMPT_TYPE_MISMATCH = "sample 'prompt' is not a str"
+        UNSPECIFIED = "dataset validation failed"
+
+    def __init__(
+        self,
+        reason: "DatasetValidationError.Reason",
+        detail: str | None = None,
+    ) -> None:
+        self.reason = reason
+        self.detail = detail
+        message = reason.value if detail is None else f"{reason.value}: {detail}"
+        super().__init__(message)
 
 
 class SetupError(CLIError):