flyte 2.0.0b23__py3-none-any.whl → 2.0.0b25__py3-none-any.whl

flyte/_task.py

@@ -20,6 +20,7 @@ from typing import (
     TypeVar,
     Union,
     cast,
+    overload,
 )
 
 from flyte._pod import PodTemplate
@@ -37,7 +38,7 @@ from ._timeout import TimeoutType
 from .models import MAX_INLINE_IO_BYTES, NativeInterface, SerializationContext
 
 if TYPE_CHECKING:
-    from flyteidl.core.tasks_pb2 import DataLoadingConfig
+    from flyteidl2.core.tasks_pb2 import DataLoadingConfig
 
     from flyte.trigger import Trigger
 
@@ -48,11 +49,12 @@ R = TypeVar("R")  # return type
 
 AsyncFunctionType: TypeAlias = Callable[P, Coroutine[Any, Any, R]]
 SyncFunctionType: TypeAlias = Callable[P, R]
-FunctionTypes: TypeAlias = Union[AsyncFunctionType, SyncFunctionType]
+FunctionTypes: TypeAlias = AsyncFunctionType | SyncFunctionType
+F = TypeVar("F", bound=FunctionTypes)
 
 
 @dataclass(kw_only=True)
-class TaskTemplate(Generic[P, R]):
+class TaskTemplate(Generic[P, R, F]):
     """
     Task template is a template for a task that can be executed. It defines various parameters for the task, which
     can be defined statically at the time of task definition or dynamically at the time of task invocation using
@@ -87,6 +89,7 @@ class TaskTemplate(Generic[P, R]):
     :param pod_template: Optional The pod template to use for the task.
     :param report: Optional Whether to report the task execution to the Flyte console, defaults to False.
     :param queue: Optional The queue to use for the task. If not provided, the default queue will be used.
+    :param debuggable: Optional Whether the task supports debugging capabilities, defaults to False.
     """
 
     name: str
@@ -107,8 +110,10 @@ class TaskTemplate(Generic[P, R]):
     pod_template: Optional[Union[str, PodTemplate]] = None
     report: bool = False
     queue: Optional[str] = None
+    debuggable: bool = False
 
     parent_env: Optional[weakref.ReferenceType[TaskEnvironment]] = None
+    parent_env_name: Optional[str] = None
     ref: bool = field(default=False, init=False, repr=False, compare=False)
     max_inline_io_bytes: int = MAX_INLINE_IO_BYTES
     triggers: Tuple[Trigger, ...] = field(default_factory=tuple)
@@ -225,6 +230,14 @@ class TaskTemplate(Generic[P, R]):
     def native_interface(self) -> NativeInterface:
         return self.interface
 
+    @overload
+    async def aio(self: TaskTemplate[P, R, SyncFunctionType], *args: P.args, **kwargs: P.kwargs) -> R: ...
+
+    @overload
+    async def aio(
+        self: TaskTemplate[P, R, AsyncFunctionType], *args: P.args, **kwargs: P.kwargs
+    ) -> Coroutine[Any, Any, R]: ...
+
     async def aio(self, *args: P.args, **kwargs: P.kwargs) -> Coroutine[Any, Any, R] | R:
         """
         The aio function allows executing "sync" tasks, in an async context. This helps with migrating v1 defined sync
@@ -248,7 +261,6 @@ class TaskTemplate(Generic[P, R]):
         :param kwargs:
         :return:
         """
-
         ctx = internal_ctx()
         if ctx.is_task_context():
             from ._internal.controllers import get_controller
@@ -273,6 +285,14 @@ class TaskTemplate(Generic[P, R]):
             # even for synchronous tasks. This is to support migration.
             return self.forward(*args, **kwargs)
 
