flyte 0.0.1b0__py3-none-any.whl

flyte/remote/_data.py

@@ -0,0 +1,155 @@
+import asyncio
+import hashlib
+import os
+import typing
+import uuid
+from base64 import b64encode
+from datetime import timedelta
+from functools import lru_cache
+from pathlib import Path
+from typing import Tuple
+
+import aiofiles
+import grpc
+import httpx
+from flyteidl.service import dataproxy_pb2
+from google.protobuf import duration_pb2
+
+from flyte._initialize import CommonInit, get_client, get_common_config, requires_client
+from flyte.errors import RuntimeSystemError
+
+_UPLOAD_EXPIRES_IN = timedelta(seconds=60)
+
+
+def get_extra_headers_for_protocol(native_url: str) -> typing.Dict[str, str]:
+    """
+    For Azure Blob Storage, we need to set certain headers for http request.
+    This is used when we work with signed urls.
+    :param native_url:
+    :return:
+    """
+    if native_url.startswith("abfs://"):
+        return {"x-ms-blob-type": "BlockBlob"}
+    return {}
+
+
+@lru_cache
+def hash_file(file_path: typing.Union[os.PathLike, str]) -> Tuple[bytes, str, int]:
+    """
+    Hash a file and produce a digest to be used as a version
+    """
+    h = hashlib.md5()
+    size = 0
+
+    with open(file_path, "rb") as file:
+        while True:
+            # Reading is buffered, so we can read smaller chunks.
+            chunk = file.read(h.block_size)
+            if not chunk:
+                break
+            h.update(chunk)
+            size += len(chunk)
+
+    return h.digest(), h.hexdigest(), size
+
+
+async def _upload_single_file(
+    cfg: CommonInit, fp: Path, verify: bool = True, basedir: str | None = None
+) -> Tuple[str, str]:
+    md5_bytes, str_digest, _ = hash_file(fp)
+    from flyte._logging import logger
+
+    try:
+        expires_in_pb = duration_pb2.Duration()
+        expires_in_pb.FromTimedelta(_UPLOAD_EXPIRES_IN)
+        client = get_client()
+        resp = await client.dataproxy_service.CreateUploadLocation(  # type: ignore
+            dataproxy_pb2.CreateUploadLocationRequest(
+                project=cfg.project,
+                domain=cfg.domain,
+                content_md5=md5_bytes,
+                filename=fp.name,
+                expires_in=expires_in_pb,
+                filename_root=basedir,
+                add_content_md5_metadata=True,
+            )
+        )
+    except grpc.aio.AioRpcError as e:
+        if e.code() == grpc.StatusCode.NOT_FOUND:
+            raise RuntimeSystemError(
+                "NotFound", f"Failed to get signed url for {fp}, please check your project and domain."
+            )
+        elif e.code() == grpc.StatusCode.PERMISSION_DENIED:
+            raise RuntimeSystemError(
+                "PermissionDenied", f"Failed to get signed url for {fp}, please check your permissions."
+            )
+        else:
+            raise RuntimeSystemError(e.code().value, f"Failed to get signed url for {fp}.")
+    except Exception as e:
+        raise RuntimeSystemError(type(e).__name__, f"Failed to get signed url for {fp}.") from e
+    logger.debug(f"Uploading to signed url {resp.signed_url} for {fp}")
+    extra_headers = get_extra_headers_for_protocol(resp.native_url)
+    extra_headers.update(resp.headers)
+    encoded_md5 = b64encode(md5_bytes)
+    content_length = fp.stat().st_size
+
+    async with aiofiles.open(str(fp), "rb") as file:
+        extra_headers.update({"Content-Length": str(content_length), "Content-MD5": encoded_md5.decode("utf-8")})
+        async with httpx.AsyncClient(verify=verify) as aclient:
+            put_resp = await aclient.put(resp.signed_url, headers=extra_headers, content=file)
+            if put_resp.status_code != 200:
+                raise RuntimeSystemError(
+                    "UploadFailed",
+                    f"Failed to upload {fp} to {resp.signed_url}, status code: {put_resp.status_code}",
+                )
+        # TODO in old code we did this
+        #             if self._config.platform.insecure_skip_verify is True
+        #             else self._config.platform.ca_cert_file_path,
+    logger.debug(f"Uploaded with digest {str_digest}, blob location is {resp.native_url}")
+    return str_digest, resp.native_url
+
+
+@requires_client
+async def upload_file(fp: Path, verify: bool = True) -> Tuple[str, str]:
+    """
+    Uploads a file to a remote location and returns the remote URI.
+
+    :param fp: The file path to upload.
+    :param verify: Whether to verify the certificate for HTTPS requests.
+    :return: A tuple containing the MD5 digest and the remote URI.
+    """
+    # This is a placeholder implementation. Replace with actual upload logic.
+    cfg = get_common_config()
+    if not fp.is_file():
+        raise ValueError(f"{fp} is not a single file, upload arg must be a single file.")
+    return await _upload_single_file(cfg, fp, verify=verify)
+
+
+@requires_client
+async def upload_dir(dir_path: Path, verify: bool = True) -> str:
+    """
+    Uploads a directory to a remote location and returns the remote URI.
+
+    :param dir_path: The directory path to upload.
+    :param verify: Whether to verify the certificate for HTTPS requests.
+    :return: The remote URI of the uploaded directory.
+    """
+    # This is a placeholder implementation. Replace with actual upload logic.
+    cfg = get_common_config()
+    if not dir_path.is_dir():
+        raise ValueError(f"{dir_path} is not a directory, upload arg must be a directory.")
+
+    prefix = uuid.uuid4().hex
+
+    files = dir_path.rglob("*")
+    uploaded_files = []
+    for file in files:
+        if file.is_file():
+            uploaded_files.append(_upload_single_file(cfg, file, verify=verify, basedir=prefix))
+
+    urls = await asyncio.gather(*uploaded_files)
+    native_url = urls[0][1]  # Assuming all files are uploaded to the same prefix
+    # native_url is of the form s3://my-s3-bucket/flytesnacks/development/{prefix}/source/empty.md
+    uri = native_url.split(prefix)[0] + "/" + prefix
+
+    return uri

