diff --git a/sdk/evaluation/azure-ai-evaluation/CHANGELOG.md b/sdk/evaluation/azure-ai-evaluation/CHANGELOG.md index 50defabe0ef2..4282ab310e4e 100644 --- a/sdk/evaluation/azure-ai-evaluation/CHANGELOG.md +++ b/sdk/evaluation/azure-ai-evaluation/CHANGELOG.md @@ -1,5 +1,17 @@ # Release History +## 1.18.12 (Unreleased) + +### Features Added + +### Breaking Changes + +### 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. + +### Other Changes + ## 1.18.11 (2026-11-03) ### Features Added @@ -17,6 +29,7 @@ - `evaluate()` computes `.binary_aggregate` over the rows whose result is `pass` or `fail`. Other rows, such as not applicable or errored rows, used to count as failures, although the mean score leaves them out. An evaluator without `pass` or `fail` rows, including a custom evaluator whose results have other values, has no `binary_aggregate`, as one without scores has no mean. - An evaluator's inputs keep the order in which its `__call__` overloads declare them, instead of the order of a Python `set`, which varies with the process's hash seed. An evaluator that sends these inputs to a Decision model unchanged, as several Microsoft Foundry catalog evaluators do, could list the same fields in a different order in each process, and the Decision model could answer the same input differently. The built-in evaluators in this package already send their Decision inputs in a fixed order. + ## 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 b2035dda5695..ecc024ce873c 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" LOGGER = logging.getLogger(__name__) _PREFIX = "azure.ai.evaluation.decision." @@ -362,6 +368,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 _provider_error_code(payload: Any) -> Optional[str]: """Return the provider's nested ``error.code`` when it is a bounded identifier, otherwise None.""" provider_error = payload.get("error") if isinstance(payload, dict) else None @@ -695,25 +720,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, @@ -766,6 +838,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 code = _provider_error_code(payload) kind = _provider_error_kind(status, payload) 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/azure/ai/evaluation/_version.py b/sdk/evaluation/azure-ai-evaluation/azure/ai/evaluation/_version.py index 03d429c451a0..a307b7b58822 100644 --- a/sdk/evaluation/azure-ai-evaluation/azure/ai/evaluation/_version.py +++ b/sdk/evaluation/azure-ai-evaluation/azure/ai/evaluation/_version.py @@ -3,4 +3,4 @@ # --------------------------------------------------------- # represents upcoming version -VERSION = "1.18.11" +VERSION = "1.18.12" 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 5761abc9e135..f91f8b0a7019 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,14 +5,18 @@ """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 os +import pickle import subprocess import sys import threading +import time +import weakref from copy import deepcopy from logging.handlers import BufferingHandler from concurrent.futures import ThreadPoolExecutor @@ -22,11 +26,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, GroundednessEvaluator +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 ( @@ -141,6 +147,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", @@ -652,6 +694,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())