flyte 2.0.0b22__py3-none-any.whl → 2.0.0b23__py3-none-any.whl

flyte/models.py

@@ -77,6 +77,37 @@ class ActionID:
         return self.new_sub_action(new_name)
 
 
+@rich.repr.auto
+@dataclass
+class PathRewrite:
+    """
+    Configuration for rewriting paths during input loading.
+    """
+
+    # If set, rewrites any path starting with this prefix to the new prefix.
+    old_prefix: str
+    new_prefix: str
+
+    def __post_init__(self):
+        if not self.old_prefix or not self.new_prefix:
+            raise ValueError("Both old_prefix and new_prefix must be non-empty strings.")
+        if self.old_prefix == self.new_prefix:
+            raise ValueError("old_prefix and new_prefix must be different.")
+
+    @classmethod
+    def from_str(cls, pattern: str) -> PathRewrite:
+        """
+        Create a PathRewrite from a string pattern of the form `old_prefix->new_prefix`.
+        """
+        parts = pattern.split("->")
+        if len(parts) != 2:
+            raise ValueError(f"Invalid path rewrite pattern: {pattern}. Expected format 'old_prefix->new_prefix'.")
+        return cls(old_prefix=parts[0], new_prefix=parts[1])
+
+    def __repr__(self) -> str:
+        return f"{self.old_prefix}->{self.new_prefix}"
+
+
 @rich.repr.auto
 @dataclass(frozen=True, kw_only=True)
 class RawDataPath:
@@ -86,6 +117,7 @@ class RawDataPath:
     """
 
     path: str
+    path_rewrite: Optional[PathRewrite] = None
 
     @classmethod
     def from_local_folder(cls, local_folder: str | pathlib.Path | None = None) -> RawDataPath:

flyte/remote/__init__.py

@@ -7,12 +7,15 @@ __all__ = [
     "ActionDetails",
     "ActionInputs",
     "ActionOutputs",
+    "Phase",
     "Project",
     "Run",
     "RunDetails",
     "Secret",
     "SecretTypes",
     "Task",
+    "Trigger",
+    "User",
     "create_channel",
     "upload_dir",
     "upload_file",
@@ -22,6 +25,8 @@ from ._action import Action, ActionDetails, ActionInputs, ActionOutputs
 from ._client.auth import create_channel
 from ._data import upload_dir, upload_file
 from ._project import Project
-from ._run import Run, RunDetails
+from ._run import Phase, Run, RunDetails
 from ._secret import Secret, SecretTypes
 from ._task import Task
+from ._trigger import Trigger
+from ._user import User

flyte/remote/_client/_protocols.py

@@ -1,12 +1,12 @@
 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 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
+from flyte._protos.workflow import run_logs_service_pb2, run_service_pb2, task_service_pb2, trigger_service_pb2
 
 
 class MetadataServiceProtocol(Protocol):
@@ -131,3 +131,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/controlplane.py

@@ -15,19 +15,26 @@ if "GRPC_VERBOSITY" not in os.environ:
 #### Has to be before grpc
 
 import grpc
-from flyteidl.service import admin_pb2_grpc, dataproxy_pb2_grpc
+from flyteidl.service import admin_pb2_grpc, dataproxy_pb2_grpc, identity_pb2_grpc
 
 from flyte._protos.secret import secret_pb2_grpc
-from flyte._protos.workflow import run_logs_service_pb2_grpc, run_service_pb2_grpc, task_service_pb2_grpc
+from flyte._protos.workflow import (
+    run_logs_service_pb2_grpc,
+    run_service_pb2_grpc,
+    task_service_pb2_grpc,
+    trigger_service_pb2_grpc,
+)
 
 from ._protocols import (
     DataProxyService,
+    IdentityService,
     MetadataServiceProtocol,
     ProjectDomainService,
     RunLogsService,
     RunService,
     SecretService,
     TaskService,
+    TriggerService,
 )
 from .auth import create_channel
 
@@ -38,7 +45,6 @@ class ClientSet:
         channel: grpc.aio.Channel,
         endpoint: str,
         insecure: bool = False,
-        data_proxy_channel: grpc.aio.Channel | None = None,
         **kwargs,
     ):
         self.endpoint = endpoint
@@ -50,6 +56,8 @@ class ClientSet:
         self._dataproxy = dataproxy_pb2_grpc.DataProxyServiceStub(channel=channel)
         self._log_service = run_logs_service_pb2_grpc.RunLogsServiceStub(channel=channel)
         self._secrets_service = secret_pb2_grpc.SecretServiceStub(channel=channel)
+        self._identity_service = identity_pb2_grpc.IdentityServiceStub(channel=channel)
+        self._trigger_service = trigger_service_pb2_grpc.TriggerServiceStub(channel=channel)
 
     @classmethod
     async def for_endpoint(cls, endpoint: str, *, insecure: bool = False, **kwargs) -> ClientSet:
@@ -105,5 +113,13 @@ class ClientSet:
     def secrets_service(self) -> SecretService:
         return self._secrets_service
 
+    @property
+    def identity_service(self) -> IdentityService:
+        return self._identity_service
+
+    @property
+    def trigger_service(self) -> TriggerService:
+        return self._trigger_service
+
     async def close(self, grace: float | None = None):
         return await self._channel.close(grace=grace)

flyte/remote/_run.py

@@ -7,6 +7,7 @@ import grpc
 import rich.repr
 
 from flyte._initialize import ensure_client, get_client, get_common_config
+from flyte._logging import logger
 from flyte._protos.common import identifier_pb2, list_pb2
 from flyte._protos.workflow import run_definition_pb2, run_service_pb2
 from flyte.syncify import syncify
@@ -16,6 +17,11 @@ from ._action import _action_details_rich_repr, _action_rich_repr
 from ._common import ToJSONMixin
 from ._console import get_run_url
 
+# @kumare3 is sadpanda, because we have to create a mirror of phase types here, because protobuf phases are ghastly
+Phase = Literal[
+    "queued", "waiting_for_resources", "initializing", "running", "succeeded", "failed", "aborted", "timed_out"
+]
+
 
 @dataclass
 class Run(ToJSONMixin):
@@ -40,14 +46,16 @@ class Run(ToJSONMixin):
     @classmethod
     async def listall(
         cls,
-        filters: str | None = None,
+        in_phase: Tuple[Phase] | None = None,
+        created_by_subject: str | None = None,
         sort_by: Tuple[str, Literal["asc", "desc"]] | None = None,
         limit: int = 100,
     ) -> AsyncIterator[Run]:
         """
         Get all runs for the current project and domain.
 
