From 853a32418253d96a91408c782096a67d973dae92 Mon Sep 17 00:00:00 2001 From: Ahmad Qardahji Date: Tue, 3 Nov 2026 13:32:19 -0800 Subject: [PATCH] Reuse credential tokens for Decision and chat requests Decision requests, and Azure OpenAI chat model requests that authenticate with Microsoft Entra ID, reuse a credential's token until five minutes before it expires, as azure-core's bearer token policy does, instead of requesting a token for every request. Concurrent requests that share a synchronous credential wait for one token request on a worker thread without blocking their event loops; a credential whose get_token is a coroutine function, or that implements AsyncTokenCredential but not the synchronous context manager protocol, is asked on the caller's event loop. Only an authentication failure is shared with the requests that waited on it, and only with requests of the same kind, so a chat failure is never replayed as a Decision failure or the reverse. A token rejected with HTTP 401 is dropped. Tokens are kept per scope, so Decision keeps the Cognitive Services scope that the hosted evaluation service relies on, and the cache never keeps a credential alive. --- .../azure-ai-evaluation/CHANGELOG.md | 4 + .../_evaluators/_common/_model_runners.py | 91 +- .../_legacy/_common/_async_token_provider.py | 64 +- .../_legacy/_common/_token_cache.py | 273 ++++++ .../ai/evaluation/_legacy/prompty/_prompty.py | 107 ++- .../unittests/test_decision_model_support.py | 601 +++++++++++- .../unittests/test_prompty_token_reuse.py | 891 ++++++++++++++++++ 7 files changed, 1985 insertions(+), 46 deletions(-) create mode 100644 sdk/evaluation/azure-ai-evaluation/azure/ai/evaluation/_legacy/_common/_token_cache.py create mode 100644 sdk/evaluation/azure-ai-evaluation/tests/unittests/test_prompty_token_reuse.py diff --git a/sdk/evaluation/azure-ai-evaluation/CHANGELOG.md b/sdk/evaluation/azure-ai-evaluation/CHANGELOG.md index a69fabca1246..dec4b7d1c529 100644 --- a/sdk/evaluation/azure-ai-evaluation/CHANGELOG.md +++ b/sdk/evaluation/azure-ai-evaluation/CHANGELOG.md @@ -9,6 +9,10 @@ - The binary evaluators (task adherence, task completion, tool call success, tool input accuracy, tool output utilization and tool selection) accept a `probability_threshold` keyword. For Decision models it sets the probability at or above which the score is 1 (default 0.75) and overrides the value in `model_config`; chat models ignore it with a warning. - Conversation inputs are handled as before with a Decision model: evaluators that accept them score each user and assistant turn separately and aggregate the results. The catalog's whole-conversation rubrics are not included. +### Bugs Fixed + +- Decision requests, and Azure OpenAI chat model requests that authenticate with Microsoft Entra ID, reuse a credential's token until five minutes before it expires, as azure-core's bearer token policy does, instead of requesting a token for every request. Concurrent evaluations that share a synchronous credential wait for one token request, without blocking their event loops. A credential whose `get_token` is a coroutine function, or that implements `AsyncTokenCredential` but not the synchronous context manager protocol, is asked for its token on the caller's event loop, and any other credential on a worker thread, so a plain `get_token` that needs the running event loop must belong to an `AsyncTokenCredential`. With developer credentials such as `AzureCliCredential`, which start a process per request, concurrent evaluations were slow and could fail. Requests that waited on a token request that failed authentication share its failure instead of each requesting a token again; after any other failure, such as a token service outage, each waiting request makes its own attempt, so a transient failure fails only the request that hit it. Decision and chat requests behave the same way. A token rejected with HTTP 401 is dropped and requested again; a token whose request began before the rejection is not kept for reuse, and requests that waited on it use it rather than each requesting another one. The default credential that a chat evaluator creates when no `credential` is passed is shared the same way. Tokens are kept per scope, and the cache never keeps a credential alive. A token without an expiry, or that expires within five minutes, is not kept, but the requests that waited for it use it. + ## 1.18.10 (2026-10-22) ### Features Added diff --git a/sdk/evaluation/azure-ai-evaluation/azure/ai/evaluation/_evaluators/_common/_model_runners.py b/sdk/evaluation/azure-ai-evaluation/azure/ai/evaluation/_evaluators/_common/_model_runners.py index e2d34742c296..f6d2e22a9695 100644 --- a/sdk/evaluation/azure-ai-evaluation/azure/ai/evaluation/_evaluators/_common/_model_runners.py +++ b/sdk/evaluation/azure-ai-evaluation/azure/ai/evaluation/_evaluators/_common/_model_runners.py @@ -16,14 +16,17 @@ import math import os import re +import threading import time from typing import Any, Callable, Dict, List, Optional from urllib.parse import urlsplit +from azure.core.exceptions import ClientAuthenticationError from azure.core.pipeline.policies import AsyncRedirectPolicy, AsyncRetryPolicy, CustomHookPolicy from azure.core.pipeline.transport import AioHttpTransport from azure.ai.evaluation._exceptions import ErrorBlame, ErrorCategory, ErrorTarget, EvaluationException +from azure.ai.evaluation._legacy._common._token_cache import TOKENS as _TOKENS, is_async_credential from ..._common.utils import construct_prompty_model_config, validate_model_config from ..._http_utils import get_async_http_client @@ -44,6 +47,9 @@ def value(self) -> str: DECISION_MODEL_TYPE = "decision" _DECISION_PATH = "/providers/microsoft/v1/systemone" +# Decision endpoints also accept the chat scope, https://ai.azure.com/.default, but the hosted evaluation service +# relies on this audience to call Decision with its exchanged or managed identity token rather than the caller's own +# token, so Decision requests keep this scope and their own cached token. _DECISION_TOKEN_SCOPE = "https://cognitiveservices.azure.com/.default" _DECISION_RETRY_STATUS_CODES = [408, 429, 500, 502, 503, 504, 529] LOGGER = logging.getLogger(__name__) @@ -340,6 +346,25 @@ def _is_finite_number(value: Any) -> bool: return False +def _token_failure(error: BaseException) -> Optional[Dict[str, Any]]: + """Describe an authentication failure for the requests that waited on it, or return None for any other failure. + + As on the chat path, only an authentication failure is shared: any other failure, such as a token service outage, + is not, so each waiting request makes its own attempt and one transient failure fails only the request that hit + it. The description never keeps the error, whose traceback references the credential. + """ + if not isinstance(error, ClientAuthenticationError): + return None + message = "Could not acquire a token for the decision model." + return { + "message": message, + "internal_message": message, + "target": ErrorTarget.EVALUATE, + "category": ErrorCategory.FAILED_EXECUTION, + "blame": ErrorBlame.USER_ERROR, + } + + def _probability(value: Any) -> float: if not _is_finite_number(value): raise ValueError("Decision probabilities must be finite numbers.") @@ -648,25 +673,72 @@ async def _await_request(self, content: bytes, deadline: float, execution: Decis task.cancel() await asyncio.gather(task, return_exceptions=True) - async def _get_token(self) -> str: - if inspect.iscoroutinefunction(self._credential.get_token): - return (await self._credential.get_token(_DECISION_TOKEN_SCOPE)).token - token = await asyncio.to_thread(self._credential.get_token, _DECISION_TOKEN_SCOPE) - if inspect.isawaitable(token): - token = await token - return token.token + async def _get_token(self, deadline: float) -> Optional[str]: + """Return a bearer token, or None when the deadline passed while waiting for another request's token. + + A credential whose ``get_token`` is a coroutine function, or that implements ``AsyncTokenCredential`` but not + the synchronous context manager protocol, is asked on this event loop; any other credential is asked on a + worker thread, once for concurrent requests. + + :param float deadline: Absolute ``time.monotonic()`` instant by which the request must complete, not a + duration or a wall-clock time; the caller derives it from the request timeout. + :return: The bearer token, or None when the deadline passed while this request waited for a token that + another request was acquiring. + :rtype: Optional[str] + """ + credential = self._credential + token = _TOKENS.cached(credential, _DECISION_TOKEN_SCOPE) + if token is None: + # An acquisition that begins before a service rejects a token may return that token; the generation it + # began under lets the cache refuse it. + generation = _TOKENS.generation(credential, _DECISION_TOKEN_SCOPE) + if is_async_credential(credential): + # An asynchronous credential's token request may need this event loop. + token = credential.get_token(_DECISION_TOKEN_SCOPE) + if inspect.isawaitable(token): + token = await token + _TOKENS.store(credential, _DECISION_TOKEN_SCOPE, token, generation) + else: + abandoned = threading.Event() + try: + # A request waiting for another request's token holds a default-executor thread while it waits; + # past the executor's size, further requests queue for a thread before they reach the cache. + token = await asyncio.to_thread( + _TOKENS.fetch, + credential, + _DECISION_TOKEN_SCOPE, + abandoned=abandoned, + deadline=deadline, + describe_failure=_token_failure, + ) + finally: + abandoned.set() + if token is None and time.monotonic() >= deadline: + return None + if inspect.isawaitable(token): + token = await token + _TOKENS.store(credential, _DECISION_TOKEN_SCOPE, token, generation) + value = getattr(token, "token", None) + if not isinstance(value, str): + raise _failed("Could not acquire a token for the decision model.", ErrorBlame.USER_ERROR) + return value async def _post(self, content: bytes, *, deadline: float, execution: DecisionExecution) -> Any: headers = {"x-ms-useragent": self._user_agent, "Content-Type": "application/json"} + token = None if self._api_key: headers["api-key"] = self._api_key else: try: - headers["Authorization"] = f"Bearer {await self._get_token()}" + token = await self._get_token(deadline) except EvaluationException: raise except Exception: # pylint: disable=broad-except raise _failed("Could not acquire a token for the decision model.", ErrorBlame.USER_ERROR) from None + if token is None: + execution.outcome = "deadline" + raise _failed("The decision model request timed out.") + headers["Authorization"] = "Bearer " + token policies: Dict[str, Any] = { "custom_hook_policy": execution, @@ -718,6 +790,9 @@ async def _post(self, content: bytes, *, deadline: float, execution: DecisionExe raise _failed("The decision model returned invalid JSON.") from None payload = None if status != 200: + if status == 401 and token is not None: + # A rejected token is not reused. + _TOKENS.discard(self._credential, _DECISION_TOKEN_SCOPE, token) blame = ErrorBlame.USER_ERROR if status in (400, 401, 403, 404, 422) else ErrorBlame.SYSTEM_ERROR provider_error = payload.get("error") if isinstance(payload, dict) else None code = provider_error.get("code") if isinstance(provider_error, dict) else None diff --git a/sdk/evaluation/azure-ai-evaluation/azure/ai/evaluation/_legacy/_common/_async_token_provider.py b/sdk/evaluation/azure-ai-evaluation/azure/ai/evaluation/_legacy/_common/_async_token_provider.py index 4509604921c7..71f76db8b7f3 100644 --- a/sdk/evaluation/azure-ai-evaluation/azure/ai/evaluation/_legacy/_common/_async_token_provider.py +++ b/sdk/evaluation/azure-ai-evaluation/azure/ai/evaluation/_legacy/_common/_async_token_provider.py @@ -4,6 +4,7 @@ import logging import os +import threading from typing import Any, AsyncContextManager, Optional from azure.core.credentials import AccessToken, TokenCredential @@ -27,11 +28,13 @@ def __init__(self, *, base_url: Optional[str] = None, **kwargs: Any) -> None: """Initialize the AsyncAzureTokenProvider.""" self._credential: Optional[TokenCredential] = None self._env_client: Optional[AzureEnvironmentClient] = AzureEnvironmentClient(base_url=base_url, **kwargs) + self._credential_lock = threading.Lock() async def close(self) -> None: - credential, client = self._credential, self._env_client - self._credential = None - self._env_client = None + with self._credential_lock: + credential, client = self._credential, self._env_client + self._credential = None + self._env_client = None try: close_credential = getattr(credential, "close", None) if close_credential is not None: @@ -40,31 +43,58 @@ async def close(self) -> None: if client is not None: await client.close() - async def get_token( - self, - *scopes: str, - claims: Optional[str] = None, - tenant_id: Optional[str] = None, - enable_cae: bool = False, - **kwargs: Any, - ) -> AccessToken: - if self._credential is None: - self._credential = await self._initialize_async(self._env_client) + async def get_credential(self) -> TokenCredential: + """Return the synchronous credential that supplies this provider's tokens, creating it on first use. + + Concurrent first requests all use the first credential created, so they can share its token. + + :return: The credential. + :rtype: ~azure.core.credentials.TokenCredential + """ + credential = self._credential + if credential is None: + created = await self._initialize_async(self._env_client) + with self._credential_lock: + closed = self._env_client is None + # A provider closed while the credential was created keeps no credential. + if self._credential is None and not closed: + self._credential = created + credential = self._credential + if credential is not created: + close_created = getattr(created, "close", None) + if close_created is not None: + close_created() + if closed: + raise EvaluationException( + f"{AsyncAzureTokenProvider.__name__} instance has already been closed.", + target=ErrorTarget.UNKNOWN, + category=ErrorCategory.INVALID_VALUE, + blame=ErrorBlame.USER_ERROR, + ) - if self._credential is None: + if credential is None: raise EvaluationException( f"{self.__class__.__name__} could not determine the credential to use.", target=ErrorTarget.UNKNOWN, category=ErrorCategory.INVALID_VALUE, blame=ErrorBlame.SYSTEM_ERROR, ) + return credential - return self._credential.get_token(*scopes, claims=claims, tenant_id=tenant_id, enable_cae=enable_cae, **kwargs) + async def get_token( + self, + *scopes: str, + claims: Optional[str] = None, + tenant_id: Optional[str] = None, + enable_cae: bool = False, + **kwargs: Any, + ) -> AccessToken: + credential = await self.get_credential() + return credential.get_token(*scopes, claims=claims, tenant_id=tenant_id, enable_cae=enable_cae, **kwargs) async def __aenter__(self) -> "AsyncAzureTokenProvider": try: - if self._credential is None: - self._credential = await self._initialize_async(self._env_client) + await self.get_credential() except BaseException: try: await self.close() diff --git a/sdk/evaluation/azure-ai-evaluation/azure/ai/evaluation/_legacy/_common/_token_cache.py b/sdk/evaluation/azure-ai-evaluation/azure/ai/evaluation/_legacy/_common/_token_cache.py new file mode 100644 index 000000000000..be74ab288085 --- /dev/null +++ b/sdk/evaluation/azure-ai-evaluation/azure/ai/evaluation/_legacy/_common/_token_cache.py @@ -0,0 +1,273 @@ +# --------------------------------------------------------- +# Copyright (c) Microsoft Corporation. All rights reserved. +# --------------------------------------------------------- + +"""Reuse a caller credential's access tokens across model requests. + +Developer credentials such as ``AzureCliCredential`` start a process for every token request, so evaluations that +requested a token for each model call were slow and, when concurrent, could fail. A token is kept per credential and +scope until shortly before it expires, as azure-core's bearer token policy does. Credentials are keyed by identity +through a weak reference and never kept alive; a credential that cannot be weakly referenced, or a token without a +numeric expiry, is not cached. A token that a service rejected is dropped, and a token whose acquisition began before +that rejection is not stored, so a rejected token never re-enters the cache through a slower request. Requests that +waited on an acquisition whose token was not cached, because of such a rejection or because of its expiry, use that +token instead of each requesting another one. + +azure-identity's ``get_bearer_token_provider`` is not used because it keeps one token per provider instance with no +coordination between concurrent requests, so concurrent first requests each call the credential; because its +wrapper never sees a service's 401 response, so a rejected token is served until it expires; and because its +asynchronous form binds its lock to one event loop while each synchronous evaluator call runs on a new loop. +Sharing one token across evaluators, waiting with a deadline and under cancellation, and dropping a rejected token +need the coordination here. +""" + +import contextlib +import inspect +import math +import threading +import time +import weakref +from typing import Any, Callable, Dict, NamedTuple, Optional, Tuple + +from azure.core.credentials_async import AsyncTokenCredential +from azure.core.exceptions import ClientAuthenticationError + +from azure.ai.evaluation._exceptions import EvaluationException + +# A token is reused until this many seconds before it expires, like azure-core's bearer token policy. +TOKEN_REFRESH_SECONDS = 300 +# How often a request waiting for another request's token checks whether it was abandoned. +TOKEN_WAIT_INTERVAL_SECONDS = 0.05 + + +def is_async_credential(credential: Any) -> bool: + """Whether a credential's token is requested on the caller's event loop rather than on a worker thread. + + That is a credential whose ``get_token`` is a coroutine function, or one that implements the asynchronous + credential protocol but not the synchronous context manager protocol: its ``get_token`` may return an awaitable + that needs the loop it was called on. Any other credential is asked on a worker thread, so a plain ``get_token`` + that needs the caller's running loop must belong to an ``AsyncTokenCredential``. + """ + return inspect.iscoroutinefunction(getattr(credential, "get_token", None)) or ( + isinstance(credential, AsyncTokenCredential) and not isinstance(credential, contextlib.AbstractContextManager) + ) + + +def _is_fresh(token: Any) -> bool: + expires_on = getattr(token, "expires_on", None) + if isinstance(expires_on, bool) or not isinstance(expires_on, (int, float)): + return False + try: + return ( + isinstance(getattr(token, "token", None), str) + and math.isfinite(expires_on) + and expires_on - time.time() > TOKEN_REFRESH_SECONDS + ) + except OverflowError: + return False + + +def _has_value(token: Any) -> bool: + """Whether a token has a string value to send. Its expiry is not checked: a credential may report none, or 0.""" + return isinstance(getattr(token, "token", None), str) + + +def describe_credential_failure(error: BaseException) -> Optional[Dict[str, Any]]: + """Describe a credential's authentication failure by its message, or return None for any other failure. + + Requests that waited on an authentication failure share it as a ``ClientAuthenticationError``. Any other failure, + such as a token service outage, is not shared, so each waiting request makes its own attempt and raises the + credential's own error type. The description never keeps the error or its traceback. + """ + if isinstance(error, ClientAuthenticationError): + return {"message": str(error) or type(error).__name__} + return None + + +class _SharedFailure(NamedTuple): + """A failed token request's description and the attempt count after it, kept for requests of the same kind.""" + + attempt: int + description: Dict[str, Any] + + +class _CredentialTokens: + """One credential's cached token for one scope, the lock for its token requests, the latest failure that each + kind of request described, and the count of rejections that invalidates acquisitions begun before them.""" + + def __init__(self, reference: Any) -> None: + self.reference = reference + self.token: Any = None + self.lock = threading.Lock() + self.attempts = 0 + # Keyed by the failure formatter and exception type of the requests that share the failure, so requests that + # describe or raise failures differently never raise each other's. Cleared when a token request succeeds. + self.failures: Dict[Tuple[Callable[..., Any], Callable[..., BaseException]], _SharedFailure] = {} + # The latest token request's token when it was not cached, because a rejection kept it out or because of its + # expiry, for the requests that waited on that request. It is never cached. + self.unstored: Any = None + # Incremented when a service rejects a token. An acquisition that began before the latest rejection may + # return the rejected token, so a token is stored only under the generation its acquisition began with. + self.generation = 0 + + +class CredentialTokenCache: + """Share each caller credential's token for a scope across requests until shortly before it expires. + + Concurrent requests that share a synchronous credential wait for one token request. A caller can also ask that + requests which waited on a failed token request share its failure instead of each requesting a token again; only + requests that describe and raise failures the same way share one. A token a service rejected is discarded; the + credential's next token is cached again, even when it is the same token, while an acquisition that began before + the rejection is never stored. Requests that waited on an acquisition whose token was not cached use that token + instead of each requesting another one. + """ + + def __init__(self) -> None: + self._entries: Dict[Tuple[int, str], _CredentialTokens] = {} + # Reentrant because a weak reference callback can run while a lookup holds the lock. + self._guard = threading.RLock() + + def _entry(self, credential: Any, scope: str, create: bool = False) -> Optional[_CredentialTokens]: + key = (id(credential), scope) + with self._guard: + entry = self._entries.get(key) + if entry is not None and entry.reference() is credential: + return entry + if not create: + return None + + def forget(reference: Any) -> None: + with self._guard: + current = self._entries.get(key) + if current is not None and current.reference is reference: + del self._entries[key] + + try: + entry = _CredentialTokens(weakref.ref(credential, forget)) + except TypeError: + return None + self._entries[key] = entry + return entry + + def cached(self, credential: Any, scope: str) -> Any: + entry = self._entry(credential, scope) + with self._guard: + token = entry.token if entry is not None else None + return token if _is_fresh(token) else None + + def generation(self, credential: Any, scope: str) -> int: + """Return the rejection generation to pass to ``store`` for a token acquisition that begins now.""" + entry = self._entry(credential, scope) + with self._guard: + return entry.generation if entry is not None else 0 + + def store(self, credential: Any, scope: str, token: Any, generation: Optional[int] = None) -> bool: + """Keep a fresh token for reuse and return whether it was kept. + + ``generation`` is what ``generation`` returned when the acquisition began. A token acquired before a later + rejection is not stored, because it may be the rejected token; without a generation the token is stored. + """ + entry = self._entry(credential, scope, create=True) if _is_fresh(token) else None + if entry is not None: + with self._guard: + if generation is None or generation == entry.generation: + entry.token = token + return True + return False + + def discard(self, credential: Any, scope: str, token: str) -> None: + """Stop reusing a token a service rejected, and invalidate acquisitions that began before the rejection.""" + entry = self._entry(credential, scope, create=True) + if entry is not None: + with self._guard: + entry.generation += 1 + if entry.token is not None and entry.token.token == token: + entry.token = None + if entry.unstored is not None and entry.unstored.token == token: + entry.unstored = None + + def fetch( + self, + credential: Any, + scope: str, + *, + abandoned: Optional[threading.Event] = None, + deadline: Optional[float] = None, + describe_failure: Optional[Callable[[BaseException], Optional[Dict[str, Any]]]] = None, + failure_type: Callable[..., BaseException] = EvaluationException, + ) -> Any: + """Request a synchronous credential's token once for concurrent requests. + + Returns None, without requesting a token, when the request was abandoned or its monotonic ``deadline`` passed + while it waited for another request's token. With ``describe_failure``, a request that waited on a failed + request of its own kind, one that passed the same ``describe_failure`` and ``failure_type``, raises + ``failure_type`` built from that request's description instead of requesting a token again. A None + description is not shared, and a request never raises a failure that another kind of request described, so + each keeps its own error type and classification; with no failure of its own kind, it makes its own attempt. + The description must not keep the error, whose traceback references the credential, and ``describe_failure`` + and ``failure_type`` must be module-level functions or classes: they key the shared failures, so a bound method + or closure would keep what it references, such as the credential, alive. Requests that waited on a request + whose token was not cached use that token. + """ + entry = self._entry(credential, scope, create=True) + if entry is None: + if (abandoned is not None and abandoned.is_set()) or ( + deadline is not None and time.monotonic() >= deadline + ): + return None + return credential.get_token(scope) + with self._guard: + queued_after = entry.attempts + # Wait in short steps so an abandoned request stops waiting promptly. + while True: + remaining = math.inf if deadline is None else deadline - time.monotonic() + if (abandoned is not None and abandoned.is_set()) or remaining <= 0: + return None + if entry.lock.acquire(timeout=min(remaining, TOKEN_WAIT_INTERVAL_SECONDS)): + break + try: + token = self.cached(credential, scope) + if ( + token is not None + or (abandoned is not None and abandoned.is_set()) + or (deadline is not None and time.monotonic() >= deadline) + ): + return token + kind = (describe_failure, failure_type) + with self._guard: + waited = entry.attempts != queued_after + failure = entry.failures.get(kind) if waited and describe_failure is not None else None + unstored = entry.unstored if waited else None + if failure is not None and failure.attempt > queued_after: + raise failure_type(**failure.description) + if _has_value(unstored): + # Issued for the request this one waited on and not rejected since it was shared, so use it rather + # than request another token. + return unstored + with self._guard: + generation = entry.generation + try: + token = credential.get_token(scope) + except Exception as error: + description = describe_failure(error) if describe_failure is not None else None + with self._guard: + entry.attempts += 1 + if description is not None: + entry.failures[kind] = _SharedFailure(entry.attempts, description) + else: + entry.failures.pop(kind, None) + entry.unstored = None + raise + # Decide whether the token is cached and publish the outcome together, so a rejection cannot land between + # them and be undone. + with self._guard: + stored = inspect.isawaitable(token) or self.store(credential, scope, token, generation) + entry.attempts += 1 + entry.failures.clear() + entry.unstored = token if not stored and _has_value(token) else None + return token + finally: + entry.lock.release() + + +TOKENS = CredentialTokenCache() diff --git a/sdk/evaluation/azure-ai-evaluation/azure/ai/evaluation/_legacy/prompty/_prompty.py b/sdk/evaluation/azure-ai-evaluation/azure/ai/evaluation/_legacy/prompty/_prompty.py index 87876d997321..fa1ad1b89b1e 100644 --- a/sdk/evaluation/azure-ai-evaluation/azure/ai/evaluation/_legacy/prompty/_prompty.py +++ b/sdk/evaluation/azure-ai-evaluation/azure/ai/evaluation/_legacy/prompty/_prompty.py @@ -3,17 +3,19 @@ # --------------------------------------------------------- import asyncio +import inspect import re +import threading from logging import Logger from os import PathLike from pathlib import Path -from typing import Any, AsyncGenerator, Awaitable, Dict, Final, List, Mapping, Optional, Sequence, Tuple, Union, cast +from typing import Any, AsyncGenerator, Dict, Final, List, Mapping, Optional, Sequence, Tuple, Union, cast -from openai import AsyncAzureOpenAI, AsyncOpenAI, NotGiven, OpenAIError -from openai.lib.azure import AsyncAzureADTokenProvider +from openai import AsyncAzureOpenAI, AsyncOpenAI, AuthenticationError, NotGiven, OpenAIError from azure.core.credentials import TokenCredential from azure.core.credentials_async import AsyncTokenCredential +from azure.core.exceptions import ClientAuthenticationError from azure.ai.evaluation._exceptions import ErrorTarget from azure.ai.evaluation._constants import DefaultOpenEncoding, TokenScope @@ -41,6 +43,7 @@ from azure.ai.evaluation._constants import DEFAULT_MAX_COMPLETION_TOKENS_REASONING_MODELS from azure.ai.evaluation._legacy._common._logging import get_logger from azure.ai.evaluation._legacy._common._async_token_provider import AsyncAzureTokenProvider +from azure.ai.evaluation._legacy._common._token_cache import TOKENS, describe_credential_failure, is_async_credential from azure.ai.evaluation._user_agent import UserAgentSingleton PROMPTY_EXTENSION: Final[str] = ".prompty" @@ -296,6 +299,7 @@ async def __call__( # pylint: disable=docstring-keyword-should-match-keyword-on default_headers = {"User-Agent": UserAgentSingleton().value} api_client: Union[AsyncAzureOpenAI, AsyncOpenAI, AsyncByoProjectResponsesClient] + token_provider: Optional[_ReusedTokenProvider] = None if is_byo: # BYO judge model (connection/deployment): route chat.completions through the Foundry # Responses API, which resolves the connection and auth (API key / managed identity / OAuth2). @@ -311,15 +315,15 @@ async def __call__( # pylint: disable=docstring-keyword-should-match-keyword-on else: connection = Connection.parse_from_config(configuration) if isinstance(connection, AzureOpenAIConnection): + if not connection.api_key: + token_provider = self.get_token_provider(self._token_credential) api_client = AsyncAzureOpenAI( azure_endpoint=connection.azure_endpoint, api_key=connection.api_key, azure_deployment=connection.azure_deployment, api_version=connection.api_version, max_retries=max_retries, - azure_ad_token_provider=( - self.get_token_provider(self._token_credential) if not connection.api_key else None - ), + azure_ad_token_provider=token_provider, default_headers=default_headers, ) elif isinstance(connection, OpenAIConnection): @@ -335,11 +339,17 @@ async def __call__( # pylint: disable=docstring-keyword-should-match-keyword-on f"'{type(connection).__name__}' is not a supported connection type.", target=ErrorTarget.EVAL_RUN ) - response: OpenAIChatResponseType = await self._send_with_retries( - api_client=api_client, - params=params, - timeout=timeout, - ) + try: + response: OpenAIChatResponseType = await self._send_with_retries( + api_client=api_client, + params=params, + timeout=timeout, + ) + except WrappedOpenAIError as error: + if token_provider is not None and isinstance(error.__cause__, AuthenticationError): + # A rejected token is not reused. + token_provider.discard() + raise return await format_llm_response( response=response, @@ -431,18 +441,75 @@ async def _send_with_retries( retry += 1 @staticmethod - def get_token_provider(cred: Union[TokenCredential, AsyncTokenCredential]) -> AsyncAzureADTokenProvider: + def get_token_provider(cred: Union[TokenCredential, AsyncTokenCredential]) -> "_ReusedTokenProvider": """Get the token provider for the prompty. :param Union[TokenCredential, AsyncTokenCredential] cred: The Azure authentication credential. - :return: The token provider if a credential is provided, otherwise None. - :rtype: Optional[AsyncAzureADTokenProvider] + :return: A token provider for the chat client that reuses the credential's token. + :rtype: _ReusedTokenProvider """ + return _ReusedTokenProvider(cred, TokenScope.COGNITIVE_SERVICES_MANAGEMENT) - async def _wrapper() -> str: - token = cred.get_token(TokenScope.COGNITIVE_SERVICES_MANAGEMENT) - if isinstance(token, Awaitable): - token = await token - return token.token - return _wrapper +class _ReusedTokenProvider: + """Supply a credential's token to a chat client, reusing it until shortly before it expires. + + Concurrent requests that share a synchronous credential wait for one token request, on a worker thread so the + event loop is not blocked. An asynchronous credential, one whose ``get_token`` is a coroutine function or that + implements ``AsyncTokenCredential`` but not the synchronous context manager protocol, is asked on the caller's + event loop instead; a plain ``get_token`` that needs the running loop must therefore belong to an + ``AsyncTokenCredential``. The default token provider is represented by the synchronous credential it creates. The + request that made a failed token request raises the credential's error. When that failure was an authentication + failure, requests that waited on it raise a ``ClientAuthenticationError`` with the same message instead of each + requesting a token again; after any other failure, each waiting request makes its own attempt. A token the service + rejected is discarded, and an acquisition that began before the rejection is not stored. + """ + + def __init__(self, credential: Union[TokenCredential, AsyncTokenCredential], scope: str) -> None: + self._credential = credential + self._scope = scope + self._source: Any = credential + self._token: Optional[str] = None + + async def __call__(self) -> str: + credential: Any = self._credential + if isinstance(credential, AsyncAzureTokenProvider): + # Its tokens come from a synchronous credential, so concurrent requests can share one token request. + credential = await credential.get_credential() + self._source = credential + token = TOKENS.cached(credential, self._scope) + if token is None: + # An acquisition that begins before a service rejects a token may return that token; the generation it + # began under lets the cache refuse it. + generation = TOKENS.generation(credential, self._scope) + if is_async_credential(credential): + # An asynchronous credential's token request may need this event loop. + token = credential.get_token(self._scope) + if inspect.isawaitable(token): + token = await token + TOKENS.store(credential, self._scope, token, generation) + else: + abandoned = threading.Event() + try: + # A request waiting for another request's token holds a default-executor thread while it waits; + # past the executor's size, further requests queue for a thread before they reach the cache. + token = await asyncio.to_thread( + TOKENS.fetch, + credential, + self._scope, + abandoned=abandoned, + describe_failure=describe_credential_failure, + failure_type=ClientAuthenticationError, + ) + finally: + abandoned.set() + if inspect.isawaitable(token): + token = await token + TOKENS.store(credential, self._scope, token, generation) + self._token = token.token + return self._token + + def discard(self) -> None: + """Stop reusing the token this provider last supplied, which the service rejected.""" + if self._token is not None: + TOKENS.discard(self._source, self._scope, self._token) diff --git a/sdk/evaluation/azure-ai-evaluation/tests/unittests/test_decision_model_support.py b/sdk/evaluation/azure-ai-evaluation/tests/unittests/test_decision_model_support.py index 49c2b863c71e..dfec6c743e7d 100644 --- a/sdk/evaluation/azure-ai-evaluation/tests/unittests/test_decision_model_support.py +++ b/sdk/evaluation/azure-ai-evaluation/tests/unittests/test_decision_model_support.py @@ -5,12 +5,16 @@ """Offline tests for shared decision-model runners, base selection, and evaluation strategy.""" import asyncio +import gc import hashlib import json import logging import math +import pickle import sys import threading +import time +import weakref from copy import deepcopy from logging.handlers import BufferingHandler from concurrent.futures import ThreadPoolExecutor @@ -19,11 +23,13 @@ import pytest from openai.types.chat import ChatCompletion +from azure.core.credentials import AccessToken from azure.core.pipeline.policies import AsyncRedirectPolicy, AsyncRetryPolicy from azure.core.pipeline.transport import AsyncHttpTransport, HttpResponse -from azure.core.exceptions import ServiceRequestError +from azure.core.exceptions import ClientAuthenticationError, ServiceRequestError from azure.ai.evaluation import CoherenceEvaluator, FluencyEvaluator +from azure.ai.evaluation._constants import TokenScope from azure.ai.evaluation._evaluators._common import EvaluatorBase, ModelEvaluatorBase, PromptyEvaluatorBase from azure.ai.evaluation._evaluators._common import _model_runners from azure.ai.evaluation._evaluators._common._decision_strategy import ( @@ -138,6 +144,42 @@ def run(runner, **overrides): return asyncio.run(runner(**arguments)) +# Only bounds a hung test; no assertion depends on how long anything takes. +WAIT_SECONDS = 10 + + +def watch_token_queue(credential, requests): + """Return an event set once ``requests`` threads have queued for the credential's Decision token request. + + Each request reads the cache's attempt count before it queues, so a token request released after this event + completes while every request is queued, whatever the thread scheduling. + """ + entry = _model_runners._TOKENS._entry(credential, _model_runners._DECISION_TOKEN_SCOPE, create=True) + lock, queued, all_queued = entry.lock, set(), threading.Event() + + class QueueWatchingLock: + def acquire(self, *args, **kwargs): + queued.add(threading.get_ident()) + if len(queued) >= requests: + all_queued.set() + return lock.acquire(*args, **kwargs) + + def release(self): + lock.release() + + entry.lock = QueueWatchingLock() + return all_queued + + +class PicklableCredential: + def __init__(self): + self.requested = 0 + + def get_token(self, *_, **__): + self.requested += 1 + return AccessToken("picklable", int(time.time()) + 3600) + + class TestDecisionModelRunnerConfiguration: @pytest.mark.parametrize( "changes", @@ -541,6 +583,563 @@ def test_unsupported_question_type(self, client): assert not client.calls +class TestDecisionTokenReuse: + """A credential's token serves repeated and concurrent requests until shortly before it expires.""" + + CONFIG = {key: value for key, value in DECISION_CONFIG.items() if key != "api_key"} + + @staticmethod + def access_token(value, lifetime=3600): + return AccessToken(value, int(time.time()) + lifetime) + + @staticmethod + def authorizations(client): + return [call["headers"]["Authorization"] for call in client.calls] + + @pytest.mark.parametrize("asynchronous", [False, True]) + def test_fresh_token_is_reused_across_requests_and_runners(self, client, asynchronous): + client.response = FakeResponse(body=json.dumps(level_payload(4)).encode()) + token = self.access_token("shared-token") + credential = Mock() + credential.get_token = AsyncMock(return_value=token) if asynchronous else Mock(return_value=token) + first, second = DecisionModelRunner(self.CONFIG, credential), DecisionModelRunner(self.CONFIG, credential) + for runner in (first, second, first): + run(runner) + credential.get_token.assert_called_once_with("https://cognitiveservices.azure.com/.default") + assert self.authorizations(client) == ["Bearer shared-token"] * 3 + + @pytest.mark.parametrize("lifetime", [None, 60], ids=["no-expiry", "near-expiry"]) + def test_tokens_without_a_usable_expiry_are_requested_each_time(self, client, lifetime): + client.response = FakeResponse(body=json.dumps(level_payload(4)).encode()) + token = SimpleNamespace(token="per-request") if lifetime is None else self.access_token("per-request", lifetime) + credential = Mock(get_token=Mock(return_value=token)) + runner = DecisionModelRunner(self.CONFIG, credential) + run(runner) + run(runner) + assert credential.get_token.call_count == 2 + + def test_concurrent_requests_share_one_synchronous_token_request(self, client): + client.response = FakeResponse(body=json.dumps(level_payload(4)).encode()) + release = threading.Event() + requested = [] + + def get_token(*scopes, **_): + requested.append(scopes) + release.wait() + return self.access_token("single-request") + + credential = Mock(get_token=Mock(side_effect=get_token)) + all_queued = watch_token_queue(credential, 6) + runners = [DecisionModelRunner(self.CONFIG, credential) for _ in range(6)] + # Like evaluate(), each runner awaits on its own thread and event loop. + with ThreadPoolExecutor(max_workers=6) as pool: + try: + requests = [pool.submit(run, runner, timeout=2 * WAIT_SECONDS) for runner in runners] + assert all_queued.wait(WAIT_SECONDS) + finally: + release.set() + for request in requests: + request.result(2 * WAIT_SECONDS) + assert len(requested) == 1 + assert self.authorizations(client) == ["Bearer single-request"] * 6 + + @pytest.mark.parametrize("expiry", ["none", "near", "past", "zero"]) + def test_concurrent_requests_share_a_token_that_is_not_cached(self, client, expiry): + client.response = FakeResponse(body=json.dumps(level_payload(4)).encode()) + release = threading.Event() + requested = [] + + def get_token(*_, **__): + requested.append(1) + release.wait() + expires_on = {"none": None, "near": time.time() + 120, "past": time.time() - 1, "zero": 0}[expiry] + return SimpleNamespace(token="uncached", expires_on=expires_on) + + credential = Mock(get_token=Mock(side_effect=get_token)) + all_queued = watch_token_queue(credential, 6) + with ThreadPoolExecutor(max_workers=6) as pool: + try: + requests = [ + pool.submit(run, DecisionModelRunner(self.CONFIG, credential), timeout=2 * WAIT_SECONDS) + for _ in range(6) + ] + assert all_queued.wait(WAIT_SECONDS) + finally: + release.set() + for request in requests: + request.result(2 * WAIT_SECONDS) + assert len(requested) == 1 + assert self.authorizations(client) == ["Bearer uncached"] * 6 + assert _model_runners._TOKENS.cached(credential, _model_runners._DECISION_TOKEN_SCOPE) is None + + def test_decision_tokens_are_kept_apart_from_chat_tokens(self, client): + client.response = FakeResponse(body=json.dumps(level_payload(4)).encode()) + credential = Mock(get_token=Mock(return_value=self.access_token("decision"))) + run(DecisionModelRunner(self.CONFIG, credential)) + credential.get_token.assert_called_once_with("https://cognitiveservices.azure.com/.default") + # Chat requests key the cache with their own scope, so they never reuse a Decision token. + assert _model_runners._TOKENS.cached(credential, TokenScope.COGNITIVE_SERVICES_MANAGEMENT) is None + + def test_rejected_token_is_not_reused(self, client): + client.response = FakeResponse(status_code=401, body=b"{}") + tokens = iter([self.access_token("rejected"), self.access_token("renewed")]) + credential = Mock(get_token=Mock(side_effect=lambda *_, **__: next(tokens))) + runner = DecisionModelRunner(self.CONFIG, credential) + with pytest.raises(EvaluationException, match="HTTP 401"): + run(runner) + client.response = FakeResponse(body=json.dumps(level_payload(4)).encode()) + run(runner) + run(runner) + assert credential.get_token.call_count == 2 + assert self.authorizations(client) == ["Bearer rejected", "Bearer renewed", "Bearer renewed"] + + @pytest.mark.parametrize("status", [403, 429, 500]) + def test_a_token_is_kept_after_other_errors(self, client, status): + client.response = FakeResponse(status_code=status, body=b"{}") + credential = Mock(get_token=Mock(return_value=self.access_token("kept"))) + runner = DecisionModelRunner(self.CONFIG, credential, max_retries=0) + with pytest.raises(EvaluationException, match=f"HTTP {status}"): + run(runner) + client.response = FakeResponse(body=json.dumps(level_payload(4)).encode()) + run(runner) + credential.get_token.assert_called_once() + assert self.authorizations(client) == ["Bearer kept"] * 2 + + def test_a_token_the_credential_reissues_after_a_rejection_is_cached_again(self, client): + client.response = FakeResponse(status_code=401, body=b"{}") + credential = Mock(get_token=Mock(return_value=self.access_token("same"))) + runner = DecisionModelRunner(self.CONFIG, credential) + with pytest.raises(EvaluationException, match="HTTP 401"): + run(runner) + assert _model_runners._TOKENS.cached(credential, _model_runners._DECISION_TOKEN_SCOPE) is None + # The service accepts the token now, for example after a role assignment propagated. The credential + # re-issues the same token: it is requested once more, then reused again instead of once per request. + client.response = FakeResponse(body=json.dumps(level_payload(4)).encode()) + for _ in range(3): + run(runner) + assert credential.get_token.call_count == 2 + assert self.authorizations(client) == ["Bearer same"] * 4 + assert _model_runners._TOKENS.cached(credential, _model_runners._DECISION_TOKEN_SCOPE).token == "same" + + def test_an_acquisition_that_began_before_two_rejections_is_not_stored(self): + scope = _model_runners._DECISION_TOKEN_SCOPE + + async def scenario(): + entered, gate = asyncio.Event(), asyncio.Event() + calls = [] + + class AsyncCredential: + async def get_token(self, *_, **__): + calls.append(1) + call = len(calls) + if call == 1: + entered.set() + await gate.wait() + # The delayed first acquisition and the second return token A; later ones return token B. + return TestDecisionTokenReuse.access_token("token-A" if call <= 2 else "token-B") + + credential = AsyncCredential() + runner = DecisionModelRunner(self.CONFIG, credential) + deadline = time.monotonic() + WAIT_SECONDS + pending = asyncio.ensure_future(runner._get_token(deadline)) + await entered.wait() + # Two later acquisitions complete first, and the service rejects each token with a 401. + for expected in ("token-A", "token-B"): + assert await runner._get_token(deadline) == expected + _model_runners._TOKENS.discard(credential, scope, expected) + gate.set() + # The delayed acquisition returns a token the service already rejected; it must not be cached. + assert await pending == "token-A" + assert _model_runners._TOKENS.cached(credential, scope) is None + # The next acquisition is cached again. + assert await runner._get_token(deadline) == "token-B" + assert _model_runners._TOKENS.cached(credential, scope).token == "token-B" + + asyncio.run(scenario()) + + def test_credential_without_weak_references_is_used_without_caching(self, client): + client.response = FakeResponse(body=json.dumps(level_payload(4)).encode()) + requested = [] + credential = SimpleNamespace(get_token=lambda *scopes, **_: requested.append(scopes) or self.access_token("t")) + runner = DecisionModelRunner(self.CONFIG, credential) + run(runner) + run(runner) + assert len(requested) == 2 + + def test_unhashable_credential_is_cached_by_identity(self, client): + from dataclasses import dataclass + + @dataclass + class Credential: + requested: int = 0 + + def get_token(self, *_, **__): + self.requested += 1 + return TestDecisionTokenReuse.access_token("by-identity") + + client.response = FakeResponse(body=json.dumps(level_payload(4)).encode()) + credential = Credential() + runner = DecisionModelRunner(self.CONFIG, credential) + run(runner) + run(runner) + assert credential.requested == 1 + + def test_equal_credentials_keep_their_own_tokens(self, client): + class Credential: + def __init__(self, token): + self.token = token + + def __eq__(self, other): + return isinstance(other, Credential) + + def __hash__(self): + return 0 + + def get_token(self, *_, **__): + return TestDecisionTokenReuse.access_token(self.token) + + client.response = FakeResponse(body=json.dumps(level_payload(4)).encode()) + first, second = Credential("first-identity"), Credential("second-identity") + assert first == second + run(DecisionModelRunner(self.CONFIG, first)) + run(DecisionModelRunner(self.CONFIG, second)) + assert self.authorizations(client) == ["Bearer first-identity", "Bearer second-identity"] + + def test_request_whose_deadline_passes_while_waiting_does_not_request_a_token(self, client): + client.response = FakeResponse(body=json.dumps(level_payload(4)).encode()) + entered, release = threading.Event(), threading.Event() + requested = [] + + def get_token(*_, **__): + requested.append(1) + entered.set() + release.wait() + raise RuntimeError("Azure CLI is not logged in") + + credential = Mock(get_token=Mock(side_effect=get_token)) + outcomes = {} + + def request(name, timeout): + try: + run(DecisionModelRunner(self.CONFIG, credential), timeout=timeout) + except EvaluationException as error: + outcomes[name] = str(error) + + leader = threading.Thread(target=request, args=("leader", 2 * WAIT_SECONDS)) + waiter = threading.Thread(target=request, args=("waiter", 0.3)) + try: + leader.start() + assert entered.wait(WAIT_SECONDS) + # The first request holds the credential's token request until the second request's deadline has passed. + waiter.start() + waiter.join(WAIT_SECONDS) + waiter_finished = not waiter.is_alive() + finally: + release.set() + leader.join(WAIT_SECONDS) + assert waiter_finished and not leader.is_alive() + assert len(requested) == 1 + assert "timed out" in outcomes["waiter"] + assert "Could not acquire a token" in outcomes["leader"] + + @pytest.mark.parametrize( + "failure,token_requests,blame,message", + [ + pytest.param( + lambda: ClientAuthenticationError("Please run 'az login' (private-detail)"), + 1, + ErrorBlame.USER_ERROR, + "Could not acquire a token", + id="authentication-shared-at-once", + ), + pytest.param( + lambda: EvaluationException( + message="Service-owned Decision credential acquisition failed (RuntimeError, status 5xx).", + blame=ErrorBlame.SYSTEM_ERROR, + category=ErrorCategory.FAILED_EXECUTION, + target=ErrorTarget.EVALUATE, + ), + 4, + ErrorBlame.SYSTEM_ERROR, + "acquisition failed", + id="other-not-shared-each-request-tries", + ), + ], + ) + def test_requests_waiting_on_a_failed_token_request_share_its_failure( + self, client, failure, token_requests, blame, message + ): + client.response = FakeResponse(body=json.dumps(level_payload(4)).encode()) + entered, release = threading.Event(), threading.Event() + requested = [] + + def get_token(*_, **__): + requested.append(1) + entered.set() + release.wait() + raise failure() + + credential = Mock(get_token=Mock(side_effect=get_token)) + all_queued = watch_token_queue(credential, 4) + errors = {} + + def request(name): + try: + run(DecisionModelRunner(self.CONFIG, credential), timeout=2 * WAIT_SECONDS) + except EvaluationException as error: + errors[name] = error + + leader = threading.Thread(target=request, args=("leader",)) + waiters = [threading.Thread(target=request, args=(f"waiter-{index}",)) for index in range(3)] + try: + leader.start() + assert entered.wait(WAIT_SECONDS) + for waiter in waiters: + waiter.start() + assert all_queued.wait(WAIT_SECONDS) + finally: + release.set() + for thread in [leader, *waiters]: + thread.join(WAIT_SECONDS) + # Queued requests share an authentication failure instead of each starting another request; any other failure is + # not shared, so each request makes its own attempt and keeps its own typed error. + assert len(requested) == token_requests and len(errors) == 4 + assert all(error.blame == blame and message in str(error) for error in errors.values()) + assert all("private-detail" not in str(error) for error in errors.values()) + credential.get_token.side_effect = lambda *_, **__: self.access_token("recovered") + run(DecisionModelRunner(self.CONFIG, credential)) + assert credential.get_token.call_count == token_requests + 1 and self.authorizations(client) == [ + "Bearer recovered" + ] + + def test_requests_waiting_on_a_transient_token_failure_recover_with_one_more_request(self, client): + client.response = FakeResponse(body=json.dumps(level_payload(4)).encode()) + entered, release = threading.Event(), threading.Event() + requested = [] + + def get_token(*_, **__): + requested.append(1) + if len(requested) == 1: + entered.set() + release.wait() + raise RuntimeError("token service unavailable") + return self.access_token("recovered") + + credential = Mock(get_token=Mock(side_effect=get_token)) + all_queued = watch_token_queue(credential, 4) + outcomes = {} + + def request(name): + try: + run(DecisionModelRunner(self.CONFIG, credential), timeout=2 * WAIT_SECONDS) + outcomes[name] = "completed" + except EvaluationException as error: + outcomes[name] = str(error) + + leader = threading.Thread(target=request, args=("leader",)) + waiters = [threading.Thread(target=request, args=(f"waiter-{index}",)) for index in range(3)] + try: + leader.start() + assert entered.wait(WAIT_SECONDS) + for waiter in waiters: + waiter.start() + assert all_queued.wait(WAIT_SECONDS) + finally: + release.set() + for thread in [leader, *waiters]: + thread.join(WAIT_SECONDS) + # Only the request whose token request failed fails; the next request recovers the token for the others. + assert len(requested) == 2 + assert "Could not acquire a token" in outcomes["leader"] + assert [outcomes[f"waiter-{index}"] for index in range(3)] == ["completed"] * 3 + assert self.authorizations(client) == ["Bearer recovered"] * 3 + + def test_requests_waiting_on_a_token_kept_out_of_the_cache_use_it(self, client): + client.response = FakeResponse(body=json.dumps(level_payload(4)).encode()) + entered, release = threading.Event(), threading.Event() + requested = [] + + def get_token(*_, **__): + requested.append(1) + entered.set() + release.wait() + return self.access_token("issued") + + credential = Mock(get_token=Mock(side_effect=get_token)) + all_queued = watch_token_queue(credential, 4) + scope = _model_runners._DECISION_TOKEN_SCOPE + with ThreadPoolExecutor(max_workers=4) as pool: + try: + requests = [ + pool.submit(run, DecisionModelRunner(self.CONFIG, credential), timeout=2 * WAIT_SECONDS) + for _ in range(4) + ] + assert entered.wait(WAIT_SECONDS) and all_queued.wait(WAIT_SECONDS) + # Another request's 401 lands while the token is requested, so the token is not cached. + _model_runners._TOKENS.discard(credential, scope, "rejected-earlier") + finally: + release.set() + for request in requests: + request.result(2 * WAIT_SECONDS) + assert len(requested) == 1 + assert self.authorizations(client) == ["Bearer issued"] * 4 + assert _model_runners._TOKENS.cached(credential, scope) is None + # A request that starts afterwards requests its own token, which is cached again. + run(DecisionModelRunner(self.CONFIG, credential)) + assert len(requested) == 2 + assert _model_runners._TOKENS.cached(credential, scope).token == "issued" + + def test_a_token_rejected_before_its_request_resumes_is_not_stored_again(self, client, monkeypatch): + client.response = FakeResponse(body=json.dumps(level_payload(4)).encode()) + credential = Mock(get_token=Mock(return_value=self.access_token("shared"))) + fetch = _model_runners._TOKENS.fetch + scope = _model_runners._DECISION_TOKEN_SCOPE + + def fetch_then_reject(*args, **kwargs): + token = fetch(*args, **kwargs) + # Another request's 401 discards the token before this request's event loop resumes. + _model_runners._TOKENS.discard(credential, scope, token.token) + return token + + monkeypatch.setattr(_model_runners._TOKENS, "fetch", fetch_then_reject) + run(DecisionModelRunner(self.CONFIG, credential)) + assert _model_runners._TOKENS.cached(credential, scope) is None + + def test_failed_token_request_does_not_keep_the_credential_alive(self, client): + class Credential: + def get_token(self, *_, **__): + raise RuntimeError("Azure CLI is not logged in") + + credential = Credential() + with pytest.raises(EvaluationException, match="Could not acquire a token"): + run(DecisionModelRunner(self.CONFIG, credential)) + reference = weakref.ref(credential) + del credential + gc.collect() + assert reference() is None + + def test_cancelled_request_stops_waiting_for_another_token_request(self, client): + client.response = FakeResponse(body=json.dumps(level_payload(4)).encode()) + entered, release = threading.Event(), threading.Event() + requested = [] + + def get_token(*_, **__): + requested.append(1) + entered.set() + release.wait() + return self.access_token("late") + + credential = Mock(get_token=Mock(side_effect=get_token)) + all_queued = watch_token_queue(credential, 2) + cancelled = threading.Event() + outcomes = {} + + def waiter(): + try: + run( + DecisionModelRunner(self.CONFIG, credential, cancellation_requested=cancelled.is_set), + timeout=2 * WAIT_SECONDS, + ) + except EvaluationException as error: + outcomes["waiter"] = str(error) + + leader = threading.Thread( + target=lambda: run(DecisionModelRunner(self.CONFIG, credential), timeout=2 * WAIT_SECONDS) + ) + waiting = threading.Thread(target=waiter) + try: + leader.start() + assert entered.wait(WAIT_SECONDS) + waiting.start() + # Cancel only once the second request waits for the first request's token. + assert all_queued.wait(WAIT_SECONDS) + cancelled.set() + # The first request's token request stays blocked until ``finally``, so the cancelled request can only + # return by giving up its wait. + waiting.join(WAIT_SECONDS) + still_waiting = waiting.is_alive() + finally: + release.set() + leader.join(WAIT_SECONDS) + waiting.join(WAIT_SECONDS) + assert not still_waiting and "cancelled" in outcomes["waiter"] + assert len(requested) == 1 + + def test_cache_does_not_keep_the_credential_alive(self, client): + client.response = FakeResponse(body=json.dumps(level_payload(4)).encode()) + token = self.access_token("caller-owned") + + class Credential: + def get_token(self, *_, **__): + return token + + credential = Credential() + run(DecisionModelRunner(self.CONFIG, credential)) + assert _model_runners._TOKENS.cached(credential, _model_runners._DECISION_TOKEN_SCOPE) is token + reference = weakref.ref(credential) + del credential + gc.collect() + assert reference() is None + + def test_a_runner_still_pickles_after_its_token_is_cached(self, client): + client.response = FakeResponse(body=json.dumps(level_payload(4)).encode()) + credential = PicklableCredential() + runner = DecisionModelRunner(self.CONFIG, credential) + run(runner) + copy = pickle.loads(pickle.dumps(runner)) + # The copy's credential is another object, so it requests its own token once and then reuses it. + run(copy) + run(copy) + run(runner) + assert credential.requested == 1 + assert copy._credential.requested == 2 + assert self.authorizations(client) == ["Bearer picklable"] * 4 + + @pytest.mark.parametrize("asynchronous", [False, True]) + @pytest.mark.parametrize( + "token", + [None, SimpleNamespace(token=None, expires_on=None), SimpleNamespace(token=123, expires_on=None)], + ids=["none", "no-token", "non-string-token"], + ) + def test_a_credential_that_returns_no_usable_token_is_a_user_error(self, client, token, asynchronous): + credential = Mock() + credential.get_token = AsyncMock(return_value=token) if asynchronous else Mock(return_value=token) + with pytest.raises(EvaluationException, match="Could not acquire a token") as raised: + run(DecisionModelRunner(self.CONFIG, credential)) + assert raised.value.blame == ErrorBlame.USER_ERROR + assert not client.calls + + def test_a_token_wait_ends_at_the_deadline_without_requesting_a_token(self): + credential = Mock(get_token=Mock(return_value=self.access_token("unused"))) + scope = _model_runners._DECISION_TOKEN_SCOPE + entry = _model_runners._TOKENS._entry(credential, scope, create=True) + runner = DecisionModelRunner(self.CONFIG, credential) + # Another request holds the credential's token request throughout. + assert entry.lock.acquire(timeout=WAIT_SECONDS) + try: + assert _model_runners._TOKENS.fetch(credential, scope, deadline=time.monotonic() + 0.1) is None + assert asyncio.run(runner._get_token(time.monotonic() + 0.1)) is None + finally: + entry.lock.release() + credential.get_token.assert_not_called() + + def test_a_credential_without_weak_references_is_not_asked_after_the_deadline(self): + requested = [] + credential = SimpleNamespace(get_token=lambda *scopes, **_: requested.append(scopes)) + scope = _model_runners._DECISION_TOKEN_SCOPE + + assert _model_runners._TOKENS.fetch(credential, scope, deadline=time.monotonic() - 1) is None + assert requested == [] + + def test_a_request_whose_token_wait_reached_the_deadline_times_out(self, client, monkeypatch): + runner = DecisionModelRunner(self.CONFIG, Mock(get_token=Mock())) + monkeypatch.setattr(runner, "_get_token", AsyncMock(return_value=None)) + with pytest.raises(EvaluationException, match="timed out") as raised: + run(runner) + assert raised.value.blame == ErrorBlame.SYSTEM_ERROR + assert not client.calls + + class TestAnswerValidation: def test_noul(self): assert validate_noul_answer({"type": "noul", "noul": 0.25}) == 0.25 diff --git a/sdk/evaluation/azure-ai-evaluation/tests/unittests/test_prompty_token_reuse.py b/sdk/evaluation/azure-ai-evaluation/tests/unittests/test_prompty_token_reuse.py new file mode 100644 index 000000000000..ad568952ce2f --- /dev/null +++ b/sdk/evaluation/azure-ai-evaluation/tests/unittests/test_prompty_token_reuse.py @@ -0,0 +1,891 @@ +# --------------------------------------------------------- +# Copyright (c) Microsoft Corporation. All rights reserved. +# --------------------------------------------------------- + +"""The chat flow's token provider reuses a caller credential's token, kept apart from Decision tokens.""" + +import asyncio +import gc +import threading +import time +import weakref +from pathlib import Path +from types import SimpleNamespace +from unittest.mock import MagicMock + +import httpx +import pytest +from azure.core.credentials import AccessToken +from azure.core.exceptions import ClientAuthenticationError +from openai import AuthenticationError, InternalServerError + +import azure.ai.evaluation._evaluators._coherence as coherence +from azure.ai.evaluation._evaluators._common import _model_runners +from azure.ai.evaluation._common.utils import construct_prompty_model_config, validate_model_config +from azure.ai.evaluation._constants import TokenScope +from azure.ai.evaluation._exceptions import EvaluationException +from azure.ai.evaluation._legacy._common._async_token_provider import AsyncAzureTokenProvider +from azure.ai.evaluation._legacy._common._token_cache import TOKENS, describe_credential_failure, is_async_credential +from azure.ai.evaluation._legacy.prompty import AsyncPrompty +from azure.ai.evaluation._legacy.prompty._exceptions import WrappedOpenAIError + +CHAT_SCOPE = TokenScope.COGNITIVE_SERVICES_MANAGEMENT +DECISION_SCOPE = _model_runners._DECISION_TOKEN_SCOPE +COHERENCE_PROMPTY = Path(coherence.__file__).parent / "coherence.prompty" +CHAT_MODEL = construct_prompty_model_config( + validate_model_config({"azure_endpoint": "https://example.openai.azure.com", "azure_deployment": "judge"}), + "2024-10-21", + "test-agent", +) +# Only bounds a hung test; no assertion depends on how long anything takes. +WAIT_SECONDS = 10 + + +def access_token(value): + return AccessToken(value, int(time.time()) + 3600) + + +def request_token(credential): + return asyncio.run(AsyncPrompty.get_token_provider(credential)()) + + +def join_all(threads): + for thread in threads: + thread.join(WAIT_SECONDS) + assert not any(thread.is_alive() for thread in threads) + + +def watch_token_queue(credential, requests): + """Return an event set once ``requests`` threads have queued for the credential's chat token request. + + Each request reads the cache's attempt count before it queues, so a token request released after this event + completes while every request is queued, whatever the thread scheduling. + """ + entry = TOKENS._entry(credential, CHAT_SCOPE, create=True) + lock, queued, all_queued = entry.lock, set(), threading.Event() + + class QueueWatchingLock: + def acquire(self, *args, **kwargs): + queued.add(threading.get_ident()) + if len(queued) >= requests: + all_queued.set() + return lock.acquire(*args, **kwargs) + + def release(self): + lock.release() + + entry.lock = QueueWatchingLock() + return all_queued + + +class Credential: + """A synchronous credential that records each token request and answers from a queue of results. + + With ``gate``, each token request waits until the test sets the event; tests set it in ``finally``. + """ + + def __init__(self, *results, gate=None): + self.scopes = [] + self._results = list(results) + self._gate = gate + self._lock = threading.Lock() + + def get_token(self, *scopes, **_): + with self._lock: + self.scopes.append(scopes) + result = self._results.pop(0) if len(self._results) > 1 else self._results[0] + if self._gate is not None: + self._gate.wait() + if isinstance(result, Exception): + raise result + return result + + +@pytest.mark.unittest +class TestChatTokenReuse: + def test_a_fresh_token_is_reused_across_requests_and_flows(self): + credential = Credential(access_token("chat")) + + assert [request_token(credential) for _ in range(3)] == ["chat"] * 3 + assert asyncio.run(AsyncPrompty.get_token_provider(credential)()) == "chat" + assert credential.scopes == [(CHAT_SCOPE,)] + + def test_concurrent_requests_wait_for_one_token_request(self): + release = threading.Event() + credential = Credential(access_token("shared"), gate=release) + all_queued = watch_token_queue(credential, 8) + tokens = [] + threads = [threading.Thread(target=lambda: tokens.append(request_token(credential))) for _ in range(8)] + try: + for thread in threads: + thread.start() + assert all_queued.wait(WAIT_SECONDS) + finally: + release.set() + join_all(threads) + + assert tokens == ["shared"] * 8 + assert credential.scopes == [(CHAT_SCOPE,)] + + @pytest.mark.parametrize("expiry", ["none", "near", "past", "zero"]) + def test_concurrent_requests_share_a_token_that_is_not_cached(self, expiry): + expires_on = {"none": None, "near": time.time() + 120, "past": time.time() - 1, "zero": 0}[expiry] + issued = SimpleNamespace(token="uncached", expires_on=expires_on) + release = threading.Event() + credential = Credential(issued, gate=release) + all_queued = watch_token_queue(credential, 6) + tokens = [] + threads = [threading.Thread(target=lambda: tokens.append(request_token(credential))) for _ in range(6)] + try: + for thread in threads: + thread.start() + assert all_queued.wait(WAIT_SECONDS) + finally: + release.set() + join_all(threads) + + assert tokens == ["uncached"] * 6 + assert len(credential.scopes) == 1 + assert TOKENS.cached(credential, CHAT_SCOPE) is None + # A request that did not wait on that token request makes its own. + assert request_token(credential) == "uncached" + assert len(credential.scopes) == 2 + + def test_a_token_without_a_value_is_not_given_to_waiting_requests(self): + release = threading.Event() + credential = Credential(SimpleNamespace(token=None, expires_on=None), gate=release) + all_queued = watch_token_queue(credential, 3) + outcomes = [] + threads = [threading.Thread(target=lambda: outcomes.append(request_token(credential))) for _ in range(3)] + try: + for thread in threads: + thread.start() + assert all_queued.wait(WAIT_SECONDS) + finally: + release.set() + join_all(threads) + + # Each request asks the credential itself and returns its value, as before. + assert outcomes == [None] * 3 + assert len(credential.scopes) == 3 + + def test_waiting_for_another_requests_token_does_not_block_the_event_loop(self): + entered, release = threading.Event(), threading.Event() + + class SlowCredential(Credential): + def get_token(self, *scopes, **kwargs): + entered.set() + return super().get_token(*scopes, **kwargs) + + credential = SlowCredential(access_token("shared"), gate=release) + all_queued = watch_token_queue(credential, 2) + leader = threading.Thread(target=lambda: request_token(credential)) + # Releases the token request if waiting blocked the loop, so a regression fails instead of hanging. + watchdog = threading.Timer(WAIT_SECONDS, release.set) + + async def scenario(): + waiting = asyncio.ensure_future(AsyncPrompty.get_token_provider(credential)()) + assert await asyncio.to_thread(all_queued.wait, WAIT_SECONDS) + # The loop ran this coroutine while the request waits for the first request's token. + assert not waiting.done() and not release.is_set() + release.set() + return await waiting + + try: + watchdog.start() + leader.start() + assert entered.wait(WAIT_SECONDS) + assert asyncio.run(scenario()) == "shared" + finally: + release.set() + watchdog.cancel() + join_all([leader]) + assert len(credential.scopes) == 1 + + def test_chat_and_decision_tokens_are_kept_apart(self): + credential = Credential(access_token("chat")) + TOKENS.store(credential, DECISION_SCOPE, access_token("decision")) + + assert request_token(credential) == "chat" + assert credential.scopes == [(CHAT_SCOPE,)] + assert TOKENS.cached(credential, DECISION_SCOPE).token == "decision" + + def test_a_failed_token_request_raises_the_credential_error_and_is_retried(self): + credential = Credential(ClientAuthenticationError("Please run az login."), access_token("recovered")) + + with pytest.raises(ClientAuthenticationError, match="az login"): + request_token(credential) + assert request_token(credential) == "recovered" + assert len(credential.scopes) == 2 + + def test_requests_queued_behind_a_failed_token_request_share_its_failure(self): + release = threading.Event() + credential = Credential( + ClientAuthenticationError("Please run az login."), access_token("recovered"), gate=release + ) + all_queued = watch_token_queue(credential, 4) + outcomes = [] + + def request(): + try: + outcomes.append(request_token(credential)) + except ClientAuthenticationError as error: + outcomes.append(str(error)) + + threads = [threading.Thread(target=request) for _ in range(4)] + try: + for thread in threads: + thread.start() + assert all_queued.wait(WAIT_SECONDS) + finally: + release.set() + join_all(threads) + + assert outcomes == ["Please run az login."] * 4 + assert len(credential.scopes) == 1 + # A request made after the failure tries again. + assert request_token(credential) == "recovered" + + def test_other_failures_are_not_shared_so_each_request_keeps_the_error_type(self): + class TokenServiceOutage(Exception): + pass + + release = threading.Event() + credential = Credential(TokenServiceOutage("The token service is unavailable."), gate=release) + all_queued = watch_token_queue(credential, 3) + outcomes = [] + + def request(): + try: + outcomes.append(request_token(credential)) + except Exception as error: # pylint: disable=broad-except + outcomes.append(type(error).__name__) + + threads = [threading.Thread(target=request) for _ in range(3)] + try: + for thread in threads: + thread.start() + assert all_queued.wait(WAIT_SECONDS) + finally: + release.set() + join_all(threads) + + assert outcomes == ["TokenServiceOutage"] * 3 + assert len(credential.scopes) == 3 + + def test_requests_waiting_on_a_token_kept_out_of_the_cache_use_it(self): + entered, release = threading.Event(), threading.Event() + + class SlowCredential(Credential): + def get_token(self, *scopes, **kwargs): + entered.set() + release.wait() + return super().get_token(*scopes, **kwargs) + + credential = SlowCredential(access_token("issued")) + all_queued = watch_token_queue(credential, 4) + tokens = [] + threads = [threading.Thread(target=lambda: tokens.append(request_token(credential))) for _ in range(4)] + try: + for thread in threads: + thread.start() + assert entered.wait(WAIT_SECONDS) and all_queued.wait(WAIT_SECONDS) + # Another request's 401 lands while the token is requested, so the token is not cached. + TOKENS.discard(credential, CHAT_SCOPE, "rejected-earlier") + finally: + release.set() + join_all(threads) + + assert tokens == ["issued"] * 4 + assert len(credential.scopes) == 1 + assert TOKENS.cached(credential, CHAT_SCOPE) is None + # A request that starts afterwards requests its own token, which is cached again. + assert request_token(credential) == "issued" + assert len(credential.scopes) == 2 + assert TOKENS.cached(credential, CHAT_SCOPE).token == "issued" + + def test_a_rejected_token_is_not_given_to_waiting_requests(self): + credential = Credential(access_token("unused")) + entry = TOKENS._entry(credential, CHAT_SCOPE, create=True) + entry.unstored = access_token("issued") + + TOKENS.discard(credential, CHAT_SCOPE, "other") + assert entry.unstored.token == "issued" + TOKENS.discard(credential, CHAT_SCOPE, "issued") + assert entry.unstored is None + + def test_a_rejection_cannot_land_between_caching_and_the_hand_off_to_waiting_requests(self, monkeypatch): + credential = Credential(access_token("issued")) + store = TOKENS.store + guard_free = [] + + def probe_guard(): + # A rejection from another request needs this lock. + acquired = TOKENS._guard.acquire(blocking=False) + if acquired: + TOKENS._guard.release() + guard_free.append(acquired) + + def store_then_probe(*args, **kwargs): + # A rejection while the token was requested keeps it out of the cache for waiting requests. + TOKENS.discard(credential, CHAT_SCOPE, "rejected-earlier") + stored = store(*args, **kwargs) + probe = threading.Thread(target=probe_guard) + probe.start() + join_all([probe]) + return stored + + monkeypatch.setattr(TOKENS, "store", store_then_probe) + assert request_token(credential) == "issued" + + # The decision to cache and the hand-off happen while the lock is held, so no rejection lands between them. + assert guard_free == [False] + assert TOKENS._entry(credential, CHAT_SCOPE).unstored.token == "issued" + # A rejection afterwards still reaches the handed-off token. + TOKENS.discard(credential, CHAT_SCOPE, "issued") + assert TOKENS._entry(credential, CHAT_SCOPE).unstored is None + + def test_requests_waiting_on_a_failed_token_request_do_not_get_an_older_token(self): + class TokenServiceOutage(Exception): + pass + + entered, first_release, second_release = threading.Event(), threading.Event(), threading.Event() + calls = [] + + class StagedCredential(Credential): + def get_token(self, *scopes, **kwargs): + calls.append(1) + if len(calls) == 1: + entered.set() + first_release.wait() + elif len(calls) == 2: + second_release.wait() + return super().get_token(*scopes, **kwargs) + + credential = StagedCredential( + access_token("issued"), TokenServiceOutage("The token service is unavailable."), access_token("recovered") + ) + results = {} + + def request(name): + try: + results[name] = request_token(credential) + except Exception as error: # pylint: disable=broad-except + results[name] = type(error).__name__ + + # A rejection while the first token is requested keeps that token out of the cache. + first = threading.Thread(target=request, args=("first",)) + try: + first.start() + assert entered.wait(WAIT_SECONDS) + TOKENS.discard(credential, CHAT_SCOPE, "rejected-earlier") + finally: + first_release.set() + join_all([first]) + assert results["first"] == "issued" and TOKENS.cached(credential, CHAT_SCOPE) is None + + # The next token request fails; the requests waiting on it make their own attempt. + all_queued = watch_token_queue(credential, 3) + names = ["second", "waiter-1", "waiter-2"] + threads = [threading.Thread(target=request, args=(name,)) for name in names] + try: + for thread in threads: + thread.start() + assert all_queued.wait(WAIT_SECONDS) + finally: + second_release.set() + join_all(threads) + + assert sorted(results[name] for name in names) == ["TokenServiceOutage", "recovered", "recovered"] + assert len(calls) == 3 + + def test_a_token_without_an_expiry_is_requested_each_time(self): + credential = Credential(SimpleNamespace(token="no-expiry", expires_on=None)) + + assert [request_token(credential) for _ in range(2)] == ["no-expiry"] * 2 + assert len(credential.scopes) == 2 + + def test_an_asynchronous_credentials_token_is_reused(self): + class AsyncCredential: + def __init__(self): + self.calls = 0 + + async def get_token(self, *_, **__): + self.calls += 1 + return access_token("async") + + credential = AsyncCredential() + + assert [request_token(credential) for _ in range(3)] == ["async"] * 3 + assert credential.calls == 1 + + def test_the_cache_does_not_keep_the_credential_alive(self): + credential = Credential(access_token("caller-owned")) + assert request_token(credential) == "caller-owned" + assert TOKENS.cached(credential, CHAT_SCOPE) is not None + + reference = weakref.ref(credential) + del credential + gc.collect() + assert reference() is None + + @pytest.mark.parametrize("default_provider", [False, True], ids=["caller credential", "default provider"]) + @pytest.mark.parametrize( + "error_type, status, tokens_sent", + [(AuthenticationError, 401, ["rejected", "replacement"]), (InternalServerError, 500, ["rejected", "rejected"])], + ) + def test_only_a_token_rejected_with_401_is_requested_again( + self, monkeypatch, error_type, status, tokens_sent, default_provider + ): + credential = Credential(access_token("rejected"), access_token("replacement")) + if default_provider: + provider = AsyncAzureTokenProvider() + provider._credential = credential + token_credential = provider + else: + token_credential = credential + flow = AsyncPrompty.load(source=COHERENCE_PROMPTY, model=CHAT_MODEL, token_credential=token_credential) + sent = [] + + async def send(self, *, api_client, params, timeout): + sent.append(await api_client._azure_ad_token_provider()) + response = httpx.Response(status, request=httpx.Request("POST", "https://example.openai.azure.com")) + error = error_type("request failed", response=response, body=None) + raise WrappedOpenAIError(error=error) from error + + monkeypatch.setattr(AsyncPrompty, "_send_with_retries", send) + for _ in range(2): + with pytest.raises(WrappedOpenAIError): + asyncio.run(flow(query="When does the museum open?", response="At noon.")) + + assert sent == tokens_sent + assert len(credential.scopes) == len(set(tokens_sent)) + + def test_the_default_provider_reuses_its_token(self): + provider = AsyncAzureTokenProvider() + inner = Credential(access_token("default")) + provider._credential = inner + + assert [request_token(provider) for _ in range(3)] == ["default"] * 3 + assert inner.scopes == [(CHAT_SCOPE,)] + asyncio.run(provider.close()) + + def test_concurrent_requests_share_the_default_providers_token_request(self): + release = threading.Event() + provider = AsyncAzureTokenProvider() + inner = Credential(access_token("default"), gate=release) + provider._credential = inner + all_queued = watch_token_queue(inner, 6) + tokens = [] + threads = [threading.Thread(target=lambda: tokens.append(request_token(provider))) for _ in range(6)] + try: + for thread in threads: + thread.start() + assert all_queued.wait(WAIT_SECONDS) + finally: + release.set() + join_all(threads) + + assert tokens == ["default"] * 6 + assert inner.scopes == [(CHAT_SCOPE,)] + asyncio.run(provider.close()) + + def test_the_default_provider_creates_its_credential_once(self, monkeypatch): + inner = Credential(access_token("created")) + created = [] + + async def initialize(client): + created.append(client) + return inner + + monkeypatch.setattr(AsyncAzureTokenProvider, "_initialize_async", staticmethod(initialize)) + provider = AsyncAzureTokenProvider() + + assert [request_token(provider) for _ in range(3)] == ["created"] * 3 + assert len(created) == 1 and inner.scopes == [(CHAT_SCOPE,)] + # Its own get_token still asks the credential it created. + assert asyncio.run(provider.get_token(CHAT_SCOPE)).token == "created" + assert len(inner.scopes) == 2 + asyncio.run(provider.close()) + + def test_concurrent_first_requests_share_one_default_credential_and_token_request(self, monkeypatch): + requests = 4 + barrier = threading.Barrier(requests) + created, closed = [], [] + + class ClosableCredential(Credential): + def close(self): + closed.append(self) + + async def initialize(client): + # Every first request creates a credential at the same time. + await asyncio.to_thread(barrier.wait, WAIT_SECONDS) + credential = ClosableCredential(access_token("default")) + created.append(credential) + return credential + + monkeypatch.setattr(AsyncAzureTokenProvider, "_initialize_async", staticmethod(initialize)) + provider = AsyncAzureTokenProvider() + tokens = [] + threads = [threading.Thread(target=lambda: tokens.append(request_token(provider))) for _ in range(requests)] + for thread in threads: + thread.start() + join_all(threads) + + assert tokens == ["default"] * requests + # The first credential created is kept and the others are closed, so one token request serves every row. + assert len(created) == requests and len(closed) == requests - 1 + assert sum(len(credential.scopes) for credential in created) == 1 + asyncio.run(provider.close()) + + def test_a_credential_created_while_the_default_provider_closes_is_closed(self, monkeypatch): + closed = [] + + class ClosableCredential(Credential): + def close(self): + closed.append(self) + + created = ClosableCredential(access_token("late")) + + async def scenario(): + entered, gate = asyncio.Event(), asyncio.Event() + + async def initialize(client): + entered.set() + await gate.wait() + return created + + monkeypatch.setattr(AsyncAzureTokenProvider, "_initialize_async", staticmethod(initialize)) + provider = AsyncAzureTokenProvider() + pending = asyncio.ensure_future(provider.get_credential()) + await entered.wait() + await provider.close() + gate.set() + with pytest.raises(EvaluationException, match="already been closed"): + await pending + assert provider._credential is None + + asyncio.run(scenario()) + assert closed == [created] + + def test_a_token_rejected_while_another_request_awaits_it_is_not_stored_again(self): + async def scenario(): + entered, gate = asyncio.Event(), asyncio.Event() + calls = [] + + class AsyncCredential: + async def get_token(self, *_, **__): + calls.append(1) + if len(calls) == 1: + entered.set() + await gate.wait() + return access_token("rejected") + + credential = AsyncCredential() + waiting = AsyncPrompty.get_token_provider(credential) + rejected = AsyncPrompty.get_token_provider(credential) + pending = asyncio.ensure_future(waiting()) + # The first request misses the cache and awaits its token. + await entered.wait() + assert await rejected() == "rejected" + # The service rejects that token with a 401 before the first request resumes. + rejected.discard() + gate.set() + assert await pending == "rejected" + assert TOKENS.cached(credential, CHAT_SCOPE) is None + # A different token from the credential is cached again. + TOKENS.store(credential, CHAT_SCOPE, access_token("new")) + assert TOKENS.cached(credential, CHAT_SCOPE).token == "new" + + asyncio.run(scenario()) + + def test_a_token_the_credential_reissues_after_a_rejection_is_cached_again(self): + credential = Credential(access_token("same")) + + async def reject_once(): + provider = AsyncPrompty.get_token_provider(credential) + assert await provider() == "same" + # The service rejects the token with a 401, for example while a role assignment still propagates. + provider.discard() + + asyncio.run(reject_once()) + assert TOKENS.cached(credential, CHAT_SCOPE) is None + # The credential re-issues the same token: it is requested once more, then reused again. + assert [request_token(credential) for _ in range(3)] == ["same"] * 3 + assert len(credential.scopes) == 2 + assert TOKENS.cached(credential, CHAT_SCOPE).token == "same" + + def test_an_acquisition_that_began_before_two_rejections_is_not_stored(self): + async def scenario(): + entered, gate = asyncio.Event(), asyncio.Event() + calls = [] + + class AsyncCredential: + async def get_token(self, *_, **__): + calls.append(1) + call = len(calls) + if call == 1: + entered.set() + await gate.wait() + # The delayed first acquisition and the second return token A; later ones return token B. + return access_token("token-A" if call <= 2 else "token-B") + + credential = AsyncCredential() + pending = asyncio.ensure_future(AsyncPrompty.get_token_provider(credential)()) + await entered.wait() + # Two later acquisitions complete first, and the service rejects each token with a 401. + for expected in ("token-A", "token-B"): + provider = AsyncPrompty.get_token_provider(credential) + assert await provider() == expected + provider.discard() + gate.set() + # The delayed acquisition returns a token the service already rejected; it must not be cached. + assert await pending == "token-A" + assert TOKENS.cached(credential, CHAT_SCOPE) is None + # The next acquisition is cached again. + assert await AsyncPrompty.get_token_provider(credential)() == "token-B" + assert TOKENS.cached(credential, CHAT_SCOPE).token == "token-B" + + asyncio.run(scenario()) + + +DECISION_CONFIG = { + "type": "decision", + "azure_endpoint": "https://example.services.ai.azure.com", + "azure_deployment": "Decision-1", +} + + +def request_decision_token(credential): + runner = _model_runners.DecisionModelRunner(DECISION_CONFIG, credential) + return asyncio.run(runner._get_token(time.monotonic() + 2 * WAIT_SECONDS)) + + +class GatedCredential(Credential): + """A ``Credential`` whose token requests wait until the test sets ``release``, and that reports the first one.""" + + def __init__(self, *results): + self.entered, self.release = threading.Event(), threading.Event() + super().__init__(*results, gate=self.release) + + def get_token(self, *scopes, **kwargs): + self.entered.set() + return super().get_token(*scopes, **kwargs) + + +def queue_behind(credential, leader, *waiters): + """Return the outcomes of ``leader``'s token request and of each of ``waiters``, all queued behind it.""" + requests = [leader, *waiters] + all_queued = watch_token_queue(credential, len(requests)) + outcomes = [None] * len(requests) + + def request(index): + try: + outcomes[index] = requests[index](credential) + except Exception as error: # pylint: disable=broad-except + outcomes[index] = error + + threads = [threading.Thread(target=request, args=(index,)) for index in range(len(requests))] + try: + threads[0].start() + assert credential.entered.wait(WAIT_SECONDS) + for thread in threads[1:]: + thread.start() + assert all_queued.wait(WAIT_SECONDS) + finally: + credential.release.set() + join_all(threads) + return outcomes + + +def fetch_as(describe_failure, failure_type): + """Return a request for the credential's chat-scope token that describes and raises failures its own way.""" + + def request(credential): + return TOKENS.fetch(credential, CHAT_SCOPE, describe_failure=describe_failure, failure_type=failure_type).token + + return request + + +# Two kinds of request that describe and raise failures differently, as chat and Decision requests do. +AUTHENTICATION_KIND = fetch_as(describe_credential_failure, ClientAuthenticationError) +EVALUATION_KIND = fetch_as(_model_runners._token_failure, EvaluationException) +EVALUATION_MESSAGE = "Could not acquire a token for the decision model." + + +@pytest.mark.unittest +class TestFailuresAreSharedWithinOneKindOfRequest: + """Requests that describe failures differently share a credential's tokens, but not each other's failures.""" + + @pytest.mark.parametrize("leader", ["authentication", "evaluation"]) + def test_both_kinds_of_request_share_one_token_request(self, leader): + kinds = {"authentication": AUTHENTICATION_KIND, "evaluation": EVALUATION_KIND} + waiter = "evaluation" if leader == "authentication" else "authentication" + credential = GatedCredential(access_token("shared")) + + assert queue_behind(credential, kinds[leader], kinds[waiter]) == ["shared", "shared"] + assert len(credential.scopes) == 1 + + @pytest.mark.parametrize("retried", ["recovered", "failed again"]) + def test_an_evaluation_failure_is_not_replayed_to_another_kind_of_request(self, retried): + retry = access_token("recovered") if retried == "recovered" else ClientAuthenticationError("Run az login.") + credential = GatedCredential(ClientAuthenticationError("Please run az login."), retry) + + evaluation, authentication = queue_behind(credential, EVALUATION_KIND, AUTHENTICATION_KIND) + + assert isinstance(evaluation, ClientAuthenticationError) + # The other request asks the credential itself rather than raising the evaluation description as its own + # authentication error. + if retried == "recovered": + assert authentication == "recovered" + else: + assert isinstance(authentication, ClientAuthenticationError) and str(authentication) == "Run az login." + assert len(credential.scopes) == 2 + + @pytest.mark.parametrize("retried", ["recovered", "failed again"]) + def test_an_authentication_failure_is_not_replayed_to_another_kind_of_request(self, retried): + retry = access_token("recovered") if retried == "recovered" else ClientAuthenticationError("Run az login.") + credential = GatedCredential(ClientAuthenticationError("Please run az login."), retry) + + authentication, evaluation = queue_behind(credential, AUTHENTICATION_KIND, EVALUATION_KIND) + + assert isinstance(authentication, ClientAuthenticationError) and "az login" in str(authentication) + # The other request asks the credential itself rather than raising an unclassified EvaluationException. + if retried == "recovered": + assert evaluation == "recovered" + else: + assert isinstance(evaluation, ClientAuthenticationError) and str(evaluation) == "Run az login." + assert len(credential.scopes) == 2 + + def test_evaluation_failures_are_shared_whichever_request_wakes_first(self): + credential = GatedCredential(ClientAuthenticationError("Please run az login.")) + + evaluation, authentication, other_evaluation = queue_behind( + credential, EVALUATION_KIND, AUTHENTICATION_KIND, EVALUATION_KIND + ) + + assert isinstance(evaluation, ClientAuthenticationError) + # The other kind of request makes its own attempt and raises the credential's error. The request of the same + # kind shares the first failure even when the other kind's attempt ends first. + assert isinstance(authentication, ClientAuthenticationError) and str(authentication) == "Please run az login." + assert isinstance(other_evaluation, EvaluationException) and EVALUATION_MESSAGE in str(other_evaluation) + assert len(credential.scopes) == 2 + + def test_authentication_failures_are_shared_whichever_request_wakes_first(self): + credential = GatedCredential(ClientAuthenticationError("Please run az login.")) + + authentication, evaluation, other_authentication = queue_behind( + credential, AUTHENTICATION_KIND, EVALUATION_KIND, AUTHENTICATION_KIND + ) + + assert isinstance(authentication, ClientAuthenticationError) + assert isinstance(evaluation, ClientAuthenticationError) and not isinstance(evaluation, EvaluationException) + assert isinstance(other_authentication, ClientAuthenticationError) + assert str(other_authentication) == "Please run az login." + assert len(credential.scopes) == 2 + + def test_a_successful_token_request_ends_sharing_of_earlier_failures(self): + credential = Credential(ClientAuthenticationError("Please run az login."), access_token("recovered")) + + with pytest.raises(ClientAuthenticationError): + EVALUATION_KIND(credential) + entry = TOKENS._entry(credential, CHAT_SCOPE) + assert list(entry.failures) == [(_model_runners._token_failure, EvaluationException)] + + assert AUTHENTICATION_KIND(credential) == "recovered" + # A request that waited on the failure but wakes after this success gets the token, not the earlier failure. + assert entry.failures == {} + + +class LoopBoundAsyncCredential: + """An asynchronous credential whose ``get_token`` is a plain function that needs the caller's event loop.""" + + def __init__(self): + self.calls = 0 + + def get_token(self, *_, **__): + self.calls += 1 + loop = asyncio.get_running_loop() + + async def request(): + assert asyncio.get_running_loop() is loop + return access_token("loop-bound") + + return request() + + async def close(self): + pass + + async def __aenter__(self): + return self + + async def __aexit__(self, *_): + pass + + +class DualProtocolCredential(GatedCredential): + """A synchronous credential that implements both the synchronous and the asynchronous context manager protocols.""" + + def close(self): + pass + + def __enter__(self): + return self + + def __exit__(self, *_): + pass + + async def __aenter__(self): + return self + + async def __aexit__(self, *_): + pass + + +@pytest.mark.unittest +class TestAsynchronousCredentialRequests: + @pytest.mark.parametrize("kind", ["chat", "decision"]) + def test_an_asynchronous_credential_requests_its_token_on_the_callers_event_loop(self, kind): + make_request = {"chat": request_token, "decision": request_decision_token}[kind] + credential = LoopBoundAsyncCredential() + + assert [make_request(credential) for _ in range(2)] == ["loop-bound"] * 2 + assert credential.calls == 1 + + def test_a_synchronous_credential_with_both_context_manager_protocols_is_asked_on_a_worker_thread(self): + credential = DualProtocolCredential(access_token("dual")) + + # Both requests wait for one token request on a worker thread; on the event loop each would ask the credential. + assert queue_behind(credential, request_token, request_token) == ["dual", "dual"] + assert len(credential.scopes) == 1 + + def test_a_plain_get_token_outside_the_asynchronous_protocol_is_asked_on_a_worker_thread(self): + class MinimalWrapper: + def get_token(self, *_, **__): + return asyncio.ensure_future(asyncio.sleep(0, access_token("minimal"))) + + # A plain get_token that needs the running event loop must belong to an AsyncTokenCredential. + assert not is_async_credential(MinimalWrapper()) + + def test_synchronous_identity_credentials_request_tokens_on_a_worker_thread(self): + from azure.identity import ( + AzureCliCredential, + ClientSecretCredential, + DefaultAzureCredential, + ManagedIdentityCredential, + ) + from azure.identity.aio import AzureCliCredential as AsyncAzureCliCredential + + synchronous = [ + AzureCliCredential(), + ClientSecretCredential("tenant", "client", "secret"), + DefaultAzureCredential(exclude_shared_token_cache_credential=True), + ManagedIdentityCredential(), + Credential(access_token("sync")), + MagicMock(get_token=MagicMock(return_value=access_token("mock")), close=MagicMock()), + DualProtocolCredential(access_token("dual")), + ] + + assert not any(is_async_credential(credential) for credential in synchronous) + assert is_async_credential(AsyncAzureCliCredential()) + assert is_async_credential(LoopBoundAsyncCredential())