flyte 2.0.0b13__py3-none-any.whl → 2.0.0b30__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

@@ -132,7 +132,9 @@ class CustomError(RuntimeUserError):
         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.
         """
-        return cls(e.__class__.__name__, str(e))
+        new_exc = cls(e.__class__.__name__, str(e))
+        new_exc.__cause__ = e
+        return new_exc
 
 
 class NotInTaskContextError(RuntimeUserError):
@@ -172,7 +174,7 @@ class RuntimeDataValidationError(RuntimeUserError):
 
     def __init__(self, var: str, e: Exception | str, task_name: str = ""):
         super().__init__(
-            "DataValiationError", f"In task {task_name} variable {var}, failed to serialize/deserialize because of {e}"
+            "DataValidationError", f"In task {task_name} variable {var}, failed to serialize/deserialize because of {e}"
         )
 
 
@@ -221,3 +223,21 @@ class RunAbortedError(RuntimeUserError):
 
     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

@@ -1,12 +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
+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/_container.py

@@ -2,7 +2,7 @@ import os
 import pathlib
 from typing import Any, Dict, List, Literal, Optional, Tuple, Type, Union
 
-from flyteidl.core import tasks_pb2
+from flyteidl2.core import tasks_pb2
 
 from flyte import Image, storage
 from flyte._logging import logger
@@ -83,6 +83,11 @@ class ContainerTask(TaskTemplate):
                 self._image = Image.from_debian_base()
             else:
                 self._image = Image.from_base(image)
+
+        if command and any(not isinstance(c, str) for c in command):
+            raise ValueError("All elements in the command list must be strings.")
+        if arguments and any(not isinstance(a, str) for a in arguments):
+            raise ValueError("All elements in the arguments list must be strings.")
         self._cmd = command
         self._args = arguments
         self._input_data_dir = input_data_dir

flyte/git/__init__.py

@@ -0,0 +1,3 @@
+from ._config import config_from_root
+
+__all__ = ["config_from_root"]

flyte/git/_config.py

@@ -0,0 +1,21 @@
+import pathlib
+import subprocess
+
+import flyte.config
+
+
+def config_from_root(path: pathlib.Path | str = ".flyte/config.yaml") -> flyte.config.Config | None:
+    """Get the config file from the git root directory.
+
+    By default, the config file is expected to be in `.flyte/config.yaml` in the git root directory.
+    """
+    try:
+        result = subprocess.run(["git", "rev-parse", "--show-toplevel"], check=False, capture_output=True, text=True)
+        if result.returncode != 0:
+            return None
+        root = pathlib.Path(result.stdout.strip())
+        if not (root / path).exists():
+            return None
+        return flyte.config.auto(root / path)
+    except Exception:
+        return None

flyte/io/__init__.py

@@ -7,6 +7,7 @@ of large datasets in Union.
 """
 
 __all__ = [
+    "PARQUET",
     "DataFrame",
     "DataFrameDecoder",
     "DataFrameEncoder",
@@ -17,6 +18,7 @@ __all__ = [
 ]
 
 from ._dataframe import (
+    PARQUET,
     DataFrame,
     DataFrameDecoder,
     DataFrameEncoder,

flyte/io/_dataframe/__init__.py

@@ -18,6 +18,7 @@ from flyte._logging import logger
 from flyte._utils.lazy_module import is_imported
 
 from .dataframe import (
+    PARQUET,
     DataFrame,
     DataFrameDecoder,
     DataFrameEncoder,
@@ -121,6 +122,7 @@ def lazy_import_dataframe_handler():
 
 
 __all__ = [
+    "PARQUET",
     "DataFrame",
     "DataFrameDecoder",
     "DataFrameEncoder",

flyte/io/_dataframe/basic_dfs.py

@@ -3,7 +3,7 @@ import typing
 from pathlib import Path
 from typing import TypeVar
 
-from flyteidl.core import literals_pb2, types_pb2
+from flyteidl2.core import literals_pb2, types_pb2
 from fsspec.core import split_protocol, strip_protocol
 
 import flyte.storage as storage
@@ -58,16 +58,16 @@ class PandasToCSVEncodingHandler(DataFrameEncoder):
 
         if not storage.is_remote(uri):
             Path(uri).mkdir(parents=True, exist_ok=True)
-        path = os.path.join(uri, ".csv")
+        csv_file = storage.join(uri, "data.csv")
         df = typing.cast(pd.DataFrame, dataframe.val)
         df.to_csv(
-            path,
+            csv_file,
             index=False,
-            storage_options=get_pandas_storage_options(uri=path),
+            storage_options=get_pandas_storage_options(uri=csv_file),
         )
         structured_dataset_type.format = CSV
         return literals_pb2.StructuredDataset(
-            uri=uri, metadata=literals_pb2.StructuredDatasetMetadata(structured_dataset_type)
+            uri=uri, metadata=literals_pb2.StructuredDatasetMetadata(structured_dataset_type=structured_dataset_type)
         )
 
 
@@ -83,16 +83,25 @@ class CSVToPandasDecodingHandler(DataFrameDecoder):
         uri = proto_value.uri
         columns = None
         kwargs = get_pandas_storage_options(uri=uri)
-        path = os.path.join(uri, ".csv")
+        csv_file = storage.join(uri, "data.csv")
         if current_task_metadata.structured_dataset_type and current_task_metadata.structured_dataset_type.columns:
             columns = [c.name for c in current_task_metadata.structured_dataset_type.columns]
         try:
-            return pd.read_csv(path, usecols=columns, storage_options=kwargs)
+            import io
+
+            # The pattern used here is a bit wonky because of obstore issues with csv, getting early eof error.
+            buf = io.BytesIO()
+            async for chunk in storage.get_stream(csv_file):
+                buf.write(chunk)
+            buf.seek(0)
+            df = pd.read_csv(buf)
+            return df
+
         except Exception as exc:
             if exc.__class__.__name__ == "NoCredentialsError":
                 logger.debug("S3 source detected, attempting anonymous S3 access")
                 kwargs = get_pandas_storage_options(uri=uri, anonymous=True)
-                return pd.read_csv(path, usecols=columns, storage_options=kwargs)
+                return pd.read_csv(csv_file, usecols=columns, storage_options=kwargs)
             else:
                 raise