flyte 2.0.0b32__py3-none-any.whl

flyte/connectors/_server.py

@@ -0,0 +1,197 @@
+import inspect
+import os
+import sys
+from http import HTTPStatus
+from typing import Callable, Dict, List, Tuple, Type, Union
+
+import grpc
+from flyteidl2.core.security_pb2 import Connection
+from flyteidl2.plugins.connector_pb2 import (
+    CreateTaskRequest,
+    CreateTaskResponse,
+    DeleteTaskRequest,
+    DeleteTaskResponse,
+    GetConnectorRequest,
+    GetConnectorResponse,
+    GetTaskLogsRequest,
+    GetTaskLogsResponse,
+    GetTaskMetricsRequest,
+    GetTaskMetricsResponse,
+    GetTaskRequest,
+    GetTaskResponse,
+    ListConnectorsRequest,
+    ListConnectorsResponse,
+)
+from flyteidl2.service.connector_pb2_grpc import (
+    AsyncConnectorServiceServicer,
+    ConnectorMetadataServiceServicer,
+)
+from prometheus_client import Counter, Summary
+
+from flyte._internal.runtime.convert import Inputs, convert_from_inputs_to_native
+from flyte._logging import logger
+from flyte.connectors._connector import ConnectorRegistry, FlyteConnectorNotFound, get_resource_proto
+from flyte.connectors.utils import _start_grpc_server
+from flyte.models import NativeInterface, _has_default
+from flyte.syncify import syncify
+from flyte.types import TypeEngine
+
+metric_prefix = "flyte_connector_"
+create_operation = "create"
+get_operation = "get"
+delete_operation = "delete"
+
+# Follow the naming convention. https://prometheus.io/docs/practices/naming/
+request_success_count = Counter(
+    f"{metric_prefix}requests_success_total",
+    "Total number of successful requests",
+    ["task_type", "operation"],
+)
+request_failure_count = Counter(
+    f"{metric_prefix}requests_failure_total",
+    "Total number of failed requests",
+    ["task_type", "operation", "error_code"],
+)
+request_latency = Summary(
+    f"{metric_prefix}request_latency_seconds",
+    "Time spent processing connector request",
+    ["task_type", "operation"],
+)
+input_literal_size = Summary(f"{metric_prefix}input_literal_bytes", "Size of input literal", ["task_type"])
+
+
+def _handle_exception(e: Exception, context: grpc.ServicerContext, task_type: str, operation: str):
+    if isinstance(e, FlyteConnectorNotFound):
+        error_message = f"Cannot find connector for task type: {task_type}."
+        logger.error(error_message)
+        context.set_code(grpc.StatusCode.NOT_FOUND)
+        context.set_details(error_message)
+        request_failure_count.labels(task_type=task_type, operation=operation, error_code=HTTPStatus.NOT_FOUND).inc()
+    else:
+        error_message = f"failed to {operation} {task_type} task with error:\n {e}."
+        logger.error(error_message)
+        context.set_code(grpc.StatusCode.INTERNAL)
+        context.set_details(error_message)
+        request_failure_count.labels(
+            task_type=task_type, operation=operation, error_code=HTTPStatus.INTERNAL_SERVER_ERROR
+        ).inc()
+
+
+class ConnectorService:
+    @syncify
+    @classmethod
+    async def run(cls, port: int, prometheus_port: int, worker: int, timeout: int | None, modules: List[str] | None):
+        working_dir = os.getcwd()
+        if all(os.path.realpath(path) != working_dir for path in sys.path):
+            sys.path.append(working_dir)
+        await _start_grpc_server(port, prometheus_port, worker, timeout, modules)
+
+
+def record_connector_metrics(func: Callable):
+    async def wrapper(
+        self,
+        request: Union[CreateTaskRequest, GetTaskRequest, DeleteTaskRequest],
+        context: grpc.ServicerContext,
+        *args,
+        **kwargs,
+    ):
+        if isinstance(request, CreateTaskRequest):
+            task_type = request.template.type
+            operation = create_operation
+            if request.inputs:
+                input_literal_size.labels(task_type=task_type).observe(request.inputs.ByteSize())
+        elif isinstance(request, GetTaskRequest):
+            task_type = request.task_category.name
+            operation = get_operation
+        elif isinstance(request, DeleteTaskRequest):
+            task_type = request.task_category.name
+            operation = delete_operation
+        else:
+            context.set_code(grpc.StatusCode.UNIMPLEMENTED)
+            context.set_details("Method not implemented!")
+            return None
+
+        try:
+            with request_latency.labels(task_type=task_type, operation=operation).time():
+                res = await func(self, request, context, *args, **kwargs)
+            request_success_count.labels(task_type=task_type, operation=operation).inc()
+            return res
+        except Exception as e:
+            _handle_exception(e, context, task_type, operation)
+
+    return wrapper
+
+
+def _get_connection_kwargs(request: Connection) -> Dict[str, str]:
+    kwargs = {}
+
+    for k, v in request.secrets.items():
+        kwargs[k] = v
+    for k, v in request.configs.items():
+        kwargs[k] = v
+
+    return kwargs
+
+
+class AsyncConnectorService(AsyncConnectorServiceServicer):
+    @record_connector_metrics
+    async def CreateTask(self, request: CreateTaskRequest, context: grpc.ServicerContext) -> CreateTaskResponse:
+        template = request.template
+        connector = ConnectorRegistry.get_connector(template.type, template.task_type_version)
+        logger.info(f"{connector.name} start creating the job")
+        python_interface_inputs: Dict[str, Tuple[Type, Type[_has_default] | Type[inspect._empty]]] = {
+            name: (TypeEngine.guess_python_type(lt.type), inspect.Parameter.empty)
+            for name, lt in template.interface.inputs.variables.items()
+        }
+        native_interface = NativeInterface.from_types(inputs=python_interface_inputs, outputs={})
+        native_inputs = await convert_from_inputs_to_native(native_interface, Inputs(proto_inputs=request.inputs))
+        resource_meta = await connector.create(
+            task_template=request.template,
+            inputs=native_inputs,
+            output_prefix=request.output_prefix,
+            task_execution_metadata=request.task_execution_metadata,
+            connection=_get_connection_kwargs(request.connection),
+        )
+        return CreateTaskResponse(resource_meta=resource_meta.encode())
+
+    @record_connector_metrics
+    async def GetTask(self, request: GetTaskRequest, context: grpc.ServicerContext) -> GetTaskResponse:
+        connector = ConnectorRegistry.get_connector(request.task_category.name, request.task_category.version)
+        logger.info(f"{connector.name} start checking the status of the job")
+        res = await connector.get(
+            resource_meta=connector.metadata_type.decode(request.resource_meta),
+            connection=_get_connection_kwargs(request.connection),
+        )
+        return GetTaskResponse(resource=await get_resource_proto(res))
+
+    @record_connector_metrics
+    async def DeleteTask(self, request: DeleteTaskRequest, context: grpc.ServicerContext) -> DeleteTaskResponse:
+        connector = ConnectorRegistry.get_connector(request.task_category.name, request.task_category.version)
+        logger.info(f"{connector.name} start deleting the job")
+        await connector.delete(
+            resource_meta=connector.metadata_type.decode(request.resource_meta),
+            connection=_get_connection_kwargs(request.connection),
+        )
+        return DeleteTaskResponse()
+
+    async def GetTaskMetrics(
+        self, request: GetTaskMetricsRequest, context: grpc.ServicerContext
+    ) -> GetTaskMetricsResponse:
+        connector = ConnectorRegistry.get_connector(request.task_category.name, request.task_category.version)
+        logger.info(f"{connector.name} start getting metrics of the job")
+        return await connector.get_metrics(resource_meta=connector.metadata_type.decode(request.resource_meta))
+
+    async def GetTaskLogs(self, request: GetTaskLogsRequest, context: grpc.ServicerContext) -> GetTaskLogsResponse:
+        connector = ConnectorRegistry.get_connector(request.task_category.name, request.task_category.version)
+        logger.info(f"{connector.name} start getting logs of the job")
+        return await connector.get_logs(resource_meta=connector.metadata_type.decode(request.resource_meta))
+
+
+class ConnectorMetadataService(ConnectorMetadataServiceServicer):
+    async def GetConnector(self, request: GetConnectorRequest, context: grpc.ServicerContext) -> GetConnectorResponse:
+        return GetConnectorResponse(connector=ConnectorRegistry._get_connector_metadata(request.name))
+
+    async def ListConnectors(
+        self, request: ListConnectorsRequest, context: grpc.ServicerContext
+    ) -> ListConnectorsResponse:
+        return ListConnectorsResponse(connectors=ConnectorRegistry._list_connectors())