+    @overload
+    def __call__(self: TaskTemplate[P, R, SyncFunctionType], *args: P.args, **kwargs: P.kwargs) -> R: ...
+
+    @overload
+    def __call__(
+        self: TaskTemplate[P, R, AsyncFunctionType], *args: P.args, **kwargs: P.kwargs
+    ) -> Coroutine[Any, Any, R]: ...
+
     def __call__(self, *args: P.args, **kwargs: P.kwargs) -> Coroutine[Any, Any, R] | R:
         """
         This is the entrypoint for an async function task at runtime. It will be called during an execution.
@@ -424,14 +444,15 @@ class TaskTemplate(Generic[P, R]):
 
 
 @dataclass(kw_only=True)
-class AsyncFunctionTaskTemplate(TaskTemplate[P, R]):
+class AsyncFunctionTaskTemplate(TaskTemplate[P, R, F]):
     """
     A task template that wraps an asynchronous functions. This is automatically created when an asynchronous function
     is decorated with the task decorator.
     """
 
-    func: FunctionTypes
+    func: F
     plugin_config: Optional[Any] = None  # This is used to pass plugin specific configuration
+    debuggable: bool = True
 
     def __post_init__(self):
         super().__post_init__()

flyte/_task_environment.py

@@ -15,6 +15,7 @@ from typing import (
     Tuple,
     Union,
     cast,
+    overload,
 )
 
 import rich.repr
@@ -33,7 +34,7 @@ from ._trigger import Trigger
 from .models import MAX_INLINE_IO_BYTES, NativeInterface
 
 if TYPE_CHECKING:
-    from ._task import FunctionTypes, P, R
+    from ._task import F, P, R
 
 
 @rich.repr.auto
@@ -150,9 +151,9 @@ class TaskEnvironment(Environment):
             kwargs["interruptible"] = interruptible
         return replace(self, **kwargs)
 
+    @overload
     def task(
         self,
-        _func: Callable[P, R] | None = None,
         *,
         short_name: Optional[str] = None,
         cache: CacheRequest | None = None,
@@ -165,7 +166,31 @@ class TaskEnvironment(Environment):
         max_inline_io_bytes: int = MAX_INLINE_IO_BYTES,
         queue: Optional[str] = None,
         triggers: Tuple[Trigger, ...] | Trigger = (),
-    ) -> Union[AsyncFunctionTaskTemplate[P, R], Callable[P, R]]:
+    ) -> Callable[[Callable[P, R]], AsyncFunctionTaskTemplate[P, R, Callable[P, R]]]: ...
+
+    @overload
+    def task(
+        self,
+        _func: Callable[P, R],
+        /,
+    ) -> AsyncFunctionTaskTemplate[P, R, Callable[P, R]]: ...
+
+    def task(
+        self,
+        _func: F | None = None,
+        *,
+        short_name: Optional[str] = None,
+        cache: CacheRequest | None = None,
+        retries: Union[int, RetryStrategy] = 0,
+        timeout: Union[timedelta, int] = 0,
+        docs: Optional[Documentation] = None,
+        pod_template: Optional[Union[str, PodTemplate]] = None,
+        report: bool = False,
+        interruptible: bool | None = None,
+        max_inline_io_bytes: int = MAX_INLINE_IO_BYTES,
+        queue: Optional[str] = None,
+        triggers: Tuple[Trigger, ...] | Trigger = (),
+    ) -> Callable[[F], AsyncFunctionTaskTemplate[P, R, F]] | AsyncFunctionTaskTemplate[P, R, F]:
         """
         Decorate a function to be a task.
 
