flyte 0.2.0b1__py3-none-any.whl → 2.0.0b46__py3-none-any.whl

flyte/remote/_app.py

@@ -0,0 +1,346 @@
+from __future__ import annotations
+
+from typing import AsyncIterator, Literal, Mapping, Tuple, cast
+
+import grpc
+import rich.repr
+from flyteidl2.app import app_definition_pb2, app_payload_pb2
+from flyteidl2.common import identifier_pb2, list_pb2
+
+from flyte._initialize import ensure_client, get_client, get_init_config
+from flyte._logging import logger
+from flyte.syncify import syncify
+
+from ._common import ToJSONMixin, filtering, sorting
+
+WaitFor = Literal["activated", "deactivated"]
+
+
+def _is_active(state: app_definition_pb2.Status.DeploymentStatus) -> bool:
+    return state in [
+        app_definition_pb2.Status.DeploymentStatus.DEPLOYMENT_STATUS_ACTIVE,
+        app_definition_pb2.Status.DeploymentStatus.DEPLOYMENT_STATUS_STARTED,
+    ]
+
+
+def _is_deactivated(state: app_definition_pb2.Status.DeploymentStatus) -> bool:
+    return state in [
+        app_definition_pb2.Status.DeploymentStatus.DEPLOYMENT_STATUS_UNASSIGNED,
+        app_definition_pb2.Status.DeploymentStatus.DEPLOYMENT_STATUS_STOPPED,
+    ]
+
+
+class App(ToJSONMixin):
+    pb2: app_definition_pb2.App
+
+    def __init__(self, pb2: app_definition_pb2.App):
+        self.pb2 = pb2
+
+    @property
+    def name(self) -> str:
+        """
+        Get the name of the app.
+        """
+        return self.pb2.metadata.id.name
+
+    @property
+    def revision(self) -> int:
+        """
+        Get the revision number of the app.
+        """
+        return self.pb2.metadata.revision
+
+    @property
+    def endpoint(self) -> str:
+        """
+        Get the public endpoint URL of the app.
+        """
+        return self.pb2.status.ingress.public_url
+
+    @property
+    def deployment_status(self) -> app_definition_pb2.Status.DeploymentStatus:
+        """
+        Get the deployment status of the app
+        Returns:
+
+        """
+        if len(self.pb2.status.conditions) > 0:
+            return self.pb2.status.conditions[-1].deployment_status
+        else:
+            return app_definition_pb2.Status.DeploymentStatus.DEPLOYMENT_STATUS_UNSPECIFIED
+
+    @property
+    def desired_state(self) -> app_definition_pb2.Spec.DesiredState:
+        """
+        Get the desired state of the app.
+        """
+        return self.pb2.spec.desired_state
+
+    def is_active(self) -> bool:
+        """
+        Check if the app is currently active or started.
+        """
+        return _is_active(self.deployment_status)
+
+    def is_deactivated(self) -> bool:
+        """
+        Check if the app is currently deactivated or stopped.
+        """
+        return _is_deactivated(self.deployment_status)
+
+    @property
+    def url(self) -> str:
+        """
+        Get the console URL for viewing the app.
+        """
+        client = get_client()
+        return client.console.app_url(
+            project=self.pb2.metadata.id.project,
+            domain=self.pb2.metadata.id.domain,
+            app_name=self.name,
+        )
+
+    @syncify
+    async def watch(self, wait_for: WaitFor = "activated") -> App:
+        """
+        Watch for the app to reach activated or deactivated state.
+        :param wait_for: ["activated", "deactivated"]
+
+        Returns: The app in the desired state.
+        Raises: RuntimeError if the app did not reach desired state and failed!
+        """
+
+        if wait_for == "activated" and self.is_active():
+            return self
+        elif wait_for == "deactivated" and self.is_deactivated():
+            return self
+
+        call = cast(
+            AsyncIterator[app_payload_pb2.WatchResponse],
+            get_client().app_service.Watch(
+                request=app_payload_pb2.WatchRequest(
+                    app_id=self.pb2.metadata.id,
+                )
+            ),
+        )
+        async for resp in call:
+            if resp.update_event:
+                updated_app = resp.update_event.updated_app
+                current_status = updated_app.status.conditions[-1].deployment_status
+                if current_status == app_definition_pb2.Status.DeploymentStatus.DEPLOYMENT_STATUS_FAILED:
+                    raise RuntimeError(f"App deployment for app {self.name} has failed!")
+                if wait_for == "activated":
+                    if _is_active(current_status):
+                        return App(updated_app)
+                    elif _is_deactivated(current_status):
+                        raise RuntimeError(f"App deployment for app {self.name} has failed!")
+                elif wait_for == "deactivated":
+                    if _is_deactivated(current_status):
+                        return App(updated_app)
+        raise RuntimeError(f"App deployment for app {self.name} stalled!")
+
+    async def _update(
+        self, desired_state: app_definition_pb2.Spec.DesiredState, reason: str, wait_for: WaitFor | None = None
+    ) -> App:
+        """
+        Internal method to update the app's desired state.
+
+        :param desired_state: The new desired state for the app.
+        :param reason: Reason for the update.
+        :param wait_for: Optional state to wait for after update.
+        :return: The updated app.
+        """
+        new_pb2 = app_definition_pb2.App()
+        new_pb2.CopyFrom(self.pb2)
+        new_pb2.spec.desired_state = desired_state
+        updated_app = await App.update.aio(new_pb2, reason=reason)
+        if wait_for:
+            await updated_app.watch.aio(wait_for)
+        return updated_app
+
+    @syncify
+    async def activate(self, wait: bool = False) -> App:
+        """
+        Start the app
+        :param wait: Wait for the app to reach started state
+
+        """
+        if self.is_active():
+            return self
+        return await self._update(
+            app_definition_pb2.Spec.DESIRED_STATE_STARTED,
+            "User requested to activate app from flyte-sdk",
+            "activated" if wait else None,
+        )
+
+    @syncify
+    async def deactivate(self, wait: bool = False):
+        """
+        Stop the app
+        :param wait: Wait for the app to reach the deactivated state
+        """
+        if self.is_deactivated():
+            return
+        return await self._update(
+            app_definition_pb2.Spec.DESIRED_STATE_STOPPED,
+            "User requested to deactivate app from flyte-sdk",
+            "deactivated" if wait else None,
+        )
+
+    def __rich_repr__(self) -> rich.repr.Result:
+        """
+        Rich representation of the App object for pretty printing.
+        """
+        yield "name", self.name
+        yield "revision", self.revision
+        yield "endpoint", self.endpoint
+        yield (
+            "deployment_status",
+            app_definition_pb2.Status.DeploymentStatus.Name(self.deployment_status)[len("DEPLOYMENT_STATUS_") :],
+        )
+        yield "desired_state", app_definition_pb2.Spec.DesiredState.Name(self.desired_state)[len("DESIRED_STATE_") :]
+
+    @syncify
+    @classmethod
+    async def update(cls, updated_app_proto: app_definition_pb2.App, reason: str) -> App:
+        ensure_client()
+        resp = await get_client().app_service.Update(
+            request=app_payload_pb2.UpdateRequest(
+                app=updated_app_proto,
+                reason=reason,
+            )
+        )
+        return App(pb2=resp.app)
+
+    @syncify
+    @classmethod
+    async def replace(
+        cls,
+        name: str,
+        updated_app_spec: app_definition_pb2.Spec,
+        reason: str,
+        labels: Mapping[str, str] | None = None,
+        project: str | None = None,
+        domain: str | None = None,
+    ) -> App:
+        """
+        Replace an existing app's that matches the given name, with a new spec and optionally labels.
+        :param name: Name of the new app
+        :param updated_app_spec: Updated app spec
+        :param labels: Optional labels for the new app
+        :param project: Optional project for the new app
+        :param domain: Optional domain for the new app
+        :return: A new app
+        """
+        ensure_client()
+        app = await cls.get.aio(name=name, project=project, domain=domain)
+        updated_app_spec.creator.CopyFrom(app.pb2.spec.creator)
+        new_app = app_definition_pb2.App(
+            metadata=app_definition_pb2.Meta(
+                id=app.pb2.metadata.id,
+                revision=app.revision,
+                labels=labels if labels else app.pb2.metadata.labels,
+            ),
+            spec=updated_app_spec,
+            status=app.pb2.status,
+        )
+        return await cls.update.aio(new_app, reason=reason)
+
+    @syncify
+    @classmethod
+    async def get(
+        cls,
+        name: str,
+        project: str | None = None,
+        domain: str | None = None,
+    ) -> App:
+        """
+        Get an app by name.
+
+        :param name: The name of the app.
+        :param project: The project of the app.
+        :param domain: The domain of the app.
+        :return: The app remote object.
+        """
+        ensure_client()
+        cfg = get_init_config()
+        resp = await get_client().app_service.Get(
+            request=app_payload_pb2.GetRequest(
+                app_id=app_definition_pb2.Identifier(
+                    org=cfg.org,
+                    project=project or cfg.project,
+                    domain=domain or cfg.domain,
+                    name=name,
+                ),
+            )
+        )
+        return cls(pb2=resp.app)
+
+    @syncify
+    @classmethod
+    async def listall(
+        cls,
+        created_by_subject: str | None = None,
+        sort_by: Tuple[str, Literal["asc", "desc"]] | None = None,
+        limit: int = 100,
+    ) -> AsyncIterator[App]:
+        ensure_client()
+        cfg = get_init_config()
+        i = 0
+        token = None
+        sort_pb2 = sorting(sort_by)
+        filters = filtering(created_by_subject)
+        project = None
+        if cfg.project:
+            project = identifier_pb2.ProjectIdentifier(
+                organization=cfg.org,
+                name=cfg.project,
+                domain=cfg.domain,
+            )
+        while True:
+            req = app_payload_pb2.ListRequest(
+                request=list_pb2.ListRequest(
+                    limit=min(100, limit),
+                    token=token,
+                    sort_by=sort_pb2,
+                    filters=filters,
+                ),
+                org=cfg.org,
+                project=project,
+            )
+            resp = await get_client().app_service.List(
+                request=req,
+            )
+            token = resp.token
+            for a in resp.apps:
+                i += 1
+                if i > limit:
+                    return
+                yield cls(a)
+            if not token:
+                break
+
+    @syncify
+    @classmethod
+    async def create(cls, app: app_definition_pb2.App) -> App:
+        ensure_client()
+        try:
+            resp = await get_client().app_service.Create(app_payload_pb2.CreateRequest(app=app))
+            created_app = cls(resp.app)
+            logger.info(f"Deployed app {created_app.name} with revision {created_app.revision}")
+            return created_app
+        except grpc.aio.AioRpcError as e:
+            if e.code() in [grpc.StatusCode.ABORTED, grpc.StatusCode.ALREADY_EXISTS]:
+                if e.code() == grpc.StatusCode.ALREADY_EXISTS:
+                    logger.warning(f"App {app.metadata.id.name} already exists, updating...")
+                elif e.code() == grpc.StatusCode.ABORTED:
+                    logger.warning(f"Create App {app.metadata.id.name} was aborted on server, check state!")
+                return await App.replace.aio(
+                    name=app.metadata.id.name,
+                    labels=app.metadata.labels,
+                    updated_app_spec=app.spec,
+                    reason="User requested serve from sdk",
+                    project=app.metadata.id.project,
+                    domain=app.metadata.id.domain,
+                )
+            raise

