flyte 0.0.1b0__py3-none-any.whl → 2.0.0b46__py3-none-any.whl

flyte/connectors/_connector.py

@@ -0,0 +1,330 @@
+import asyncio
+import json
+import os
+import typing
+from abc import ABC, abstractmethod
+from dataclasses import asdict, dataclass
+from typing import Any, Dict, List, Optional
+
+from flyteidl2.connector import connector_pb2
+from flyteidl2.connector.connector_pb2 import Connector as ConnectorProto
+from flyteidl2.connector.connector_pb2 import (
+    GetTaskLogsResponse,
+    GetTaskMetricsResponse,
+    TaskCategory,
+    TaskExecutionMetadata,
+)
+from flyteidl2.core import tasks_pb2
+from flyteidl2.core.execution_pb2 import TaskExecution, TaskLog
+from google.protobuf import json_format
+from google.protobuf.struct_pb2 import Struct
+
+import flyte.storage as storage
+from flyte import Secret
+from flyte._code_bundle import build_code_bundle
+from flyte._context import internal_ctx
+from flyte._deploy import build_images
+from flyte._initialize import get_init_config
+from flyte._internal.runtime import convert, io
+from flyte._internal.runtime.convert import convert_from_native_to_inputs, convert_from_native_to_outputs
+from flyte._internal.runtime.io import upload_inputs
+from flyte._internal.runtime.task_serde import get_proto_task
+from flyte._logging import logger
+from flyte._task import AsyncFunctionTaskTemplate, TaskTemplate
+from flyte.connectors.utils import _render_task_template, is_terminal_phase
+from flyte.models import CodeBundle, NativeInterface, SerializationContext
+from flyte.types._type_engine import dataclass_from_dict
+
+
+@dataclass(frozen=True)
+class ConnectorRegistryKey:
+    task_type_name: str
+    task_type_version: int
+
+
+@dataclass
+class ResourceMeta:
+    """
+    This is the metadata for the job. For example, the id of the job.
+    """
+
+    def encode(self) -> bytes:
+        """
+        Encode the resource meta to bytes.
+        """
+        return json.dumps(asdict(self)).encode("utf-8")
+
+    @classmethod
+    def decode(cls, data: bytes) -> "ResourceMeta":
+        """
+        Decode the resource meta from bytes.
+        """
+        return dataclass_from_dict(cls, json.loads(data.decode("utf-8")))
+
+
+@dataclass
+class Resource:
+    """
+    This is the output resource of the job.
+
+    Attributes
+    ----------
+        phase : TaskExecution.Phase
+            The phase of the job.
+        message : Optional[str]
+            The return message from the job.
+        log_links : Optional[List[TaskLog]]
+            The log links of the job. For example, the link to the BigQuery Console.
+        outputs : Optional[Union[LiteralMap, typing.Dict[str, Any]]]
+            The outputs of the job. If return python native types, the agent will convert them to flyte literals.
+        custom_info : Optional[typing.Dict[str, Any]]
+            The custom info of the job. For example, the job config.
+    """
+
+    phase: TaskExecution.Phase
+    message: Optional[str] = None
+    log_links: Optional[List[TaskLog]] = None
+    outputs: Optional[Dict[str, Any]] = None
+    custom_info: Optional[typing.Dict[str, Any]] = None
+
+
+class AsyncConnector(ABC):
+    """
+    This is the base class for all async connectors, and it defines the interface that all connectors must implement.
+    The connector service is responsible for invoking connectors.
+    The executor will communicate with the connector service to create tasks, get the status of tasks, and delete tasks.
+
+    All the connectors should be registered in the ConnectorRegistry.
+    Connector Service will look up the connector based on the task type and version.
+    """
+
+    name = "Async Connector"
+    task_type_name: str
+    task_type_version: int = 0
+    metadata_type: ResourceMeta
+
+    @abstractmethod
+    async def create(
+        self,
+        task_template: tasks_pb2.TaskTemplate,
+        output_prefix: str,
+        inputs: Optional[Dict[str, typing.Any]] = None,
+        task_execution_metadata: Optional[TaskExecutionMetadata] = None,
+        **kwargs,
+    ) -> ResourceMeta:
+        """
+        Return a resource meta that can be used to get the status of the task.
+        """
+        raise NotImplementedError
+
+    @abstractmethod
+    async def get(self, resource_meta: ResourceMeta, **kwargs) -> Resource:
+        """
+        Return the status of the task, and return the outputs in some cases. For example, bigquery job
+        can't write the structured dataset to the output location, so it returns the output literals to the propeller,
+        and the propeller will write the structured dataset to the blob store.
+        """
+        raise NotImplementedError
+
+    @abstractmethod
+    async def delete(self, resource_meta: ResourceMeta, **kwargs):
+        """
+        Delete the task. This call should be idempotent. It should raise an error if fails to delete the task.
+        """
+        raise NotImplementedError
+
+    async def get_metrics(self, resource_meta: ResourceMeta, **kwargs) -> GetTaskMetricsResponse:
+        """
+        Return the metrics for the task.
+        """
+        raise NotImplementedError
+
+    async def get_logs(self, resource_meta: ResourceMeta, **kwargs) -> GetTaskLogsResponse:
+        """
+        Return the metrics for the task.
+        """
+        raise NotImplementedError
+
+
+class ConnectorRegistry(object):
+    """
+    This is the registry for all connectors.
+    The connector service will look up the connector registry based on the task type and version.
+    """
+
+    _REGISTRY: typing.ClassVar[Dict[ConnectorRegistryKey, AsyncConnector]] = {}
+    _METADATA: typing.ClassVar[Dict[str, ConnectorProto]] = {}
+
+    @staticmethod
+    def register(connector: AsyncConnector, override: bool = False):
+        key = ConnectorRegistryKey(
+            task_type_name=connector.task_type_name, task_type_version=connector.task_type_version
+        )
+        if key in ConnectorRegistry._REGISTRY and override is False:
+            raise ValueError(
+                f"Duplicate connector for task type: {connector.task_type_name}"
+                f" and version: {connector.task_type_version}"
+            )
+        ConnectorRegistry._REGISTRY[key] = connector
+
+        task_category = TaskCategory(name=connector.task_type_name, version=connector.task_type_version)
+
+        if connector.name in ConnectorRegistry._METADATA:
+            connector_metadata = ConnectorRegistry._get_connector_metadata(connector.name)
+            connector_metadata.supported_task_categories.append(task_category)
+        else:
+            connector_metadata = ConnectorProto(
+                name=connector.name,
+                supported_task_categories=[task_category],
+            )
+            ConnectorRegistry._METADATA[connector.name] = connector_metadata
+
+    @staticmethod
+    def get_connector(task_type_name: str, task_type_version: int = 0) -> AsyncConnector:
+        key = ConnectorRegistryKey(task_type_name=task_type_name, task_type_version=task_type_version)
+        if key not in ConnectorRegistry._REGISTRY:
+            raise FlyteConnectorNotFound(
+                f"Cannot find connector for task type: {task_type_name} and version: {task_type_version}"
+            )
+        return ConnectorRegistry._REGISTRY[key]
+
+    @staticmethod
+    def _list_connectors() -> List[ConnectorProto]:
+        return list(ConnectorRegistry._METADATA.values())
+
+    @staticmethod
+    def _get_connector_metadata(name: str) -> ConnectorProto:
+        if name not in ConnectorRegistry._METADATA:
+            raise FlyteConnectorNotFound(f"Cannot find connector for name: {name}.")
+        return ConnectorRegistry._METADATA[name]
+
+
+class ConnectorSecretsMixin:
+    def __init__(self, secrets: Dict[str, str]):
+        # Key is the id of the secret, value is the secret name.
+        self._secrets = secrets
+
+    @property
+    def secrets(self) -> List[Secret]:
+        return [Secret(key=k, as_env_var=v) for k, v in self._secrets.items()]
+
+
+class AsyncConnectorExecutorMixin:
+    """
+    This mixin class is used to run the connector task locally, and it's only used for local execution.
+    Task should inherit from this class if the task can be run in the connector.
+    """
+
+    async def execute(self, **kwargs) -> Any:
+        task = typing.cast(TaskTemplate, self)
+        connector = ConnectorRegistry.get_connector(task.task_type, task.task_type_version)
+
+        ctx = internal_ctx()
+        tctx = internal_ctx().data.task_context
+        cfg = get_init_config()
+
+        if tctx is None:
+            raise RuntimeError("Task context is not set.")
+
+        if tctx.mode == "remote" and isinstance(self, AsyncFunctionTaskTemplate):
+            return await AsyncFunctionTaskTemplate.execute(self, **kwargs)
+
+        prefix = tctx.raw_data_path.get_random_remote_path()
+        if isinstance(self, AsyncFunctionTaskTemplate):
+            if not storage.is_remote(tctx.raw_data_path.path):
+                return await TaskTemplate.execute(self, **kwargs)
+            else:
+                local_code_bundle = await build_code_bundle(
+                    from_dir=cfg.root_dir,
+                    dryrun=True,
+                )
+                if local_code_bundle.tgz is None:
+                    raise RuntimeError("no tgz found in code bundle")
+                remote_code_path = await storage.put(
+                    local_code_bundle.tgz, prefix + "/code_bundle/" + os.path.basename(local_code_bundle.tgz)
+                )
+                sc = SerializationContext(
+                    project=tctx.action.project,
+                    domain=tctx.action.domain,
+                    org=tctx.action.org,
+                    code_bundle=CodeBundle(
+                        tgz=remote_code_path,
+                        computed_version=local_code_bundle.computed_version,
+                        destination="/opt/flyte/",
+                    ),
+                    version=tctx.version,
+                    image_cache=await build_images.aio(task.parent_env()) if task.parent_env else None,
+                    root_dir=cfg.root_dir,
+                )
+                tt = get_proto_task(task, sc)
+
+                tt = _render_task_template(tt, prefix)
+                inputs = await convert_from_native_to_inputs(task.native_interface, **kwargs)
+                inputs_uri = io.inputs_path(prefix)
+                await upload_inputs(inputs, inputs_uri)
+        else:
+            sc = SerializationContext(
+                project=tctx.action.project,
+                domain=tctx.action.domain,
+                org=tctx.action.org,
+                code_bundle=tctx.code_bundle,
+                version=tctx.version,
+                image_cache=tctx.compiled_image_cache,
+                root_dir=cfg.root_dir,
+            )
+            tt = get_proto_task(task, sc)
+
+        custom = json_format.MessageToDict(tt.custom)
+        secrets = custom["secrets"] if "secrets" in custom else {}
+        for k, v in secrets.items():
+            env_var = os.getenv(v)
+            if env_var is None:
+                raise ValueError(f"Secret {v} not found in environment.")
+            secrets[k] = env_var
+        resource_meta = await connector.create(
+            task_template=tt, output_prefix=ctx.raw_data.path, inputs=kwargs, **secrets
+        )
+        resource = Resource(phase=TaskExecution.RUNNING)
+
+        while not is_terminal_phase(resource.phase):
+            resource = await connector.get(resource_meta=resource_meta, **secrets)
+
+            if resource.log_links:
+                for link in resource.log_links:
+                    logger.info(f"{link.name}: {link.uri}")
+            await asyncio.sleep(3)
+
+        if resource.phase != TaskExecution.SUCCEEDED:
+            raise RuntimeError(f"Failed to run the task {task.name} with error: {resource.message}")
+
+        # TODO: Support abort
+        if (
+            isinstance(self, AsyncFunctionTaskTemplate)
+            and storage.is_remote(tctx.raw_data_path.path)
+            and await storage.exists(io.outputs_path(prefix))
+        ):
+            outputs = await io.load_outputs(io.outputs_path(prefix))
+            return await convert.convert_outputs_to_native(task.interface, outputs)
+
+        if resource.outputs is None:
+            return None
+        return tuple(resource.outputs.values())
+
+
+async def get_resource_proto(resource: Resource) -> connector_pb2.Resource:
+    if resource.outputs:
+        interface = NativeInterface.from_types(inputs={}, outputs={k: type(v) for k, v in resource.outputs.items()})
+        outputs = (await convert_from_native_to_outputs(tuple(resource.outputs.values()), interface)).proto_outputs
+    else:
+        outputs = None
+
+    return connector_pb2.Resource(
+        phase=resource.phase,
+        message=resource.message,
+        log_links=resource.log_links,
+        outputs=outputs,
+        custom_info=(json_format.Parse(json.dumps(resource.custom_info), Struct()) if resource.custom_info else None),
+    )
+
+
+class FlyteConnectorNotFound(ValueError): ...