flyte/connectors/utils.py

@@ -0,0 +1,135 @@
+import importlib
+from concurrent import futures
+from importlib.metadata import entry_points
+from typing import List
+
+import click
+import grpc
+from flyteidl2.core.execution_pb2 import TaskExecution
+from flyteidl2.service import connector_pb2
+from flyteidl2.service.connector_pb2_grpc import (
+    add_AsyncConnectorServiceServicer_to_server,
+    add_ConnectorMetadataServiceServicer_to_server,
+)
+from rich.console import Console
+from rich.table import Table
+
+from flyte import logger
+
+
+def is_terminal_phase(phase: TaskExecution.Phase) -> bool:
+    """
+    Return true if the phase is terminal.
+    """
+    return phase in [TaskExecution.SUCCEEDED, TaskExecution.ABORTED, TaskExecution.FAILED]
+
+
+def convert_to_flyte_phase(state: str) -> TaskExecution.Phase:
+    """
+    Convert the state from the connector to the phase in flyte.
+    """
+    state = state.lower()
+    if state in ["failed", "timeout", "timedout", "canceled", "cancelled", "skipped"]:
+        return TaskExecution.FAILED
+    if state in ["internal_error"]:
+        return TaskExecution.RETRYABLE_FAILED
+    elif state in ["done", "succeeded", "success", "completed"]:
+        return TaskExecution.SUCCEEDED
+    elif state in ["running", "terminating"]:
+        return TaskExecution.RUNNING
+    elif state in ["pending"]:
+        return TaskExecution.INITIALIZING
+    raise ValueError(f"Unrecognized state: {state}")
+
+
+async def _start_grpc_server(
+    port: int, prometheus_port: int, worker: int, timeout: int | None, modules: List[str] | None
+):
+    try:
+        from flyte.connectors._server import (
+            AsyncConnectorService,
+            ConnectorMetadataService,
+        )
+    except ImportError as e:
+        raise ImportError(
+            "Flyte connector dependencies are not installed."
+            " Please install it using `pip install flyteplugins-connector`"
+        ) from e
+
+    click.secho("🚀 Starting the connector service...")
+    _load_connectors(modules)
+    _start_http_server(prometheus_port)
+
+    print_metadata()
+
+    server = grpc.aio.server(futures.ThreadPoolExecutor(max_workers=worker))
+
+    add_AsyncConnectorServiceServicer_to_server(AsyncConnectorService(), server)
+    add_ConnectorMetadataServiceServicer_to_server(ConnectorMetadataService(), server)
+    _start_health_check_server(server, worker)
+
+    server.add_insecure_port(f"[::]:{port}")
+    await server.start()
+    await server.wait_for_termination(timeout)
+
+
+def _start_http_server(prometheus_port: int):
+    try:
+        from prometheus_client import start_http_server
+
+        click.secho("Starting up the server to expose the prometheus metrics...")
+        start_http_server(prometheus_port)
+    except ImportError as e:
+        click.secho(f"Failed to start the prometheus server with error {e}", fg="red")
+
+
+def _start_health_check_server(server: grpc.Server, worker: int):
+    try:
+        from grpc_health.v1 import health, health_pb2, health_pb2_grpc
+
+        health_servicer = health.HealthServicer(
+            experimental_non_blocking=True,
+            experimental_thread_pool=futures.ThreadPoolExecutor(max_workers=worker),
+        )
+
+        for service in connector_pb2.DESCRIPTOR.services_by_name.values():
+            health_servicer.set(service.full_name, health_pb2.HealthCheckResponse.SERVING)
+        health_servicer.set(health.SERVICE_NAME, health_pb2.HealthCheckResponse.SERVING)
+
+        health_pb2_grpc.add_HealthServicer_to_server(health_servicer, server)
+
+    except ImportError as e:
+        click.secho(f"Failed to start the health check servicer with error {e}", fg="red")
+
+
+def print_metadata():
+    from flyte.connectors import ConnectorRegistry
+
+    connectors = ConnectorRegistry._list_connectors()
+
+    table = Table(title="Connector Metadata")
+    table.add_column("Connector Name", style="cyan", no_wrap=True)
+    table.add_column("Support Task Types", style="cyan")
+
+    for connector in connectors:
+        categories = ""
+        for category in connector.supported_task_categories:
+            categories += f"{category.name} ({category.version}) "
+        table.add_row(connector.name, categories)
+
+    console = Console()
+    console.print(table)
+
+
+def _load_connectors(modules: List[str] | None):
+    plugins = entry_points(group="flyte.connectors")
+    for ep in plugins:
+        try:
+            logger.info(f"Loading connector: {ep.name}")
+            ep.load()
+        except Exception as e:
+            logger.warning(f"Failed to load type transformer {ep.name} with error: {e}")
+
+    if modules:
+        for m in modules:
+            importlib.import_module(m)

