flyte 0.0.1b0__py3-none-any.whl

union/remote/__init__.py

@@ -0,0 +1,23 @@
+"""
+Remote Entities that are accessible from the Union Server once deployed or created.
+"""
+
+__all__ = [
+    "Action",
+    "ActionDetails",
+    "Project",
+    "Run",
+    "RunDetails",
+    "Secret",
+    "Task",
+    "create_channel",
+    "upload_dir",
+    "upload_file",
+]
+
+from ._client.auth import create_channel
+from ._data import upload_dir, upload_file
+from ._project import Project
+from ._run import Action, ActionDetails, Run, RunDetails
+from ._secret import Secret
+from ._task import Task

union/remote/_client/_protocols.py

@@ -0,0 +1,129 @@
+from typing import AsyncIterator, Protocol
+
+from flyteidl.admin import project_attributes_pb2, project_pb2, version_pb2
+from flyteidl.service import dataproxy_pb2
+
+from union._protos.secret import payload_pb2
+from union._protos.workflow import run_logs_service_pb2, run_service_pb2, task_service_pb2
+
+
+class MetadataServiceProtocol(Protocol):
+    async def GetVersion(self, request: version_pb2.GetVersionRequest) -> version_pb2.GetVersionResponse: ...
+
+
+class ProjectDomainService(Protocol):
+    async def RegisterProject(
+        self, request: project_pb2.ProjectRegisterRequest
+    ) -> project_pb2.ProjectRegisterResponse: ...
+
+    async def UpdateProject(self, request: project_pb2.Project) -> project_pb2.ProjectUpdateResponse: ...
+
+    async def GetProject(self, request: project_pb2.ProjectGetRequest) -> project_pb2.Project: ...
+
+    async def ListProjects(self, request: project_pb2.ProjectListRequest) -> project_pb2.Projects: ...
+
+    async def GetDomains(self, request: project_pb2.GetDomainRequest) -> project_pb2.GetDomainsResponse: ...
+
+    async def UpdateProjectDomainAttributes(
+        self, request: project_attributes_pb2.ProjectAttributesUpdateRequest
+    ) -> project_pb2.ProjectUpdateResponse: ...
+
+    async def GetProjectDomainAttributes(
+        self, request: project_attributes_pb2.ProjectAttributesGetRequest
+    ) -> project_attributes_pb2.ProjectAttributes: ...
+
+    async def DeleteProjectDomainAttributes(
+        self, request: project_attributes_pb2.ProjectAttributesDeleteRequest
+    ) -> project_attributes_pb2.ProjectAttributesDeleteResponse: ...
+
+    async def UpdateProjectAttributes(
+        self, request: project_attributes_pb2.ProjectAttributesUpdateRequest
+    ) -> project_attributes_pb2.ProjectAttributesUpdateResponse: ...
+
+    async def GetProjectAttributes(
+        self, request: project_attributes_pb2.ProjectAttributesGetRequest
+    ) -> project_attributes_pb2.ProjectAttributes: ...
+
+    async def DeleteProjectAttributes(
+        self, request: project_attributes_pb2.ProjectAttributesDeleteRequest
+    ) -> project_attributes_pb2.ProjectAttributesDeleteResponse: ...
+
+
+class TaskService(Protocol):
+    async def DeployTask(self, request: task_service_pb2.DeployTaskRequest) -> task_service_pb2.DeployTaskResponse: ...
+
+    async def GetTaskDetails(
+        self, request: task_service_pb2.GetTaskDetailsRequest
+    ) -> task_service_pb2.GetTaskDetailsResponse: ...
+
+
+class RunService(Protocol):
+    async def CreateRun(self, request: run_service_pb2.CreateRunRequest) -> run_service_pb2.CreateRunResponse: ...
+
+    async def AbortRun(self, request: run_service_pb2.AbortRunRequest) -> run_service_pb2.AbortRunResponse: ...
+
+    async def GetRunDetails(
+        self, request: run_service_pb2.GetRunDetailsRequest
+    ) -> run_service_pb2.GetRunDetailsResponse: ...
+
+    async def WatchRunDetails(
+        self, request: run_service_pb2.WatchRunDetailsRequest
+    ) -> AsyncIterator[run_service_pb2.WatchRunDetailsResponse]: ...
+
+    async def GetActionDetails(
+        self, request: run_service_pb2.GetActionDetailsRequest
+    ) -> run_service_pb2.GetActionDetailsResponse: ...
+
+    async def WatchActionDetails(
+        self, request: run_service_pb2.WatchActionDetailsRequest
+    ) -> AsyncIterator[run_service_pb2.WatchActionDetailsResponse]: ...
+
+    async def GetActionData(
+        self, request: run_service_pb2.GetActionDataRequest
+    ) -> run_service_pb2.GetActionDataResponse: ...
+
+    async def ListRuns(self, request: run_service_pb2.ListRunsRequest) -> run_service_pb2.ListRunsResponse: ...
+
+    async def WatchRuns(
+        self, request: run_service_pb2.WatchRunsRequest
+    ) -> AsyncIterator[run_service_pb2.WatchRunsResponse]: ...
+
+    async def ListActions(self, request: run_service_pb2.ListActionsRequest) -> run_service_pb2.ListActionsResponse: ...
+
+    async def WatchActions(
+        self, request: run_service_pb2.WatchActionsRequest
+    ) -> AsyncIterator[run_service_pb2.WatchActionsResponse]: ...
+
+
+class DataProxyService(Protocol):
+    async def CreateUploadLocation(
+        self, request: dataproxy_pb2.CreateUploadLocationRequest
+    ) -> dataproxy_pb2.CreateUploadLocationResponse: ...
+
+    async def CreateDownloadLocation(
+        self, request: dataproxy_pb2.CreateDownloadLocationRequest
+    ) -> dataproxy_pb2.CreateDownloadLocationResponse: ...
+
+    async def CreateDownloadLink(
+        self, request: dataproxy_pb2.CreateDownloadLinkRequest
+    ) -> dataproxy_pb2.CreateDownloadLinkResponse: ...
+
+    async def GetData(self, request: dataproxy_pb2.GetDataRequest) -> dataproxy_pb2.GetDataResponse: ...
+
+
+class RunLogsService(Protocol):
+    async def TailLogs(
+        self, request: run_logs_service_pb2.TailLogsRequest
+    ) -> AsyncIterator[run_logs_service_pb2.TailLogsResponse]: ...
+
+
+class SecretService(Protocol):
+    async def CreateSecret(self, request: payload_pb2.CreateSecretRequest) -> payload_pb2.CreateSecretResponse: ...
+
+    async def UpdateSecret(self, request: payload_pb2.UpdateSecretRequest) -> payload_pb2.UpdateSecretResponse: ...
+
+    async def GetSecret(self, request: payload_pb2.GetSecretRequest) -> payload_pb2.GetSecretResponse: ...
+
+    async def ListSecrets(self, request: payload_pb2.ListSecretsRequest) -> payload_pb2.ListSecretsResponse: ...
+
+    async def DeleteSecret(self, request: payload_pb2.DeleteSecretRequest) -> payload_pb2.DeleteSecretResponse: ...

