Skip to main content

Authentication Internals

When communicating with a secure Flyte control plane, client requests must authenticate across both asynchronous gRPC channels and asynchronous HTTP sessions. Token lifecycles, OAuth2 grants, and network retries must occur seamlessly without dropping active RPC calls or triggering concurrent refresh stampedes.

In flyte-sdk, the authentication layer bridges gRPC interceptors and HTTP client adapters with a unified async credential management architecture located in remote._client.auth.

The Authenticator Base Architecture

The core of authentication in flyte-sdk is the Authenticator base class defined in remote._client.auth._authenticators.base. It provides an abstraction for managing token lifecycles, configuration discovery, credential storage, and concurrency control.

from flyte.remote._client.auth._authenticators.base import Authenticator, GrpcAuthMetadata
from flyte.remote._client.auth._keyring import KeyringStore, Credentials
from flyte.remote._client.auth._client_config import ClientConfig, ClientConfigStore

Credential Resolution and Storage

When an Authenticator is initialized, it attempts to load existing cached credentials from KeyringStore using the stripped endpoint string:

self._creds = credentials or KeyringStore.retrieve(endpoint)

The Credentials model (remote._client.auth._keyring.Credentials) automatically calculates a unique credential identifier id using an MD5 hash of the current access_token. This identifier is used to track token mutations across concurrent execution flows.

Configuration resolution is lazily evaluated through _resolve_config(). If a ClientConfigStore (such as RemoteClientConfigStore or StaticClientConfigStore) is provided, Authenticator fetches remote server configuration and merges it with local overrides:

async def _resolve_config(self) -> ClientConfig:
if self._resolved_config is not None:
return self._resolved_config

if self._cfg_store is None:
raise ValueError("ClientConfigStore is not set. Cannot resolve configuration.")

remote_config = await self._cfg_store.get_client_config()
self._resolved_config = (
remote_config.with_override(self._client_config) if self._client_config else remote_config
)
return self._resolved_config

Concrete Authenticator Implementations

Concrete subclasses of Authenticator implement _do_refresh_credentials() to execute grant-specific token acquisition:

  • PKCEAuthenticator: Implements the OAuth2 Authorization Code flow with Proof Key for Code Exchange (PKCE) and a local loopback server (OAuthCallbackHandler).
  • ClientCredentialsAuthenticator: Implements the OAuth2 Client Credentials flow for machine-to-machine authentication via client ID and secret.
  • DeviceCodeAuthenticator: Implements the OAuth2 Device Authorization Grant flow for interactive terminal authorization.
  • AsyncCommandAuthenticator: Generates authentication tokens dynamically by executing external CLI commands.

Concurrency and Stampede Protection

When access tokens expire, multiple concurrent RPC calls or HTTP requests may receive unauthenticated responses at the same instant. Without synchronization, this would trigger multiple simultaneous token refresh requests against the identity provider.

Authenticator.refresh_credentials prevents this using an asyncio.Lock paired with a double-checked matching pattern against creds_id:

async def refresh_credentials(self, creds_id: str | None = None):
# Fast path check before acquiring the lock
if creds_id and creds_id != self._creds_id:
return

async with self._async_lock:
# Double-check inside the lock: another coroutine may have refreshed while we waited
if creds_id and creds_id != self._creds_id:
return

try:
self._creds = await self._do_refresh_credentials()
KeyringStore.store(self._creds)
except Exception:
KeyringStore.delete(self._endpoint)
raise

self._creds_id = self._creds.id

The refresh flow behaves as follows:

  1. Callers pass the creds_id that was active when their initial request failed.
  2. If creds_id does not match self._creds_id, another coroutine has already refreshed the token; the caller skips the refresh and retries with the new credentials immediately.
  3. If creds_id matches, the coroutine acquires self._async_lock and checks again.
  4. If a refresh fails with an exception, KeyringStore.delete(self._endpoint) removes the invalid credentials from storage before the exception propagates.

gRPC Authentication Interceptors

Flyte's asynchronous gRPC channels integrate authentication via client interceptors in remote._client.auth._grpc_utils.auth_interceptor.

Interceptor Hierarchy

All auth interceptors inherit from _BaseAuthInterceptor, which holds a callable factory returning an Authenticator instance and attaches authentication metadata to gRPC call details:

class _BaseAuthInterceptor:
def __init__(self, get_authenticator: typing.Callable[[], Authenticator]):
self._get_authenticator = get_authenticator
self._authenticator: typing.Optional[Authenticator] = None

@property
def authenticator(self) -> Authenticator:
if self._authenticator is None:
self._authenticator = self._get_authenticator()
return self._authenticator

async def call_details_with_auth_metadata(
self, client_call_details: grpc.aio.ClientCallDetails
) -> typing.Tuple[grpc.aio.ClientCallDetails, str]:
auth_metadata = await self.authenticator.get_grpc_call_auth_metadata()
if auth_metadata:
return with_metadata(client_call_details, auth_metadata.pairs), auth_metadata.creds_id
else:
return client_call_details, ""

Handling Unauthenticated Responses

Flyte defines interceptors for all four gRPC communication patterns:

  • AuthUnaryUnaryInterceptor (grpc.aio.UnaryUnaryClientInterceptor)
  • AuthUnaryStreamInterceptor (grpc.aio.UnaryStreamClientInterceptor)
  • AuthStreamUnaryInterceptor (grpc.aio.StreamUnaryClientInterceptor)
  • AuthStreamStreamInterceptor (grpc.aio.StreamStreamClientInterceptor)