flyte/remote/_logs.py

@@ -0,0 +1,116 @@
+import asyncio
+from collections import deque
+from dataclasses import dataclass
+from typing import AsyncGenerator, AsyncIterator
+
+from rich.console import Console
+from rich.live import Live
+from rich.panel import Panel
+from rich.text import Text
+
+from flyte._api_commons import syncer
+from flyte._initialize import get_client, requires_client
+from flyte._protos.logs.dataplane import payload_pb2
+from flyte._protos.workflow import run_definition_pb2, run_logs_service_pb2
+
+
+def _format_line(logline: payload_pb2.LogLine, show_ts: bool) -> Text:
+    style_map = {
+        payload_pb2.LogLineOriginator.SYSTEM: "bold magenta",
+        payload_pb2.LogLineOriginator.USER: "cyan",
+        payload_pb2.LogLineOriginator.UNKNOWN: "light red",
+    }
+    style = style_map.get(logline.originator, "")
+    if "flyte" in logline.message and "flyte.errors" not in logline.message:
+        style = "dim"
+    ts = ""
+    if show_ts:
+        ts = f"[{logline.timestamp.ToDatetime().isoformat()}]"
+    return Text(f"{ts} {logline.message}", style=style)
+
+
+class AsyncLogViewer:
+    """
+    A class to view logs asynchronously in the console or terminal or jupyter notebook.
+    """
+
+    def __init__(self, log_source: AsyncIterator, max_lines: int = 30, name: str = "Logs", show_ts: bool = False):
+        self.console = Console()
+        self.log_source = log_source
+        self.max_lines = max_lines
+        self.lines: deque = deque(maxlen=max_lines + 1)
+        self.name = name
+        self.show_ts = show_ts
+        self.total_lines = 0
+
+    def _render(self):
+        log_text = Text()
+        for line in self.lines:
+            log_text.append(line)
+        return Panel(log_text, title=self.name, border_style="yellow")
+
+    async def run(self):
+        with Live(self._render(), refresh_per_second=10, console=self.console) as live:
+            try:
+                async for logline in self.log_source:
+                    formatted = _format_line(logline, show_ts=self.show_ts)
+                    self.lines.append(formatted)
+                    self.total_lines += 1
+                    live.update(self._render())
+            except asyncio.CancelledError:
+                pass
+        self.console.print(f"Scrolled {self.total_lines} lines of logs.")
+
+
+@dataclass
+class Logs:
+    @classmethod
+    @requires_client
+    @syncer.wrap
+    async def tail(
+        cls, action_id: run_definition_pb2.ActionIdentifier, attempt: int = 1
+    ) -> AsyncGenerator[payload_pb2.LogLine, None]:
+        """
+        Tail the logs for a given action ID and attempt.
+        :param action_id: The action ID to tail logs for.
+        :param attempt: The attempt number (default is 0).
+        """
+        resp = get_client().logs_service.TailLogs(
+            run_logs_service_pb2.TailLogsRequest(action_id=action_id, attempt=attempt)
+        )
+        async for log_set in resp:
+            if log_set.logs:
+                for log in log_set.logs:
+                    for line in log.lines:
+                        yield line
+
+    @classmethod
+    async def create_viewer(
+        cls,
+        action_id: run_definition_pb2.ActionIdentifier,
+        attempt: int = 1,
+        max_lines: int = 30,
+        show_ts: bool = False,
+        raw: bool = False,
+    ):
+        """
+        Create a log viewer for a given action ID and attempt.
+        :param action_id: Action ID to view logs for.
+        :param attempt: Attempt number (default is 1).
+        :param max_lines: Maximum number of lines to show if using the viewer. The logger will scroll
+           and keep only max_lines in view.
+        :param show_ts: Whether to show timestamps in the logs.
+        :param raw: if True, return the raw log lines instead of a viewer.
+        """
+        if raw:
+            console = Console()
+            async for line in cls.tail.aio(cls, action_id=action_id, attempt=attempt):
+                console.print(_format_line(line, show_ts=show_ts), end="")
+            return
+        viewer = AsyncLogViewer(
+            log_source=cls.tail.aio(cls, action_id=action_id, attempt=attempt),
+            max_lines=max_lines,
+            show_ts=show_ts,
+            name=f"{action_id.run.name}:{action_id.name}",
+        )
+        await viewer.run()