union/remote/_client/auth/__init__.py

@@ -0,0 +1,12 @@
+from union.remote._client.auth._channel import create_channel
+from union.remote._client.auth._client_config import AuthType, ClientConfig
+from union.remote._client.auth.errors import AccessTokenNotFoundError, AuthenticationError, AuthenticationPending
+
+__all__ = [
+    "AccessTokenNotFoundError",
+    "AuthType",
+    "AuthenticationError",
+    "AuthenticationPending",
+    "ClientConfig",
+    "create_channel",
+]

union/remote/_client/auth/_authenticators/base.py

@@ -0,0 +1,391 @@
+import asyncio
+import ssl
+import typing
+from abc import abstractmethod
+from http import HTTPStatus
+
+import httpx
+import pydantic
+
+from union.remote._client.auth._client_config import ClientConfig, ClientConfigStore
+from union.remote._client.auth._keyring import Credentials, KeyringStore
+
+
+class GrpcAuthMetadata(pydantic.BaseModel):
+    creds_id: str
+    pairs: typing.Dict[str, str]
+
+
+class Authenticator(object):
+    """
+    Base authenticator for all authentication flows
+    """
+
+    def __init__(
+        self,
+        endpoint: str,
+        *,
+        cfg_store: typing.Optional[ClientConfigStore] = None,
+        client_config: typing.Optional[ClientConfig] = None,
+        credentials: typing.Optional[Credentials] = None,
+        http_session: typing.Optional[httpx.AsyncClient] = None,
+        http_proxy_url: typing.Optional[str] = None,
+        verify: bool = True,
+        ca_cert_path: typing.Optional[str] = None,
+        **kwargs,
+    ):
+        """
+        Initialize the base authenticator.
+
+        :param endpoint: The endpoint URL for authentication
+        :param cfg_store: Optional client configuration store for retrieving remote configuration
+        :param client_config: Optional client configuration containing authentication settings
+        :param credentials: Optional credentials to use for authentication
+        :param http_session: Optional HTTP session to use for requests
+        :param http_proxy_url: Optional HTTP proxy URL
+        :param verify: Whether to verify SSL certificates
+        :param ca_cert_path: Optional path to CA certificate file
+        :param kwargs: Additional keyword arguments passed to get_async_session, which may include:
+            - auth: Authentication implementation to use
+            - params: Query parameters to include in request URLs
+            - headers: HTTP headers to include in requests
+            - cookies: Cookies to include in requests
+            - cert: SSL client certificate (path or tuple)
+            - http1: Whether to enable HTTP/1.1 support
+            - http2: Whether to enable HTTP/2 support
+            - proxies: Proxy configuration mapping
+            - mounts: Mounted transports for specific URL patterns
+            - timeout: Request timeout configuration
+            - follow_redirects: Whether to follow redirects
+            - limits: Connection pool limits
+            - max_redirects: Maximum number of redirects to follow
+            - event_hooks: Event hooks for request/response lifecycle
+            - base_url: Base URL to join with relative URLs
+            - transport: Transport implementation to use
+            - app: ASGI application to handle requests
+        """
+        self._endpoint = endpoint
+        self._creds = credentials or KeyringStore.retrieve(endpoint)
+        self._http_proxy_url = http_proxy_url
+        self._verify = verify
+        self._ca_cert_path = ca_cert_path
+        self._client_config = client_config
+        self._cfg_store = cfg_store
+        # Will be populated by _ensure_remote_config
+        self._resolved_config = None
+        # Lock for coroutine safety
+        self._async_lock = asyncio.Lock()
+        self._http_session = http_session or get_async_session(**kwargs)
+        # Id for tracking credential refresh state
+        self._creds_id = self._creds.id if self._creds else None
+
+    async def _resolve_config(self) -> ClientConfig:
+        """
+        Resolves and merges client configuration with remote configuration.
+
+        This method fetches the remote configuration from the cfg_store and merges it with
+        the local client_config, prioritizing local settings over remote ones.
+
+        This method is thread-safe and coroutine-safe, ensuring the remote config is fetched
+        only once regardless of concurrent access from multiple threads or coroutines.
+
+        :return: A merged ClientConfig object containing resolved configuration settings
+        """
+        # First check without locks for performance
+        if self._resolved_config is not None:
+            return self._resolved_config
+
+        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
+
+    def get_credentials(self) -> typing.Optional[Credentials]:
+        """
+        Get the current credentials.
+
+        :return: The current credentials or None if not set
+        """
+        return self._creds
+
+    def _set_credentials(self, creds: Credentials):
+        """
+        Set the credentials.
+
+        :param creds: The credentials to set
+        """
+        self._creds = creds
+
+    async def get_grpc_call_auth_metadata(self) -> typing.Optional[GrpcAuthMetadata]:
+        """
+        Fetch the authentication metadata for gRPC calls.
+
+        :return: A tuple of (header_key, header_value) or None if no credentials are available
+        """
+        creds = self.get_credentials()
+        if creds:
+            cfg = await self._resolve_config()
+            return GrpcAuthMetadata(
+                creds_id=creds.id,
+                pairs={
+                    cfg.header_key: f"Bearer {creds.access_token}",
+                },
+            )
+        return None
+
+    async def refresh_credentials(self, creds_id: str | None = None):
+        """
+        Refresh the credentials asynchronously with thread and asyncio safety.
+
+        This method implements a thread-safe and coroutine-safe credential refresh mechanism.
+        It uses a timestamp-based approach to prevent redundant credential refreshes when
+        multiple threads or coroutines attempt to refresh credentials simultaneously.
+
+        The caller should capture the current _creds_timestamp before attempting to use credentials.
+        If credential usage fails, the caller can pass that timestamp to this method.
+        If the timestamp matches the current value, a refresh is needed; otherwise,
+        another thread has already refreshed the credentials.
+
+        :param creds_id: The id of credentials when they were last accessed by the caller.
+                               If None, force a refresh regardless of id.
+        :raises: May raise authentication-related exceptions if the refresh fails
+        """
+        # If creds_id is None, force refresh
+        # If creds_id matches current value, credentials need refresh
+        # If creds_id doesn't match, another thread already refreshed credentials
+        if creds_id is not None and creds_id != self._creds_id:
+            # Credentials have been refreshed by another thread/coroutine since caller read them
+            return
+
+        # Use the async lock to ensure coroutine safety
+        async with self._async_lock:
+            # Double-check pattern to avoid unnecessary work
+            if creds_id is not None and creds_id != self._creds_id:
+                # Another thread/coroutine refreshed credentials while we were waiting for the lock
+                return
+
+            # Perform the actual credential refresh
+            try:
+                self._creds = await self._do_refresh_credentials()
+                KeyringStore.store(self._creds)
+            except Exception:
+                KeyringStore.delete(self._endpoint)
+                raise
+
+            # Update the timestamp to indicate credentials have been refreshed
+            self._creds_id = self._creds.id
+
+    @abstractmethod
+    async def _do_refresh_credentials(self) -> Credentials:
+        """
+        Perform the actual credential refresh operation.
+
+        This method must be implemented by subclasses to handle the specific authentication flow.
+        It should update the internal credentials object (_creds) with a new access token.
+
+        Implementations typically use the resolved configuration from _resolve_config() to
+        determine authentication endpoints, scopes, audience, and other parameters needed for
+        the specific authentication flow.
+
+        :raises: May raise authentication-related exceptions if the refresh fails
+        """
+        ...
+
+
+class AsyncAuthenticatedClient(httpx.AsyncClient):
+    """
+    An httpx.AsyncClient that automatically adds authentication headers to requests.
+    This class extends httpx.AsyncClient which is inherently async for network operations.
+    """
+
+    def __init__(self, authenticator: Authenticator, **kwargs):
+        """
+        Initialize the authenticated client.
+
+        :param authenticator: The authenticator to use for authentication
+        :param kwargs: Additional arguments passed to the httpx.AsyncClient constructor
+        """
+        super().__init__(**kwargs)
+        self.auth_adapter = AsyncAuthenticationHTTPAdapter(authenticator)
+        self.authenticator = authenticator
+
+    async def send(self, request: httpx.Request, **kwargs) -> httpx.Response:
+        """
+        Sends the request with added authentication headers.
+        Must be async because it performs network IO operations and may need to refresh credentials.
+        If the response returns a 401 status code, refreshes the credentials and retries the request.
+
+        :param request: The request object to send.
+        :param kwargs: Additional keyword arguments passed to the parent httpx.AsyncClient.send method, which may
+        include:
+            - auth: Authentication implementation to use for this request
+            - follow_redirects: Whether to follow redirects for this request
+            - timeout: Request timeout configuration for this request
+        :return: The response object.
+        """
+
+        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
+
+
+class AsyncAuthenticationHTTPAdapter:
+    """
+    A custom async HTTP adapter that adds authentication headers to requests of an httpx.AsyncClient.
+    This is the async equivalent of AuthenticationHTTPAdapter for requests.
+    """
+
+    def __init__(self, authenticator: Authenticator):
+        """
+        Initialize the authentication HTTP adapter.
+
+        :param authenticator: The authenticator to use for authentication
+        """
+        self.authenticator = authenticator
+
+    async def add_auth_header(self, request: httpx.Request) -> str:
+        """
+        Adds authentication headers to the request.
+        Must be async because it may call refresh_credentials which performs IO operations.
+
+        :param request: The request object to add headers to.
+        :return: The credentials ID (creds_id) used for tracking credential refresh state
+        """
+        if self.authenticator.get_credentials() is None:
+            await self.authenticator.refresh_credentials()
+
+        metadata = await self.authenticator.get_grpc_call_auth_metadata()
+        for key, value in metadata.pairs:
+            request.headers[key] = value
+        return metadata.creds_id
+
+
+def upgrade_async_session_to_proxy_authenticated(
+    http_session: httpx.AsyncClient, proxy_authenticator: typing.Optional[Authenticator] = None, **kwargs
+) -> httpx.AsyncClient:
+    """
+    Given an httpx.AsyncClient, it returns a new session that uses AsyncAuthenticationHTTPAdapter
+    to perform authentication with a proxy in front of Flyte
+
+    :param http_session: httpx.AsyncClient Precreated session
+    :param proxy_authenticator: Optional authenticator for proxy authentication
+    :param kwargs: Additional arguments passed to AsyncAuthenticatedClient, which may include:
+        - auth: Authentication implementation to use
+        - params: Query parameters to include in request URLs
+        - headers: HTTP headers to include in requests
+        - cookies: Cookies to include in requests
+        - verify: SSL verification mode (True/False/path to certificate)
+        - cert: SSL client certificate (path or tuple)
+        - http1: Whether to enable HTTP/1.1 support
+        - http2: Whether to enable HTTP/2 support
+        - proxies: Proxy configuration mapping
+        - mounts: Mounted transports for specific URL patterns
+        - timeout: Request timeout configuration
+        - follow_redirects: Whether to follow redirects
+        - limits: Connection pool limits
+        - max_redirects: Maximum number of redirects to follow
+        - event_hooks: Event hooks for request/response lifecycle
+        - base_url: Base URL to join with relative URLs
+        - transport: Transport implementation to use
+        - app: ASGI application to handle requests
+    :return: httpx.AsyncClient with authentication
+    """
+    if proxy_authenticator:
+        return AsyncAuthenticatedClient(proxy_authenticator, **kwargs)
+    else:
+        return http_session
+
+
+def get_async_session(
+    proxy_authenticator: Authenticator | None = None,
+    ca_cert_path: str | None = None,
+    verify: bool | None = None,
+    **kwargs,
+) -> httpx.AsyncClient:
+    """
+    Returns a new httpx.AsyncClient with proxy authentication if proxy_authenticator is provided.
+
+    This function creates a new httpx.AsyncClient and optionally configures it with proxy authentication
+    if a proxy authenticator is provided.
+
+    :param proxy_authenticator: Optional authenticator for proxy authentication
+    :param ca_cert_path: Optional path to CA certificate file for SSL verification
+    :param verify: Optional SSL verification mode (True/False/path to certificate)
+    :param kwargs: Additional keyword arguments passed to httpx.AsyncClient constructor and AsyncAuthenticatedClient,
+     which may include:
+        - auth: Authentication implementation to use
+        - params: Query parameters to include in request URLs
+        - headers: HTTP headers to include in requests
+        - cookies: Cookies to include in requests
+        - cert: SSL client certificate (path or tuple)
+        - http1: Whether to enable HTTP/1.1 support
+        - http2: Whether to enable HTTP/2 support
+        - proxies: Proxy configuration mapping
+        - mounts: Mounted transports for specific URL patterns
+        - timeout: Request timeout configuration
+        - follow_redirects: Whether to follow redirects
+        - limits: Connection pool limits
+        - max_redirects: Maximum number of redirects to follow
+        - event_hooks: Event hooks for request/response lifecycle
+        - base_url: Base URL to join with relative URLs
+        - transport: Transport implementation to use
+        - app: ASGI application to handle requests
+        - proxy_env: Environment variables for proxy command
+        - proxy_timeout: Timeout for proxy command execution
+        - header_key: Header key to use for authentication
+        - endpoint: The endpoint URL for authentication
+        - client_id: Client ID for authentication
+        - client_secret: Client secret for authentication
+        - scopes: List of scopes to request during authentication
+        - audience: Audience for the token
+        - http_proxy_url: HTTP proxy URL
+    :return: An httpx.AsyncClient instance, optionally configured with proxy authentication
+    """
+
+    # Extract known httpx.AsyncClient parameters from kwargs
+    client_kwargs = {
+        k: v
+        for k, v in kwargs.items()
+        if k
+        in [
+            "auth",
+            "params",
+            "headers",
+            "cookies",
+            "verify",
+            "cert",
+            "http1",
+            "http2",
+            "proxies",
+            "mounts",
+            "timeout",
+            "follow_redirects",
+            "limits",
+            "max_redirects",
+            "event_hooks",
+            "base_url",
+            "transport",
+            "app",
+        ]
+    }
+
+    if ca_cert_path:
+        verify = ssl.create_default_context(capath=ca_cert_path)
+
+    if verify is not None:
+        client_kwargs["verify"] = verify
+
+    http_session = httpx.AsyncClient(**client_kwargs)
+    if proxy_authenticator:
+        http_session = upgrade_async_session_to_proxy_authenticated(
+            http_session, proxy_authenticator=proxy_authenticator, **kwargs
+        )
+    return http_session