@@ -189,13 +214,13 @@ class TaskEnvironment(Environment):
 
         :return: A TaskTemplate that can be used to deploy the task.
         """
-        from ._task import P, R
+        from ._task import F, P, R
 
         if self.reusable is not None:
             if pod_template is not None:
                 raise ValueError("Cannot set pod_template when environment is reusable.")
 
-        def decorator(func: FunctionTypes) -> AsyncFunctionTaskTemplate[P, R]:
+        def decorator(func: F) -> AsyncFunctionTaskTemplate[P, R, F]:
             short = short_name or func.__name__
             task_name = self.name + "." + func.__name__
 
@@ -209,7 +234,7 @@ class TaskEnvironment(Environment):
             if self.plugin_config is not None:
                 from flyte.extend import TaskPluginRegistry
 
-                task_template_class: type[AsyncFunctionTaskTemplate[P, R]] | None = TaskPluginRegistry.find(
+                task_template_class: type[AsyncFunctionTaskTemplate[P, R, F]] | None = TaskPluginRegistry.find(
                     config_type=type(self.plugin_config)
                 )
                 if task_template_class is None:
@@ -218,9 +243,9 @@ class TaskEnvironment(Environment):
                         f"Please register a plugin using flyte.extend.TaskPluginRegistry.register() api."
                     )
             else:
-                task_template_class = AsyncFunctionTaskTemplate[P, R]
+                task_template_class = AsyncFunctionTaskTemplate[P, R, F]
 
-            task_template_class = cast(type[AsyncFunctionTaskTemplate[P, R]], task_template_class)
+            task_template_class = cast(type[AsyncFunctionTaskTemplate[P, R, F]], task_template_class)
             tmpl = task_template_class(
                 func=func,
                 name=task_name,
@@ -235,6 +260,7 @@ class TaskEnvironment(Environment):
                 secrets=self.secrets,
                 pod_template=pod_template or self.pod_template,
                 parent_env=weakref.ref(self),
+                parent_env_name=self.name,
                 interface=NativeInterface.from_callable(func),
                 report=report,
                 short_name=short,
@@ -248,8 +274,8 @@ class TaskEnvironment(Environment):
             return tmpl
 
         if _func is None:
-            return cast(AsyncFunctionTaskTemplate, decorator)
-        return cast(AsyncFunctionTaskTemplate, decorator(_func))
+            return cast(Callable[[F], AsyncFunctionTaskTemplate[P, R, F]], decorator)
+        return cast(AsyncFunctionTaskTemplate[P, R, F], decorator(_func))
 
     @property
     def tasks(self) -> Dict[str, TaskTemplate]:
@@ -286,4 +312,5 @@ class TaskEnvironment(Environment):
         for t in tasks:
             env._tasks[t.name] = t
             t.parent_env = weakref.ref(env)
+            t.parent_env_name = name
         return env

flyte/_utils/module_loader.py

@@ -5,8 +5,6 @@ import sys
 from pathlib import Path
 from typing import List, Tuple
 
-from rich.progress import BarColumn, Progress, TextColumn, TimeElapsedColumn, TimeRemainingColumn
-
 import flyte.errors
 
 
@@ -18,6 +16,8 @@ def load_python_modules(path: Path, recursive: bool = False) -> Tuple[List[str],
     :param recursive: If True, load modules recursively from subdirectories
     :return: List of loaded module names, and list of file paths that failed to load
     """
+    from rich.progress import BarColumn, Progress, TextColumn, TimeElapsedColumn, TimeRemainingColumn
+
     loaded_modules = []
     failed_paths = []
 

flyte/_version.py

@@ -28,7 +28,7 @@ version_tuple: VERSION_TUPLE
 commit_id: COMMIT_ID
 __commit_id__: COMMIT_ID
 
-__version__ = version = '2.0.0b23'
-__version_tuple__ = version_tuple = (2, 0, 0, 'b23')
+__version__ = version = '2.0.0b25'
+__version_tuple__ = version_tuple = (2, 0, 0, 'b25')
 
-__commit_id__ = commit_id = 'g52dc01de2'
+__commit_id__ = commit_id = 'gf704e4b75'

flyte/cli/_common.py

@@ -7,6 +7,7 @@ import os
 import sys
 from abc import abstractmethod
 from dataclasses import dataclass, replace
+from functools import lru_cache
 from pathlib import Path
 from types import MappingProxyType, ModuleType
 from typing import Any, Dict, Iterable, List, Literal, Optional
@@ -20,6 +21,7 @@ from rich.pretty import pretty_repr
 from rich.table import Table
 from rich.traceback import Traceback
 
+import flyte.config
 import flyte.errors
 from flyte.config import Config
 
@@ -112,13 +114,19 @@ class CLIConfig:
         """
         return replace(self, **kwargs)
 
-    def init(self, project: str | None = None, domain: str | None = None, root_dir: str | None = None):
+    def init(
+        self,
+        project: str | None = None,
+        domain: str | None = None,
+        root_dir: str | None = None,
+        images: tuple[str, ...] | None = None,
+    ):
         from flyte.config._config import TaskConfig
 
         task_cfg = TaskConfig(
             org=self.org or self.config.task.org,
-            project=project or self.config.task.project,
-            domain=domain or self.config.task.domain,
+            project=project if project is not None else self.config.task.project,
+            domain=domain if domain is not None else self.config.task.domain,
         )
 
         kwargs: Dict[str, Any] = {}
@@ -132,7 +140,7 @@ class CLIConfig:
 
         updated_config = self.config.with_params(platform_cfg, task_cfg)
 
-        flyte.init_from_config(updated_config, log_level=self.log_level, root_dir=root_dir)
+        flyte.init_from_config(updated_config, log_level=self.log_level, root_dir=root_dir, images=images)
 
 
 class InvokeBaseMixin:
@@ -355,7 +363,7 @@ def _table_format(table: Table, vals: Iterable[Any]) -> Table:
         if headers is None:
             headers = [k for k, _ in o]
             for h in headers:
-                table.add_column(h.capitalize())
+                table.add_column(h.capitalize(), no_wrap=True if "name" in h.casefold() else False)
         table.add_row(*[str(v) for _, v in o])
     return table
 
@@ -375,6 +383,7 @@ def format(title: str, vals: Iterable[Any], of: OutputFormat = "table") -> Table
                     header_style=HEADER_STYLE,
                     show_header=True,
                     border_style=PREFERRED_BORDER_COLOR,
+                    expand=True,
                 ),
                 vals,
             )
@@ -408,3 +417,36 @@ def get_console() -> Console:
     Get a console that is configured to use colors if the terminal supports it.
     """
     return Console(color_system="auto", force_terminal=True, width=120)