-        :param filters: The filters to apply to the project list.
+        :param in_phase: Filter runs by one or more phases.
+        :param created_by_subject: Filter runs by the subject that created them. (this is not username, but the subject)
         :param sort_by: The sorting criteria for the project list, in the format (field, order).
         :param limit: The maximum number of runs to return.
         :return: An iterator of runs.
@@ -59,6 +67,36 @@ class Run(ToJSONMixin):
             key=sort_by[0],
             direction=(list_pb2.Sort.ASCENDING if sort_by[1] == "asc" else list_pb2.Sort.DESCENDING),
         )
+        filters = []
+        if in_phase:
+            phases = [str(run_definition_pb2.Phase.Value(f"PHASE_{p.upper()}")) for p in in_phase]
+            logger.debug(f"Fetching run phases: {phases}")
+            if len(phases) > 1:
+                filters.append(
+                    list_pb2.Filter(
+                        function=list_pb2.Filter.Function.VALUE_IN,
+                        field="phase",
+                        values=phases,
+                    ),
+                )
+            else:
+                filters.append(
+                    list_pb2.Filter(
+                        function=list_pb2.Filter.Function.EQUAL,
+                        field="phase",
+                        values=phases[0],
+                    ),
+                )
+        if created_by_subject:
+            logger.debug(f"Fetching runs created by: {created_by_subject}")
+            filters.append(
+                list_pb2.Filter(
+                    function=list_pb2.Filter.Function.EQUAL,
+                    field="created_by",
+                    values=[created_by_subject],
+                ),
+            )
+
         cfg = get_common_config()
         i = 0
         while True:
@@ -66,6 +104,7 @@ class Run(ToJSONMixin):
                 limit=min(100, limit),
                 token=token,
                 sort_by=sort_pb2,
