flyte 2.0.0b22__py3-none-any.whl → 2.0.0b30__py3-none-any.whl

flyte/connectors/_connector.py

@@ -0,0 +1,270 @@
+import asyncio
+import json
+import typing
+from abc import ABC, abstractmethod
+from dataclasses import asdict, dataclass
+from typing import Any, Dict, List, Optional
+
+from flyteidl2.core import tasks_pb2
+from flyteidl2.core.execution_pb2 import TaskExecution, TaskLog
+from flyteidl2.plugins import connector_pb2
+from flyteidl2.plugins.connector_pb2 import Connector as ConnectorProto
+from flyteidl2.plugins.connector_pb2 import (
+    GetTaskLogsResponse,
+    GetTaskMetricsResponse,
+    TaskCategory,
+    TaskExecutionMetadata,
+)
+from google.protobuf import json_format
+from google.protobuf.struct_pb2 import Struct
+
+from flyte import Secret
+from flyte._context import internal_ctx
+from flyte._initialize import get_init_config
+from flyte._internal.runtime.convert import convert_from_native_to_outputs
+from flyte._internal.runtime.task_serde import get_proto_task
+from flyte._logging import logger
+from flyte._task import TaskTemplate
+from flyte.connectors.utils import is_terminal_phase
+from flyte.models import 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.")
+
+        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)
+        resource_meta = await connector.create(task_template=tt, output_prefix=ctx.raw_data.path, inputs=kwargs)
+        resource = Resource(phase=TaskExecution.RUNNING)
+
+        while not is_terminal_phase(resource.phase):
+            resource = await connector.get(resource_meta=resource_meta)
+
+            if resource.log_links:
+                for link in resource.log_links:
+                    logger.info(f"{link.name}: {link.uri}")
+            await asyncio.sleep(1)
+
+        if resource.phase != TaskExecution.SUCCEEDED:
+            raise RuntimeError(f"Failed to run the task {task.name} with error: {resource.message}")
+
+        # TODO: Support abort
+
+        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)
+    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,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

@@ -174,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}"
         )
 
 
@@ -232,3 +232,12 @@ class SlowDownError(RuntimeUserError):
 
     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