[feat] Initial multi-turn enabling

[tianmu-li]

What does this PR do?

Initial implementation for multi-turn benchmarks.
Known limitations:

1. Datasets with tool calls are not yet tested.
2. Turn n+1 sent to the server contains turn n from server, instead of turn n from dataset (server output in turn n is not dropped). This is closer to model usage in reality but doesn't fully replay dataset.

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 ✍️ ✅

[Copilot]

Pull request overview

Adds first-pass multi-turn benchmarking support by introducing conversation-aware dataset loading, scheduling, request construction, and event metadata propagation to enable turn sequencing and optional global concurrency limiting.

Changes:

• Introduces MultiTurnDataset, ConversationManager, and MultiTurnScheduler to support multi-turn conversation issuance (parallel/sequential; poisson fallback).
• Propagates conversation metadata (conversation_id, turn_number) through Query/transport and into events.db.
• Updates OpenAI adapters to accept pre-built messages arrays; adds docs, examples, and extensive unit/integration tests.

Reviewed changes

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

Show a summary per file

File	Description
tests/unit/load_generator/test_sample.py	Adds a unit test for error handling in multi-turn completion handling.
tests/unit/load_generator/test_multi_turn_scheduler.py	New unit tests covering multi-turn scheduling modes, blocking, and concurrency control.
tests/unit/load_generator/test_multi_turn_load_generator.py	New unit tests for multi-turn dataset param forwarding behavior.
tests/unit/load_generator/test_multi_turn_conversation_manager.py	New unit tests for conversation state management and concurrency behavior.
tests/unit/load_generator/test_load_generator.py	Updates existing tests to use dict-like sample data consistent with new load path.
tests/unit/dataset_manager/test_multi_turn_dataset.py	Adds comprehensive tests for MultiTurnDataset validation/metadata/sample loading.
tests/integration/test_multi_turn.py	Adds end-to-end integration tests for multi-turn modes, history, and concurrency.
src/inference_endpoint/openai/openai_msgspec_adapter.py	Adds support for messages input in request construction.
src/inference_endpoint/openai/openai_adapter.py	Adds support for messages input in request construction.
src/inference_endpoint/metrics/recorder.py	Extends event rows with conversation_id and turn_number fields and plumbs them through record_event.
src/inference_endpoint/load_generator/scheduler.py	Adds MultiTurnScheduler and blocking sentinels; supports optional global concurrency limiting.
src/inference_endpoint/load_generator/sample.py	Adds conversation-aware completion handling and ConversationSample.
src/inference_endpoint/load_generator/load_generator.py	Builds multi-turn requests using conversation history and forwards generation params.
src/inference_endpoint/load_generator/conversation_manager.py	New conversation state manager with turn sequencing, failure handling, and completion tracking.
src/inference_endpoint/load_generator/init.py	Exports new multi-turn types.
src/inference_endpoint/endpoint_client/worker.py	Preserves/merges query metadata into results/errors for round-tripping.
src/inference_endpoint/endpoint_client/http.py	Adds query_metadata to InFlightRequest.
src/inference_endpoint/endpoint_client/http_sample_issuer.py	Attaches conversation metadata to Query.metadata when issuing multi-turn samples.
src/inference_endpoint/dataset_manager/multi_turn_dataset.py	New dataset implementation for multi-turn JSONL with validation and param forwarding.
src/inference_endpoint/dataset_manager/factory.py	Selects MultiTurnDataset + skips adapter compatibility transforms when multi-turn is enabled.
src/inference_endpoint/dataset_manager/init.py	Exports MultiTurnDataset.
src/inference_endpoint/core/types.py	Adds metadata field to Query for internal metadata transport.
src/inference_endpoint/config/schema.py	Adds MULTI_TURN load pattern + multi-turn configuration models/enums.
MULTI_TURN_QUICKSTART.md	Adds a multi-turn quickstart guide.
HYBRID_SCHEDULER_GUIDE.md	Adds documentation for hybrid turn sequencing + concurrency limiting.
examples/09_MultiTurn/README.md	Adds multi-turn example documentation.
examples/09_MultiTurn/multi_turn_with_concurrency.yaml	Adds example config with concurrency limiting.
examples/09_MultiTurn/multi_turn_benchmark.yaml	Adds example baseline config.
examples/09_MultiTurn/customer_support_conversations.jsonl	Adds an example multi-turn dataset.

