flyte 2.0.0b32__py3-none-any.whl

flyte/remote/_client/controlplane.py

@@ -0,0 +1,128 @@
+from __future__ import annotations
+
+import os
+
+# Set environment variables for gRPC, this reduces log spew and avoids unnecessary warnings
+# before importing grpc
+if "GRPC_VERBOSITY" not in os.environ:
+    os.environ["GRPC_VERBOSITY"] = "ERROR"
+    os.environ["GRPC_CPP_MIN_LOG_LEVEL"] = "ERROR"
+    # Disable fork support (stops "skipping fork() handlers")
+    os.environ["GRPC_ENABLE_FORK_SUPPORT"] = "0"
+    # Reduce absl/glog verbosity
+    os.environ["GLOG_minloglevel"] = "2"
+    os.environ["ABSL_LOG"] = "0"
+#### Has to be before grpc
+
+import grpc
+from flyteidl.service import admin_pb2_grpc, dataproxy_pb2_grpc, identity_pb2_grpc
+from flyteidl2.app import app_service_pb2_grpc
+from flyteidl2.secret import secret_pb2_grpc
+from flyteidl2.task import task_service_pb2_grpc
+from flyteidl2.trigger import trigger_service_pb2_grpc
+from flyteidl2.workflow import run_logs_service_pb2_grpc, run_service_pb2_grpc
+
+from ._protocols import (
+    AppService,
+    DataProxyService,
+    IdentityService,
+    MetadataServiceProtocol,
+    ProjectDomainService,
+    RunLogsService,
+    RunService,
+    SecretService,
+    TaskService,
+    TriggerService,
+)
+from .auth import create_channel
+
+
+class ClientSet:
+    def __init__(
+        self,
+        channel: grpc.aio.Channel,
+        endpoint: str,
+        insecure: bool = False,
+        **kwargs,
+    ):
+        self.endpoint = endpoint
+        self.insecure = insecure
+        self._channel = channel
+        self._admin_client = admin_pb2_grpc.AdminServiceStub(channel=channel)
+        self._task_service = task_service_pb2_grpc.TaskServiceStub(channel=channel)
+        self._app_service = app_service_pb2_grpc.AppServiceStub(channel=channel)
+        self._run_service = run_service_pb2_grpc.RunServiceStub(channel=channel)
+        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:
+        return cls(
+            await create_channel(endpoint, None, insecure=insecure, **kwargs), endpoint, insecure=insecure, **kwargs
+        )
+
+    @classmethod
+    async def for_api_key(cls, api_key: str, *, insecure: bool = False, **kwargs) -> ClientSet:
+        from flyte.remote._client.auth._auth_utils import decode_api_key
+
+        # Parsing the API key is done in create_channel, but cleaner to redo it here rather than getting create_channel
+        # to return the endpoint
+        endpoint, _, _, _ = decode_api_key(api_key)
+
+        return cls(
+            await create_channel(None, api_key, insecure=insecure, **kwargs), endpoint, insecure=insecure, **kwargs
+        )
+
+    @classmethod
+    async def for_serverless(cls) -> ClientSet:
+        raise NotImplementedError
+
+    @classmethod
+    async def from_env(cls) -> ClientSet:
+        raise NotImplementedError
+
+    @property
+    def metadata_service(self) -> MetadataServiceProtocol:
+        return self._admin_client
+
+    @property
+    def project_domain_service(self) -> ProjectDomainService:
+        return self._admin_client
+
+    @property
+    def task_service(self) -> TaskService:
+        return self._task_service
+
+    @property
+    def app_service(self) -> AppService:
+        return self._app_service
+
+    @property
+    def run_service(self) -> RunService:
+        return self._run_service
+
+    @property
+    def dataproxy_service(self) -> DataProxyService:
+        return self._dataproxy
+
+    @property
+    def logs_service(self) -> RunLogsService:
+        return self._log_service
+
+    @property
+    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/_common.py

@@ -0,0 +1,30 @@
+import json
+
+from google.protobuf.json_format import MessageToDict, MessageToJson
+
+
+class ToJSONMixin:
+    """
+    A mixin class that provides a method to convert an object to a JSON-serializable dictionary.
+    """
+
+    def to_dict(self) -> dict:
+        """
+        Convert the object to a JSON-serializable dictionary.
+
+        Returns:
+            dict: A dictionary representation of the object.
+        """
+        if hasattr(self, "pb2"):
+            return MessageToDict(self.pb2)
+        else:
+            return {k: v for k, v in self.__dict__.items() if not k.startswith("_")}
+
+    def to_json(self) -> str:
+        """
+        Convert the object to a JSON string.
+
+        Returns:
+            str: A JSON string representation of the object.
+        """
+        return MessageToJson(self.pb2) if hasattr(self, "pb2") else json.dumps(self.to_dict())

