feat: pre-compute ISL token counts for multi-turn dataset-history mode

- Add _precompute_isl_for_multi_turn() in execute.py: runs
  apply_chat_template(messages, tokenize=True, add_generation_prompt=True)
  once per client turn at setup time and stores results in
  sample["input_tokens"], hitting the IslTrigger sync fast path
  (len(token_ids)) with zero hot-path cost.
- Add _extract_prompt_text() in session.py: refactors inline message
  content extraction to handle list-form multimodal content safely,
  fixing a crash when content is a list (e.g. vision/tool-call messages).
- Add unit tests for both helpers and two integration tests covering
  target_concurrency cap enforcement and pipeline exception propagation.

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>

Author: tianmu-li

src/inference_endpoint/commands/benchmark/execute.py

@@ -288,6 +288,42 @@ def _load_datasets(
     return dataloader, accuracy_datasets, eval_configs
 
 
+def _precompute_isl_for_multi_turn(
+    dataloader: MultiTurnDataset, tokenizer_name: str
+) -> None:
+    """Tokenize pre-built message lists and store token counts in each sample.
+
+    Runs apply_chat_template once per client turn so the hot-path IslTrigger
+    sync path (len(token_ids)) is used instead of on-the-fly text tokenization.
+    Only affects dataset-history turns; live-history turns override 'messages'
+    at runtime so the stored input_tokens are stale (acceptable approximation).
+    """
+    # Local import: optional dependency, circular-import avoidance (consistent
+    # with _annotate_response_token_counts in this file).
+    from transformers import AutoTokenizer  # noqa: PLC0415
+
+    tokenizer = AutoTokenizer.from_pretrained(tokenizer_name)
+    skipped = 0
+    for sample in dataloader.data or []:
+        messages = sample.get("messages")
+        if not messages:
+            continue
+        try:
+            token_ids: list[int] = tokenizer.apply_chat_template(
+                messages,
+                tokenize=True,
+                add_generation_prompt=True,
+            )
+            sample["input_tokens"] = token_ids
+        except Exception:  # template errors vary by model; skip gracefully
+            skipped += 1
+    if skipped:
+        logger.warning(
+            "ISL pre-computation: %d turn(s) skipped (apply_chat_template failed)",
+            skipped,
+        )
+
+
 def setup_benchmark(config: BenchmarkConfig, test_mode: TestMode) -> BenchmarkContext:
     """Load tokenizer, dataset, create scheduler, setup report dir."""
     # CPU affinity
@@ -317,6 +353,10 @@ def setup_benchmark(config: BenchmarkConfig, test_mode: TestMode) -> BenchmarkCo
     # Datasets
     dataloader, accuracy_datasets, eval_configs = _load_datasets(config, report_dir)
 
+    if isinstance(dataloader, MultiTurnDataset) and tokenizer_name is not None:
+        logger.info("Pre-computing ISL token counts for multi-turn dataset…")
+        _precompute_isl_for_multi_turn(dataloader, tokenizer_name)
+
     # Setup runtime settings using factory method
     rt_settings = RuntimeSettings.from_config(config, dataloader.num_samples())
 

src/inference_endpoint/load_generator/session.py

@@ -47,6 +47,26 @@
 _WARMUP_ENABLED = os.environ.get("ENABLE_WARMUP") == "1"
 
 
+def _extract_prompt_text(messages: list[Any]) -> str | None:
+    """Join text content from an OpenAI messages list; handles list-form multimodal content."""
+    parts: list[str] = []
+    for m in messages:
+        if not isinstance(m, dict):
+            continue
+        c = m.get("content")
+        if isinstance(c, str) and c:
+            parts.append(c)
+        elif isinstance(c, list):
+            parts.extend(
+                p["text"]
+                for p in c
+                if isinstance(p, dict)
+                and p.get("type") == "text"
+                and isinstance(p.get("text"), str)
+            )
+    return "\n".join(parts) if parts else None
+
+
 # ---------------------------------------------------------------------------
 # Phase configuration
 # ---------------------------------------------------------------------------
@@ -204,12 +224,7 @@ def issue(
             token_ids = data.get("input_tokens") or data.get("token_ids")
             prompt_text = data.get("prompt")
             if prompt_text is None and "messages" in data:
-                parts: list[str] = [
-                    m["content"]
-                    for m in data["messages"]
-                    if isinstance(m, dict) and m.get("content")
-                ]
-                prompt_text = "\n".join(parts) if parts else None
+                prompt_text = _extract_prompt_text(data["messages"])
             prompt_data = PromptData(
                 text=prompt_text,
                 token_ids=tuple(token_ids) if token_ids is not None else None,

tests/integration/test_multi_turn.py

@@ -77,6 +77,7 @@ def _make_dataset(rows: list[dict]) -> MultiTurnDataset:
 def _make_strategy(
     ds: MultiTurnDataset,
     use_dataset_history: bool = True,
+    target_concurrency: int | None = None,
 ) -> MultiTurnStrategy:
     mt_cfg = MultiTurnConfig(
         turn_timeout_s=10.0,
@@ -86,6 +87,7 @@ def _make_strategy(
         conversation_manager=ConversationManager(),
         dataset_metadata=ds.conversation_metadata,
         multi_turn_config=mt_cfg,
+        target_concurrency=target_concurrency,
     )
 
 
@@ -600,6 +602,114 @@ async def test_concurrent_conversations_stress(echo_server):
     assert len(responses) == expected_client_turns
 
 
+@pytest.mark.integration
+@pytest.mark.asyncio
+async def test_multi_turn_active_conversations_respects_target_concurrency(echo_server):
+    num_convs = 20
+    rows = []
+    for i in range(num_convs):
+        conv_id = f"cap_conv_{i}"
+        rows += [
+            {
+                "conversation_id": conv_id,
+                "turn": 1,
+                "role": "user",
+                "content": f"Q1-{i}",
+            },
+            {
+                "conversation_id": conv_id,
+                "turn": 2,
+                "role": "assistant",
+                "content": f"A1-{i}",
+            },
+            {
+                "conversation_id": conv_id,
+                "turn": 3,
+                "role": "user",
+                "content": f"Q2-{i}",
+            },
+        ]
+
+    ds = _make_dataset(rows)
+    strategy = _make_strategy(ds, target_concurrency=4)
+    responses: dict = {}
+
+    observed_max: list[int] = []
+    orig_on_sample_complete = strategy.on_sample_complete
+
+    def tracked_on_sample_complete(result) -> None:
+        observed_max.append(len(strategy._active_iters))
+        orig_on_sample_complete(result)
+
+    strategy.on_sample_complete = tracked_on_sample_complete
+
+    await _run_session(echo_server.url, ds, strategy, responses)
+
+    assert len(responses) == num_convs * 2  # 2 client turns per conversation
+    assert max(observed_max, default=0) <= 4
+
+
+@pytest.mark.integration
+@pytest.mark.asyncio
+async def test_multi_turn_pipeline_exception_propagates(echo_server):
+    rows = [
+        {"conversation_id": "err_c1", "turn": 1, "role": "user", "content": "Q1"},
+        {"conversation_id": "err_c1", "turn": 2, "role": "assistant", "content": "A1"},
+        {"conversation_id": "err_c1", "turn": 3, "role": "user", "content": "Q2"},
+    ]
+    ds = _make_dataset(rows)
+    strategy = _make_strategy(ds)
+
+    call_count = 0
+    orig_issue_next_turn = strategy._issue_next_turn
+
+    def failing_issue_next_turn(*args, **kwargs):
+        nonlocal call_count
+        call_count += 1
+        if call_count >= 2:
+            raise RuntimeError("injected pipeline error")
+        return orig_issue_next_turn(*args, **kwargs)
+
+    strategy._issue_next_turn = failing_issue_next_turn
+
+    loop = asyncio.get_running_loop()
+    http_config = HTTPClientConfig(
+        endpoint_urls=[urljoin(echo_server.url, "/v1/chat/completions")],
+        warmup_connections=0,
+        num_workers=2,
+    )
+    http_client = await HTTPEndpointClient.create(http_config, loop)
+    issuer = HttpClientSampleIssuer(http_client)
+
+    try:
+        session = BenchmarkSession(
+            issuer=issuer,
+            event_publisher=_NoOpPublisher(),
+            loop=loop,
+            on_sample_complete=strategy.on_sample_complete,
+        )
+        rt = RuntimeSettings(
+            metrics.Throughput(1000),
+            [metrics.Throughput(1000)],
+            min_duration_ms=0,
+            max_duration_ms=30_000,
+            n_samples_from_dataset=ds.num_samples(),
+            n_samples_to_issue=ds.num_samples(),
+            min_sample_count=1,
+            rng_sched=random.Random(42),
+            rng_sample_index=random.Random(42),
+            load_pattern=LoadPattern(type=LoadPatternType.MAX_THROUGHPUT),
+        )
+        phase = PhaseConfig("perf", rt, ds, PhaseType.PERFORMANCE, strategy=strategy)
+
+        with pytest.raises(RuntimeError, match="injected pipeline error"):
+            await asyncio.wait_for(session.run([phase]), timeout=30.0)
+
+        assert strategy._inflight == {}
+    finally:
+        await http_client.shutdown_async()
+
+
 @pytest.mark.integration
 @pytest.mark.asyncio
 async def test_tools_field_forwarded_to_endpoint(echo_server):

tests/unit/commands/test_precompute_isl.py

@@ -0,0 +1,106 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Unit tests for _precompute_isl_for_multi_turn."""
+
+from unittest.mock import MagicMock, patch
+
+import pytest
+from inference_endpoint.commands.benchmark.execute import _precompute_isl_for_multi_turn
+
+
+def _make_dataloader(samples: list[dict]) -> MagicMock:
+    dl = MagicMock()
+    dl.data = samples
+    return dl
+
+
+class TestPrecomputeIslForMultiTurn:
+    @pytest.mark.unit
+    def test_sets_input_tokens_for_samples_with_messages(self):
+        samples = [
+            {"messages": [{"role": "user", "content": "hello"}]},
+            {"messages": [{"role": "user", "content": "world"}]},
+        ]
+        dataloader = _make_dataloader(samples)
+        mock_tokenizer = MagicMock()
+        mock_tokenizer.apply_chat_template.side_effect = lambda msgs, **_: list(
+            range(len(msgs) * 3)
+        )
+
+        with patch("transformers.AutoTokenizer") as mock_cls:
+            mock_cls.from_pretrained.return_value = mock_tokenizer
+            _precompute_isl_for_multi_turn(dataloader, "test-model")
+
+        for sample in samples:
+            assert "input_tokens" in sample
+            assert isinstance(sample["input_tokens"], list)
+
+    @pytest.mark.unit
+    def test_leaves_samples_without_messages_untouched(self):
+        samples = [
+            {"prompt": "no messages here"},
+            {"input_tokens": [1, 2, 3]},
+        ]
+        dataloader = _make_dataloader(samples)
+        mock_tokenizer = MagicMock()
+
+        with patch("transformers.AutoTokenizer") as mock_cls:
+            mock_cls.from_pretrained.return_value = mock_tokenizer
+            _precompute_isl_for_multi_turn(dataloader, "test-model")
+
+        mock_tokenizer.apply_chat_template.assert_not_called()
+        assert "input_tokens" not in samples[0]
+        assert samples[1]["input_tokens"] == [1, 2, 3]
+
+    @pytest.mark.unit
+    def test_skips_failed_template_calls_with_warning(self, caplog):
+        samples = [
+            {"messages": [{"role": "user", "content": "good"}]},
+            {"messages": [{"role": "user", "content": "bad"}]},
+        ]
+        dataloader = _make_dataloader(samples)
+
+        def side_effect(msgs, **_):
+            if msgs[0]["content"] == "bad":
+                raise ValueError("template error")
+            return [10, 20, 30]
+
+        mock_tokenizer = MagicMock()
+        mock_tokenizer.apply_chat_template.side_effect = side_effect
+
+        with patch("transformers.AutoTokenizer") as mock_cls:
+            mock_cls.from_pretrained.return_value = mock_tokenizer
+            with caplog.at_level("WARNING"):
+                _precompute_isl_for_multi_turn(dataloader, "test-model")
+
+        assert "input_tokens" in samples[0]
+        assert "input_tokens" not in samples[1]
+        assert "1 turn(s) skipped" in caplog.text
+
+    @pytest.mark.unit
+    def test_add_generation_prompt_true(self):
+        samples = [{"messages": [{"role": "user", "content": "hi"}]}]
+        dataloader = _make_dataloader(samples)
+        mock_tokenizer = MagicMock()
+        mock_tokenizer.apply_chat_template.return_value = [1, 2, 3]
+
+        with patch("transformers.AutoTokenizer") as mock_cls:
+            mock_cls.from_pretrained.return_value = mock_tokenizer
+            _precompute_isl_for_multi_turn(dataloader, "test-model")
+
+        _, kwargs = mock_tokenizer.apply_chat_template.call_args
+        assert kwargs.get("add_generation_prompt") is True
+        assert kwargs.get("tokenize") is True

tests/unit/load_generator/test_async_session.py

@@ -39,6 +39,7 @@
     PhaseResult,
     PhaseType,
     SessionResult,
+    _extract_prompt_text,
 )
 from inference_endpoint.metrics.metric import Throughput
 
@@ -882,3 +883,53 @@ def test_perf_results_filter(self, enable_warmup):
         assert len(sr.perf_results) == 2
         assert len(sr.accuracy_results) == 1
         assert sr.perf_results[0].name == "perf1"
+
+
+@pytest.mark.unit
+class TestExtractPromptText:
+    def test_string_content_extracted(self):
+        messages = [
+            {"role": "user", "content": "Hello"},
+            {"role": "assistant", "content": "Hi"},
+        ]
+        assert _extract_prompt_text(messages) == "Hello\nHi"
+
+    def test_multimodal_list_content_text_parts_extracted(self):
+        messages = [
+            {
+                "role": "user",
+                "content": [
+                    {"type": "text", "text": "Describe this image"},
+                    {"type": "image_url"},
+                ],
+            }
+        ]
+        assert _extract_prompt_text(messages) == "Describe this image"
+
+    def test_mixed_string_and_list_content(self):
+        messages = [
+            {"role": "system", "content": "You are helpful"},
+            {
+                "role": "user",
+                "content": [
+                    {"type": "text", "text": "What is this?"},
+                    {"type": "image_url"},
+                ],
+            },
+        ]
+        assert _extract_prompt_text(messages) == "You are helpful\nWhat is this?"
+
+    def test_none_content_skipped(self):
+        messages = [
+            {"role": "assistant", "content": None},
+            {"role": "user", "content": "Hello"},
+        ]
+        assert _extract_prompt_text(messages) == "Hello"
+
+    def test_list_content_with_no_text_parts_returns_none(self):
+        messages = [{"role": "user", "content": [{"type": "image_url"}]}]
+        assert _extract_prompt_text(messages) is None
+
+    def test_non_dict_messages_skipped(self):
+        messages = ["not a dict", {"role": "user", "content": "Valid"}]
+        assert _extract_prompt_text(messages) == "Valid"