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

flyte/remote/_trigger.py

@@ -0,0 +1,368 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+from functools import cached_property
+from typing import AsyncIterator
+
+import grpc.aio
+from flyteidl2.common import identifier_pb2, list_pb2
+from flyteidl2.task import common_pb2, task_definition_pb2
+from flyteidl2.trigger import trigger_definition_pb2, trigger_service_pb2
+
+import flyte
+from flyte._initialize import ensure_client, get_client, get_init_config
+from flyte._internal.runtime import trigger_serde
+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_init_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:
+        """
+        Name of the trigger.
+        """
+        return self.id.name.name
+
+    @property
+    def id(self) -> identifier_pb2.TriggerIdentifier:
+        """
+        Identifier for the trigger.
+        """
+        return self.pb2.id
+
+    @property
+    def task_name(self) -> str:
+        """
+        Name of the associated task
+        """
+        return self.pb2.id.name.task_name
+
+    @property
+    def automation_spec(self) -> common_pb2.TriggerAutomationSpec:
+        """
+        Get the automation specification for the trigger (e.g., schedule, event).
+        """
+        return self.pb2.automation_spec
+
+    @property
+    def metadata(self) -> trigger_definition_pb2.TriggerMetadata:
+        """
+        Get the metadata for the trigger.
+        """
+        return self.pb2.metadata
+
+    @property
+    def status(self) -> trigger_definition_pb2.TriggerStatus:
+        """
+        Get the current status of the trigger.
+        """
+        return self.pb2.status
+
+    @property
+    def is_active(self) -> bool:
+        """
+        Check if the trigger is currently active.
+        """
+        return self.pb2.spec.active
+
+    @cached_property
+    def trigger(self) -> trigger_definition_pb2.Trigger:
+        """
+        Get the trigger protobuf object constructed from the details.
+        """
+        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):
+    """
+    Represents a trigger in the Flyte platform.
+    """
+
+    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_init_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(
+                    name=identifier_pb2.TriggerName(
+                        name=trigger.name,
+                        task_name=task_name,
+                        org=cfg.org,
+                        project=cfg.project,
+                        domain=cfg.domain,
+                    ),
+                    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.aio(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_init_config()
+        token = None
+        task_name_id = None
+        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:
+            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_init_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_init_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:
+        """
+        Get the unique identifier for the trigger.
+        """
+        return self.pb2.id
+
+    @property
+    def name(self) -> str:
+        """
+        Get the name of the trigger.
+        """
+        return self.id.name.name
+
+    @property
+    def task_name(self) -> str:
+        """
+        Get the name of the task associated with this trigger.
+        """
+        return self.id.name.task_name
+
+    @property
+    def automation_spec(self) -> common_pb2.TriggerAutomationSpec:
+        """
+        Get the automation specification for the trigger.
+        """
+        return self.pb2.automation_spec
+
+    @property
+    def url(self) -> str:
+        """
+        Get the console URL for viewing the trigger.
+        """
+        client = get_client()
+        return client.console.trigger_url(
+            project=self.pb2.id.name.project,
+            domain=self.pb2.id.name.domain,
+            task_name=self.pb2.id.name.task_name,
+            trigger_name=self.name,
+        )
+
+    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:
+        """
+        Check if the trigger is currently active.
+        """
+        return self.pb2.active
+
+    def _rich_automation(self, automation: common_pb2.TriggerAutomationSpec):
+        """
+        Generate rich representation fields for the automation specification.
+        """
+        if automation.type == common_pb2.TriggerAutomationSpec.type.TYPE_NONE:
+            yield "none", None
+        elif automation.type == common_pb2.TriggerAutomationSpec.type.TYPE_SCHEDULE:
+            if automation.schedule.cron is not None:
+                yield "cron", automation.schedule.cron
+            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):
+        """
+        Rich representation of the Trigger object for pretty printing.
+        """
+        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,43 @@
+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):
+    """
+    Represents a user in the Flyte platform.
+    """
+
+    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:
+        """
+        Get the subject identifier of the user.
+        """
+        return self.pb2.subject
+
+    def name(self) -> str:
+        """
+        Get the name of the user.
+        """
+        return self.pb2.name

flyte/report/_report.py

@@ -4,10 +4,9 @@ import string
 from dataclasses import dataclass, field
 from typing import TYPE_CHECKING, Dict, List, Union
 
-from flyte._api_commons import syncer
-from flyte._internal.runtime import io
 from flyte._logging import logger
 from flyte._tools import ipython_check
+from flyte.syncify import syncify
 
 if TYPE_CHECKING:
     from IPython.core.display import HTML
@@ -112,7 +111,7 @@ def get_tab(name: str, /, create_if_missing: bool = True) -> Tab:
     return report.get_tab(name, create_if_missing=create_if_missing)
 
 
-@syncer.wrap
+@syncify
 async def log(content: str, do_flush: bool = False):
     """
     Log content to the main tab. The content should be a valid HTML string, but not a complete HTML document,
@@ -126,13 +125,14 @@ async def log(content: str, do_flush: bool = False):
         await flush.aio()
 
 
-@syncer.wrap
+@syncify
 async def flush():
     """
     Flush the report.
     """
     import flyte.storage as storage
     from flyte._context import internal_ctx
+    from flyte._internal.runtime import io
 
     if not internal_ctx().is_task_context():
         return
@@ -145,11 +145,15 @@ async def flush():
     assert report_html is not None
     assert isinstance(report_html, str)
     report_path = io.report_path(internal_ctx().data.task_context.output_path)
-    final_path = await storage.put_stream(report_html.encode("utf-8"), to_path=report_path)
+    content_types = {
+        "Content-Type": "text/html",  # For s3
+        "content_type": "text/html",  # For gcs
+    }
+    final_path = await storage.put_stream(report_html.encode("utf-8"), to_path=report_path, attributes=content_types)
     logger.debug(f"Report flushed to {final_path}")
 
 
-@syncer.wrap
+@syncify
 async def replace(content: str, do_flush: bool = False):
     """
     Get the report. Replaces the content of the main tab.

flyte/storage/__init__.py

@@ -1,24 +1,36 @@
 __all__ = [
+    "ABFS",
+    "GCS",
+    "S3",
+    "Storage",
+    "exists",
+    "exists_sync",
     "get",
+    "get_configured_fsspec_kwargs",
     "get_random_local_directory",
     "get_random_local_path",
     "get_stream",
     "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,
     get_random_local_path,
     get_stream,
     get_underlying_filesystem,
     is_remote,
     join,
+    open,
     put,
     put_stream,
 )