diff --git a/src/inference_endpoint/commands/benchmark/execute.py b/src/inference_endpoint/commands/benchmark/execute.py index af7e9d6c..b8a8d5a9 100644 --- a/src/inference_endpoint/commands/benchmark/execute.py +++ b/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 diff --git a/src/inference_endpoint/load_generator/agentic_inference_strategy.py b/src/inference_endpoint/load_generator/agentic_inference_strategy.py index fb7b7edb..3f4e6262 100644 --- a/src/inference_endpoint/load_generator/agentic_inference_strategy.py +++ b/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. diff --git a/tests/unit/commands/test_benchmark.py b/tests/unit/commands/test_benchmark.py index 6dac045b..0045de54 100644 --- a/tests/unit/commands/test_benchmark.py +++ b/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 diff --git a/tests/unit/load_generator/test_agentic_inference_strategy.py b/tests/unit/load_generator/test_agentic_inference_strategy.py index fabc00a0..3962c559 100644 --- a/tests/unit/load_generator/test_agentic_inference_strategy.py +++ b/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,