flyte/remote/_auth_metadata.py

@@ -0,0 +1,42 @@
+from contextlib import contextmanager
+from typing import Tuple
+
+
+@contextmanager
+def auth_metadata(*kv: Tuple[str, str]):
+    """
+    This context manager allows you to pass contextualized auth metadata downstream to the Flyte authentication system.
+
+    This is only useful if flyte.init_passthrough() has been called.
+
+    Example:
+    ```python
+
+    flyte.init_passthrough("my-endpoint")
+
+    ...
+
+    with auth_metadata((key1, value1), (key2, value2)):
+        ...
+    ```
+
+    Args:
+        *kv: Tuple of auth metadata key/value pairs.
+    """
+    if not kv:
+        raise ValueError("No auth metadata provided.")
+    from flyte._context import internal_ctx
+
+    ctx = internal_ctx()
+    with ctx.new_metadata(kv):
+        yield
+
+
+def get_auth_metadata() -> Tuple[Tuple[str, str], ...]:
+    """
+    Returns auth metadata as tuple.
+    """
+    from flyte._context import internal_ctx
+
+    ctx = internal_ctx()
+    return ctx.data.metadata or ()

flyte/remote/_client/_protocols.py

@@ -1,13 +1,15 @@
 from typing import AsyncIterator, Protocol
 
 from flyteidl.admin import project_attributes_pb2, project_pb2, version_pb2
