Feat/tool sequences

[tianmu-li]

What does this PR do?

Updated multi-turn implementation for #232. Added tool sequencing, fixed scheduler for concurrent requests.

Type of change

• Bug fix
• New feature
• Documentation update
• Refactor/cleanup

Related issues

Testing

• Tests added/updated
• All tests pass locally
• Manual testing completed

Checklist

• Code follows project style
• Pre-commit hooks pass
• Documentation updated (if needed)

[github-actions]

MLCommons CLA bot All contributors have signed the MLCommons CLA ✍️ ✅

[gemini-code-assist]

Code Review

This pull request introduces a comprehensive multi-turn conversation benchmarking framework, including a new MultiTurnScheduler, ConversationManager, and MultiTurnDataset. These additions enable benchmarking of conversational AI workloads with turn sequencing, conversation history management, and optional concurrency control. My review identified potential issues regarding the usage of sentinel objects in the scheduler and the robustness of the timeout logic in the conversation manager.

[arekay-nv]

Partial review - will complete later.

[arekay-nv]

⚠️ Superseded. This comment was posted while a pending review was blocking inline-comment delivery. The review has since been re-posted with all 14 inline comments attached: see review #4184417968 and the updated summary. The findings are unchanged.

[arekay-nv]

Review Council — Multi-AI Code Review

Reviewed by: Claude (Codex review failed with a CLI config error — invalid 'features' requirement 'browser_use' from cloud requirements) | Depth: thorough

Found 15 issues (0 critical, 2 high, 5 medium, 8 low). 14 posted as inline comments, 1 in summary table only (line outside diff hunk).

[arekay-nv]

Review Council — Multi-AI Code Review

Reviewed by: Claude | Depth: thorough | Commit: 0a7ad37

14 inline comments posted in review #4184417968. 1 finding (openai_adapter.py:131) is summary-only because the line falls outside the diff hunk.

Found 15 issues (0 critical, 2 high, 5 medium, 8 low). Codex review failed with a CLI config error (invalid features requirement browser_use from cloud requirements) — Claude-only review.

🟠 Must Fix (high)

Issues that will cause incorrect behavior users will hit in normal usage.

#	Location	Category	Summary
1	src/inference_endpoint/openai/openai_msgspec_adapter.py:222	bug	from_endpoint_response violates the QueryResult.metadata type contract by passing metadata=None when no…
2	src/inference_endpoint/load_generator/multi_turn_strategy.py:137	error-handling	await asyncio.gather(*tasks, return_exceptions=True) silently swallows every exception raised by a _conv_pipeline coroutine. The result…

🟡 Should Fix (medium)

Real issues that trigger under specific conditions, or design flaws that will compound.

#	Location	Category	Summary
3	src/inference_endpoint/load_generator/multi_turn_strategy.py:164	bug	When a turn times out waiting for the previous turn's response, the pipeline does state.failed_turns += 1 but never increments…
4	src/inference_endpoint/load_generator/session.py:206	data-integrity	PromptData.text = data.get("prompt") and token_ids = data.get("input_tokens") or data.get("token_ids") — neither is set on multi-turn…
5	src/inference_endpoint/dataset_manager/multi_turn_dataset.py:212	bug	system_prompts_by_conv is keyed by str(conv_id) (line 212), but pre_built_messages_by_key and current_turn_messages_by_key are…
6	src/inference_endpoint/load_generator/multi_turn_strategy.py:180	bug	In live-history mode (use_dataset_history=False), the per-turn tool message reuses the dataset's hardcoded tool_call_id (e.g.…
7	src/inference_endpoint/load_generator/conversation_manager.py:142	error-handling	mark_turn_complete and mark_turn_failed raise KeyError if conversation_id is missing. These are invoked from…

🔵 Consider (low)

Valid improvements that could be follow-ups.