+
+
+def parse_images(cfg: Config, values: tuple[str, ...] | None) -> None:
+    """
+    Parse image values and update the config.
+
+    Args:
+        cfg: The Config object to write images to
+        values: List of image strings in format "imagename=imageuri" or just "imageuri"
+    """
+    if values is None:
+        return
+    for value in values:
+        if "=" in value:
+            image_name, image_uri = value.split("=", 1)
+            cfg.image.image_refs[image_name] = image_uri
+        else:
+            # If no name specified, use "default" as the name
+            cfg.image.image_refs["default"] = value
+
+
+@lru_cache()
+def initialize_config(
+    ctx: click.Context, project: str, domain: str, root_dir: str | None = None, images: tuple[str, ...] | None = None
+):
+    obj: CLIConfig | None = ctx.obj
+    if obj is None:
+        import flyte.config
+
+        obj = CLIConfig(flyte.config.auto(), ctx)
+
+    obj.init(project, domain, root_dir, images)
+    return obj

flyte/cli/_create.py

@@ -79,6 +79,10 @@ def secret(
     """
     from flyte.remote import Secret
 
+    # todo: remove this hack when secrets creation more easily distinguishes between org and project/domain level
+    #   (and domain level) secrets
+    project = "" if project is None else project
+    domain = "" if domain is None else domain
     cfg.init(project, domain)
     if from_file:
         with open(from_file, "rb") as f:

flyte/cli/_deploy.py

@@ -189,6 +189,14 @@ class EnvPerFileGroup(common.ObjectsPerFileGroup):
     def _filter_objects(self, module: ModuleType) -> Dict[str, Any]:
         return {k: v for k, v in module.__dict__.items() if isinstance(v, flyte.Environment)}
 
+    def list_commands(self, ctx):
+        common.initialize_config(ctx, self.deploy_args.project, self.deploy_args.domain, self.deploy_args.root_dir)
+        return super().list_commands(ctx)
+
+    def get_command(self, ctx, obj_name):
+        common.initialize_config(ctx, self.deploy_args.project, self.deploy_args.domain, self.deploy_args.root_dir)
+        return super().get_command(ctx, obj_name)
+
     def _get_command_for_obj(self, ctx: click.Context, obj_name: str, obj: Any) -> click.Command:
         obj = cast(flyte.Environment, obj)
         return DeployEnvCommand(

flyte/cli/_get.py

@@ -246,6 +246,10 @@ def secret(
     """
     Get a list of all secrets, or details of a specific secret by name.
     """
+    if project is None:
+        project = ""
+    if domain is None:
+        domain = ""
     cfg.init(project=project, domain=domain)
 
     console = common.get_console()

flyte/cli/_params.py

@@ -15,9 +15,9 @@ from typing import get_args
 import rich_click as click
 import yaml
 from click import Parameter
-from flyteidl.core.interface_pb2 import Variable
-from flyteidl.core.literals_pb2 import Literal
-from flyteidl.core.types_pb2 import BlobType, LiteralType, SimpleType
+from flyteidl2.core.interface_pb2 import Variable
+from flyteidl2.core.literals_pb2 import Literal
+from flyteidl2.core.types_pb2 import BlobType, LiteralType, SimpleType
 from google.protobuf.json_format import MessageToDict
 from mashumaro.codecs.json import JSONEncoder
 
@@ -505,7 +505,7 @@ def to_click_option(
     This handles converting workflow input types to supported click parameters with callbacks to initialize
     the input values to their expected types.
     """
-    from flyteidl.core.types_pb2 import SimpleType
+    from flyteidl2.core.types_pb2 import SimpleType
 
     if input_name != input_name.lower():
         # Click does not support uppercase option names: https://github.com/pallets/click/issues/837

flyte/cli/_run.py

@@ -15,7 +15,7 @@ from .._code_bundle._utils import CopyFiles
 from .._task import TaskTemplate
 from ..remote import Run
 from . import _common as common
-from ._common import CLIConfig
+from ._common import CLIConfig, initialize_config
 from ._params import to_click_option
 
 RUN_REMOTE_CMD = "deployed-task"
@@ -43,7 +43,7 @@ def _list_tasks(
 ) -> list[str]:
     import flyte.remote
 
-    _initialize_config(ctx, project, domain)
+    common.initialize_config(ctx, project, domain)
     return [task.name for task in flyte.remote.Task.listall(by_task_name=by_task_name, by_task_env=by_task_env)]
 
 
@@ -108,6 +108,17 @@ class RunArguments:
             )
         },
     )
