flyte 0.0.1b0__py3-none-any.whl

flyte/remote/_secret.py

@@ -0,0 +1,132 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import AsyncGenerator, Iterator, Literal, Union
+
+import rich.repr
+
+from flyte._api_commons import syncer
+from flyte._initialize import get_client, get_common_config, requires_client
+from flyte._protos.secret import definition_pb2, payload_pb2
+
+SecretTypes = Literal["regular", "image_pull"]
+
+
+@dataclass
+class Secret:
+    pb2: definition_pb2.Secret
+
+    @classmethod
+    @requires_client
+    @syncer.wrap
+    async def create(cls, name: str, value: Union[str, bytes], type: SecretTypes = "regular"):
+        cfg = get_common_config()
+        secret_type = (
+            definition_pb2.SecretType.SECRET_TYPE_GENERIC
+            if type == "regular"
+            else definition_pb2.SecretType.SECRET_TYPE_IMAGE_PULL_SECRET
+        )
+
+        if isinstance(value, str):
+            secret = definition_pb2.SecretSpec(
+                type=secret_type,
+                string_value=value,
+            )
+        else:
+            secret = definition_pb2.SecretSpec(
+                type=secret_type,
+                binary_value=value,
+            )
+        await get_client().secrets_service.CreateSecret(  # type: ignore
+            request=payload_pb2.CreateSecretRequest(
+                id=definition_pb2.SecretIdentifier(
+                    organization=cfg.org,
+                    project=cfg.project,
+                    domain=cfg.domain,
+                    name=name,
+                ),
+                secret_spec=secret,
+            ),
+        )
+
+    @classmethod
+    @requires_client
+    @syncer.wrap
+    async def get(cls, name: str) -> Secret:
+        cfg = get_common_config()
+        resp = await get_client().secrets_service.GetSecret(
+            request=payload_pb2.GetSecretRequest(
+                id=definition_pb2.SecretIdentifier(
+                    organization=cfg.org,
+                    project=cfg.project,
+                    domain=cfg.domain,
+                    name=name,
+                )
+            )
+        )
+        return Secret(pb2=resp.secret)
+
+    @classmethod
+    @requires_client
+    @syncer.wrap
+    async def listall(cls, limit: int = 100) -> Union[Iterator[Secret], AsyncGenerator[Secret, None]]:
+        cfg = get_common_config()
+        token = None
+        while True:
+            resp = await get_client().secrets_service.ListSecrets(  # type: ignore
+                request=payload_pb2.ListSecretsRequest(
+                    organization=cfg.org,
+                    project=cfg.project,
+                    domain=cfg.domain,
+                    token=token,
+                    limit=limit,
+                ),
+            )
+            token = resp.token
+            for r in resp.secrets:
+                yield cls(r)
+            if not token:
+                break
+
+    @classmethod
+    @requires_client
+    @syncer.wrap
+    async def delete(cls, name):
+        cfg = get_common_config()
+        await get_client().secrets_service.DeleteSecret(  # type: ignore
+            request=payload_pb2.DeleteSecretRequest(
+                id=definition_pb2.SecretIdentifier(
+                    organization=cfg.org,
+                    project=cfg.project,
+                    domain=cfg.domain,
+                    name=name,
+                )
+            )
+        )
+
+    @property
+    def name(self) -> str:
+        return self.pb2.id.name
+
+    @property
+    def type(self) -> str:
+        if self.pb2.secret_metadata.type == definition_pb2.SecretType.SECRET_TYPE_GENERIC:
+            return "regular"
+        elif self.pb2.secret_metadata.type == definition_pb2.SecretType.SECRET_TYPE_IMAGE_PULL_SECRET:
+            return "image_pull"
+        raise ValueError("unknown type")
+
+    def __rich_repr__(self) -> rich.repr.Result:
+        yield "project", self.pb2.id.project or "-"
+        yield "domain", self.pb2.id.domain or "-"
+        yield "name", self.name
+        yield "type", self.type
+        yield "created_time", self.pb2.secret_metadata.created_time.ToDatetime().isoformat()
+        yield "status", definition_pb2.OverallStatus.Name(self.pb2.secret_metadata.secret_status.overall_status)
+        yield (
+            "cluster_status",
+            {
+                s.cluster.name: definition_pb2.SecretPresenceStatus.Name(s.presence_status)
+                for s in self.pb2.secret_metadata.secret_status.cluster_status
+            },
+        )