#	Location	Category	Summary
8	src/inference_endpoint/config/runtime_settings.py:200	design	self.load_pattern.type.value == "multi_turn" compares the enum value as a string literal instead of the enum:…
9	src/inference_endpoint/config/schema.py:253	design	MultiTurnConfig uses model_config = {"extra": "forbid"} (raw dict) and is missing frozen=True, while every other config model in this…
10	src/inference_endpoint/dataset_manager/multi_turn_dataset.py:222	bug	When iterating prior_rows to build pre-built messages, a tool row whose tool_results field is an empty list (e.g. tool_results: [])…
11	tests/integration/test_multi_turn.py:559	testing	The "tool use" integration tests verify that tools, tool_calls, and tool_results are forwarded to the endpoint and that all client…
12	src/inference_endpoint/dataset_manager/multi_turn_dataset.py:142	performance	_validate_conversation_structure, _validate_turn_numbering, and _build_metadata each call self.dataframe.groupby("conversation_id")…
13	src/inference_endpoint/load_generator/conversation_manager.py:45	concurrency	turn_done: asyncio.Event = field(default_factory=asyncio.Event) is constructed at ConversationState instantiation time, before…
14	src/inference_endpoint/load_generator/multi_turn_strategy.py:95	design	self._inflight: dict[str, str] = {} is mutated from two execution contexts (the per-conversation pipeline tasks at line 190, and…
15	src/inference_endpoint/openai/openai_adapter.py:131 (summary-only)	api-contract	OpenAIAdapter.from_endpoint_response (the non-msgspec adapter, used when callers explicitly select it) does not extract tool_calls,…

[tianmu-li]

Hi @arekay-nv, I've addressed the comments. Appreciate it if you could take another look.

[Copilot]

Pull request overview

This PR updates the benchmarking system to support multi-turn conversational workloads (including tool-calling sequences), adds a dedicated multi-turn dataset format + conversion/validation tooling, and wires a new multi-turn load strategy into the benchmarking session and OpenAI adapters.

Changes:

• Add MultiTurnDataset (flat-row JSONL format) with validation, metadata precomputation, and adapter-default handling for per-turn parameters/tools.
• Add MultiTurnStrategy + ConversationManager to enforce per-conversation turn sequencing with optional global concurrency limiting, and integrate it into BenchmarkSession.
• Extend OpenAI request/response handling for messages, tools, tool-call metadata, and streaming tool-call accumulation; add extensive unit/integration tests and multi-turn docs/examples/scripts.

Reviewed changes

Copilot reviewed 43 out of 44 changed files in this pull request and generated 5 comments.

Show a summary per file

File	Description
tests/unit/openai/test_openai_adapter.py	New unit tests for OpenAIAdapter tool serialization and tools forwarding.
tests/unit/openai/test_msgspec_adapter.py	New unit tests for msgspec OpenAI adapter tool-call fields and message dict conversion.
tests/unit/load_generator/test_multi_turn_strategy.py	New unit tests for turn sequencing, concurrency semaphore behavior, and metadata propagation.
tests/unit/load_generator/test_multi_turn_conversation_manager.py	New unit tests for conversation state bookkeeping and event gating.
tests/unit/dataset_manager/test_transforms.py	Add coverage for the new AddDefaultColumns transform.
tests/unit/dataset_manager/test_multi_turn_dataset.py	Comprehensive tests for MultiTurnDataset, including tool sequences and metadata correctness.
tests/unit/core/test_types.py	Add tests for QueryResult.with_metadata() and Query.metadata round-tripping.
tests/unit/config/test_schema.py	Add tests for multi-turn config validation and multi-turn sample counting logic.
tests/integration/test_multi_turn.py	End-to-end integration tests exercising dataset-history/live-history modes and tool-use conversations.
src/inference_endpoint/openai/types.py	Extend msgspec OpenAI types to include tool_calls, tool_call_id, and tools.
src/inference_endpoint/openai/openai_msgspec_adapter.py	Support messages input, tool-call fields, tools forwarding, and richer response metadata.
src/inference_endpoint/openai/openai_adapter.py	Support messages input, tools forwarding, and return richer response metadata.
src/inference_endpoint/openai/accumulator.py	Accumulate streamed tool_calls + finish_reason into final QueryResult.metadata.
src/inference_endpoint/load_generator/strategy.py	Extend PhaseIssuerProtocol.issue() to accept an optional data_override.
src/inference_endpoint/load_generator/session.py	Allow injecting a per-phase strategy; support data overrides in sample issuance.
src/inference_endpoint/load_generator/multi_turn_strategy.py	New multi-turn strategy implementing per-conversation sequencing + global concurrency limiting.
src/inference_endpoint/load_generator/conversation_manager.py	New synchronous conversation state manager used by multi-turn strategy.
src/inference_endpoint/endpoint_client/worker.py	Propagate Query.metadata through requests and merge into results.
src/inference_endpoint/endpoint_client/http.py	Add query_metadata field to InFlightRequest.
src/inference_endpoint/endpoint_client/adapter_protocol.py	Generalize SSE decoding/parse APIs to return adapter-specific chunk objects.
src/inference_endpoint/dataset_manager/transforms.py	Add AddDefaultColumns (fill-missing-only) transform.
src/inference_endpoint/dataset_manager/multi_turn_dataset.py	New multi-turn dataset implementation with tool-sequence handling and metadata building.
src/inference_endpoint/dataset_manager/factory.py	Select MultiTurnDataset when dataset config includes multi_turn; skip prompt-based transforms for it.
src/inference_endpoint/dataset_manager/init.py	Export MultiTurnDataset and AddDefaultColumns.
src/inference_endpoint/core/types.py	Add Query.metadata and QueryResult.with_metadata().
src/inference_endpoint/config/templates/online_template_full.yaml	Expose multi_turn dataset block and multi_turn load pattern option in template.
src/inference_endpoint/config/templates/online_template.yaml	Expose multi_turn load pattern option in template.
src/inference_endpoint/config/templates/offline_template_full.yaml	Expose multi_turn dataset block and load pattern option in template.
src/inference_endpoint/config/templates/concurrency_template_full.yaml	Expose multi_turn dataset block and load pattern option in template.
src/inference_endpoint/config/templates/concurrency_template.yaml	Expose multi_turn load pattern option in template.
src/inference_endpoint/config/schema.py	Add multi-turn schema objects and cross-validate dataset.multi_turn ↔ load_pattern.type.
src/inference_endpoint/config/runtime_settings.py	Make multi-turn sample count issue all dataset client turns (min-sample-count aware).
src/inference_endpoint/commands/benchmark/execute.py	Instantiate and wire MultiTurnStrategy automatically when using MultiTurnDataset.
scripts/validate_jsonl_schema.py	New CLI script to validate multi-turn JSONL rows against schema.
scripts/multi_turn_dataset_schema.json	New JSON Schema for multi-turn flat-row JSONL datasets.
scripts/convert_agentic_snapshot.py	New conversion+verification script from snapshot-style agentic datasets to flat-row JSONL.
examples/09_MultiTurn/multi_turn_with_concurrency.yaml	Example config: multi-turn with global concurrency limiting.
examples/09_MultiTurn/multi_turn_benchmark.yaml	Example config: basic multi-turn benchmark.
examples/09_MultiTurn/datasets/.gitkeep	Placeholder for converted example datasets.
examples/09_MultiTurn/customer_support_conversations.jsonl	Example multi-turn dataset.
examples/09_MultiTurn/agentic_workflow_benchmark.yaml	Example config for converted agentic workflow dataset.
examples/09_MultiTurn/agentic_coding_benchmark.yaml	Example config for converted agentic coding dataset.
examples/09_MultiTurn/README.md	Multi-turn feature documentation and agentic conversion guidance.
docs/MULTI_TURN_QUICKSTART.md	Quickstart guide for running multi-turn benchmarks.

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

[arekay-nv]

There is some confusion around the issue policy, lets sync today.

[Copilot]

Pull request overview

Copilot reviewed 43 out of 44 changed files in this pull request and generated 2 comments.

Comments suppressed due to low confidence (2)

src/inference_endpoint/endpoint_client/adapter_protocol.py:134

• parse_sse_chunk now appends whatever decode_sse_message returns, including None (e.g., when an SSE message has no choices). This contradicts the docstring (“Silently ignores non-content SSE messages”) and forces downstream accumulators to defensively handle None. Consider filtering out None return values here (and/or handling exceptions per JSON doc) so call sites only see meaningful chunk objects.

    def parse_sse_chunk(cls, buffer: bytes, end_pos: int) -> list[Any]:
        """
        Parse SSE chunk and extract all chunk objects.

        Extracts JSON documents from SSE stream and decodes them to chunk objects.
        Silently ignores non-content SSE messages (role, finish_reason, etc).

        Args:
            buffer: Byte buffer containing SSE data
            end_pos: End position in buffer to parse up to

        Returns:
            List of chunk objects extracted from the SSE chunk
        """
        json_docs = cls.SSE_DATA_PATTERN.findall(buffer[:end_pos])
        parsed_contents = []

        try:
            for json_doc in json_docs:
                content = cls.decode_sse_message(json_doc)
                parsed_contents.append(content)
        except Exception:
            # Normal for non-content SSE messages (role, finish_reason, etc)
            pass

        return parsed_contents

src/inference_endpoint/core/types.py:242

• The Query docstring’s gc=False note only mentions data/headers, but Query now also has a mutable metadata dict. To avoid future misuse (and to match the more explicit QueryResult guidance), consider updating this note to include metadata as well and/or adding an AT-RISK (gc=False) warning that data/metadata/headers must not be mutated to introduce cycles.

    Attributes:
        id: Unique identifier for this query (auto-generated UUID).
        data: Request payload as a dictionary (typically contains prompt, model, etc.).
        metadata: Internal metadata that round-trips through transport (e.g., conversation_id).
        headers: HTTP headers to include in the request (e.g., authorization).
        created_at: Timestamp when query was created (seconds since epoch).

    Example:
        >>> query = Query(
        ...     data={"prompt": "Hello", "model": "Qwen/Qwen3-8B", "max_tokens": 100},
        ...     headers={"Authorization": "Bearer token123"},
        ... )

    Note:
        gc=False: Safe because data/headers are simple key-value pairs without cycles.
        Do NOT store self-referential or cyclic structures in data/headers fields.

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

[tianmu-li]

Hi @arekay-nv, I've addressed the comments. Main change is refactoring to an event-based model and limiting number of active conversations to concurrency. Appreciate it if you could take another look.

[arekay-nv]

Almost there - the only issue is the stickiness of the client-turns.

[arekay-nv]

@viraatc can you check the perf implications - there might be non-negligible overhead for non-agentic/multi-turn workloads.

[arekay-nv]

conversation_metadata: dict[str, Any] is load-bearing across two files but has no schema. Consumers in multi_turn_strategy.py read metadata["samples"], .get("system_prompts_by_conv", {}), .get("current_turn_messages_by_key", {}), and .get("pre_built_messages_by_key", {}) — typos return silent empty dicts. Worse, this dict is two-phase-initialized: _build_metadata() populates entries without sample_index, and load() mutates them in place to add it (see line 451). Any consumer that reads metadata between __init__ and load() gets a different shape — tests at test_multi_turn_dataset.py:188 already rely on the partial shape, locking in the temporal coupling.

Three concrete risks:

1. Stringly-typed (conv_id, turn) tuple keys — multi_turn_strategy.py:201 does .get((conv_id, turn)). If conv_id is anything but str (e.g. an int from JSON), it silently returns None → empty messages.
2. Samples are silently dropped at line 450 if their post-transform key isn't in key_to_sample_index — no log, no count.
3. The aligned-indices invariant between samples[i]["sample_index"] and self.data is tribal knowledge.

Fix: replace with a typed ConversationMetadata dataclass (frozen, slots) produced once by load(). This eliminates the two-phase shape, makes key types explicit, and gives consumers attribute access with type-checker coverage. This is the single highest-impact refactor available in the PR.

[arekay-nv]

Integration tests don't validate metrics for multi-turn. This file has no assertion that the metrics aggregator produces ISL/OSL/TPOT values for a multi-turn run. Given the PR adds _precompute_isl_for_multi_turn specifically to populate input_tokens for the hot-path IslTrigger, and OslTrigger/TpotTrigger now have _extract_message paths for tool-call responses (metrics_table.py:312-407), at least one end-to-end test should assert that:

• IslTrigger reports the precomputed token count (not the text-fallback count)
• OslTrigger/TpotTrigger produce a value for a streamed tool-call response
• For a SSE response with tool_calls only (no text content), the TTFT sentinel chunk fires

Without this, a regression in either trigger goes unnoticed.

[arekay-nv]

Additional findings outside the PR diff

These can't be posted inline because they reference lines that aren't part of this change, but they're directly worsened or surfaced by it.

src/inference_endpoint/endpoint_client/adapter_protocol.py:126-132 — broad except Exception: pass swallows tool-call decode failures. The block is pre-existing, but its blast radius has grown dramatically in this PR: decode_sse_message now goes through msgspec.json.decode(..., type=SSEMessage) with an expanded SSEDelta / SSEChoice schema that includes tool_calls: list[dict[str, Any]] | None. The catch now silently hides:

• msgspec.DecodeError from any malformed SSE payload
• msgspec.ValidationError when a streamed tool-call partial doesn't conform (server sends tool_calls as an object, non-string id, schema drift)
• The comment claims this is "normal for non-content SSE messages (role, finish_reason, etc)" — but the new accumulator handles those explicitly, so the rationale is stale and the catch now hides real bugs.

A loop-level try means one decode failure mid-batch silently truncates the chunk list — surviving documents in the same SSE flush are lost. Tool-call accumulation can drop the suffix of a streamed response, producing invalid JSON in conversation history without any log. Recommend: narrow the catch to msgspec.DecodeError/ValidationError, log at warning with query_id and the offending payload prefix, and don't break out of the loop on a single bad doc — continue past it.

src/inference_endpoint/load_generator/session.py:465-472 — inflight count drifts negative after timed-out late response. _handle_timeout in MultiTurnStrategy removes the entry from MultiTurnStrategy._inflight, but phase_issuer.uuid_to_index retains the query id. If the late response eventually arrives, _handle_response still hits the query_id in phase_issuer.uuid_to_index branch and decrements phase_issuer.inflight — potentially driving it below zero. The <= 0 check in the drain logic still fires (already <= 0), so no symptom in normal runs, but the invariant violation will silently mask future bugs. Recommend: clear uuid_to_index[query_id] from _handle_timeout, or guard the decrement.

tests/unit/core/test_types.py — missing class-level @pytest.mark.unit markers. TestErrorData (line 36), TestQuerySerialization (50), TestQueryResultSerialization (189), TestStreamChunkSerialization (470), TestQueryResultWorkerPatterns (592), TestTextAfterFirstChunk (693), TestMixedTypeSerialization (747) lack the marker. Only TestTextAfterFirstChunkToolCalls near line 896 has it. This breaks pytest -m unit filtering — these tests would not run in a -m unit CI lane.

Test gaps worth adding (the existing tests catch them but not as explicit cases):

• _handle_timeout failure bookkeeping (state.failed_turns increments for the timed-out turn and every remaining un-issued turn, state.is_complete() becomes True, _fill_slot() runs)
• phase_issuer.issue() returning None mid-stream — the FakePhaseIssuer in test_multi_turn_strategy.py defines stop_after but no test ever passes a non-None value, so the session-stop branch in multi_turn_strategy.py:219-223 is untested
• KeyError paths in on_sample_complete and ConversationManager.mark_turn_*
• _expand_tool_results validation paths (tool_call_id is None, content is None) and the "empty tool_results list" warning
• _PROPAGATED_KEYS = {"model", "max_completion_tokens", "max_new_tokens", "stream", "tools"} propagation from first user row — current tests supply tools on every relevant row, so propagation logic isn't exercised
• Accumulator: parallel tool_calls at different indices (only idx=0 tested), out-of-order indices (currently sorted at finalize, never tested), partial without index key (defaults to 0)
• SSE streaming end-to-end with real tool-call deltas — the echo server doesn't emit tool_calls deltas, so the new SSEDelta.tool_calls decoding path is only covered by unit tests, not the real transport pipeline

Type tightening — ChatMessage.role: str could be Literal["user", "assistant", "system", "tool"]. msgspec accepts Literal types and would reject unknown values on decode — a free encapsulation win for the wire-format type. The dataset validator already enforces this at multi_turn_dataset.py:154-159; the API boundary doesn't.

---

Generated by a multi-agent review pass (code-reviewer, pr-test-analyzer, silent-failure-hunter, comment-analyzer, type-design-analyzer). See inline comments for the in-diff findings.

[Copilot]

Pull request overview

Copilot reviewed 56 out of 57 changed files in this pull request and generated 22 comments.

[Copilot]

Pull request overview

Copilot reviewed 54 out of 55 changed files in this pull request and generated 12 comments.

[viraatc]

looks great thanks for the MR!
( completed first look will be doing another pass )

[viraatc]

we intentionally put try-except outside the for-loop for perf reasons.
this has large impact in streaming mode.

[tianmu-li]

Moved try-except to outside the for-loop

[viraatc]

adapter code should prioritize performance, wondering if its possible to avoid the isinstance checks and propagate the error directly at runtime.

[tianmu-li]

Fixed. isinstance check is removed.

[viraatc]

Q: why do we need have mutable content? does chat-message with None content make sense?

[tianmu-li]

Yes, tool-dispatching assistant message needs to carry content=null.

[Copilot]

Pull request overview

Copilot reviewed 56 out of 57 changed files in this pull request and generated 15 comments.

Comments suppressed due to low confidence (3)

tests/unit/load_generator/test_multi_turn_strategy.py:110

• This monkey-patched issue method replaces FakePhaseIssuer.issue but also lacks the conversation_id and turn keyword parameters that MultiTurnStrategy passes, so this test will fail with TypeError on the first issued turn.
  tests/unit/load_generator/test_multi_turn_strategy.py:174
• MultiTurnStrategy passes conversation_id and turn into issue(); this TimedIssuer signature does not accept those keywords, so test_turn_ordering_enforced fails before it can validate ordering.
  tests/unit/load_generator/test_multi_turn_strategy.py:393
• MultiTurnStrategy calls issue with conversation_id and turn keyword arguments. This ErrorIssuer signature rejects those keywords, so the test raises a TypeError instead of the intended simulated pipeline error.

[Copilot]

Pull request overview

Copilot reviewed 54 out of 55 changed files in this pull request and generated 11 comments.