Comments suppressed due to low confidence (1)

src/inference_endpoint/load_generator/load_generator.py:189

• load_sample_data() records LOADGEN_DATA_LOAD using a default sample_uuid="placeholder", and __next__() always passes that placeholder. This makes all data-load events share the same UUID, breaking correlation of per-sample timelines in events.db (and any downstream metrics that join events by sample_uuid). Generate the sample UUID first and pass the real UUID into load_sample_data (or let load_sample_data return data without recording, and record the event once a real UUID exists).

    def load_sample_data(
        self, sample_index: int, sample_uuid: str = "placeholder"
    ) -> Any:
        """Load sample data from dataloader and record event.

        Helper method that loads sample data and records the data load event
        for accurate timing measurements.

        Args:
            sample_index: Index of sample in dataset.
            sample_uuid: UUID of the sample being created (default: "placeholder").

        Returns:
            Sample data loaded from dataloader (format depends on dataset).
        """
        sample_data = self.dataloader.load_sample(sample_index)
        EventRecorder.record_event(
            SessionEvent.LOADGEN_DATA_LOAD,
            time.monotonic_ns(),
            sample_uuid=sample_uuid,
        )
        return sample_data

---

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

[Copilot]

In __next__(), sample data is loaded and the load event is recorded with sample_uuid="placeholder", then a new Sample/ConversationSample is created with a different UUID. This guarantees a UUID mismatch between LOADGEN_DATA_LOAD and subsequent events for the same sample. Please ensure the same UUID is used across all events for a sample (e.g., create the Sample first to obtain uuid, then load data and set it, or allow passing a pre-generated uuid into ConversationSample).

[Copilot]

This test’s assertions don’t match the current multi-turn error-handling behavior: SampleEventHandler.query_result_complete() calls conversation_manager.mark_turn_failed(conv_id) when result.error is set, which advances current_turn, clears pending_user_turn, and appends an assistant error placeholder to message_history. As written, the test should fail. Either update the test expectations to reflect mark_turn_failed semantics, or change the handler to not advance the conversation on error (but then scheduling can deadlock unless timeouts/skip logic handles it).

[Copilot]

test_multi_turn_scheduler_timeout_handling is checking for a skipped turn == 2, but the provided multi_turn_dataset_metadata contains only user turns (1, 3, 5). As a result, first_conv_turn2_found will always be false and the test doesn’t actually verify the skip behavior. Consider asserting that the first conversation’s next user turn (e.g., turn 3) is skipped after a timeout, and set multi_turn_config.turn_timeout_s to a small value in this test to avoid a multi-second unit test runtime.

[Copilot]

The docstring states that role can be "user", "assistant", or "system", but _validate_conversation_structure() enforces strict alternation of only "user"/"assistant" and will reject any "system" role rows. Either update the documentation to match the implemented schema (system prompt via separate system field), or extend validation/processing to support role == "system" rows explicitly.

Suggested change

	- role: Speaker role ("user", "assistant", or "system")
	- content: Message content
	Optional columns:
	- system: System prompt (only for first user turn)
	- role: Speaker role ("user" or "assistant")
	- content: Message content
	Optional columns:
	- system: System prompt associated with the conversation (typically set on the first user turn)

[Copilot]

set_conversation_manager() enables multi-turn behavior via global state on the singleton handler, but there’s no corresponding way to unset/reset it. If the process runs multiple benchmark sessions, multi-turn state may leak into subsequent runs. Consider adding a clear_conversation_manager() / reset() API (and calling it from session teardown) so single-turn runs aren’t affected by a previously configured conversation manager.

Suggested change

	def clear_conversation_manager(self) -> None:
	"""Disable multi-turn mode by clearing the conversation manager.
	This should be called during session teardown to avoid leaking
	conversation state between benchmark runs.
	"""
	self.conversation_manager = None

