Skip to content

Commit 4301bb8

Browse files
authored
Add shuffle to agentic workload (#501)
Shuffles agentic conversation start order using a seeded without-replacement permutation over conversation indices, matching the default every other perf workload already uses (agentic previously ran in dataset encounter order). Intra-conversation turn order is unchanged; ordering is deterministic via dataloader_random_seed (default 42). Note: agentic perf results are no longer bit-comparable to pre-change runs, with no opt-out to restore encounter order. Signed-off-by: Rashid Kaleem <230885705+arekay-nv@users.noreply.github.com>
1 parent 47cc5c8 commit 4301bb8

10 files changed

Lines changed: 311 additions & 17 deletions

File tree

‎docs/config/DESIGN.md‎

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -146,9 +146,9 @@ has not fully delegated runtime construction to rulesets yet.
146146
**Reproducibility via explicit seeds**
147147

148148
`RuntimeSettings` contains two seeded `Random` instances: one for scheduler timing jitter
149-
(`rng_sched`) and one for dataset sample ordering (`rng_sample_index`). These make the runtime
150-
configuration reproducible in principle, but the original seed values are not currently persisted
151-
to the report output.
149+
(`rng_sched`) and one for dataset sample ordering (`rng_sample_index`, including agentic
150+
conversation start order). These make the runtime configuration reproducible in principle, but
151+
the original seed values are not currently persisted to the report output.
152152

153153
## Integration Points
154154

‎docs/load_generator/DESIGN.md‎

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -491,7 +491,10 @@ gates the _timing_ of the next turn, it does not alter its prompt. Concurrency i
491491
**active conversations** (`_target_concurrency`), not in-flight requests.
492492

493493
- `execute(phase_issuer)` seeds up to `_target_concurrency` initial conversations, then awaits
494-
completion of all of them.
494+
completion of all of them. Conversation start order is a `WithoutReplacementSampleOrder` over
495+
conversation indices when `rng_sample_index` is set (`dataloader_random_seed`, same RNG as
496+
non-agentic sample shuffle); intra-conversation turn order is unchanged. Without an RNG,
497+
conversations start in dataset encounter order.
495498
- On each response, `on_sample_complete(result)` synchronously routes it to its conversation and
496499
calls `_issue_next_turn()` → `_issue_turn_now()`, which stores the new `query_id → conversation_id`
497500
mapping and issues the turn with zero event-loop delay (or an optional per-turn delay). When a

‎src/inference_endpoint/commands/benchmark/execute.py‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -754,11 +754,15 @@ def _build_agentic_strategy(
754754
if perf_ds_cfg is not None:
755755
agentic_cfg = perf_ds_cfg.agentic_inference
756756
assert ctx.dataloader.conversation_metadata is not None
757+
rng_sample_index = (
758+
ctx.rt_settings.rng_sample_index if ctx.rt_settings is not None else None
759+
)
757760
return AgenticInferenceStrategy(
758761
conversation_manager=ConversationManager(),
759762
dataset_metadata=ctx.dataloader.conversation_metadata,
760763
agentic_inference_config=agentic_cfg,
761764
target_concurrency=ctx.config.settings.load_pattern.target_concurrency,
765+
rng_sample_index=rng_sample_index,
762766
)
763767

764768

‎src/inference_endpoint/config/runtime_settings.py‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -110,7 +110,7 @@ class RuntimeSettings:
110110
"""Random number generator for scheduler"""
111111

112112
rng_sample_index: random.Random
113-
"""Random number generator for sample indexing"""
113+
"""Random number generator for sample indexing / agentic conversation shuffle"""
114114

115115
load_pattern: LoadPattern | None
116116
"""Load pattern configuration"""

‎src/inference_endpoint/config/schema.py‎

Lines changed: 6 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -635,7 +635,12 @@ def _parse_duration_suffix(cls, v: object) -> object:
635635
cyclopts.Parameter(alias="--num-samples", help="Sample count override"),
636636
] = Field(None, gt=0)
637637
scheduler_random_seed: int = Field(42, description="Scheduler RNG seed")
638-
dataloader_random_seed: int = Field(42, description="Dataloader RNG seed")
638+
dataloader_random_seed: int = Field(
639+
42,
640+
description=(
641+
"Dataloader RNG seed (sample shuffle; agentic conversation order)"
642+
),
643+
)
639644

640645
@model_validator(mode="after")
641646
def _validate_durations(self) -> Self:

‎src/inference_endpoint/config/templates/concurrency_template_full.yaml‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -55,7 +55,7 @@ settings:
5555
max_issue_duration_ms: null # Cap on performance-phase issuing in ms; reaching it ends the phase normally (None = no cap)
5656
n_samples_to_issue: null # Sample count override
5757
scheduler_random_seed: 42 # Scheduler RNG seed
58-
dataloader_random_seed: 42 # Dataloader RNG seed
58+
dataloader_random_seed: 42 # Dataloader RNG seed (sample shuffle; agentic conversation order)
5959
load_pattern:
6060
type: concurrency # Load pattern type | options: max_throughput, poisson, concurrency, agentic_inference, burst, step
6161
target_qps: null # Target QPS

‎src/inference_endpoint/config/templates/offline_template_full.yaml‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -55,7 +55,7 @@ settings:
5555
max_issue_duration_ms: null # Cap on performance-phase issuing in ms; reaching it ends the phase normally (None = no cap)
5656
n_samples_to_issue: null # Sample count override
5757
scheduler_random_seed: 42 # Scheduler RNG seed
58-
dataloader_random_seed: 42 # Dataloader RNG seed
58+
dataloader_random_seed: 42 # Dataloader RNG seed (sample shuffle; agentic conversation order)
5959
load_pattern:
6060
type: max_throughput # Load pattern type | offline only: max_throughput
6161
target_qps: null # Target QPS

‎src/inference_endpoint/config/templates/online_template_full.yaml‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -55,7 +55,7 @@ settings:
5555
max_issue_duration_ms: null # Cap on performance-phase issuing in ms; reaching it ends the phase normally (None = no cap)
5656
n_samples_to_issue: null # Sample count override
5757
scheduler_random_seed: 42 # Scheduler RNG seed
58-
dataloader_random_seed: 42 # Dataloader RNG seed
58+
dataloader_random_seed: 42 # Dataloader RNG seed (sample shuffle; agentic conversation order)
5959
load_pattern:
6060
type: poisson # Load pattern type | options: max_throughput, poisson, concurrency, agentic_inference, burst, step
6161
target_qps: 10.0 # Target QPS

‎src/inference_endpoint/load_generator/agentic_inference_strategy.py‎

Lines changed: 50 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -18,6 +18,7 @@
1818
import asyncio
1919
import hashlib
2020
import logging
21+
import random
2122
import time
2223
from collections import defaultdict
2324
from typing import Any
@@ -28,6 +29,11 @@
2829
from ..dataset_manager.agentic_inference_dataset import ConversationMetadata
2930
from ..exceptions import InputValidationError
3031
from .conversation_manager import ConversationManager
32+
from .sample_order import (
33+
SampleOrder,
34+
SequentialSampleOrder,
35+
WithoutReplacementSampleOrder,
36+
)
3137
from .strategy import PhaseIssuerProtocol
3238

3339
logger = logging.getLogger(__name__)
@@ -45,10 +51,14 @@ class AgenticInferenceStrategy:
4551
"""Event-driven agentic inference strategy. Completion of each turn triggers the next.
4652
4753
execute() seeds the first N conversations (issues turn 1 for each), then
48-
awaits _all_done. on_sample_complete() is called synchronously from the
49-
receive coroutine for each response — it issues the next turn immediately
50-
(zero event-loop iterations between response and next issuance), or starts
51-
a new conversation when the current one finishes all turns.
54+
awaits _all_done. Conversation start order uses the same without-replacement
55+
shuffle as non-agentic datasets when rng_sample_index is provided
56+
(RuntimeSettings.rng_sample_index / dataloader_random_seed); turns inside a
57+
conversation stay in dataset order. on_sample_complete() is called
58+
synchronously from the receive coroutine for each response — it issues the
59+
next turn immediately (zero event-loop iterations between response and next
60+
issuance), or starts a new conversation when the current one finishes all
61+
turns.
5262
5363
At most target_concurrency conversations are active simultaneously. When
5464
target_concurrency is None, all conversations start at once.
@@ -76,6 +86,7 @@ def __init__(
7686
dataset_metadata: ConversationMetadata,
7787
agentic_inference_config: AgenticInferenceConfig | None = None,
7888
target_concurrency: int | None = None,
89+
rng_sample_index: random.Random | None = None,
7990
):
8091
"""Initialize agentic inference strategy.
8192
@@ -85,6 +96,9 @@ def __init__(
8596
agentic_inference_config: Agentic inference conversation configuration.
8697
target_concurrency: Maximum number of simultaneously active conversations.
8798
None means all conversations run concurrently.
99+
rng_sample_index: Dataloader RNG used to shuffle conversation start
100+
order (same seed as non-agentic WithoutReplacementSampleOrder).
101+
None keeps dataset encounter order.
88102
"""
89103
self._conv_manager = conversation_manager
90104
self._dataset_metadata = dataset_metadata
@@ -105,6 +119,7 @@ def __init__(
105119
else _DEFAULT_TURN_TIMEOUT_S
106120
)
107121
self._target_concurrency = target_concurrency
122+
self._rng_sample_index = rng_sample_index
108123
self._enable_salt = (
109124
agentic_inference_config.enable_salt
110125
if agentic_inference_config is not None
@@ -121,6 +136,7 @@ def __init__(
121136

122137
# Event-driven state — populated in execute().
123138
self._base_convs: list[tuple[str, ConversationTurns]] = []
139+
self._conversation_order: SampleOrder | None = None
124140
self._active_iters: dict[str, ActiveConversationState] = {}
125141
self._timeout_handles: dict[str, asyncio.TimerHandle] = {}
126142
self._delay_handles: dict[str, asyncio.TimerHandle] = {}
@@ -161,6 +177,7 @@ async def execute(self, phase_issuer: PhaseIssuerProtocol) -> int:
161177
(conv_id, sorted(turns, key=lambda x: x[1]))
162178
for conv_id, turns in conv_samples.items()
163179
]
180+
self._conversation_order = self._make_conversation_order()
164181
self._validate_salt_system_prompts()
165182
n_to_start = self._initial_conversations_to_start()
166183
try:
@@ -210,12 +227,39 @@ def _has_trajectory_budget(self) -> bool:
210227
def _has_more_conversation_instances(self) -> bool:
211228
return bool(self._base_convs and self._has_trajectory_budget())
212229

230+
def _make_conversation_order(self) -> SampleOrder | None:
231+
"""Conversation-index iterator; shuffle when a dataloader RNG is set.
232+
233+
Only permutation orders are valid here. `_next_conversation_instance`
234+
derives `repeat_id` as `started_count // n_convs + 1`, which assumes the
235+
order yields each conversation index exactly once per block of `n_convs`
236+
draws. `WithoutReplacementSampleOrder` and `SequentialSampleOrder` both
237+
satisfy this. Do NOT swap in `WithReplacementSampleOrder` (or any order
238+
that can repeat an index within a block) without first reworking
239+
`repeat_id`: a conversation recurring within a block collides on
240+
`(source_id, repeat_id)`, which duplicates the anti-cache salt (server
241+
prefix-cache hits corrupt throughput) and the `logical_id` key in
242+
`_active_iters` (conversation tracking overwrites).
243+
"""
244+
n_convs = len(self._base_convs)
245+
if n_convs == 0:
246+
return None
247+
if self._rng_sample_index is not None:
248+
return WithoutReplacementSampleOrder(
249+
n_samples_in_dataset=n_convs,
250+
rng=self._rng_sample_index,
251+
)
252+
return SequentialSampleOrder(n_samples_in_dataset=n_convs)
253+
213254
def _next_conversation_instance(self) -> ConversationInstance | None:
214255
if not self._has_more_conversation_instances():
215256
return None
216257

217-
source_index = self._started_trajectory_count % len(self._base_convs)
258+
assert self._conversation_order is not None
259+
source_index = next(self._conversation_order)
218260
source_id, turns = self._base_convs[source_index]
261+
# Unique (source_id, repeat_id) relies on _conversation_order being a
262+
# permutation every n_convs draws (see _make_conversation_order).
219263
instance_id = self._started_trajectory_count // len(self._base_convs) + 1
220264
logical_id = (
221265
source_id if instance_id == 1 else f"{source_id}__repeat_{instance_id}"
@@ -396,7 +440,7 @@ def _messages_with_trajectory_salt(
396440
conversation_id.encode("utf-8"), digest_size=2
397441
).hexdigest()
398442
message["content"] = (
399-
f"[salt: {repeat_salt}]\n\n" f"{content}\n\n" f"[salt: {conv_salt}]"
443+
f"[salt: {repeat_salt}]\n\n{content}\n\n[salt: {conv_salt}]"
400444
)
401445
return salted_messages
402446
raise InputValidationError(

0 commit comments

Comments
 (0)