In AuthUnaryUnaryInterceptor, the interceptor adds auth metadata, invokes the continuation, and catches grpc.aio.AioRpcError. When either grpc.StatusCode.UNAUTHENTICATED or grpc.StatusCode.UNKNOWN is encountered (as certain proxies and ingresses convert unauthenticated statuses to unknown errors), it triggers refresh_credentials(creds_id=creds_id) and retries the RPC call:

class AuthUnaryUnaryInterceptor(_BaseAuthInterceptor, grpc.aio.UnaryUnaryClientInterceptor):
async def intercept_unary_unary(
self,
continuation: typing.Callable,
client_call_details: ClientCallDetails,
request: typing.Any,
):
updated_call_details, creds_id = await self.call_details_with_auth_metadata(client_call_details)
try:
return await (await continuation(updated_call_details, request))
except grpc.aio.AioRpcError as e:
if e.code() == grpc.StatusCode.UNAUTHENTICATED or e.code() == grpc.StatusCode.UNKNOWN:
await self.authenticator.refresh_credentials(creds_id=creds_id)
updated_call_details, _ = await self.call_details_with_auth_metadata(client_call_details)
return await (await continuation(updated_call_details, request))
else:
raise e

Streaming Calls with UnaryStreamCall

Streaming RPCs (AuthUnaryStreamInterceptor and AuthStreamStreamInterceptor) wrap invocations in a custom UnaryStreamCall object. This delegates response iteration to response_iterator(), which handles error recovery mid-stream or during call initiation:

async def response_iterator(self) -> typing.AsyncIterator[ResponseType]:
call_details, creds_id = await self._parent_interceptor.call_details_with_auth_metadata(self._call_details)
self._call = await self._continuation(call_details, self._request)
try:
async for response in self._call:
yield response
except grpc.aio.AioRpcError as e:
if e.code() == grpc.StatusCode.UNAUTHENTICATED or e.code() == grpc.StatusCode.UNKNOWN:
await self._authenticator.refresh_credentials(creds_id=creds_id)
updated_call_details, _ = await self._parent_interceptor.call_details_with_auth_metadata(call_details)
self._call = await self._continuation(updated_call_details, self._request)
async for response in self._call:
yield response
else:
raise e

HTTP Authentication Layer

For HTTP-based operations—such as token exchanges, metadata discovery, and control plane HTTP endpoints—flyte-sdk provides AsyncAuthenticationHTTPAdapter and AsyncAuthenticatedClient in remote._client.auth._authenticators.base.

from http import HTTPStatus
import httpx
from flyte.remote._client.auth._authenticators.base import (
AsyncAuthenticatedClient,
AsyncAuthenticationHTTPAdapter,
Authenticator,
)

Header Injection via AsyncAuthenticationHTTPAdapter

AsyncAuthenticationHTTPAdapter inspects the Authenticator instance. If no credentials exist, it triggers an initial refresh before requesting metadata:

class AsyncAuthenticationHTTPAdapter:
def __init__(self, authenticator: Authenticator):
self.authenticator = authenticator

async def add_auth_header(self, request: httpx.Request) -> typing.Optional[str]:
if self.authenticator.get_credentials() is None:
await self.authenticator.refresh_credentials()

metadata = await self.authenticator.get_grpc_call_auth_metadata()
if metadata is None:
return None
for key, value in metadata.pairs.keys():
request.headers[key] = value
return metadata.creds_id

Request Execution and Retries via AsyncAuthenticatedClient

AsyncAuthenticatedClient extends httpx.AsyncClient. In its send() method, it injects authentication headers, executes the request, and intercepts 401 UNAUTHORIZED responses to refresh credentials and retry once:

class AsyncAuthenticatedClient(httpx.AsyncClient):
def __init__(self, authenticator: Authenticator, **kwargs):
super().__init__(**kwargs)
self.auth_adapter = AsyncAuthenticationHTTPAdapter(authenticator)
self.authenticator = authenticator

async def send(self, request: httpx.Request, **kwargs) -> httpx.Response:
creds_id = await self.auth_adapter.add_auth_header(request)
response = await super().send(request, **kwargs)

if response.status_code == HTTPStatus.UNAUTHORIZED:
await self.authenticator.refresh_credentials(creds_id=creds_id)
await self.auth_adapter.add_auth_header(request)
response = await super().send(request, **kwargs)

return response

Channel and Interceptor Assembly

In remote._client.auth._authenticators.factory, flyte-sdk packages these components when establishing connections. The factory functions create_auth_interceptors and create_proxy_auth_interceptors construct the interceptor lists using lazy authenticator factories:

def create_auth_interceptors(in_channel: grpc.aio.Channel, endpoint: str, **kwargs) -> list:
def authenticator_factory() -> Authenticator:
return get_async_authenticator(
endpoint=endpoint,
cfg_store=RemoteClientConfigStore(in_channel),
**kwargs,
)

return [
AuthUnaryUnaryInterceptor(authenticator_factory),
AuthUnaryStreamInterceptor(authenticator_factory),
AuthStreamUnaryInterceptor(authenticator_factory),
AuthStreamStreamInterceptor(authenticator_factory),
]

When building gRPC channels via remote._client.auth._channel.create_channel or initializing control plane connections via remote._client.controlplane.ClientSet, these interceptor chains ensure all outbound calls remain authenticated with automatic credential recovery across gRPC and HTTP transports.