[gemini-code-assist]

Code Review

This pull request introduces a comprehensive multi-turn conversation benchmarking feature, including a new MultiTurnScheduler with parallel and sequential modes, a ConversationManager for state tracking, and a MultiTurnDataset for handling structured conversational data. The implementation also adds concurrency control to multi-turn workloads and updates the event recording system to track conversation-specific metadata. Valid feedback from the review identifies a syntax error and a broken link in the documentation, as well as an opportunity to improve the robustness of parameter forwarding logic by using an allow-list. Additionally, a unit test was found to have incorrect assertions regarding how conversation state should advance when a turn fails.

[gemini-code-assist]

The Python code example contains invalid syntax timeout=300s. This should be timeout=300 to be valid Python, as the s suffix is not recognized. This could be confusing for users trying to understand or use the example.

Suggested change

	if sample.turn > 1:
	conversation_manager.wait_for_turn_ready(
	conv_id,
	turn,
	timeout=300s
	)
	timeout=300

[gemini-code-assist]

This line contains a reference to MULTI_TURN_IMPLEMENTATION_SUMMARY.md, but this file does not appear to be included in the pull request. This results in a broken link in the documentation.

[gemini-code-assist]

The logic for forwarding generation parameters uses an exclude_fields set. This approach can be brittle; if new conversational metadata fields are added to the sample data in the future, they must also be added to this exclusion list to prevent them from being incorrectly forwarded as generation parameters.

A more robust approach would be to use an allow-list of known generation parameters, such as the GENERATION_PARAMS set already defined in src/inference_endpoint/dataset_manager/multi_turn_dataset.py. This would ensure only valid generation parameters are forwarded.

Suggested change

	exclude_fields = {"conversation_id", "turn", "role", "content", "system"}
	for key, value in sample_data_raw.items():
	if key not in exclude_fields and value is not None:
	request_data[key] = value
	from inference_endpoint.dataset_manager.multi_turn_dataset import GENERATION_PARAMS
	# Forward all generation parameters from sample_data_raw
	for key, value in sample_data_raw.items():
	if key in GENERATION_PARAMS and value is not None:
	request_data[key] = value

[gemini-code-assist]

The test test_query_result_complete_does_not_advance_conversation_on_error has incorrect assertions and a misleading name. When a turn fails, SampleEventHandler calls ConversationManager.mark_turn_failed, which correctly advances the conversation state to prevent sequential mode from getting stuck. This test should verify that behavior.

The assertions should check that current_turn is advanced, pending_user_turn is cleared, and an error placeholder is added to the message history.

def test_query_result_complete_advances_conversation_on_error(
    record_event_mock, record_exception_mock, clean_sample_event_hooks
):
    record_event_mock.return_value = None
    record_exception_mock.return_value = None

    manager = ConversationManager()
    state = manager.get_or_create("conv_001", None)
    manager.mark_turn_issued("conv_001", 1, "Hello")
    SampleEventHandler.set_conversation_manager(manager)

    result = QueryResult(
        id="123",
        metadata={"conversation_id": "conv_001", "turn_number": 1},
        error=ErrorData(error_type="TimeoutError", error_message="request failed"),
    )

    SampleEventHandler.query_result_complete(result)

    # After failure, the turn should be marked as complete to unblock sequential mode
    assert state.current_turn == 2
    assert state.pending_user_turn is None
    assert len(state.message_history) == 2
    assert "[ERROR:" in state.message_history[1]["content"]
    record_exception_mock.assert_called_once()

[Copilot]

Pull request overview

Copilot reviewed 30 out of 30 changed files in this pull request and generated 1 comment.

Comments suppressed due to low confidence (1)

src/inference_endpoint/load_generator/sample.py:207