-from flyteidl.service import dataproxy_pb2
+from flyteidl.service import dataproxy_pb2, identity_pb2
+from flyteidl2.app import app_payload_pb2
+from flyteidl2.secret import payload_pb2
+from flyteidl2.task import task_service_pb2
+from flyteidl2.trigger import trigger_service_pb2
+from flyteidl2.workflow import run_logs_service_pb2, run_service_pb2
 from grpc.aio import UnaryStreamCall
 from grpc.aio._typing import RequestType
 
-from flyte._protos.secret import payload_pb2
-from flyte._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: ...
@@ -58,6 +60,28 @@ class TaskService(Protocol):
         self, request: task_service_pb2.GetTaskDetailsRequest
     ) -> task_service_pb2.GetTaskDetailsResponse: ...
 
+    async def ListTasks(self, request: task_service_pb2.ListTasksRequest) -> task_service_pb2.ListTasksResponse: ...
+
+
+class AppService(Protocol):
+    async def Create(self, request: app_payload_pb2.CreateRequest) -> app_payload_pb2.CreateResponse: ...
+
+    async def Get(self, request: app_payload_pb2.GetRequest) -> app_payload_pb2.GetResponse: ...
+
+    async def Update(self, request: app_payload_pb2.UpdateRequest) -> app_payload_pb2.UpdateResponse: ...
+
+    async def UpdateStatus(
+        self, request: app_payload_pb2.UpdateStatusRequest
+    ) -> app_payload_pb2.UpdateStatusResponse: ...
+
+    async def Delete(self, request: app_payload_pb2.DeleteRequest) -> app_payload_pb2.DeleteResponse: ...
+
+    async def List(self, request: app_payload_pb2.ListRequest) -> app_payload_pb2.ListResponse: ...
+
+    async def Watch(self, request: app_payload_pb2.WatchRequest) -> app_payload_pb2.WatchResponse: ...
+
+    async def Lease(self, request: app_payload_pb2.LeaseRequest) -> app_payload_pb2.LeaseResponse: ...
+
 
 class RunService(Protocol):
     async def CreateRun(self, request: run_service_pb2.CreateRunRequest) -> run_service_pb2.CreateRunResponse: ...
