From 3d5ec423b90e0691f3924c3ef2715db74e3d25da Mon Sep 17 00:00:00 2001 From: Rabah B Date: Tue, 15 Sep 2026 23:36:17 -0400 Subject: [PATCH 1/2] Add configurable AOAI output page size Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> --- ...uation-aoai-output-page-size-2026-09-15.md | 8 + sdk/evaluation/azure-ai-evaluation/api.md | 1 + .../ai/evaluation/_evaluate/_evaluate.py | 33 +++- .../ai/evaluation/_evaluate/_evaluate_aoai.py | 21 +- .../test_aoai_evaluation_pagination.py | 183 +++++++++++++++++- 5 files changed, 240 insertions(+), 6 deletions(-) create mode 100644 .chronus/changes/rboubchir-azure-ai-evaluation-aoai-output-page-size-2026-09-15.md diff --git a/.chronus/changes/rboubchir-azure-ai-evaluation-aoai-output-page-size-2026-09-15.md b/.chronus/changes/rboubchir-azure-ai-evaluation-aoai-output-page-size-2026-09-15.md new file mode 100644 index 000000000000..f4c21ca89cb7 --- /dev/null +++ b/.chronus/changes/rboubchir-azure-ai-evaluation-aoai-output-page-size-2026-09-15.md @@ -0,0 +1,8 @@ +--- +changeKind: feature +packages: + - azure-ai-evaluation +--- + +Added the `aoai_output_items_page_size` option to `evaluate` for configuring the number of native Azure OpenAI +grader output items requested per HTTP response page. The default remains 100, and all result pages are fetched. diff --git a/sdk/evaluation/azure-ai-evaluation/api.md b/sdk/evaluation/azure-ai-evaluation/api.md index 88566d82dbc0..0fc4d1bca4a5 100644 --- a/sdk/evaluation/azure-ai-evaluation/api.md +++ b/sdk/evaluation/azure-ai-evaluation/api.md @@ -3,6 +3,7 @@ namespace azure.ai.evaluation def azure.ai.evaluation.evaluate( *, + aoai_output_items_page_size: int = 100, azure_ai_project: Optional[Union[str, AzureAIProject]] = ..., data: Union[str, PathLike], evaluation_name: Optional[str] = ..., diff --git a/sdk/evaluation/azure-ai-evaluation/azure/ai/evaluation/_evaluate/_evaluate.py b/sdk/evaluation/azure-ai-evaluation/azure/ai/evaluation/_evaluate/_evaluate.py index 76cbb33098a6..9b070983e865 100644 --- a/sdk/evaluation/azure-ai-evaluation/azure/ai/evaluation/_evaluate/_evaluate.py +++ b/sdk/evaluation/azure-ai-evaluation/azure/ai/evaluation/_evaluate/_evaluate.py @@ -59,6 +59,8 @@ from ._batch_run.batch_clients import BatchClient, BatchClientRun from ._evaluate_aoai import ( + _DEFAULT_AOAI_OUTPUT_ITEMS_PAGE_SIZE, + _MAX_AOAI_OUTPUT_ITEMS_PAGE_SIZE, _begin_aoai_evaluation, _split_evaluators_and_grader_configs, _get_evaluation_run_results, @@ -818,6 +820,25 @@ def _rename_columns_conditionally(df: pd.DataFrame) -> pd.DataFrame: return df +def _validate_aoai_output_items_page_size(aoai_output_items_page_size: int) -> None: + if ( + not isinstance(aoai_output_items_page_size, int) + or isinstance(aoai_output_items_page_size, bool) + or not 1 <= aoai_output_items_page_size <= _MAX_AOAI_OUTPUT_ITEMS_PAGE_SIZE + ): + msg = ( + "'aoai_output_items_page_size' must be an integer between 1 and " + f"{_MAX_AOAI_OUTPUT_ITEMS_PAGE_SIZE}, inclusive." + ) + raise EvaluationException( + message=msg, + internal_message=msg, + target=ErrorTarget.EVALUATE, + category=ErrorCategory.INVALID_VALUE, + blame=ErrorBlame.USER_ERROR, + ) + + def evaluate( *, data: Union[str, os.PathLike], @@ -828,6 +849,7 @@ def evaluate( azure_ai_project: Optional[Union[str, AzureAIProject]] = None, output_path: Optional[Union[str, os.PathLike]] = None, fail_on_evaluator_errors: bool = False, + aoai_output_items_page_size: int = 100, tags: Optional[Dict[str, str]] = None, **kwargs, ) -> EvaluationResult: @@ -861,6 +883,10 @@ def evaluate( Defaults to false, which means that evaluations will continue regardless of failures. If such failures occur, metrics may be missing, and evidence of failures can be found in the evaluation's logs. :paramtype fail_on_evaluator_errors: bool + :keyword aoai_output_items_page_size: The maximum number of native Azure OpenAI grader output items requested + per HTTP response page. Defaults to 100. This does not limit response bytes, request latency, or the number + of dataset rows evaluated; all output-item result pages are fetched. + :paramtype aoai_output_items_page_size: int :keyword tags: A dictionary of tags to be added to the evaluation run for tracking and organization purposes. Keys and values must be strings. For more information about tag limits, see: https://learn.microsoft.com/en-us/azure/machine-learning/resource-limits-capacity?view=azureml-api-2#runs @@ -890,6 +916,7 @@ def evaluate( https://{resource_name}.services.ai.azure.com/api/projects/{project_name} """ try: + _validate_aoai_output_items_page_size(aoai_output_items_page_size) user_agent: Optional[str] = kwargs.get("user_agent") with UserAgentSingleton().add_useragent_product(user_agent) if user_agent else contextlib.nullcontext(): results = _evaluate( @@ -901,6 +928,7 @@ def evaluate( azure_ai_project=azure_ai_project, output_path=output_path, fail_on_evaluator_errors=fail_on_evaluator_errors, + aoai_output_items_page_size=aoai_output_items_page_size, tags=tags, **kwargs, ) @@ -971,6 +999,7 @@ def _evaluate( # pylint: disable=too-many-locals,too-many-statements azure_ai_project: Optional[Union[str, AzureAIProject]] = None, output_path: Optional[Union[str, os.PathLike]] = None, fail_on_evaluator_errors: bool = False, + aoai_output_items_page_size: int = _DEFAULT_AOAI_OUTPUT_ITEMS_PAGE_SIZE, tags: Optional[Dict[str, str]] = None, **kwargs, ) -> EvaluationResult: @@ -1048,7 +1077,9 @@ def _evaluate( # pylint: disable=too-many-locals,too-many-statements # Retrieve OAI eval run results if needed. if need_get_oai_results: try: - aoai_results, aoai_metrics = _get_evaluation_run_results(eval_run_info_list) # type: ignore + aoai_results, aoai_metrics = _get_evaluation_run_results( + eval_run_info_list, aoai_output_items_page_size + ) # type: ignore # Post build TODO: add equivalent of _print_summary(per_evaluator_results) here # Combine results if both evaluators and graders are present diff --git a/sdk/evaluation/azure-ai-evaluation/azure/ai/evaluation/_evaluate/_evaluate_aoai.py b/sdk/evaluation/azure-ai-evaluation/azure/ai/evaluation/_evaluate/_evaluate_aoai.py index 6d44320f9f9a..de60f05bd334 100644 --- a/sdk/evaluation/azure-ai-evaluation/azure/ai/evaluation/_evaluate/_evaluate_aoai.py +++ b/sdk/evaluation/azure-ai-evaluation/azure/ai/evaluation/_evaluate/_evaluate_aoai.py @@ -22,6 +22,8 @@ TClient = TypeVar("TClient", ProxyClient, CodeClient) LOGGER = logging.getLogger(__name__) +_DEFAULT_AOAI_OUTPUT_ITEMS_PAGE_SIZE = 100 +_MAX_AOAI_OUTPUT_ITEMS_PAGE_SIZE = 100 # Precompiled regex for extracting data paths from mapping expressions of the form # ${data.some.dotted.path}. Compiled once at import time to avoid repeated @@ -258,7 +260,10 @@ def _combine_item_schemas(data_source_config: Dict[str, Any], kwargs: Dict[str, data_source_config["item_schema"]["required"].append(key) -def _get_evaluation_run_results(all_run_info: List[OAIEvalRunCreationInfo]) -> Tuple[pd.DataFrame, Dict[str, Any]]: +def _get_evaluation_run_results( + all_run_info: List[OAIEvalRunCreationInfo], + aoai_output_items_page_size: int = _DEFAULT_AOAI_OUTPUT_ITEMS_PAGE_SIZE, +) -> Tuple[pd.DataFrame, Dict[str, Any]]: """ Get the results of an OAI evaluation run, formatted in a way that is easy for the rest of the evaluation pipeline to consume. This method accepts a list of eval run information, and will combine the @@ -267,6 +272,8 @@ def _get_evaluation_run_results(all_run_info: List[OAIEvalRunCreationInfo]) -> T :param all_run_info: A list of evaluation run information that contains the needed values to retrieve the results of the evaluation run. :type all_run_info: List[OAIEvalRunCreationInfo] + :param aoai_output_items_page_size: The maximum number of output items to request per page. + :type aoai_output_items_page_size: int :return: A tuple containing the results of the evaluation run as a dataframe, and a dictionary of metrics calculated from the evaluation run. :rtype: Tuple[pd.DataFrame, Dict[str, Any]] @@ -278,7 +285,7 @@ def _get_evaluation_run_results(all_run_info: List[OAIEvalRunCreationInfo]) -> T output_df = pd.DataFrame() for idx, run_info in enumerate(all_run_info): LOGGER.info(f"AOAI: Fetching results for run {idx + 1}/{len(all_run_info)} (ID: {run_info['eval_run_id']})...") - cur_output_df, cur_run_metrics = _get_single_run_results(run_info) + cur_output_df, cur_run_metrics = _get_single_run_results(run_info, aoai_output_items_page_size) output_df = pd.concat([output_df, cur_output_df], axis=1) run_metrics.update(cur_run_metrics) @@ -288,6 +295,7 @@ def _get_evaluation_run_results(all_run_info: List[OAIEvalRunCreationInfo]) -> T def _get_single_run_results( run_info: OAIEvalRunCreationInfo, + aoai_output_items_page_size: int = _DEFAULT_AOAI_OUTPUT_ITEMS_PAGE_SIZE, ) -> Tuple[pd.DataFrame, Dict[str, Any]]: """ Get the results of an OAI evaluation run, formatted in a way that is easy for the rest of the evaluation @@ -296,6 +304,8 @@ def _get_single_run_results( :param run_info: The evaluation run information that contains the needed values to retrieve the results of the evaluation run. :type run_info: OAIEvalRunCreationInfo + :param aoai_output_items_page_size: The maximum number of output items to request per page. + :type aoai_output_items_page_size: int :return: A tuple containing the results of the evaluation run as a dataframe, and a dictionary of metrics calculated from the evaluation run. :rtype: Tuple[pd.DataFrame, Dict[str, Any]] @@ -349,10 +359,13 @@ def _get_single_run_results( LOGGER.info(f"AOAI: Collecting output items for run {run_info['eval_run_id']} with pagination...") all_results: List[Any] = [] next_cursor: Optional[str] = None - limit = 100 # Max allowed by API while True: - list_kwargs = {"eval_id": run_info["eval_group_id"], "run_id": run_info["eval_run_id"], "limit": limit} + list_kwargs = { + "eval_id": run_info["eval_group_id"], + "run_id": run_info["eval_run_id"], + "limit": aoai_output_items_page_size, + } if next_cursor is not None: list_kwargs["after"] = next_cursor diff --git a/sdk/evaluation/azure-ai-evaluation/tests/unittests/test_aoai_evaluation_pagination.py b/sdk/evaluation/azure-ai-evaluation/tests/unittests/test_aoai_evaluation_pagination.py index ed4f74173dfa..55d6cce1ab4f 100644 --- a/sdk/evaluation/azure-ai-evaluation/tests/unittests/test_aoai_evaluation_pagination.py +++ b/sdk/evaluation/azure-ai-evaluation/tests/unittests/test_aoai_evaluation_pagination.py @@ -12,7 +12,8 @@ OAIEvalRunCreationInfo, _wait_for_run_conclusion, ) -from azure.ai.evaluation._exceptions import EvaluationException +from azure.ai.evaluation import evaluate +from azure.ai.evaluation._exceptions import ErrorBlame, ErrorCategory, ErrorTarget, EvaluationException class MockOutputItem: @@ -35,6 +36,107 @@ def __init__(self, data: List[MockOutputItem], has_more: bool = False): class TestAOAIPagination: """Test pagination functionality in AOAI evaluation results""" + @pytest.mark.parametrize( + "evaluate_kwargs, expected_page_size", + [ + ({}, 100), + ({"aoai_output_items_page_size": 3}, 3), + ], + ) + def test_evaluate_propagates_page_size_to_every_native_grader_run(self, evaluate_kwargs, expected_page_size): + clients = [Mock(), Mock()] + input_data = pd.DataFrame([{"query": "hi"}]) + run_info = [ + OAIEvalRunCreationInfo( + client=clients[0], + eval_group_id="group-1", + eval_run_id="run-1", + grader_name_map={"grader-1": "first_grader"}, + expected_rows=1, + ), + OAIEvalRunCreationInfo( + client=clients[1], + eval_group_id="group-2", + eval_run_id="run-2", + grader_name_map={"grader-2": "second_grader"}, + expected_rows=1, + ), + ] + completed_runs = [] + for grader_id in ("grader-1", "grader-2"): + completed_run = Mock() + completed_run.status = "completed" + completed_run.per_testing_criteria_results = [Mock(testing_criteria=grader_id, passed=1, failed=0)] + completed_runs.append(completed_run) + + clients[0].evals.runs.output_items.list.return_value = MockOutputItemsList( + data=[ + MockOutputItem( + id="item-0", + datasource_item_id=0, + results=[{"name": "grader-1", "passed": True, "score": 0.9}], + ) + ] + ) + clients[1].evals.runs.output_items.list.return_value = MockOutputItemsList( + data=[ + MockOutputItem( + id="item-0", + datasource_item_id=0, + results=[{"name": "grader-2", "passed": True, "score": 0.8}], + ) + ] + ) + + with patch( + "azure.ai.evaluation._evaluate._evaluate._preprocess_data", + return_value={ + "column_mapping": {}, + "evaluators": {}, + "graders": {"first_grader": object(), "second_grader": object()}, + "input_data_df": input_data, + "target_run": None, + "batch_run_client": None, + "batch_run_data": None, + }, + ), patch( + "azure.ai.evaluation._evaluate._evaluate._begin_aoai_evaluation", + return_value=run_info, + ), patch( + "azure.ai.evaluation._evaluate._evaluate_aoai._wait_for_run_conclusion", + side_effect=completed_runs, + ), patch( + "azure.ai.evaluation._evaluate._evaluate._map_names_to_builtins", + return_value={}, + ): + result = evaluate(data="unused.jsonl", evaluators={"grader": object()}, **evaluate_kwargs) + + assert len(result["rows"]) == 1 + assert result["metrics"] == {"first_grader.pass_rate": 1.0, "second_grader.pass_rate": 1.0} + for client in clients: + client.evals.runs.output_items.list.assert_called_once() + assert client.evals.runs.output_items.list.call_args.kwargs["limit"] == expected_page_size + + @pytest.mark.parametrize("invalid_page_size", [True, False, "3", 3.0, 3.5, 0, -1, 101]) + def test_evaluate_rejects_invalid_page_size_before_any_evaluation_work(self, invalid_page_size): + with patch("azure.ai.evaluation._evaluate._evaluate._preprocess_data") as mock_preprocess, patch( + "azure.ai.evaluation._evaluate._evaluate.UserAgentSingleton.add_useragent_product" + ) as mock_add_user_agent: + with pytest.raises(EvaluationException) as exc_info: + evaluate( + data="unused.jsonl", + evaluators={}, + aoai_output_items_page_size=invalid_page_size, + user_agent="custom/1", + ) + + mock_preprocess.assert_not_called() + mock_add_user_agent.assert_not_called() + assert exc_info.value.blame == ErrorBlame.USER_ERROR + assert exc_info.value.category == ErrorCategory.INVALID_VALUE + assert exc_info.value.target == ErrorTarget.EVALUATE + assert "'aoai_output_items_page_size' must be an integer between 1 and 100" in str(exc_info.value) + def test_single_page_results(self): """Test handling of single page results (no pagination needed)""" # Mock client and run info @@ -250,3 +352,82 @@ def test_result_ordering_preservation(self): scores = df["outputs.test_grader.score"].tolist() expected_scores = [0, 1, 2, 3, 4, 5, 6, 7, 8, 9] assert scores == expected_scores + + def test_configured_page_size_fetches_all_results_and_final_partial_page(self): + mock_client = Mock() + run_info = OAIEvalRunCreationInfo( + client=mock_client, + eval_group_id="test-group", + eval_run_id="test-run", + grader_name_map={"grader-1": "friendly_grader"}, + expected_rows=160, + ) + mock_run_results = Mock() + mock_run_results.status = "completed" + mock_run_results.per_testing_criteria_results = [Mock(testing_criteria="grader-1", passed=120, failed=40)] + output_items = [ + MockOutputItem( + id=f"item-{position}", + datasource_item_id=position, + results=[{"name": "grader-1", "passed": position < 120, "score": position}], + ) + for position in range(160) + ] + + def list_output_items(**kwargs): + after = kwargs.get("after") + start = int(after.removeprefix("item-")) + 1 if after else 0 + end = min(start + kwargs["limit"], len(output_items)) + return MockOutputItemsList(data=output_items[start:end], has_more=end < len(output_items)) + + mock_client.evals.runs.output_items.list.side_effect = list_output_items + + with patch( + "azure.ai.evaluation._evaluate._evaluate_aoai._wait_for_run_conclusion", + return_value=mock_run_results, + ): + df, metrics = _get_single_run_results(run_info, aoai_output_items_page_size=3) + + assert mock_client.evals.runs.output_items.list.call_count == 54 + assert len(df) == 160 + assert df["outputs.friendly_grader.score"].tolist() == list(range(160)) + assert metrics == {"friendly_grader.pass_rate": 0.75} + calls = mock_client.evals.runs.output_items.list.call_args_list + assert "after" not in calls[0].kwargs + assert calls[1].kwargs["after"] == "item-2" + assert calls[-1].kwargs["after"] == "item-158" + assert all(call_.kwargs["limit"] == 3 for call_ in calls) + + def test_later_page_failure_propagates(self): + mock_client = Mock() + run_info = OAIEvalRunCreationInfo( + client=mock_client, + eval_group_id="test-group", + eval_run_id="test-run", + grader_name_map={"grader-1": "test_grader"}, + expected_rows=4, + ) + mock_run_results = Mock() + mock_run_results.status = "completed" + mock_run_results.per_testing_criteria_results = [Mock(testing_criteria="grader-1", passed=4, failed=0)] + first_page = MockOutputItemsList( + data=[ + MockOutputItem( + id=f"item-{i}", + datasource_item_id=i, + results=[{"name": "grader-1", "passed": True, "score": i}], + ) + for i in range(2) + ], + has_more=True, + ) + mock_client.evals.runs.output_items.list.side_effect = [first_page, RuntimeError("later page failed")] + + with patch( + "azure.ai.evaluation._evaluate._evaluate_aoai._wait_for_run_conclusion", + return_value=mock_run_results, + ): + with pytest.raises(RuntimeError, match="later page failed"): + _get_single_run_results(run_info, aoai_output_items_page_size=2) + + assert mock_client.evals.runs.output_items.list.call_count == 2 From 7506b37fb4846bc2fc8aebee5099d5223107dbc7 Mon Sep 17 00:00:00 2001 From: Rabah B Date: Wed, 16 Sep 2026 12:31:02 -0400 Subject: [PATCH 2/2] Add adaptive retries for AOAI output pages Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> --- ...uation-aoai-output-page-size-2026-09-15.md | 3 +- .../ai/evaluation/_evaluate/_evaluate.py | 3 +- .../ai/evaluation/_evaluate/_evaluate_aoai.py | 78 ++++- .../test_aoai_alignment_missing_rows.py | 1 + .../test_aoai_evaluation_pagination.py | 281 ++++++++++++++++-- 5 files changed, 335 insertions(+), 31 deletions(-) diff --git a/.chronus/changes/rboubchir-azure-ai-evaluation-aoai-output-page-size-2026-09-15.md b/.chronus/changes/rboubchir-azure-ai-evaluation-aoai-output-page-size-2026-09-15.md index f4c21ca89cb7..122b7ec17946 100644 --- a/.chronus/changes/rboubchir-azure-ai-evaluation-aoai-output-page-size-2026-09-15.md +++ b/.chronus/changes/rboubchir-azure-ai-evaluation-aoai-output-page-size-2026-09-15.md @@ -5,4 +5,5 @@ packages: --- Added the `aoai_output_items_page_size` option to `evaluate` for configuring the number of native Azure OpenAI -grader output items requested per HTTP response page. The default remains 100, and all result pages are fetched. +grader output items requested per HTTP response page. The default remains 100, all result pages are fetched, and +timed-out output-item requests retry the same page with progressively smaller page sizes. diff --git a/sdk/evaluation/azure-ai-evaluation/azure/ai/evaluation/_evaluate/_evaluate.py b/sdk/evaluation/azure-ai-evaluation/azure/ai/evaluation/_evaluate/_evaluate.py index 9b070983e865..99f55bbffd95 100644 --- a/sdk/evaluation/azure-ai-evaluation/azure/ai/evaluation/_evaluate/_evaluate.py +++ b/sdk/evaluation/azure-ai-evaluation/azure/ai/evaluation/_evaluate/_evaluate.py @@ -885,7 +885,8 @@ def evaluate( :paramtype fail_on_evaluator_errors: bool :keyword aoai_output_items_page_size: The maximum number of native Azure OpenAI grader output items requested per HTTP response page. Defaults to 100. This does not limit response bytes, request latency, or the number - of dataset rows evaluated; all output-item result pages are fetched. + of dataset rows evaluated; all output-item result pages are fetched. If an output-items request times out or + returns HTTP 408 or 504, the same page is retried with a smaller page size. :paramtype aoai_output_items_page_size: int :keyword tags: A dictionary of tags to be added to the evaluation run for tracking and organization purposes. Keys and values must be strings. For more information about tag limits, see: diff --git a/sdk/evaluation/azure-ai-evaluation/azure/ai/evaluation/_evaluate/_evaluate_aoai.py b/sdk/evaluation/azure-ai-evaluation/azure/ai/evaluation/_evaluate/_evaluate_aoai.py index de60f05bd334..d8cdae379b4f 100644 --- a/sdk/evaluation/azure-ai-evaluation/azure/ai/evaluation/_evaluate/_evaluate_aoai.py +++ b/sdk/evaluation/azure-ai-evaluation/azure/ai/evaluation/_evaluate/_evaluate_aoai.py @@ -5,25 +5,26 @@ import json import logging import re - -from openai import AzureOpenAI, OpenAI -import pandas as pd -from typing import Any, Callable, Dict, Tuple, TypeVar, Union, Type, Optional, TypedDict, List, cast, Set from time import sleep +from typing import Any, Callable, Dict, List, Optional, Set, Tuple, TypedDict, TypeVar, Type, Union, cast -from ._batch_run import CodeClient, ProxyClient +import pandas as pd +from openai import APIConnectionError, APIStatusError, APITimeoutError, AzureOpenAI, OpenAI +from openai._models import FinalRequestOptions # import aoai_mapping -from azure.ai.evaluation._exceptions import ErrorBlame, ErrorCategory, ErrorTarget, EvaluationException -from azure.ai.evaluation._constants import EVALUATION_PASS_FAIL_MAPPING from azure.ai.evaluation._aoai.aoai_grader import AzureOpenAIGrader from azure.ai.evaluation._common._experimental import experimental +from azure.ai.evaluation._constants import EVALUATION_PASS_FAIL_MAPPING +from azure.ai.evaluation._exceptions import ErrorBlame, ErrorCategory, ErrorTarget, EvaluationException +from ._batch_run import CodeClient, ProxyClient TClient = TypeVar("TClient", ProxyClient, CodeClient) LOGGER = logging.getLogger(__name__) _DEFAULT_AOAI_OUTPUT_ITEMS_PAGE_SIZE = 100 _MAX_AOAI_OUTPUT_ITEMS_PAGE_SIZE = 100 +_AOAI_OUTPUT_ITEMS_MAX_ATTEMPTS = 3 # Precompiled regex for extracting data paths from mapping expressions of the form # ${data.some.dotted.path}. Compiled once at import time to avoid repeated @@ -293,6 +294,64 @@ def _get_evaluation_run_results( return output_df, run_metrics +def _list_output_items_page( + client: Union[AzureOpenAI, OpenAI], + list_kwargs: Dict[str, Any], + page_size: int, +) -> Tuple[Any, int]: + """Fetch one output-items page with the OpenAI client's retry policy and a three-attempt budget. + + :param client: A scoped OpenAI client with automatic retries disabled. + :type client: Union[AzureOpenAI, OpenAI] + :param list_kwargs: Arguments identifying the evaluation run and current cursor. + :type list_kwargs: Dict[str, Any] + :param page_size: The number of output items to request. + :type page_size: int + :return: The fetched page and the page size to retain for subsequent pages. + :rtype: Tuple[Any, int] + """ + retry_options = FinalRequestOptions( + method="get", + url="/evals/runs/output_items", + max_retries=_AOAI_OUTPUT_ITEMS_MAX_ATTEMPTS - 1, + ) + + for attempt in range(_AOAI_OUTPUT_ITEMS_MAX_ATTEMPTS): + try: + return client.evals.runs.output_items.list(**list_kwargs, limit=page_size), page_size + except (APIConnectionError, APIStatusError) as error: + should_reduce_page_size = isinstance(error, APITimeoutError) + should_retry = isinstance(error, APIConnectionError) + response_headers = None + + if isinstance(error, APIStatusError): + should_retry = client._should_retry(error.response) # pylint: disable=protected-access + should_reduce_page_size = should_retry and error.status_code in (408, 504) + response_headers = error.response.headers + + if not should_retry or attempt == _AOAI_OUTPUT_ITEMS_MAX_ATTEMPTS - 1: + raise + + if should_reduce_page_size: + page_size = max(1, (page_size + 1) // 2) + + remaining_retries = _AOAI_OUTPUT_ITEMS_MAX_ATTEMPTS - attempt - 1 + delay = client._calculate_retry_timeout( # pylint: disable=protected-access + remaining_retries, + retry_options, + response_headers, + ) + LOGGER.warning( + "AOAI output-items request failed for cursor %s. Retrying with page size %d in %.2f seconds.", + list_kwargs.get("after"), + page_size, + delay, + ) + sleep(delay) + + raise RuntimeError("AOAI output-items retry loop exited unexpectedly.") + + def _get_single_run_results( run_info: OAIEvalRunCreationInfo, aoai_output_items_page_size: int = _DEFAULT_AOAI_OUTPUT_ITEMS_PAGE_SIZE, @@ -359,17 +418,18 @@ def _get_single_run_results( LOGGER.info(f"AOAI: Collecting output items for run {run_info['eval_run_id']} with pagination...") all_results: List[Any] = [] next_cursor: Optional[str] = None + page_size = aoai_output_items_page_size + output_items_client = run_info["client"].with_options(max_retries=0) while True: list_kwargs = { "eval_id": run_info["eval_group_id"], "run_id": run_info["eval_run_id"], - "limit": aoai_output_items_page_size, } if next_cursor is not None: list_kwargs["after"] = next_cursor - raw_list_results = run_info["client"].evals.runs.output_items.list(**list_kwargs) + raw_list_results, page_size = _list_output_items_page(output_items_client, list_kwargs, page_size) # Add current page results all_results.extend(raw_list_results.data) diff --git a/sdk/evaluation/azure-ai-evaluation/tests/unittests/test_aoai_alignment_missing_rows.py b/sdk/evaluation/azure-ai-evaluation/tests/unittests/test_aoai_alignment_missing_rows.py index f1eced6670bf..efb55d4c2713 100644 --- a/sdk/evaluation/azure-ai-evaluation/tests/unittests/test_aoai_alignment_missing_rows.py +++ b/sdk/evaluation/azure-ai-evaluation/tests/unittests/test_aoai_alignment_missing_rows.py @@ -32,6 +32,7 @@ def __init__(self, data, has_more=False): def test_aoai_results_preserve_order_with_unordered_output_items(caplog): """AOAI output_items can arrive unordered; results should align to row ids (0..N-1).""" mock_client = Mock() + mock_client.with_options.return_value = mock_client expected_rows = 5 run_info = OAIEvalRunCreationInfo( client=mock_client, diff --git a/sdk/evaluation/azure-ai-evaluation/tests/unittests/test_aoai_evaluation_pagination.py b/sdk/evaluation/azure-ai-evaluation/tests/unittests/test_aoai_evaluation_pagination.py index 55d6cce1ab4f..0a99492d3984 100644 --- a/sdk/evaluation/azure-ai-evaluation/tests/unittests/test_aoai_evaluation_pagination.py +++ b/sdk/evaluation/azure-ai-evaluation/tests/unittests/test_aoai_evaluation_pagination.py @@ -2,17 +2,21 @@ # Copyright (c) Microsoft Corporation. All rights reserved. # --------------------------------------------------------- -import pytest -from unittest.mock import Mock, MagicMock, patch +from typing import Any, List +from unittest.mock import MagicMock, Mock, patch + +import httpx import pandas as pd -from typing import List, Any +import pytest +from openai import APIStatusError, APITimeoutError +from azure.ai.evaluation import evaluate from azure.ai.evaluation._evaluate._evaluate_aoai import ( - _get_single_run_results, OAIEvalRunCreationInfo, + _get_single_run_results, + _list_output_items_page, _wait_for_run_conclusion, ) -from azure.ai.evaluation import evaluate from azure.ai.evaluation._exceptions import ErrorBlame, ErrorCategory, ErrorTarget, EvaluationException @@ -33,6 +37,22 @@ def __init__(self, data: List[MockOutputItem], has_more: bool = False): self.has_more = has_more +def _mock_openai_client(): + client = Mock() + client.with_options.return_value = client + client._calculate_retry_timeout.return_value = 0 + return client + + +def _api_status_error(status_code, headers=None): + response = httpx.Response( + status_code=status_code, + headers=headers, + request=httpx.Request("GET", "https://example.test/output_items"), + ) + return APIStatusError(f"status {status_code}", response=response, body=None) + + class TestAOAIPagination: """Test pagination functionality in AOAI evaluation results""" @@ -44,7 +64,7 @@ class TestAOAIPagination: ], ) def test_evaluate_propagates_page_size_to_every_native_grader_run(self, evaluate_kwargs, expected_page_size): - clients = [Mock(), Mock()] + clients = [_mock_openai_client(), _mock_openai_client()] input_data = pd.DataFrame([{"query": "hi"}]) run_info = [ OAIEvalRunCreationInfo( @@ -114,6 +134,7 @@ def test_evaluate_propagates_page_size_to_every_native_grader_run(self, evaluate assert len(result["rows"]) == 1 assert result["metrics"] == {"first_grader.pass_rate": 1.0, "second_grader.pass_rate": 1.0} for client in clients: + client.with_options.assert_called_once_with(max_retries=0) client.evals.runs.output_items.list.assert_called_once() assert client.evals.runs.output_items.list.call_args.kwargs["limit"] == expected_page_size @@ -140,7 +161,7 @@ def test_evaluate_rejects_invalid_page_size_before_any_evaluation_work(self, inv def test_single_page_results(self): """Test handling of single page results (no pagination needed)""" # Mock client and run info - mock_client = Mock() + mock_client = _mock_openai_client() run_info = OAIEvalRunCreationInfo( client=mock_client, eval_group_id="test-group", @@ -190,7 +211,7 @@ def test_single_page_results(self): def test_multi_page_results(self): """Test handling of multi-page results with pagination""" - mock_client = Mock() + mock_client = _mock_openai_client() run_info = OAIEvalRunCreationInfo( client=mock_client, eval_group_id="test-group", @@ -262,7 +283,7 @@ def test_multi_page_results(self): def test_empty_page_handling(self): """Test handling of empty pages in pagination""" - mock_client = Mock() + mock_client = _mock_openai_client() run_info = OAIEvalRunCreationInfo( client=mock_client, eval_group_id="test-group", @@ -304,7 +325,7 @@ def test_empty_page_handling(self): def test_result_ordering_preservation(self): """Test that results maintain proper ordering after pagination""" - mock_client = Mock() + mock_client = _mock_openai_client() run_info = OAIEvalRunCreationInfo( client=mock_client, eval_group_id="test-group", @@ -354,7 +375,7 @@ def test_result_ordering_preservation(self): assert scores == expected_scores def test_configured_page_size_fetches_all_results_and_final_partial_page(self): - mock_client = Mock() + mock_client = _mock_openai_client() run_info = OAIEvalRunCreationInfo( client=mock_client, eval_group_id="test-group", @@ -398,18 +419,134 @@ def list_output_items(**kwargs): assert calls[-1].kwargs["after"] == "item-158" assert all(call_.kwargs["limit"] == 3 for call_ in calls) - def test_later_page_failure_propagates(self): - mock_client = Mock() + @pytest.mark.parametrize( + "initial_page_size, expected_attempt_sizes", + [ + (100, [100, 50, 25]), + (50, [50, 25, 13]), + ], + ) + @patch("azure.ai.evaluation._evaluate._evaluate_aoai.sleep") + def test_timeout_retries_reduce_page_size(self, mock_sleep, initial_page_size, expected_attempt_sizes): + mock_client = _mock_openai_client() + request = httpx.Request("GET", "https://example.test/output_items") + mock_client.evals.runs.output_items.list.side_effect = [ + APITimeoutError(request=request), + APITimeoutError(request=request), + MockOutputItemsList(data=[]), + ] + + _, final_page_size = _list_output_items_page(mock_client, {"after": "cursor-1"}, initial_page_size) + + assert final_page_size == expected_attempt_sizes[-1] + assert [ + call_.kwargs["limit"] for call_ in mock_client.evals.runs.output_items.list.call_args_list + ] == expected_attempt_sizes + assert all( + call_.kwargs["after"] == "cursor-1" for call_ in mock_client.evals.runs.output_items.list.call_args_list + ) + assert [call_.args[0] for call_ in mock_client._calculate_retry_timeout.call_args_list] == [2, 1] + assert mock_sleep.call_count == 2 + + @patch("azure.ai.evaluation._evaluate._evaluate_aoai.sleep") + def test_minimum_page_size_retries_no_more_than_three_attempts(self, mock_sleep): + mock_client = _mock_openai_client() + request = httpx.Request("GET", "https://example.test/output_items") + mock_client.evals.runs.output_items.list.side_effect = APITimeoutError(request=request) + + with pytest.raises(APITimeoutError): + _list_output_items_page(mock_client, {"after": "cursor-1"}, 1) + + assert [call_.kwargs["limit"] for call_ in mock_client.evals.runs.output_items.list.call_args_list] == [1, 1, 1] + assert mock_sleep.call_count == 2 + + @pytest.mark.parametrize("status_code", [408, 504]) + @patch("azure.ai.evaluation._evaluate._evaluate_aoai.sleep") + def test_http_timeout_status_reduces_page_size(self, mock_sleep, status_code): + mock_client = _mock_openai_client() + mock_client._should_retry.return_value = True + mock_client.evals.runs.output_items.list.side_effect = [ + _api_status_error(status_code), + MockOutputItemsList(data=[]), + ] + + _, final_page_size = _list_output_items_page(mock_client, {"after": "cursor-1"}, 25) + + assert final_page_size == 13 + assert [call_.kwargs["limit"] for call_ in mock_client.evals.runs.output_items.list.call_args_list] == [25, 13] + mock_sleep.assert_called_once() + + @patch("azure.ai.evaluation._evaluate._evaluate_aoai.sleep") + def test_reduced_page_size_carries_forward_and_collects_all_results_once(self, mock_sleep): + mock_client = _mock_openai_client() run_info = OAIEvalRunCreationInfo( client=mock_client, eval_group_id="test-group", eval_run_id="test-run", grader_name_map={"grader-1": "test_grader"}, - expected_rows=4, + expected_rows=160, ) mock_run_results = Mock() mock_run_results.status = "completed" - mock_run_results.per_testing_criteria_results = [Mock(testing_criteria="grader-1", passed=4, failed=0)] + mock_run_results.per_testing_criteria_results = [Mock(testing_criteria="grader-1", passed=160, failed=0)] + output_items = [ + MockOutputItem( + id=f"item-{position}", + datasource_item_id=position, + results=[{"name": "grader-1", "passed": True, "score": position}], + ) + for position in range(160) + ] + request = httpx.Request("GET", "https://example.test/output_items") + call_count = 0 + + def list_output_items(**kwargs): + nonlocal call_count + call_count += 1 + if call_count == 1: + raise APITimeoutError(request=request) + if call_count == 3: + raise _api_status_error(504) + + after = kwargs.get("after") + start = int(after.removeprefix("item-")) + 1 if after else 0 + end = min(start + kwargs["limit"], len(output_items)) + return MockOutputItemsList(data=output_items[start:end], has_more=end < len(output_items)) + + mock_client._should_retry.return_value = True + mock_client.evals.runs.output_items.list.side_effect = list_output_items + + with patch( + "azure.ai.evaluation._evaluate._evaluate_aoai._wait_for_run_conclusion", + return_value=mock_run_results, + ): + df, metrics = _get_single_run_results(run_info, aoai_output_items_page_size=100) + + calls = mock_client.evals.runs.output_items.list.call_args_list + assert [call_.kwargs["limit"] for call_ in calls] == [100, 50, 50, 25, 25, 25, 25, 25] + assert "after" not in calls[0].kwargs + assert "after" not in calls[1].kwargs + assert calls[2].kwargs["after"] == "item-49" + assert calls[3].kwargs["after"] == "item-49" + assert calls[4].kwargs["after"] == "item-74" + assert len(df) == 160 + assert df["outputs.test_grader.score"].tolist() == list(range(160)) + assert metrics == {"test_grader.pass_rate": 1.0} + assert mock_sleep.call_count == 2 + + @patch("azure.ai.evaluation._evaluate._evaluate_aoai.sleep") + def test_later_page_retry_exhaustion_propagates_without_partial_result(self, mock_sleep): + mock_client = _mock_openai_client() + run_info = OAIEvalRunCreationInfo( + client=mock_client, + eval_group_id="test-group", + eval_run_id="test-run", + grader_name_map={"grader-1": "test_grader"}, + expected_rows=100, + ) + mock_run_results = Mock() + mock_run_results.status = "completed" + mock_run_results.per_testing_criteria_results = [Mock(testing_criteria="grader-1", passed=100, failed=0)] first_page = MockOutputItemsList( data=[ MockOutputItem( @@ -417,17 +554,121 @@ def test_later_page_failure_propagates(self): datasource_item_id=i, results=[{"name": "grader-1", "passed": True, "score": i}], ) - for i in range(2) + for i in range(50) ], has_more=True, ) - mock_client.evals.runs.output_items.list.side_effect = [first_page, RuntimeError("later page failed")] + request = httpx.Request("GET", "https://example.test/output_items") + final_error = APITimeoutError(request=request) + mock_client.evals.runs.output_items.list.side_effect = [ + first_page, + APITimeoutError(request=request), + APITimeoutError(request=request), + final_error, + ] with patch( "azure.ai.evaluation._evaluate._evaluate_aoai._wait_for_run_conclusion", return_value=mock_run_results, ): - with pytest.raises(RuntimeError, match="later page failed"): - _get_single_run_results(run_info, aoai_output_items_page_size=2) + with pytest.raises(APITimeoutError) as exc_info: + _get_single_run_results(run_info, aoai_output_items_page_size=50) - assert mock_client.evals.runs.output_items.list.call_count == 2 + assert exc_info.value is final_error + calls = mock_client.evals.runs.output_items.list.call_args_list + assert [call_.kwargs["limit"] for call_ in calls] == [50, 50, 25, 13] + assert all(call_.kwargs.get("after") == "item-49" for call_ in calls[1:]) + assert mock_sleep.call_count == 2 + + @patch("azure.ai.evaluation._evaluate._evaluate_aoai.sleep") + def test_other_retryable_status_preserves_page_size_and_backoff_headers(self, mock_sleep): + mock_client = _mock_openai_client() + run_info = OAIEvalRunCreationInfo( + client=mock_client, + eval_group_id="test-group", + eval_run_id="test-run", + grader_name_map={"grader-1": "test_grader"}, + expected_rows=1, + ) + mock_run_results = Mock() + mock_run_results.status = "completed" + mock_run_results.per_testing_criteria_results = [Mock(testing_criteria="grader-1", passed=1, failed=0)] + retryable_error = _api_status_error(503, headers={"Retry-After": "4", "x-should-retry": "true"}) + mock_client._should_retry.return_value = True + mock_client._calculate_retry_timeout.return_value = 4.0 + mock_client.evals.runs.output_items.list.side_effect = [ + retryable_error, + MockOutputItemsList( + data=[ + MockOutputItem( + id="item-0", + datasource_item_id=0, + results=[{"name": "grader-1", "passed": True, "score": 1}], + ) + ] + ), + ] + + with patch( + "azure.ai.evaluation._evaluate._evaluate_aoai._wait_for_run_conclusion", + return_value=mock_run_results, + ): + df, _ = _get_single_run_results(run_info, aoai_output_items_page_size=50) + + assert len(df) == 1 + assert [call_.kwargs["limit"] for call_ in mock_client.evals.runs.output_items.list.call_args_list] == [50, 50] + mock_client._should_retry.assert_called_once_with(retryable_error.response) + assert mock_client._calculate_retry_timeout.call_args.args[2] == retryable_error.response.headers + mock_sleep.assert_called_once_with(4.0) + + @patch("azure.ai.evaluation._evaluate._evaluate_aoai.sleep") + def test_non_retryable_status_propagates_immediately(self, mock_sleep): + mock_client = _mock_openai_client() + non_retryable_error = _api_status_error(503, headers={"x-should-retry": "false"}) + mock_client._should_retry.return_value = False + mock_client.evals.runs.output_items.list.side_effect = non_retryable_error + + with pytest.raises(APIStatusError) as exc_info: + _list_output_items_page(mock_client, {"after": "cursor-1"}, 50) + + assert exc_info.value is non_retryable_error + mock_client.evals.runs.output_items.list.assert_called_once_with(after="cursor-1", limit=50) + mock_client._calculate_retry_timeout.assert_not_called() + mock_sleep.assert_not_called() + + def test_output_items_fetch_disables_client_retries_without_mutating_shared_client(self): + shared_client = Mock() + output_items_client = Mock() + shared_client.with_options.return_value = output_items_client + run_info = OAIEvalRunCreationInfo( + client=shared_client, + eval_group_id="test-group", + eval_run_id="test-run", + grader_name_map={"grader-1": "test_grader"}, + expected_rows=1, + ) + mock_run_results = Mock() + mock_run_results.status = "completed" + mock_run_results.per_testing_criteria_results = [Mock(testing_criteria="grader-1", passed=1, failed=0)] + output_items_client.evals.runs.output_items.list.return_value = MockOutputItemsList( + data=[ + MockOutputItem( + id="item-0", + datasource_item_id=0, + results=[{"name": "grader-1", "passed": True, "score": 1}], + ) + ] + ) + + with patch( + "azure.ai.evaluation._evaluate._evaluate_aoai._wait_for_run_conclusion", + return_value=mock_run_results, + ) as mock_wait: + _get_single_run_results(run_info) + + mock_wait.assert_called_once_with(shared_client, "test-group", "test-run") + shared_client.with_options.assert_called_once_with(max_retries=0) + shared_client.evals.runs.output_items.list.assert_not_called() + output_items_client.evals.runs.output_items.list.assert_called_once() + shared_client.close.assert_not_called() + output_items_client.close.assert_not_called()