flyte/errors.py

@@ -0,0 +1,243 @@
+"""
+Exceptions raised by Union.
+
+These errors are raised when the underlying task execution fails, either because of a user error, system error or an
+unknown error.
+"""
+
+from typing import Literal
+
+ErrorKind = Literal["system", "unknown", "user"]
+
+
+def silence_grpc_polling_error(loop, context):
+    """
+    Suppress specific gRPC polling errors in the event loop.
+    """
+    exc = context.get("exception")
+    if isinstance(exc, BlockingIOError):
+        return  # suppress
+    loop.default_exception_handler(context)
+
+
+class BaseRuntimeError(RuntimeError):
+    """
+    Base class for all Union runtime errors. These errors are raised when the underlying task execution fails, either
+    because of a user error, system error or an unknown error.
+    """
+
+    def __init__(self, code: str, kind: ErrorKind, root_cause_message: str, worker: str | None = None):
+        super().__init__(root_cause_message)
+        self.code = code
+        self.kind = kind
+        self.worker = worker
+
+
+class InitializationError(BaseRuntimeError):
+    """
+    This error is raised when the Union system is tried to access without being initialized.
+    """
+
+
+class RuntimeSystemError(BaseRuntimeError):
+    """
+    This error is raised when the underlying task execution fails because of a system error. This could be a bug in the
+    Union system or a bug in the user's code.
+    """
+
+    def __init__(self, code: str, message: str, worker: str | None = None):
+        super().__init__(code, "system", message, worker)
+
+
+class UnionRpcError(RuntimeSystemError):
+    """
+    This error is raised when communication with the Union server fails.
+    """
+
+
+class RuntimeUserError(BaseRuntimeError):
+    """
+    This error is raised when the underlying task execution fails because of an error in the user's code.
+    """
+
+    def __init__(self, code: str, message: str, worker: str | None = None):
+        super().__init__(code, "user", message, worker)
+
+
+class RuntimeUnknownError(BaseRuntimeError):
+    """
+    This error is raised when the underlying task execution fails because of an unknown error.
+    """
+
+    def __init__(self, code: str, message: str, worker: str | None = None):
+        super().__init__(code, "unknown", message, worker)
+
+
+class OOMError(RuntimeUserError):
+    """
+    This error is raised when the underlying task execution fails because of an out-of-memory error.
+    """
+
+
+class TaskInterruptedError(RuntimeUserError):
+    """
+    This error is raised when the underlying task execution is interrupted.
+    """
+
+
+class PrimaryContainerNotFoundError(RuntimeUserError):
+    """
+    This error is raised when the primary container is not found.
+    """
+
+
+class TaskTimeoutError(RuntimeUserError):
+    """
+    This error is raised when the underlying task execution runs for longer than the specified timeout.
+    """
+
+    def __init__(self, message: str):
+        super().__init__("TaskTimeoutError", message, "user")
+
+
+class RetriesExhaustedError(RuntimeUserError):
+    """
+    This error is raised when the underlying task execution fails after all retries have been exhausted.
+    """
+
+
+class InvalidImageNameError(RuntimeUserError):
+    """
+    This error is raised when the image name is invalid.
+    """
+
+
+class ImagePullBackOffError(RuntimeUserError):
+    """
+    This error is raised when the image cannot be pulled.
+    """
+
+
+class CustomError(RuntimeUserError):
+    """
+    This error is raised when the user raises a custom error.
+    """
+
+    def __init__(self, code: str, message: str):
+        super().__init__(code, message, "user")
+
+    @classmethod
+    def from_exception(cls, e: Exception):
+        """
+        Create a CustomError from an exception. The exception's class name is used as the error code and the exception
+        message is used as the error message.
+        """
+        new_exc = cls(e.__class__.__name__, str(e))
+        new_exc.__cause__ = e
+        return new_exc
+
+
+class NotInTaskContextError(RuntimeUserError):
+    """
+    This error is raised when the user tries to access the task context outside of a task.
+    """
+
+
+class ActionNotFoundError(RuntimeError):
+    """
+    This error is raised when the user tries to access an action that does not exist.
+    """
+
+
+class ReferenceTaskError(RuntimeUserError):
+    """
+    This error is raised when the user tries to access a task that does not exist.
+    """
+
+    def __init__(self, message: str):
+        super().__init__("ReferenceTaskUsageError", message, "user")
+
+
+class LogsNotYetAvailableError(BaseRuntimeError):
+    """
+    This error is raised when the logs are not yet available for a task.
+    """
+
+    def __init__(self, message: str):
+        super().__init__("LogsNotYetAvailable", "system", message, None)
+
+
+class RuntimeDataValidationError(RuntimeUserError):
+    """
+    This error is raised when the user tries to access a resource that does not exist or is invalid.
+    """
+
+    def __init__(self, var: str, e: Exception | str, task_name: str = ""):
+        super().__init__(
+            "DataValidationError", f"In task {task_name} variable {var}, failed to serialize/deserialize because of {e}"
+        )
+
+
+class DeploymentError(RuntimeUserError):
+    """
+    This error is raised when the deployment of a task fails, or some preconditions for deployment are not met.
+    """
+
+    def __init__(self, message: str):
+        super().__init__("DeploymentError", message, "user")
+
+
+class ImageBuildError(RuntimeUserError):
+    """
+    This error is raised when the image build fails.
+    """
+
+    def __init__(self, message: str):
+        super().__init__("ImageBuildError", message, "user")
+
+
+class ModuleLoadError(RuntimeUserError):
+    """
+    This error is raised when the module cannot be loaded, either because it does not exist or because of a
+     syntax error.
+    """
+
+    def __init__(self, message: str):
+        super().__init__("ModuleLoadError", message, "user")
+
+
+class InlineIOMaxBytesBreached(RuntimeUserError):
+    """
+    This error is raised when the inline IO max bytes limit is breached.
+    This can be adjusted per task by setting max_inline_io_bytes in the task definition.
+    """
+
+    def __init__(self, message: str):
+        super().__init__("InlineIOMaxBytesBreached", message, "user")
+
+
+class RunAbortedError(RuntimeUserError):
+    """
+    This error is raised when the run is aborted by the user.
+    """
+
+    def __init__(self, message: str):
+        super().__init__("RunAbortedError", message, "user")
+
+
+class SlowDownError(RuntimeUserError):
+    """
+    This error is raised when the user tries to access a resource that does not exist or is invalid.
+    """
+
+    def __init__(self, message: str):
+        super().__init__("SlowDownError", message, "user")
+
+
+class OnlyAsyncIOSupportedError(RuntimeUserError):
+    """
+    This error is raised when the user tries to use sync IO in an async task.
+    """
+
+    def __init__(self, message: str):
+        super().__init__("OnlyAsyncIOSupportedError", message, "user")

flyte/extend.py

@@ -0,0 +1,19 @@
+from ._initialize import is_initialized
+from ._internal.imagebuild.image_builder import ImageBuildEngine
+from ._internal.runtime.entrypoints import download_code_bundle
+from ._internal.runtime.resources_serde import get_proto_resources
+from ._resources import PRIMARY_CONTAINER_DEFAULT_NAME, pod_spec_from_resources
+from ._task import AsyncFunctionTaskTemplate, TaskTemplate
+from ._task_plugins import TaskPluginRegistry
+
+__all__ = [
+    "PRIMARY_CONTAINER_DEFAULT_NAME",
+    "AsyncFunctionTaskTemplate",
+    "ImageBuildEngine",
+    "TaskPluginRegistry",
+    "TaskTemplate",
+    "download_code_bundle",
+    "get_proto_resources",
+    "is_initialized",
+    "pod_spec_from_resources",
+]

flyte/extras/__init__.py

@@ -0,0 +1,5 @@
+from ._container import ContainerTask
+
+__all__ = [
+    "ContainerTask",
+]