1818import asyncio
1919import hashlib
2020import logging
21+ import random
2122import time
2223from collections import defaultdict
2324from typing import Any
2829from ..dataset_manager .agentic_inference_dataset import ConversationMetadata
2930from ..exceptions import InputValidationError
3031from .conversation_manager import ConversationManager
32+ from .sample_order import (
33+ SampleOrder ,
34+ SequentialSampleOrder ,
35+ WithoutReplacementSampleOrder ,
36+ )
3137from .strategy import PhaseIssuerProtocol
3238
3339logger = 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