• SampleEventHandler now stores a persistent conversation_manager, but clear_hooks() only clears hook lists and leaves conversation_manager set. This can leak conversation state across benchmark sessions/tests in the same process and cause query_result_complete() to update the wrong manager. Consider clearing conversation_manager when clear_hooks(ev_type=None) is called (or ensure session teardown explicitly calls clear_conversation_manager()).

        self.conversation_manager = None

    def set_conversation_manager(self, manager):
        """Enable multi-turn mode by setting conversation manager.

        Args:
            manager: ConversationManager instance for tracking conversation state.
        """
        self.conversation_manager = manager

    def clear_conversation_manager(self) -> None:
        """Disable multi-turn mode by clearing the conversation manager.

        This should be called during session teardown to avoid leaking
        conversation state between benchmark runs.
        """
        self.conversation_manager = None

    def register_hook(
        self,
        event_type: SampleEvent,
        hook: Callable[[StreamChunk], None] | Callable[[QueryResult], None],
    ) -> None:
        if event_type == SampleEvent.FIRST_CHUNK:
            self.first_chunk_hooks.append(hook)
        elif event_type == SampleEvent.NON_FIRST_CHUNK:
            self.non_first_chunk_hooks.append(hook)
        elif event_type == SampleEvent.COMPLETE:
            self.complete_hooks.append(hook)
        else:
            raise ValueError(f"Invalid event type: {event_type}")

    def clear_hooks(self, ev_type: SampleEvent | None = None) -> None:
        if ev_type is None:
            self.first_chunk_hooks.clear()
            self.non_first_chunk_hooks.clear()
            self.complete_hooks.clear()
        elif ev_type == SampleEvent.FIRST_CHUNK:
            self.first_chunk_hooks.clear()
        elif ev_type == SampleEvent.NON_FIRST_CHUNK:
            self.non_first_chunk_hooks.clear()
        elif ev_type == SampleEvent.COMPLETE:
            self.complete_hooks.clear()

---

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

[Copilot]

wait_for_conversation_complete() uses time.time() for timeout tracking, while wait_for_turn_ready() uses time.monotonic(). Using wall-clock time can lead to incorrect timeouts if the system clock changes (NTP adjustments, etc.). Consider switching this method to time.monotonic() for elapsed/remaining timeout calculations to match the rest of the file.

Suggested change

	start_time = time.time()
	while not state.is_complete():
	remaining_timeout = None
	if timeout is not None:
	elapsed = time.time() - start_time
	start_time = time.monotonic()
	while not state.is_complete():
	remaining_timeout = None
	if timeout is not None:
	elapsed = time.monotonic() - start_time

[viraatc]

should we move this into QueryResult class?
(perhaps a new mixin for msgspec structs, similar to with-updates for pydantic models)

[viraatc]

q: is dict() necessary here?

[viraatc]

perhaps update handle_error signature to accept req directly instead of req.query_id and req.query_metadata separately?

[Copilot]

Pull request overview

Copilot reviewed 31 out of 31 changed files in this pull request and generated 3 comments.

---

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

[Copilot]

ConversationSample.slots repeats "uuid" and "data" which are already defined in the base Sample.slots. In Python this causes a slot name conflict at class creation time (import-time TypeError). Define slots with only the new attributes (e.g., conversation_id/turn_number/conversation_state) to avoid the conflict.

Suggested change

	__slots__ = ["uuid", "data", "conversation_id", "turn_number", "conversation_state"]
	__slots__ = ["conversation_id", "turn_number", "conversation_state"]

[Copilot]

wait_for_conversation_complete() reads self._conversations without holding the manager lock/condition, despite the class claiming thread-safe access. This can race with get_or_create()/updates and lead to KeyError or stale state. Consider wrapping the map access in with self._condition: (or with self._lock:) and only waiting on the event outside the lock if needed.

[Copilot]

This unit test will block for the full multi_turn_config_parallel.turn_timeout_s (5s) when the first conversation never completes, which can significantly slow the unit suite. Consider using a much smaller timeout for this test (e.g., 0.1–0.3s) and assert an upper bound as well to keep runtime predictable.

[tianmu-li]

Added option to use dataset history instead of using model output as history. This is needed (but not sufficient) for workato dataset enabling. Still keeping the option to accept model output to allow extending to agentic benchmark in the future.