diff --git a/src/inference_endpoint/commands/benchmark/execute.py b/src/inference_endpoint/commands/benchmark/execute.py index 301c1781..d65ee4a0 100644 --- a/src/inference_endpoint/commands/benchmark/execute.py +++ b/src/inference_endpoint/commands/benchmark/execute.py @@ -380,6 +380,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: diff --git a/src/inference_endpoint/dataset_manager/dataset.py b/src/inference_endpoint/dataset_manager/dataset.py index bd259d7a..6d4e160d 100644 --- a/src/inference_endpoint/dataset_manager/dataset.py +++ b/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." diff --git a/tests/unit/commands/test_benchmark.py b/tests/unit/commands/test_benchmark.py index 491dc72b..d5892371 100644 --- a/tests/unit/commands/test_benchmark.py +++ b/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).""" diff --git a/tests/unit/dataset_manager/test_salted_dataset.py b/tests/unit/dataset_manager/test_salted_dataset.py index acff8900..329b17fa 100644 --- a/tests/unit/dataset_manager/test_salted_dataset.py +++ b/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