flyte/remote/_console.py

@@ -0,0 +1,19 @@
+from urllib.parse import urlparse
+
+
+def _get_http_domain(endpoint: str, insecure: bool) -> str:
+    scheme = "http" if insecure else "https"
+    parsed = urlparse(endpoint)
+    if parsed.scheme == "dns":
+        domain = parsed.path.lstrip("/")
+    else:
+        domain = parsed.netloc or parsed.path
+    # TODO: make console url configurable
+    domain_split = domain.split(":")
+    if domain_split[0] == "localhost":
+        domain = domain if len(domain_split) > 1 else f"{domain}:8080"
+    return f"{scheme}://{domain}"
+
+
+def get_run_url(endpoint: str, insecure: bool, project: str, domain: str, run_name: str) -> str:
+    return f"{_get_http_domain(endpoint, insecure)}/v2/runs/project/{project}/domain/{domain}/{run_name}"

flyte/remote/_data.py

@@ -0,0 +1,161 @@
+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, ensure_client, get_client, get_init_config, require_project_and_domain
+from flyte.errors import InitializationError, RuntimeSystemError
+from flyte.syncify import syncify
+
+_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
+
+
+@require_project_and_domain
+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: {e.details()}"
+            )
+        elif e.code() == grpc.StatusCode.PERMISSION_DENIED:
+            raise RuntimeSystemError(
+                "PermissionDenied", f"Failed to get signed url for {fp}, please check your permissions: {e.details()}"
+            )
+        elif e.code() == grpc.StatusCode.UNAVAILABLE:
+            raise InitializationError("EndpointUnavailable", "user", "Service is unavailable.")
+        else:
+            raise RuntimeSystemError(e.code().value, f"Failed to get signed url for {fp}: {e.details()}")
+    except Exception as e:
+        raise RuntimeSystemError(type(e).__name__, f"Failed to get signed url for {fp}.") from e
+    logger.debug(f"Uploading to [link={resp.signed_url}]signed url[/link] for [link=file://{fp}]{fp}[/link]")
+    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 not in [200, 201, 204]:
+                raise RuntimeSystemError(
+                    "UploadFailed",
+                    f"Failed to upload {fp} to {resp.signed_url}, status code: {put_resp.status_code}, "
+                    f"response: {put_resp.text}",
+                )
+        # 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
+
+
+@syncify
+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.
+    ensure_client()
+    cfg = get_init_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)
+
+
+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.
+    ensure_client()
+    cfg = get_init_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,185 @@
+import asyncio
+from collections import deque
+from dataclasses import dataclass
+from typing import AsyncGenerator, AsyncIterator
+
+import grpc
+from flyteidl2.common import identifier_pb2
+from flyteidl2.logs.dataplane import payload_pb2
+from flyteidl2.workflow import run_logs_service_pb2
+from rich.console import Console
+from rich.live import Live
+from rich.panel import Panel
+from rich.text import Text
+
+from flyte._initialize import ensure_client, get_client
+from flyte._logging import logger
+from flyte._tools import ipython_check, ipywidgets_check
+from flyte.errors import LogsNotYetAvailableError
+from flyte.syncify import syncify
+
+style_map = {
+    payload_pb2.LogLineOriginator.SYSTEM: "bold magenta",
+    payload_pb2.LogLineOriginator.USER: "cyan",
+    payload_pb2.LogLineOriginator.UNKNOWN: "light red",
+}
+
+
+def _format_line(logline: payload_pb2.LogLine, show_ts: bool, filter_system: bool) -> Text | None:
+    if filter_system:
+        if logline.originator == payload_pb2.LogLineOriginator.SYSTEM:
+            return None
+    style = style_map.get(logline.originator, "")
+    if "[flyte]" in logline.message and "flyte.errors" not in logline.message:
+        if filter_system:
+            return None
+        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,
+        filter_system: bool = False,
+        panel: 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
+        self.filter_flyte = filter_system
+        self.panel = panel
+
+    def _render(self) -> Panel | Text:
+        log_text = Text()
+        for line in self.lines:
+            log_text.append(line)
+        if self.panel:
+            return Panel(log_text, title=self.name, border_style="yellow")
+        return log_text
+
+    async def run(self):
+        with Live(self._render(), refresh_per_second=20, console=self.console) as live:
+            try:
+                async for logline in self.log_source:
+                    formatted = _format_line(logline, show_ts=self.show_ts, filter_system=self.filter_flyte)
+                    if formatted:
+                        self.lines.append(formatted)
+                    self.total_lines += 1
+                    live.update(self._render())
+            except asyncio.CancelledError:
+                pass
+            except KeyboardInterrupt:
+                pass
+            except StopAsyncIteration:
+                self.console.print("[dim]Log stream ended.[/dim]")
+            except LogsNotYetAvailableError as e:
+                self.console.print(f"[red]Error:[/red] {e}")
+                live.update("")
+        self.console.print(f"Scrolled {self.total_lines} lines of logs.")
+
+
+@dataclass
+class Logs:
+    @syncify
+    @classmethod
+    async def tail(
+        cls,
+        action_id: identifier_pb2.ActionIdentifier,
+        attempt: int = 1,
+        retry: int = 5,
+    ) -> 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).
+        """
+        ensure_client()
+        retries = 0
+        while True:
+            try:
+                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
+                return
+            except asyncio.CancelledError:
+                return
+            except KeyboardInterrupt:
+                return
+            except StopAsyncIteration:
+                return
+            except grpc.aio.AioRpcError as e:
+                retries += 1
+                if retries >= retry:
+                    if e.code() == grpc.StatusCode.NOT_FOUND:
+                        raise LogsNotYetAvailableError(
+                            f"Log stream not available for action {action_id.name} in run {action_id.run.name}."
+                        )
+                else:
+                    await asyncio.sleep(2)
+
+    @classmethod
+    async def create_viewer(
+        cls,
+        action_id: identifier_pb2.ActionIdentifier,
+        attempt: int = 1,
+        max_lines: int = 30,
+        show_ts: bool = False,
+        raw: bool = False,
+        filter_system: bool = False,
+        panel: 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.
+        :param filter_system: Whether to filter log lines based on system logs.
+        :param panel: Whether to use a panel for the log viewer. only applicable if raw is False.
+        """
+        if attempt < 1:
+            raise ValueError("Attempt number must be greater than 0.")
+
+        if ipython_check():
+            if not ipywidgets_check():
+                logger.warning("IPython widgets is not available, defaulting to console output.")
+                raw = True
+
+        if raw:
+            console = Console()
+            async for line in cls.tail.aio(action_id=action_id, attempt=attempt):
+                line_text = _format_line(line, show_ts=show_ts, filter_system=filter_system)
+                if line_text:
+                    console.print(line_text, end="")
+            return
+        viewer = AsyncLogViewer(
+            log_source=cls.tail.aio(action_id=action_id, attempt=attempt),
+            max_lines=max_lines,
+            show_ts=show_ts,
+            name=f"{action_id.run.name}:{action_id.name} ({attempt})",
+            filter_system=filter_system,
+            panel=panel,
+        )
+        await viewer.run()

flyte/remote/_project.py

@@ -0,0 +1,88 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import AsyncIterator, Iterator, Literal, Tuple, Union
+
+import rich.repr
+from flyteidl.admin import common_pb2, project_pb2
+
+from flyte._initialize import ensure_client, get_client
+from flyte.syncify import syncify
+
+from ._common import ToJSONMixin
+
+
+# TODO Add support for orgs again
+@dataclass
+class Project(ToJSONMixin):
+    """
+    A class representing a project in the Union API.
+    """
+
+    pb2: project_pb2.Project
+
+    @syncify
+    @classmethod
+    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).
+        """
+        ensure_client()
+        service = get_client().project_domain_service  # type: ignore
+        resp = await service.GetProject(
+            project_pb2.ProjectGetRequest(
+                id=name,
+                # org=org,
+            )
+        )
+        return cls(resp)
+
+    @syncify
+    @classmethod
+    async def listall(
+        cls,
+        filters: str | None = None,
+        sort_by: Tuple[str, Literal["asc", "desc"]] | None = None,
+    ) -> Union[AsyncIterator[Project], Iterator[Project]]:
+        """
+        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.
+        """
+        ensure_client()
+        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,
+        )