union/remote/_client/auth/_authenticators/client_credentials.py

@@ -0,0 +1,73 @@
+from union._logging import logger
+from union.remote._client.auth import _token_client as token_client
+from union.remote._client.auth._authenticators.base import Authenticator
+from union.remote._client.auth._keyring import Credentials
+
+
+class ClientCredentialsAuthenticator(Authenticator):
+    """
+    This Authenticator uses ClientId and ClientSecret to authenticate
+    """
+
+    def __init__(
+        self,
+        client_id: str,
+        client_credentials_secret: str,
+        **kwargs,
+    ):
+        """
+        Initialize the client credentials authenticator.
+
+        :param client_id: The client ID for authentication
+        :param client_credentials_secret: The client secret for authentication
+        :param kwargs: Additional keyword arguments passed to the base Authenticator
+
+        **Keyword Arguments passed to base Authenticator**:
+        :param endpoint: The endpoint URL for authentication (required)
+        :param cfg_store: Optional client configuration store for retrieving remote configuration
+        :param client_config: Optional client configuration containing authentication settings
+        :param credentials: Optional credentials to use for authentication
+        :param http_session: Optional HTTP session to use for requests
+        :param http_proxy_url: Optional HTTP proxy URL
+        :param verify: Whether to verify SSL certificates (default: True)
+        :param ca_cert_path: Optional path to CA certificate file
+        :param scopes: List of scopes to request during authentication
+        :param audience: Audience for the token
+        """
+        if not client_id or not client_credentials_secret:
+            raise ValueError("both client_id and client_credentials_secret are required.")
+        self._client_id = client_id
+        self._client_credentials_secret = client_credentials_secret
+        super().__init__(**kwargs)
+
+    async def _do_refresh_credentials(self) -> Credentials:
+        """
+        Refreshes the authentication credentials using client credentials flow.
+
+        This function is used by the _handle_rpc_error() decorator, depending on the AUTH_MODE config object.
+        This handler is meant for SDK use-cases of auth (like pyflyte, or when users call SDK functions that require
+        access to Admin, like when waiting for another workflow to complete from within a task). This function uses
+        basic auth, which means the credentials for basic auth must be present from wherever this code is running.
+        """
+        cfg = await self._resolve_config()
+
+        # Note that unlike the Pkce flow, the client ID does not come from Admin.
+        logger.debug(f"Basic authorization flow with client id {self._client_id} scope {cfg.scopes}")
+        authorization_header = token_client.get_basic_authorization_header(
+            self._client_id, self._client_credentials_secret
+        )
+
+        token, refresh_token, expires_in = await token_client.get_token(
+            token_endpoint=cfg.token_endpoint,
+            authorization_header=authorization_header,
+            http_proxy_url=self._http_proxy_url,
+            verify=self._verify,
+            scopes=cfg.scopes,
+            audience=cfg.audience,
+            http_session=self._http_session,
+        )
+
+        logger.info("Retrieved new token, expires in {}".format(expires_in))
+        return Credentials(
+            for_endpoint=self._endpoint, access_token=token, refresh_token=refresh_token, expires_in=expires_in
+        )