+    image: List[str] = field(
+        default_factory=list,
+        metadata={
+            "click.option": click.Option(
+                ["--image"],
+                type=str,
+                multiple=True,
+                help="Image to be used in the run. Format: imagename=imageuri. Can be specified multiple times.",
+            )
+        },
+    )
 
     @classmethod
     def from_dict(cls, d: Dict[str, Any]) -> RunArguments:
@@ -130,7 +141,9 @@ class RunTaskCommand(click.RichCommand):
         super().__init__(obj_name, *args, **kwargs)
 
     def invoke(self, ctx: click.Context):
-        obj: CLIConfig = _initialize_config(ctx, self.run_args.project, self.run_args.domain, self.run_args.root_dir)
+        obj: CLIConfig = initialize_config(
+            ctx, self.run_args.project, self.run_args.domain, self.run_args.root_dir, tuple(self.run_args.image) or None
+        )
 
         async def _run():
             import flyte
@@ -205,6 +218,14 @@ class TaskPerFileGroup(common.ObjectsPerFileGroup):
     def _filter_objects(self, module: ModuleType) -> Dict[str, Any]:
         return {k: v for k, v in module.__dict__.items() if isinstance(v, TaskTemplate)}
 
+    def list_commands(self, ctx):
+        common.initialize_config(ctx, self.run_args.project, self.run_args.domain, self.run_args.root_dir)
+        return super().list_commands(ctx)
+
+    def get_command(self, ctx, obj_name):
+        common.initialize_config(ctx, self.run_args.project, self.run_args.domain, self.run_args.root_dir)
+        return super().get_command(ctx, obj_name)
+
     def _get_command_for_obj(self, ctx: click.Context, obj_name: str, obj: Any) -> click.Command:
         obj = cast(TaskTemplate, obj)
         return RunTaskCommand(
@@ -224,10 +245,11 @@ class RunReferenceTaskCommand(click.RichCommand):
         super().__init__(*args, **kwargs)
 
     def invoke(self, ctx: click.Context):
-        obj: CLIConfig = _initialize_config(ctx, self.run_args.project, self.run_args.domain)
+        obj: CLIConfig = common.initialize_config(
+            ctx, self.run_args.project, self.run_args.domain, self.run_args.root_dir, tuple(self.run_args.image) or None
+        )
 
         async def _run():
-            import flyte
             import flyte.remote
 
             task = flyte.remote.Task.get(self.task_name, version=self.version, auto_version="latest")
@@ -262,7 +284,7 @@ class RunReferenceTaskCommand(click.RichCommand):
         import flyte.remote
         from flyte._internal.runtime.types_serde import transform_native_to_typed_interface
 
-        _initialize_config(ctx, self.run_args.project, self.run_args.domain)
+        common.initialize_config(ctx, self.run_args.project, self.run_args.domain)
 
         task = flyte.remote.Task.get(self.task_name, auto_version="latest")
         task_details = task.fetch()
@@ -408,7 +430,6 @@ class TaskFiles(common.FileGroup):
 
     def get_command(self, ctx, cmd_name):
         run_args = RunArguments.from_dict(ctx.params)
-
         if cmd_name == RUN_REMOTE_CMD:
             return ReferenceTaskGroup(
                 name=cmd_name,
@@ -455,6 +476,28 @@ Flyte environment:
 flyte run --local hello.py my_task --arg1 value1 --arg2 value2
 ```
 
+You can provide image mappings with `--image` flag. This allows you to specify
+the image URI for the task environment during CLI execution without changing
+the code. Any images defined with `Image.from_ref_name("name")` will resolve to the
+corresponding URIs you specify here.
+
+```bash
+flyte run hello.py my_task --image my_image=ghcr.io/myorg/my-image:v1.0
+```
+
+If the image name is not provided, it is regarded as a default image and will
+be used when no image is specified in TaskEnvironment:
+
+```bash
+flyte run hello.py my_task --image ghcr.io/myorg/default-image:latest
+```
+
+You can specify multiple image arguments:
+
+```bash
+flyte run hello.py my_task --image ghcr.io/org/default:latest --image gpu=ghcr.io/org/gpu:v2.0
+```
+
 To run tasks that you've already deployed to Flyte, use the {RUN_REMOTE_CMD} command:
 
 ```bash

flyte/cli/_update.py

@@ -15,7 +15,7 @@ def update():
 @update.command("trigger", cls=common.CommandBase)
 @click.argument("name", type=str)
 @click.argument("task_name", type=str)
-@click.option("--activate/--deactivate", default=None, required=True, help="Activate or deactivate the trigger.")
+@click.option("--activate/--deactivate", required=True, help="Activate or deactivate the trigger.")
 @click.pass_obj
 def trigger(cfg: common.CLIConfig, name: str, task_name: str, activate: bool, project: str | None, domain: str | None):
     """
@@ -31,6 +31,7 @@ def trigger(cfg: common.CLIConfig, name: str, task_name: str, activate: bool, pr
     """
     cfg.init(project, domain)
     console = common.get_console()
-    with console.status(f"Updating trigger {name} for task {task_name} to {'active' if activate else 'deactive'}..."):
+    to_state = "active" if activate else "deactivate"
+    with console.status(f"Updating trigger {name} for task {task_name} to {to_state}..."):
         remote.Trigger.update(name, task_name, activate)
-    console.print("Trigger updated.")
+    console.print(f"Trigger updated and is set to [fuchsia]{to_state}[/fuchsia]")

flyte/config/_config.py

@@ -148,6 +148,7 @@ class ImageConfig(object):
     """
 
     builder: str | None = None
+    image_refs: typing.Dict[str, str] = field(default_factory=dict)
 
     @classmethod
     def auto(cls, config_file: typing.Optional[typing.Union[str, ConfigFile]] = None) -> "ImageConfig":
@@ -159,6 +160,7 @@ class ImageConfig(object):
         config_file = get_config_file(config_file)
         kwargs: typing.Dict[str, typing.Any] = {}
         kwargs = set_if_exists(kwargs, "builder", _internal.Image.BUILDER.read(config_file))
+        kwargs = set_if_exists(kwargs, "image_refs", _internal.Image.IMAGE_REFS.read(config_file))
         return ImageConfig(**kwargs)
 
 

flyte/config/_internal.py

@@ -70,3 +70,4 @@ class Image(object):
     """
 
     BUILDER = ConfigEntry(YamlConfigEntry("image.builder"))
+    IMAGE_REFS = ConfigEntry(YamlConfigEntry("image.image_refs"))

flyte/config/_reader.py

@@ -138,10 +138,10 @@ class ConfigFile(object):
 def _config_path_from_git_root() -> pathlib.Path | None:
     from flyte.git import config_from_root
 
-    try:
-        return config_from_root().source
-    except RuntimeError:
+    config = config_from_root()
+    if config is None:
         return None
+    return config.source
 
 
 def resolve_config_path() -> pathlib.Path | None:

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}"
         )
 
 

flyte/extend.py

@@ -1,4 +1,6 @@
 from ._initialize import is_initialized
+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_plugins import TaskPluginRegistry
@@ -7,6 +9,8 @@ __all__ = [
     "PRIMARY_CONTAINER_DEFAULT_NAME",
     "AsyncFunctionTaskTemplate",
     "TaskPluginRegistry",
+    "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/_config.py

@@ -4,16 +4,18 @@ import subprocess
 import flyte.config
 
 
-def config_from_root(path: pathlib.Path | str = ".flyte/config.yaml") -> flyte.config.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.
     """
-
-    result = subprocess.run(["git", "rev-parse", "--show-toplevel"], check=False, capture_output=True, text=True)
-    if result.returncode != 0:
-        raise RuntimeError(f"Failed to get git root directory: {result.stderr}")
-    root = pathlib.Path(result.stdout.strip())
-    if not (root / path).exists():
-        raise RuntimeError(f"Config file {root / path} does not exist")
-    return flyte.config.auto(root / path)
+    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