@@ -129,3 +153,37 @@ class SecretService(Protocol):
     async def ListSecrets(self, request: payload_pb2.ListSecretsRequest) -> payload_pb2.ListSecretsResponse: ...
 
     async def DeleteSecret(self, request: payload_pb2.DeleteSecretRequest) -> payload_pb2.DeleteSecretResponse: ...
+
+
+class IdentityService(Protocol):
+    async def UserInfo(self, request: identity_pb2.UserInfoRequest) -> identity_pb2.UserInfoResponse: ...
+
+
+class TriggerService(Protocol):
+    async def DeployTrigger(
+        self, request: trigger_service_pb2.DeployTriggerRequest
+    ) -> trigger_service_pb2.DeployTriggerResponse: ...
+
+    async def GetTriggerDetails(
+        self, request: trigger_service_pb2.GetTriggerDetailsRequest
+    ) -> trigger_service_pb2.GetTriggerDetailsResponse: ...
+
+    async def GetTriggerRevisionDetails(
+        self, request: trigger_service_pb2.GetTriggerRevisionDetailsRequest
+    ) -> trigger_service_pb2.GetTriggerRevisionDetailsResponse: ...
+
+    async def ListTriggers(
+        self, request: trigger_service_pb2.ListTriggersRequest
+    ) -> trigger_service_pb2.ListTriggersResponse: ...
+
+    async def GetTriggerRevisionHistory(
+        self, request: trigger_service_pb2.GetTriggerRevisionHistoryRequest
+    ) -> trigger_service_pb2.GetTriggerRevisionHistoryResponse: ...
+
+    async def UpdateTriggers(
+        self, request: trigger_service_pb2.UpdateTriggersRequest
+    ) -> trigger_service_pb2.UpdateTriggersResponse: ...
+
+    async def DeleteTriggers(
+        self, request: trigger_service_pb2.DeleteTriggersRequest
+    ) -> trigger_service_pb2.DeleteTriggersResponse: ...