flyte/connectors/_server.py

@@ -0,0 +1,194 @@
+import inspect
+import os
+import sys
+from http import HTTPStatus
+from typing import Callable, Dict, List, Tuple, Type, Union
+
+import grpc
+from flyteidl2.connector.connector_pb2 import (
+    CreateTaskRequest,
+    CreateTaskResponse,
+    DeleteTaskRequest,
+    DeleteTaskResponse,
+    GetConnectorRequest,
+    GetConnectorResponse,
+    GetTaskLogsRequest,
+    GetTaskLogsResponse,
+    GetTaskMetricsRequest,
+    GetTaskMetricsResponse,
+    GetTaskRequest,
+    GetTaskResponse,
+    ListConnectorsRequest,
+    ListConnectorsResponse,
+)
+from flyteidl2.connector.service_pb2_grpc import AsyncConnectorServiceServicer, ConnectorMetadataServiceServicer
+from flyteidl2.core.security_pb2 import Connection
+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,
+            **_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),
+            **_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),
+            **_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,159 @@
+import importlib
+from concurrent import futures
+from importlib.metadata import entry_points
+from typing import List
+
+import click
+import grpc
+from flyteidl2.connector import service_pb2
+from flyteidl2.connector.service_pb2_grpc import (
+    add_AsyncConnectorServiceServicer_to_server,
+    add_ConnectorMetadataServiceServicer_to_server,
+)
+from flyteidl2.core.execution_pb2 import TaskExecution
+from flyteidl2.core.tasks_pb2 import TaskTemplate
+from rich.console import Console
+from rich.table import Table
+
+import flyte
+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 service_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 connector '{ep.name}' with error: {e}")
+
+    if modules:
+        for m in modules:
+            importlib.import_module(m)
+
+
+def _render_task_template(tt: TaskTemplate, file_prefix: str) -> TaskTemplate:
+    if tt.container is None:
+        return tt
+    args = tt.container.args
+    ctx = flyte.ctx()
+    for i in range(len(args)):
+        tt.container.args[i] = args[i].replace("{{.input}}", f"{file_prefix}/inputs.pb")
+        tt.container.args[i] = args[i].replace("{{.outputPrefix}}", f"{file_prefix}")
+        tt.container.args[i] = args[i].replace("{{.rawOutputDataPrefix}}", f"{file_prefix}/raw_output")
+        tt.container.args[i] = args[i].replace("{{.checkpointOutputPrefix}}", f"{file_prefix}/checkpoint_output")
+        tt.container.args[i] = args[i].replace("{{.prevCheckpointPrefix}}", f"{file_prefix}/prev_checkpoint")
+        tt.container.args[i] = args[i].replace("{{.runName}}", ctx.action.run_name if ctx else "test-run")
+        tt.container.args[i] = args[i].replace("{{.actionName}}", "a1")
+
+    # Add additional required args
+    tt.container.args[1:1] = ["--run-base-dir", f"{file_prefix}/base_dir"]
+    tt.container.args[1:1] = ["--org", "test-org"]
+    tt.container.args[1:1] = ["--project", "test-project"]
+    tt.container.args[1:1] = ["--domain", "test-domain"]
+    return tt