flyte/remote/_task.py

@@ -0,0 +1,227 @@
+from __future__ import annotations
+
+import functools
+from dataclasses import dataclass
+from threading import Lock
+from typing import Any, Callable, Coroutine, Dict, Literal, Optional, Union
+
+import rich.repr
+
+import flyte
+import flyte.errors
+from flyte._api_commons import syncer
+from flyte._context import internal_ctx
+from flyte._datastructures import NativeInterface
+from flyte._initialize import get_client, get_common_config
+from flyte._protos.workflow import task_definition_pb2, task_service_pb2
+
+
+class LazyEntity:
+    """
+    Fetches the entity when the entity is called or when the entity is retrieved.
+    The entity is derived from RemoteEntity so that it behaves exactly like the mimicked entity.
+    """
+
+    def __init__(self, name: str, getter: Callable[..., Coroutine[Any, Any, Task]], *args, **kwargs):
+        self._task: Optional[Task] = None
+        self._getter = getter
+        self._name = name
+        self._mutex = Lock()
+
+    @property
+    def name(self) -> str:
+        return self._name
+
+    @syncer.wrap
+    async def fetch(self) -> Task:
+        """
+        Forwards all other attributes to task, causing the task to be fetched!
+        """
+        with self._mutex:
+            if self._task is None:
+                self._task = await self._getter()
+            if self._task is None:
+                raise RuntimeError(f"Error downloading the task {self._name}, (check original exception...)")
+            return self._task
+
+    async def __call__(self, *args, **kwargs):
+        """
+        Forwards the call to the underlying task. The entity will be fetched if not already present
+        """
+        tk = await self.fetch.aio(self)
+        return await tk(*args, **kwargs)
+
+    def __repr__(self) -> str:
+        return str(self)
+
+    def __str__(self) -> str:
+        return f"Future for task with name {self._name}"
+
+
+AutoVersioning = Literal["latest", "current"]
+
+
+@dataclass
+class Task:
+    pb2: task_definition_pb2.TaskDetails
+
+    @classmethod
+    def get(cls, name: str, version: str | None = None, auto_version: AutoVersioning | None = None) -> LazyEntity:
+        """
+        Get a task by its ID or name. If both are provided, the ID will take precedence.
+
+        Either version or auto_version are required parameters.
+
+        :param uri: The URI of the task. If provided, do not provide the rest of the parameters.
+        :param name: The name of the task.
+        :param version: The version of the task.
+        :param auto_version: If set to "latest", the latest-by-time ordered from now, version of the task will be used.
+         If set to "current", the version will be derived from the callee tasks context. This is useful if you are
+         deploying all environments with the same version. If auto_version is current, you can only access the task from
+         within a task context.
+        """
+
+        if version is None and auto_version is None:
+            raise ValueError("Either version or auto_version must be provided.")
+
+        if version is None and auto_version not in ["latest", "current"]:
+            raise ValueError("auto_version must be either 'latest' or 'current'.")
+
+        async def deferred_get(_version: str | None, _auto_version: AutoVersioning | None) -> Task:
+            if _version is None:
+                if _auto_version == "latest":
+                    raise NotImplementedError("auto_version=latest is not yet implemented.")
+                elif _auto_version == "current":
+                    ctx = flyte.ctx()
+                    if ctx is None:
+                        raise ValueError("auto_version=current can only be used within a task context.")
+                    _version = ctx.version
+            cfg = get_common_config()
+            task_id = task_definition_pb2.TaskIdentifier(
+                org=cfg.org,
+                project=cfg.project,
+                domain=cfg.domain,
+                name=name,
+                version=_version,
+            )
+            resp = await get_client().task_service.GetTaskDetails(
+                task_service_pb2.GetTaskDetailsRequest(
+                    task_id=task_id,
+                )
+            )
+            return cls(resp.details)
+
+        return LazyEntity(
+            name=name, getter=functools.partial(deferred_get, _version=version, _auto_version=auto_version)
+        )
+
+    @property
+    def name(self) -> str:
+        """
+        The name of the task.
+        """
+        return self.pb2.task_id.name
+
+    @property
+    def version(self) -> str:
+        """
+        The version of the task.
+        """
+        return self.pb2.task_id.version
+
+    @property
+    def task_type(self) -> str:
+        """
+        The type of the task.
+        """
+        return self.pb2.spec.task_template.type
+
+    @functools.cached_property
+    def interface(self) -> NativeInterface:
+        """
+        The interface of the task.
+        """
+        import flyte.types as types
+
+        return types.guess_interface(self.pb2.spec.task_template.interface)
+
+    @property
+    def cache(self) -> flyte.Cache:
+        """
+        The cache policy of the task.
+        """
+        return flyte.Cache(
+            behavior="enabled" if self.pb2.spec.task_template.metadata.discoverable else "disable",
+            version_override=self.pb2.spec.task_template.metadata.discovery_version,
+            serialize=self.pb2.spec.task_template.metadata.cache_serializable,
+            ignored_inputs=tuple(self.pb2.spec.task_template.metadata.cache_ignore_input_vars),
+        )
+
+    @property
+    def secrets(self):
+        """
+        The secrets of the task.
+        """
+        return [s.key for s in self.pb2.spec.task_template.security_context.secrets]
+
+    @property
+    def resources(self):
+        """
+        The resources of the task.
+        """
+        if self.pb2.spec.task_template.container is None:
+            return ()
+        return (
+            self.pb2.spec.task_template.container.resources.requests,
+            self.pb2.spec.task_template.container.resources.limits,
+        )
+
+    async def __call__(self, *args, **kwargs):
+        """
+        Forwards the call to the underlying task. The entity will be fetched if not already present
+        """
+        ctx = internal_ctx()
+        if ctx.is_task_context():
+            # If we are in a task context, that implies we are executing a Run.
+            # In this scenario, we should submit the task to the controller.
+            # We will also check if we are not initialized, It is not expected to be not initialized
+            from flyte._internal.controllers import get_controller
+
+            controller = await get_controller()
+            if controller:
+                return await controller.submit_task_ref(self.pb2, *args, **kwargs)
+        raise flyte.errors
+
+    def override(
+        self,
+        *,
+        local: Optional[bool] = None,
+        ref: Optional[bool] = None,
+        resources: Optional[flyte.Resources] = None,
+        cache: flyte.CacheRequest = "auto",
+        retries: Union[int, flyte.RetryStrategy] = 0,
+        timeout: Optional[flyte.TimeoutType] = None,
+        reusable: Union[flyte.ReusePolicy, Literal["auto"], None] = None,
+        env: Optional[Dict[str, str]] = None,
+        secrets: Optional[flyte.SecretRequest] = None,
+        **kwargs: Any,
+    ) -> Task:
+        raise NotImplementedError
+
+    def __rich_repr__(self) -> rich.repr.Result:
+        """
+        Rich representation of the task.
+        """
+        yield "project", self.pb2.task_id.project
+        yield "domain", self.pb2.task_id.domain
+        yield "name", self.name
+        yield "version", self.version
+        yield "task_type", self.task_type
+        yield "cache", self.cache
+        yield "interface", self.name + str(self.interface)
+        yield "secrets", self.secrets
+        yield "resources", self.resources
+
+
+if __name__ == "__main__":
+    tk = Task.get(name="example_task")

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,178 @@
+import html
+import pathlib
+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
+
+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)
+
+
+@syncer.wrap
+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()
+
+
+@syncer.wrap
+async def flush():
+    """
+    Flush the report.
+    """
+    import flyte.storage as storage
+    from flyte._context import internal_ctx
+
+    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)
+    final_path = await storage.put_stream(report_html.encode("utf-8"), to_path=report_path)
+    logger.debug(f"Report flushed to {final_path}")
+
+
+@syncer.wrap
+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,24 @@
+__all__ = [
+    "get",
+    "get_random_local_directory",
+    "get_random_local_path",
+    "get_stream",
+    "get_underlying_filesystem",
+    "is_remote",
+    "join",
+    "put",
+    "put_stream",
+    "put_stream",
+]
+
+from ._storage import (
+    get,
+    get_random_local_directory,
+    get_random_local_path,
+    get_stream,
+    get_underlying_filesystem,
+    is_remote,
+    join,
+    put,
+    put_stream,
+)