flyte/remote/_client/auth/_auth_utils.py

@@ -0,0 +1,19 @@
+from __future__ import annotations
+
+import base64
+from typing import Literal
+
+
+def decode_api_key(encoded_str: str) -> tuple[str, str, str, str | Literal["None"]]:
+    """Decode encoded base64 string into app credentials. endpoint, client_id, client_secret, org"""
+    # Split with maxsplit=3 to handle endpoints with colons (e.g., dns:///endpoint.com)
+    parts = base64.b64decode(encoded_str.encode("utf-8")).decode("utf-8").split(":", 3)
+    if len(parts) != 4:
+        raise ValueError(f"Invalid API key format. Expected 4 parts separated by ':', got {len(parts)}")
+
+    endpoint, client_id, client_secret, org = parts
+    # For consistency, let's make sure org is always a non-empty string
+    if not org:
+        org = "None"
+
+    return endpoint, client_id, client_secret, org

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

@@ -34,6 +34,7 @@ class Authenticator(object):
         http_proxy_url: typing.Optional[str] = None,
         verify: bool = True,
         ca_cert_path: typing.Optional[str] = None,
+        default_header_key: str = "authorization",
         **kwargs,
     ):
         """
@@ -80,6 +81,7 @@ class Authenticator(object):
         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
+        self._default_header_key = default_header_key
 
     async def _resolve_config(self) -> ClientConfig:
         """
