flyte 2.0.0b32__py3-none-any.whl

flyte/remote/_trigger.py

@@ -0,0 +1,306 @@
+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:
+        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_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:
+        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.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):
+        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/report/__init__.py

@@ -0,0 +1,3 @@
+from ._report import Report, current_report, flush, get_tab, log, replace
+
+__all__ = ["Report", "current_report", "flush", "get_tab", "log", "replace"]

flyte/report/_report.py

@@ -0,0 +1,182 @@
+import html
+import pathlib
+import string
+from dataclasses import dataclass, field
+from typing import TYPE_CHECKING, Dict, List, Union
+
+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
+
+_MAIN_TAB_NAME = "main"
+
+
+@dataclass
+class Tab:
+    name: str
+    content: List[str] = field(default_factory=list, init=False)
+
+    def log(self, content: str):
+        """
+        Add content to the tab.
+        The content should be a valid HTML string, but not a complete HTML document, as it will be inserted into a div.
+
+        :param content: The content to add.
+        """
+        self.content.append(content)
+
+    def replace(self, content: str):
+        """
+        Replace the content of the tab.
+        The content should be a valid HTML string, but not a complete HTML document, as it will be inserted into a div.
+
+        :param content: The content to replace.
+        """
+        self.content = [content]
+
+    def get_html(self) -> str:
+        """
+        Get the HTML representation of the tab.
+
+        :return: The HTML representation of the tab.
+        """
+        return "\n".join(self.content)
+
+
+@dataclass
+class Report:
+    name: str
+    tabs: Dict[str, Tab] = field(default_factory=dict)
+    template_path: pathlib.Path = field(default_factory=lambda: pathlib.Path(__file__).parent / "_template.html")
+
+    def __post_init__(self):
+        self.tabs[_MAIN_TAB_NAME] = Tab(_MAIN_TAB_NAME)
+
+    def get_tab(self, name: str, create_if_missing: bool = True) -> Tab:
+        """
+        Get a tab by name. If the tab does not exist, create it.
+
+        :param name: The name of the tab.
+        :param create_if_missing: Whether to create the tab if it does not exist.
+        :return: The tab.
+        """
+        if name not in self.tabs:
+            if create_if_missing:
+                self.tabs[name] = Tab(name)
+            else:
+                raise ValueError(f"Tab {name} does not exist.")
+        return self.tabs[name]
+
+    def get_final_report(self) -> Union[str, "HTML"]:
+        """
+        Get the final report as a string.
+
+        :return: The final report.
+        """
+        tabs = {n: t.get_html() for n, t in self.tabs.items()}
+        nav_htmls = []
+        body_htmls = []
+
+        for key, value in tabs.items():
+            nav_htmls.append(f'<li onclick="handleLinkClick(this)">{html.escape(key)}</li>')
+            # Can not escape here because this is HTML. Escaping it will present the HTML as text.
+            # The renderer must ensure that the HTML is safe.
+            body_htmls.append(f"<div>{value}</div>")
+
+        template = string.Template(self.template_path.open("r").read())
+
+        raw_html = template.substitute(NAV_HTML="".join(nav_htmls), BODY_HTML="".join(body_htmls))
+        if ipython_check():
+            try:
+                from IPython.core.display import HTML
+
+                return HTML(raw_html)
+            except ImportError:
+                ...
+        return raw_html
+
+
+def get_tab(name: str, /, create_if_missing: bool = True) -> Tab:
+    """
+    Get a tab by name. If the tab does not exist, create it.
+
+    :param name: The name of the tab.
+    :param create_if_missing: Whether to create the tab if it does not exist.
+    :return: The tab.
+    """
+    report = current_report()
+    return report.get_tab(name, create_if_missing=create_if_missing)
+
+
+@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,
+     as it will be inserted into a div.
+
+    :param content: The content to log.
+    :param do_flush: flush the report after logging.
+    """
+    get_tab(_MAIN_TAB_NAME).log(content)
+    if do_flush:
+        await flush.aio()
+
+
+@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
+
+    report = internal_ctx().get_report()
+    if report is None:
+        return
+
+    report_html = report.get_final_report()
+    assert report_html is not None
+    assert isinstance(report_html, str)
+    report_path = io.report_path(internal_ctx().data.task_context.output_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}")
+
+
+@syncify
+async def replace(content: str, do_flush: bool = False):
+    """
+    Get the report. Replaces the content of the main tab.
+
+    :return: The report.
+    """
+    report = current_report()
+    if report is None:
+        return
+    report.get_tab(_MAIN_TAB_NAME).replace(content)
+    if do_flush:
+        await flush.aio()
+
+
+def current_report() -> Report:
+    """
+    Get the current report. This is a dummy report if not in a task context.
+
+    :return: The current report.
+    """
+    from flyte._context import internal_ctx
+
+    report = internal_ctx().get_report()
+    if report is None:
+        report = Report("dummy")
+    return report

flyte/report/_template.html

@@ -0,0 +1,124 @@
+<!doctype html>
+<html lang="">
+<head>
+    <meta charset="utf-8">
+    <title>User Content</title>
+    <meta name="viewport" content="width=device-width, initial-scale=1">
+    <link href="https://fonts.googleapis.com/css?family=Lato:300,400,700%7COpen+Sans:400,700" rel="stylesheet">
+    <style>
+        ol, ul {
+            list-style: none;
+        }
+
+        table {
+            border-collapse: collapse;
+            border-spacing: 0;
+        }
+
+        #flyte-frame-nav {
+            display: flex;
+            width: 100%;
+        }
+
+        #flyte-frame-tabs {
+            display: flex;
+            width: 100%;
+            justify-content: center;
+            margin-block: 0;
+            padding-inline-start: 0;
+        }
+
+        #flyte-frame-tabs li {
+            cursor: pointer;
+            padding: 8px;
+            margin: 0;
+            margin-right: 12px;
+            font-size: 14px;
+            line-height: 20px;
+            font-weight: 700;
+            font-style: normal;
+            font-family: Open Sans, helvetica, arial, sans-serif;
+            color: #666666;
+            width: 126px;
+            text-align: center;
+        }
+
+        #flyte-frame-tabs li:last-child {
+            margin-right: 0;
+        }
+
+        #flyte-frame-tabs li.active {
+            border-bottom: 4px solid rgb(163, 26, 255);
+            color: #333333;
+        }
+
+        #flyte-frame-container {
+            width: auto;
+        }
+
+        #flyte-frame-container > div {
+            display: None;
+        }
+
+        #flyte-frame-container > div.active {
+            display: block;
+            padding: 2rem 2rem;
+        }
+
+    </style>
+
+</head>
+<body>
+    <nav id="flyte-frame-nav">
+        <ul id="flyte-frame-tabs">
+            $NAV_HTML
+        </ul>
+    </nav>
+    <div id="flyte-frame-container">
+        $BODY_HTML
+    </div>
+</body>
+<script>
+    const setTabs = index => {
+        const container = document.getElementById('flyte-frame-tabs')
+        for (let i = 0; i < container.children.length; i++) {
+            const tabIndex = container.children[i].getAttribute('link_index')
+            if (tabIndex === index) {
+                container.children[i].classList.add('active')
+            } else {
+                container.children[i].className = ''
+            }
+        }
+    }
+    const setContent = index => {
+        const container = document.getElementById('flyte-frame-container')
+        for (let i = 0; i < container.children.length; i++) {
+            const tabIndex = container.children[i].getAttribute('link_index')
+            if (tabIndex === index) {
+                container.children[i].classList.add('active')
+            } else {
+                container.children[i].className = ''
+            }
+        }
+    }
+    const setLinkIndex = index => {
+        setTabs(index)
+        setContent(index)
+    }
+    const handleLinkClick = e => {
+        const linkIndex = e.getAttribute('link_index');
+        setLinkIndex(linkIndex)
+    }
+
+    const tabs = document.getElementById('flyte-frame-tabs');
+    const containers = document.getElementById('flyte-frame-container');
+    for(var i = 0; i < tabs.children.length; i++) {
+        if (i === 0) {
+            tabs.children[i].classList.add('active')
+            containers.children[i].classList.add('active')
+        }
+        tabs.children[i].setAttribute("link_index", i+1)
+        containers.children[i].setAttribute("link_index", i+1)
+    }
+</script>
+</html>

flyte/storage/__init__.py

@@ -0,0 +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",
+]
+
+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,
+)