+                filters=filters,
             )
             resp = await get_client().run_service.ListRuns(
                 run_service_pb2.ListRunsRequest(
@@ -225,6 +264,7 @@ class Run(ToJSONMixin):
         """
         Rich representation of the Run object.
         """
+        yield "url", f"[blue bold][link={self.url}]link[/link][/blue bold]"
         yield from _action_rich_repr(self.pb2.action)
 
     def __repr__(self) -> str:

flyte/remote/_task.py

@@ -99,6 +99,7 @@ AutoVersioning = Literal["latest", "current"]
 class TaskDetails(ToJSONMixin):
     pb2: task_definition_pb2.TaskDetails
     max_inline_io_bytes: int = 10 * 1024 * 1024  # 10 MB
+    overriden_queue: Optional[str] = None
 
     @classmethod
     def get(
@@ -284,6 +285,13 @@ class TaskDetails(ToJSONMixin):
             f"Reference tasks [{self.name}] cannot be executed locally, only remotely."
         )
 
+    @property
+    def queue(self) -> Optional[str]:
+        """
+        The queue to use for the task.
+        """
+        return self.overriden_queue
+
     def override(
         self,
         *,
@@ -295,6 +303,7 @@ class TaskDetails(ToJSONMixin):
         secrets: Optional[flyte.SecretRequest] = None,
         max_inline_io_bytes: Optional[int] = None,
         cache: Optional[flyte.Cache] = None,
+        queue: Optional[str] = None,
         **kwargs: Any,
     ) -> TaskDetails:
         if len(kwargs) > 0:
@@ -342,7 +351,11 @@ class TaskDetails(ToJSONMixin):
             md.cache_serializable = cache.serialize
             md.cache_ignore_input_vars[:] = list(cache.ignored_inputs or ())
 
-        return TaskDetails(pb2, max_inline_io_bytes=max_inline_io_bytes or self.max_inline_io_bytes)
+        return TaskDetails(
+            pb2,
+            max_inline_io_bytes=max_inline_io_bytes or self.max_inline_io_bytes,
+            overriden_queue=queue,
+        )
 
     def __rich_repr__(self) -> rich.repr.Result:
         """

flyte/remote/_trigger.py

@@ -0,0 +1,308 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+from functools import cached_property
+from typing import AsyncIterator
+
+import grpc.aio
+
+import flyte
+from flyte._initialize import ensure_client, get_client, get_common_config
+from flyte._internal.runtime import trigger_serde
+from flyte._protos.common import identifier_pb2, list_pb2
+from flyte._protos.workflow import common_pb2, task_definition_pb2, trigger_definition_pb2, trigger_service_pb2
+from flyte.syncify import syncify
+
+from ._common import ToJSONMixin
+from ._task import Task, TaskDetails
+
+
+@dataclass
+class TriggerDetails(ToJSONMixin):
+    pb2: trigger_definition_pb2.TriggerDetails
+
+    @syncify
+    @classmethod
+    async def get(cls, *, name: str, task_name: str) -> TriggerDetails:
+        """
+        Retrieve detailed information about a specific trigger by its name.
+        """
+        ensure_client()
+        cfg = get_common_config()
+        resp = await get_client().trigger_service.GetTriggerDetails(
+            request=trigger_service_pb2.GetTriggerDetailsRequest(
+                name=identifier_pb2.TriggerName(
+                    task_name=task_name,
+                    name=name,
+                    org=cfg.org,
+                    project=cfg.project,
+                    domain=cfg.domain,
+                ),
+            )
+        )
+        return cls(pb2=resp.trigger)
+
+    @property
+    def name(self) -> str:
+        return self.id.name.name
+
+    @property
+    def id(self) -> identifier_pb2.TriggerIdentifier:
+        return self.pb2.id
+
+    @property
+    def task_name(self) -> str:
+        return self.pb2.id.name.task_name
+
+    @property
+    def automation_spec(self) -> common_pb2.TriggerAutomationSpec:
+        return self.pb2.automation_spec
+
+    @property
+    def metadata(self) -> trigger_definition_pb2.TriggerMetadata:
+        return self.pb2.metadata
+
+    @property
+    def status(self) -> trigger_definition_pb2.TriggerStatus:
+        return self.pb2.status
+
+    @property
+    def is_active(self) -> bool:
+        return self.pb2.spec.active
+
+    @cached_property
+    def trigger(self) -> trigger_definition_pb2.Trigger:
+        return trigger_definition_pb2.Trigger(
+            id=self.pb2.id,
+            automation_spec=self.automation_spec,
+            metadata=self.metadata,
+            status=self.status,
+            active=self.is_active,
+        )
+
+
+@dataclass
+class Trigger(ToJSONMixin):
+    pb2: trigger_definition_pb2.Trigger
+    details: TriggerDetails | None = None
+
+    @syncify
+    @classmethod
+    async def create(
+        cls,
+        trigger: flyte.Trigger,
+        task_name: str,
+        task_version: str | None = None,
+    ) -> Trigger:
+        """
+        Create a new trigger in the Flyte platform.
+
+        :param trigger: The flyte.Trigger object containing the trigger definition.
+        :param task_name: Optional name of the task to associate with the trigger.
+        """
+        ensure_client()
+        cfg = get_common_config()
+
+        # Fetch the task to ensure it exists and to get its input definitions
+        try:
+            lazy = (
+                Task.get(name=task_name, version=task_version)
+                if task_version
+                else Task.get(name=task_name, auto_version="latest")
+            )
+            task: TaskDetails = await lazy.fetch.aio()
+
+            task_trigger = await trigger_serde.to_task_trigger(
+                t=trigger,
+                task_name=task_name,
+                task_inputs=task.pb2.spec.task_template.interface.inputs,
+                task_default_inputs=list(task.pb2.spec.default_inputs),
+            )
+
+            resp = await get_client().trigger_service.DeployTrigger(
+                request=trigger_service_pb2.DeployTriggerRequest(
+                    id=identifier_pb2.TriggerIdentifier(
+                        name=identifier_pb2.TriggerName(
+                            name=trigger.name,
+                            task_name=task_name,
+                            org=cfg.org,
+                            project=cfg.project,
+                            domain=cfg.domain,
+                        ),
+                        revision=1,
+                    ),
+                    spec=trigger_definition_pb2.TriggerSpec(
+                        active=task_trigger.spec.active,
+                        inputs=task_trigger.spec.inputs,
+                        run_spec=task_trigger.spec.run_spec,
+                        task_version=task.version,
+                    ),
+                    automation_spec=task_trigger.automation_spec,
+                )
+            )
+
+            details = TriggerDetails(pb2=resp.trigger)
+
+            return cls(pb2=details.trigger, details=details)
+        except grpc.aio.AioRpcError as e:
+            if e.code() == grpc.StatusCode.NOT_FOUND:
+                raise ValueError(f"Task {task_name}:{task_version or 'latest'} not found") from e
+            raise
+
+    @syncify
+    @classmethod
+    async def get(cls, *, name: str, task_name: str) -> TriggerDetails:
+        """
+        Retrieve a trigger by its name and associated task name.
+        """
+        return await TriggerDetails.get(name=name, task_name=task_name)
+
+    @syncify
+    @classmethod
+    async def listall(
+        cls, task_name: str | None = None, task_version: str | None = None, limit: int = 100
+    ) -> AsyncIterator[Trigger]:
+        """
+        List all triggers associated with a specific task or all tasks if no task name is provided.
+        """
+        ensure_client()
+        cfg = get_common_config()
+        token = None
+        # task_name_id = None  TODO: implement listing by task name only
+        project_id = None
+        task_id = None
+        if task_name and task_version:
+            task_id = task_definition_pb2.TaskIdentifier(
+                name=task_name,
+                project=cfg.project,
+                domain=cfg.domain,
+                org=cfg.org,
+                version=task_version,
+            )
+        # elif task_name:  TODO: implement listing by task name only
+        #     task_name_id = task_definition_pb2.TaskName(
+        #         name=task_name,
+        #         project=cfg.project,
+        #         domain=cfg.domain,
+        #         org=cfg.org,
+        #     )
+        else:
+            project_id = identifier_pb2.ProjectIdentifier(
+                organization=cfg.org,
+                domain=cfg.domain,
+                name=cfg.project,
+            )
+
+        while True:
+            resp = await get_client().trigger_service.ListTriggers(
+                request=trigger_service_pb2.ListTriggersRequest(
+                    project_id=project_id,
+                    task_id=task_id,
+                    # task_name=task_name_id,
+                    request=list_pb2.ListRequest(
+                        limit=limit,
+                        token=token,
+                    ),
+                )
+            )
+            token = resp.token
+            for r in resp.triggers:
+                yield cls(r)
+            if not token:
+                break
+
+    @syncify
+    @classmethod
+    async def update(cls, name: str, task_name: str, active: bool):
+        """
+        Pause a trigger by its name and associated task name.
+        """
+        ensure_client()
+        cfg = get_common_config()
+        await get_client().trigger_service.UpdateTriggers(
+            request=trigger_service_pb2.UpdateTriggersRequest(
+                names=[
+                    identifier_pb2.TriggerName(
+                        org=cfg.org,
+                        project=cfg.project,
+                        domain=cfg.domain,
+                        name=name,
+                        task_name=task_name,
+                    )
+                ],
+                active=active,
+            )
+        )
+
+    @syncify
+    @classmethod
+    async def delete(cls, name: str, task_name: str):
+        """
+        Delete a trigger by its name.
+        """
+        ensure_client()
+        cfg = get_common_config()
+        await get_client().trigger_service.DeleteTriggers(
+            request=trigger_service_pb2.DeleteTriggersRequest(
+                names=[
+                    identifier_pb2.TriggerName(
+                        org=cfg.org,
+                        project=cfg.project,
+                        domain=cfg.domain,
+                        name=name,
+                        task_name=task_name,
+                    )
+                ],
+            )
+        )
+
+    @property
+    def id(self) -> identifier_pb2.TriggerIdentifier:
+        return self.pb2.id
+
+    @property
+    def name(self) -> str:
+        return self.id.name.name
+
+    @property
+    def task_name(self) -> str:
+        return self.id.name.task_name
+
+    @property
+    def automation_spec(self) -> common_pb2.TriggerAutomationSpec:
+        return self.pb2.automation_spec
+
+    async def get_details(self) -> TriggerDetails:
+        """
+        Get detailed information about this trigger.
+        """
+        if not self.details:
+            details = await TriggerDetails.get.aio(name=self.pb2.id.name.name)
+            self.details = details
+        return self.details
+
+    @property
+    def is_active(self) -> bool:
+        return self.pb2.active
+
+    def _rich_automation(self, automation: common_pb2.TriggerAutomationSpec):
+        if automation.type == common_pb2.TriggerAutomationSpec.TYPE_NONE:
+            yield "none", None
+        elif automation.type == common_pb2.TriggerAutomationSpec.TYPE_SCHEDULE:
+            if automation.schedule.cron_expression is not None:
+                yield "cron", automation.schedule.cron_expression
+            elif automation.schedule.rate is not None:
+                r = automation.schedule.rate
+                yield (
+                    "fixed_rate",
+                    (
+                        f"Every [{r.value}] {r.unit} starting at "
+                        f"{r.start_time.ToDatetime() if automation.HasField('start_time') else 'now'}"
+                    ),
+                )
+
+    def __rich_repr__(self):
+        yield "task_name", self.task_name
+        yield "name", self.name
+        yield from self._rich_automation(self.pb2.automation_spec)
+        yield "auto_activate", self.is_active

flyte/remote/_user.py

@@ -0,0 +1,33 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+
+from flyteidl.service import identity_pb2
+from flyteidl.service.identity_pb2 import UserInfoResponse
+
+from .._initialize import ensure_client, get_client
+from ..syncify import syncify
+from ._common import ToJSONMixin
+
+
+@dataclass
+class User(ToJSONMixin):
+    pb2: UserInfoResponse
+
+    @syncify
+    @classmethod
+    async def get(cls) -> User:
+        """
+        Fetches information about the currently logged in user.
+        Returns: A User object containing details about the user.
+        """
+        ensure_client()
+
+        resp = await get_client().identity_service.UserInfo(identity_pb2.UserInfoRequest())
+        return cls(resp)
+
+    def subject(self) -> str:
+        return self.pb2.subject
+
+    def name(self) -> str:
+        return self.pb2.name

flyte/storage/__init__.py

@@ -3,6 +3,8 @@ __all__ = [
     "GCS",
     "S3",
     "Storage",
+    "exists",
+    "exists_sync",
     "get",
     "get_configured_fsspec_kwargs",
     "get_random_local_directory",
@@ -11,13 +13,15 @@ __all__ = [
     "get_underlying_filesystem",
     "is_remote",
     "join",
+    "open",
     "put",
     "put_stream",
-    "put_stream",
 ]
 
 from ._config import ABFS, GCS, S3, Storage
 from ._storage import (
+    exists,
+    exists_sync,
     get,
     get_configured_fsspec_kwargs,
     get_random_local_directory,
@@ -26,6 +30,7 @@ from ._storage import (
     get_underlying_filesystem,
     is_remote,
     join,
+    open,
     put,
     put_stream,
 )