flyte/remote/_project.py

@@ -0,0 +1,86 @@
+from __future__ import annotations
+
+import typing
+from dataclasses import dataclass
+from typing import AsyncGenerator, Literal, Tuple
+
+import rich.repr
+from flyteidl.admin import common_pb2, project_pb2
+
+from flyte._api_commons import syncer
+from flyte._initialize import get_client, get_common_config, requires_client
+
+
+@dataclass
+class Project:
+    """
+    A class representing a project in the Union API.
+    """
+
+    _pb2: project_pb2.Project
+
+    @classmethod
+    @requires_client
+    @syncer.wrap
+    async def get(cls, name: str, org: str | None = None) -> Project:
+        """
+        Get a run by its ID or name. If both are provided, the ID will take precedence.
+
+        :param name: The name of the project.
+        :param org: The organization of the project (if applicable).
+        """
+        service = get_client().project_domain_service  # type: ignore
+        resp = await service.GetProject(
+            project_pb2.ProjectGetRequest(
+                id=name,
+                org=org,
+            )
+        )
+        return cls(resp)
+
+    @classmethod
+    @requires_client
+    @syncer.wrap
+    async def listall(
+        cls,
+        filters: str | None = None,
+        sort_by: Tuple[str, Literal["asc", "desc"]] | None = None,
+    ) -> typing.Union[typing.Iterator[Project], AsyncGenerator[Project, None]]:
+        """
+        Get a run by its ID or name. If both are provided, the ID will take precedence.
+
+        :param filters: The filters to apply to the project list.
+        :param sort_by: The sorting criteria for the project list, in the format (field, order).
+        :return: An iterator of projects.
+        """
+        token = None
+        sort_by = sort_by or ("created_at", "asc")
+        sort_pb2 = common_pb2.Sort(
+            key=sort_by[0], direction=common_pb2.Sort.ASCENDING if sort_by[1] == "asc" else common_pb2.Sort.DESCENDING
+        )
+        org = get_common_config().org
+        while True:
+            resp = await get_client().project_domain_service.ListProjects(  # type: ignore
+                project_pb2.ProjectListRequest(
+                    limit=100,
+                    token=token,
+                    filters=filters,
+                    sort_by=sort_pb2,
+                    org=org,
+                )
+            )
+            token = resp.token
+            for p in resp.projects:
+                yield cls(p)
+            if not token:
+                break
+
+    def __rich_repr__(self) -> rich.repr.Result:
+        yield "name", self._pb2.name
+        yield "id", self._pb2.id
+        yield "description", self._pb2.description
+        yield "state", project_pb2.Project.ProjectState.Name(self._pb2.state)
+        yield (
+            "labels",
+            ", ".join([f"{k}: {v}" for k, v in self._pb2.labels.values.items()]) if self._pb2.labels else None,
+        )