@@ -131,10 +133,14 @@ class Authenticator(object):
         """
         creds = self.get_credentials()
         if creds:
-            cfg = await self._resolve_config()
+            header_key = self._default_header_key
+            if self._resolved_config is not None:
+                # We only resolve the config during authentication flow, to avoid unnecessary network calls
+                # and usually the header_key is consistent.
+                header_key = self._resolved_config.header_key
             return GrpcAuthMetadata(
                 creds_id=creds.id,
-                pairs=Metadata((cfg.header_key, f"Bearer {creds.access_token}")),
+                pairs=Metadata((header_key, f"Bearer {creds.access_token}")),
             )
         return None
 

flyte/remote/_client/auth/_authenticators/device_code.py

@@ -1,4 +1,4 @@
-import click
+from rich import print as rich_print
 
 from flyte._logging import logger
 from flyte.remote._client.auth import _token_client as token_client
@@ -81,7 +81,7 @@ class DeviceCodeAuthenticator(Authenticator):
                     for_endpoint=self._endpoint,
                 )
             except (AuthenticationError, AuthenticationPending):
-                logger.warning("Failed to refresh token. Kicking off a full authorization flow.")
+                logger.warning("Logging in...")
 
         """Fall back to device flow"""
         resp = await token_client.get_device_code(
@@ -94,10 +94,9 @@ class DeviceCodeAuthenticator(Authenticator):
 
         full_uri = f"{resp.verification_uri}?user_code={resp.user_code}"
         text = (
-            f"To Authenticate, navigate in a browser to the following URL: "
-            f"{click.style(full_uri, fg='blue', underline=True)}"
+            f"To Authenticate, navigate in a browser to the following URL: [blue link={full_uri}]{full_uri}[/blue link]"
         )
-        click.secho(text)
+        rich_print(text)
         try:
             token, refresh_token, expires_in = await token_client.poll_token_endpoint(
                 resp,

flyte/remote/_client/auth/_authenticators/factory.py

@@ -169,6 +169,10 @@ def get_async_authenticator(
             from flyte.remote._client.auth._authenticators.device_code import DeviceCodeAuthenticator
 
             return DeviceCodeAuthenticator(endpoint=endpoint, cfg_store=cfg_store, verify=verify, **kwargs)
+        case "Passthrough":
+            from flyte.remote._client.auth._authenticators.passthrough import PassthroughAuthenticator
+
+            return PassthroughAuthenticator(endpoint=endpoint, **kwargs)
         case _:
             raise ValueError(
                 f"Invalid auth mode [{auth_type}] specified. Please update the creds config to use a valid value"

flyte/remote/_client/auth/_authenticators/passthrough.py

@@ -0,0 +1,79 @@
+import typing
+
+from grpc.aio import Metadata
+
+from flyte.remote._client.auth._authenticators.base import Authenticator, GrpcAuthMetadata
+from flyte.remote._client.auth._keyring import Credentials
+
+
+class PassthroughAuthenticator(Authenticator):
+    """
+    Passthrough authenticator that extracts headers from the context and passes them
+    to the gRPC calls without performing any authentication flow.
+
+    This authenticator is used when you want to pass custom authentication metadata
+    using the flyte.remote.auth_metadata() context manager.
+    """
+
+    def __init__(self, endpoint: str, **kwargs):
+        """
+        Initialize the passthrough authenticator.
+
+        We knowingly skip calling super, as that initializes a bunch of things that are not needed.!
+
+        :param endpoint: The endpoint URL
+        :param kwargs: Additional arguments (ignored for passthrough auth)
+        """
+        # Don't call parent __init__ to avoid unnecessary setup for passthrough auth
+        self._endpoint = endpoint
+        # We don't need credentials, config store, or HTTP session for passthrough
+        # We will create dummy creds
+        self._creds = Credentials(
+            access_token="passthrough",
+            for_endpoint=self._endpoint,
+        )
+        self._creds_id: str = "passthrough"
+
+    def refresh_credentials(self, creds_id: str | None = None):
+        return
+
+    def get_credentials(self) -> typing.Optional[Credentials]:
+        """
+        Passthrough authenticator doesn't use traditional credentials.
+        Returns a dummy credential to signal that metadata is available.
+        """
+        # Return a dummy credential so the interceptor knows to call get_grpc_call_auth_metadata
+        return self._creds
+
+    async def get_grpc_call_auth_metadata(self) -> typing.Optional[GrpcAuthMetadata]:
+        """
+        Fetch the authentication metadata from the context.
+
+        :return: GrpcAuthMetadata with the metadata from the context, or None if no metadata is available
+        """
+        # Lazy import to avoid circular dependencies
+        from flyte.remote._auth_metadata import get_auth_metadata
+
+        # Get metadata from context
+        metadata_tuples = get_auth_metadata()
+
+        if not metadata_tuples:
+            return None
+
+        return GrpcAuthMetadata(
+            creds_id=self._creds_id,
+            pairs=Metadata(*metadata_tuples),
+        )
+
+    async def _do_refresh_credentials(self) -> Credentials:
+        """
+        Passthrough authenticator doesn't need to refresh credentials.
+        This method should never be called in practice.
+        """
+        if self._creds is None:
+            # Just to satisfy mypy
+            return Credentials(
+                access_token="passthrough",
+                for_endpoint=self._endpoint,
+            )
+        return self._creds