patch the dataset sample progress bar

src/inference_endpoint/commands/benchmark/execute.py

@@ -98,6 +98,7 @@
 )
 from inference_endpoint.load_generator.agentic_inference_strategy import (
     AgenticInferenceStrategy,
+    expected_agentic_sample_count,
 )
 from inference_endpoint.load_generator.conversation_manager import ConversationManager
 from inference_endpoint.load_generator.session import (
@@ -546,6 +547,26 @@ def setup_benchmark(
                 sample_order=audit_run_spec.sample_order,
             )
         total_samples = rt_settings.total_samples_to_issue()
+        # WAR (TODO: clean up the logic of agentic issue mode and single issue mode)
+        # calculate the correct number of samples to issue for agentic inference dataset
+        if isinstance(dataloader, AgenticInferenceDataset):
+            perf_ds_cfg = next(
+                (d for d in config.datasets if d.type == DatasetType.PERFORMANCE),
+                None,
+            )
+            agentic_cfg = (
+                perf_ds_cfg.agentic_inference if perf_ds_cfg is not None else None
+            )
+            assert dataloader.conversation_metadata is not None
+            total_samples = expected_agentic_sample_count(
+                dataloader.conversation_metadata,
+                (
+                    agentic_cfg.num_trajectories_to_issue
+                    if agentic_cfg is not None
+                    else None
+                ),
+                rt_settings.rng_sample_index,
+            )
 
     total_samples += sum(
         ec.dataset.num_samples() * ec.dataset.repeats

src/inference_endpoint/load_generator/agentic_inference_strategy.py

@@ -47,6 +47,39 @@
 ConversationInstance = tuple[str, str, ConversationTurns, int]
 
 
+def expected_agentic_sample_count(
+    dataset_metadata: ConversationMetadata,
+    num_trajectories_to_issue: int | None,
+    rng_sample_index: random.Random | None = None,
+) -> int:
+    """Return the client-turn count for the requested trajectory instances."""
+    turn_counts = list(dataset_metadata.client_turns_per_conversation.values())
+    if not turn_counts:
+        return 0
+
+    trajectory_count = (
+        num_trajectories_to_issue
+        if num_trajectories_to_issue is not None
+        else len(turn_counts)
+    )
+    if rng_sample_index is None:
+        order: SampleOrder = SequentialSampleOrder(
+            n_samples_in_dataset=len(turn_counts)
+        )
+    else:
+        # Mirror the scheduler order so a partial trajectory budget counts the
+        # same conversations without advancing the RNG used for actual
+        # issuance.
+        rng_copy = random.Random()
+        rng_copy.setstate(rng_sample_index.getstate())
+        order = WithoutReplacementSampleOrder(
+            n_samples_in_dataset=len(turn_counts),
+            rng=rng_copy,
+        )
+
+    return sum(turn_counts[next(order)] for _ in range(trajectory_count))
+
+
 class AgenticInferenceStrategy:
     """Event-driven agentic inference strategy. Completion of each turn triggers the next.
 

tests/unit/commands/test_benchmark.py

@@ -85,6 +85,9 @@
 )
 from inference_endpoint.config.utils import cli_error_formatter as _error_formatter
 from inference_endpoint.core.types import APIType, QueryResult
+from inference_endpoint.dataset_manager.agentic_inference_dataset import (
+    AgenticInferenceDataset,
+)
 from inference_endpoint.dataset_manager.dataset import Dataset
 from inference_endpoint.dataset_manager.predefined.swe_bench import SWEBench
 from inference_endpoint.endpoint_client.config import HTTPClientConfig
@@ -3161,6 +3164,117 @@ def test_accuracy_dataset_named_performance_is_counted_by_type(
 
         assert ctx.total_samples == 5
 
+    @pytest.mark.unit
+    def test_agentic_trajectory_repeats_are_included_in_total_samples(
+        self, tmp_path, _base_patches
+    ):
+        dataset = AgenticInferenceDataset(
+            pd.DataFrame(
+                [
+                    {
+                        "conversation_id": "c1",
+                        "turn": 1,
+                        "role": "user",
+                        "content": "c1 first user turn",
+                    },
+                    {
+                        "conversation_id": "c1",
+                        "turn": 2,
+                        "role": "assistant",
+                        "content": "c1 assistant turn",
+                    },
+                    {
+                        "conversation_id": "c1",
+                        "turn": 3,
+                        "role": "user",
+                        "content": "c1 second user turn",
+                    },
+                    {
+                        "conversation_id": "c2",
+                        "turn": 1,
+                        "role": "user",
+                        "content": "c2 only user turn",
+                    },
+                    {
+                        "conversation_id": "c3",
+                        "turn": 1,
+                        "role": "user",
+                        "content": "c3 first user turn",
+                    },
+                    {
+                        "conversation_id": "c3",
+                        "turn": 2,
+                        "role": "assistant",
+                        "content": "c3 first assistant turn",
+                    },
+                    {
+                        "conversation_id": "c3",
+                        "turn": 3,
+                        "role": "user",
+                        "content": "c3 second user turn",
+                    },
+                    {
+                        "conversation_id": "c3",
+                        "turn": 4,
+                        "role": "assistant",
+                        "content": "c3 second assistant turn",
+                    },
+                    {
+                        "conversation_id": "c3",
+                        "turn": 5,
+                        "role": "user",
+                        "content": "c3 third user turn",
+                    },
+                ]
+            )
+        )
+        # fake dataset, 3 trajectory:
+        # traj1: user1 -> assistant1 -> user2 ==== totally 2 issued samples
+        # traj2: user1 ==== totally 1 issued sample
+        # traj3: user1 -> assistant1 -> user2 -> assistant2 -> user3 ==== totally 3 issued samples
+        dataset.load()
+        load_pattern = LoadPattern(
+            type=LoadPatternType.AGENTIC_INFERENCE, target_concurrency=2
+        )
+        config = OnlineConfig(
+            endpoint_config={"endpoints": ["http://x"]},
+            model_params={"name": "test-model"},
+            datasets=[
+                {
+                    "name": "agentic",
+                    "type": "performance",
+                    "path": "unused.jsonl",
+                    "agentic_inference": {"num_trajectories_to_issue": 5},
+                }
+            ],
+            settings=OnlineSettings(load_pattern=load_pattern),
+            report_dir=str(tmp_path),
+        )
+        rt_settings = RuntimeSettings(
+            metric_target=Throughput(10.0),
+            reported_metrics=[Throughput(10.0)],
+            min_issue_duration_ms=0,
+            max_issue_duration_ms=None,
+            n_samples_from_dataset=6,
+            n_samples_to_issue=None,
+            min_sample_count=1,
+            rng_sched=random.Random(0),
+            # Seed 5 produces c1, c2, c3 for each sequential dataset pass.
+            rng_sample_index=random.Random(5),
+            load_pattern=load_pattern,
+        )
+
+        ctx = self._setup(
+            config,
+            TestMode.PERF,
+            (dataset, []),
+            rt_settings,
+        )
+
+        # One complete pass contributes 2 + 1 + 3 = 6 client-turn samples.
+        # Issuing five trajectories repeats c1 and c2, adding 2 + 1 more.
+        assert ctx.total_samples == 9
+
 
 class TestReportConfigSecretRedaction:
     @pytest.mark.unit

tests/unit/load_generator/test_agentic_inference_strategy.py

@@ -31,6 +31,7 @@
 from inference_endpoint.exceptions import InputValidationError
 from inference_endpoint.load_generator.agentic_inference_strategy import (
     AgenticInferenceStrategy,
+    expected_agentic_sample_count,
 )
 from inference_endpoint.load_generator.conversation_manager import ConversationManager
 
@@ -172,6 +173,36 @@ def _trajectory_start_order(issuer: RecordingPhaseIssuer) -> list[str]:
     return starts
 
 
+@pytest.mark.unit
+def test_expected_sample_count_includes_trajectory_repeats():
+    metadata = _make_dataset_metadata({"conv1": [1, 2], "conv2": [1, 2, 3]})
+
+    # Different trajectory issue counts result in different numbers of samples:
+    # the two trajectories contribute two and three client-turn samples per pass.
+    assert expected_agentic_sample_count(_make_dataset_metadata({}), None) == 0
+    assert expected_agentic_sample_count(metadata, None) == 5
+    assert expected_agentic_sample_count(metadata, 1) == 2
+    assert expected_agentic_sample_count(metadata, 2) == 5
+    assert expected_agentic_sample_count(metadata, 3) == 7
+    assert expected_agentic_sample_count(metadata, 4) == 10
+    assert expected_agentic_sample_count(metadata, 5) == 12
+    assert expected_agentic_sample_count(metadata, 6) == 15
+
+
+@pytest.mark.unit
+def test_expected_sample_count_does_not_advance_scheduler_rng():
+    metadata = _make_dataset_metadata(
+        {"conv1": [1], "conv2": [1, 2], "conv3": [1, 2, 3]}
+    )
+    rng = random.Random(42)
+    original_state = rng.getstate()
+
+    expected = expected_agentic_sample_count(metadata, 2, rng)
+
+    assert expected == 3
+    assert rng.getstate() == original_state
+
+
 async def _drain_recording_strategy(
     strategy: AgenticInferenceStrategy,
     issuer: RecordingPhaseIssuer,