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

flyte/_internal/imagebuild/remote_builder.py

@@ -16,6 +16,7 @@ import flyte.errors
 from flyte import Image, remote
 from flyte._code_bundle._utils import tar_strip_file_attributes
 from flyte._image import (
+    _BASE_REGISTRY,
     AptPackages,
     Architecture,
     Commands,
@@ -24,6 +25,7 @@ from flyte._image import (
     Env,
     PipOption,
     PipPackages,
+    PoetryProject,
     PythonWheels,
     Requirements,
     UVProject,
@@ -31,14 +33,14 @@ from flyte._image import (
     WorkDir,
 )
 from flyte._internal.imagebuild.image_builder import ImageBuilder, ImageChecker
-from flyte._internal.imagebuild.utils import copy_files_to_context
+from flyte._internal.imagebuild.utils import copy_files_to_context, get_and_list_dockerignore
 from flyte._internal.runtime.task_serde import get_security_context
 from flyte._logging import logger
 from flyte._secret import Secret
 from flyte.remote import ActionOutputs, Run
 
 if TYPE_CHECKING:
-    from flyte._protos.imagebuilder import definition_pb2 as image_definition_pb2
+    from flyteidl2.imagebuilder import definition_pb2 as image_definition_pb2
 
 IMAGE_TASK_NAME = "build-image"
 IMAGE_TASK_PROJECT = "system"
@@ -68,10 +70,11 @@ class RemoteImageChecker(ImageChecker):
         image_name = f"{repository.split('/')[-1]}:{tag}"
 
         try:
+            from flyteidl2.imagebuilder import definition_pb2 as image_definition__pb2
+            from flyteidl2.imagebuilder import payload_pb2 as image_payload__pb2
+            from flyteidl2.imagebuilder import service_pb2_grpc as image_service_pb2_grpc
+
             from flyte._initialize import _get_init_config
-            from flyte._protos.imagebuilder import definition_pb2 as image_definition__pb2
-            from flyte._protos.imagebuilder import payload_pb2 as image_payload__pb2
-            from flyte._protos.imagebuilder import service_pb2_grpc as image_service_pb2_grpc
 
             cfg = _get_init_config()
             if cfg is None:
@@ -96,7 +99,7 @@ class RemoteImageBuilder(ImageBuilder):
         return [RemoteImageChecker]
 
     async def build_image(self, image: Image, dry_run: bool = False) -> str:
-        from flyte._protos.workflow import run_definition_pb2
+        from flyteidl2.workflow import run_definition_pb2
 
         image_name = f"{image.name}:{image._final_tag}"
         spec, context = await _validate_configuration(image)
@@ -110,10 +113,15 @@ class RemoteImageBuilder(ImageBuilder):
         ).override.aio(secrets=_get_build_secrets_from_image(image))
 
         logger.warning("[bold blue]🐳 Submitting a new build...[/bold blue]")
+        if image.registry and image.registry != _BASE_REGISTRY:
+            target_image = f"{image.registry}/{image_name}"
+        else:
+            # Use the default system registry in the backend.
+            target_image = image_name
         run = cast(
             Run,
             await flyte.with_runcontext(project=IMAGE_TASK_PROJECT, domain=IMAGE_TASK_DOMAIN).run.aio(
-                entity, spec=spec, context=context, target_image=image_name
+                entity, spec=spec, context=context, target_image=target_image
             ),
         )
         logger.warning(f"⏳ Waiting for build to finish at: [bold cyan link={run.url}]{run.url}[/bold cyan link]")
@@ -180,7 +188,7 @@ async def _validate_configuration(image: Image) -> Tuple[str, Optional[str]]:
 
 
 def _get_layers_proto(image: Image, context_path: Path) -> "image_definition_pb2.ImageSpec":
-    from flyte._protos.imagebuilder import definition_pb2 as image_definition_pb2
+    from flyteidl2.imagebuilder import definition_pb2 as image_definition_pb2
 
     if image.dockerfile is not None:
         raise flyte.errors.ImageBuildError(
@@ -190,7 +198,7 @@ def _get_layers_proto(image: Image, context_path: Path) -> "image_definition_pb2
     layers = []
     for layer in image._layers:
         secret_mounts = None
-        pip_options = None
+        pip_options = image_definition_pb2.PipOptions()
 
         if isinstance(layer, PipOption):
             pip_options = image_definition_pb2.PipOptions(
@@ -256,12 +264,18 @@ def _get_layers_proto(image: Image, context_path: Path) -> "image_definition_pb2
                 if "tool.uv.index" in line:
                     raise ValueError("External sources are not supported in pyproject.toml")
 
-            if layer.extra_args and "--no-install-project" in layer.extra_args:
+            if layer.project_install_mode == "dependencies_only":
                 # Copy pyproject itself
                 pyproject_dst = copy_files_to_context(layer.pyproject, context_path)
+                if pip_options.extra_args:
+                    if "--no-install-project" not in pip_options.extra_args:
+                        pip_options.extra_args += " --no-install-project"
+                else:
+                    pip_options.extra_args = " --no-install-project"
             else:
                 # Copy the entire project
-                pyproject_dst = copy_files_to_context(layer.pyproject.parent, context_path)
+                docker_ignore_patterns = get_and_list_dockerignore(image)
+                pyproject_dst = copy_files_to_context(layer.pyproject.parent, context_path, docker_ignore_patterns)
 
             uv_layer = image_definition_pb2.Layer(
                 uv_project=image_definition_pb2.UVProject(
@@ -272,6 +286,27 @@ def _get_layers_proto(image: Image, context_path: Path) -> "image_definition_pb2
                 )
             )
             layers.append(uv_layer)
+        elif isinstance(layer, PoetryProject):
+            for line in layer.pyproject.read_text().splitlines():
+                if "tool.poetry.source" in line:
+                    raise ValueError("External sources are not supported in pyproject.toml")
+
+            if layer.extra_args and "--no-root" in layer.extra_args:
+                # Copy pyproject itself
+                pyproject_dst = copy_files_to_context(layer.pyproject, context_path)
+            else:
+                # Copy the entire project
+                pyproject_dst = copy_files_to_context(layer.pyproject.parent, context_path)
+
+            poetry_layer = image_definition_pb2.Layer(
+                poetry_project=image_definition_pb2.PoetryProject(
+                    pyproject=str(pyproject_dst.relative_to(context_path)),
+                    poetry_lock=str(copy_files_to_context(layer.poetry_lock, context_path).relative_to(context_path)),
+                    extra_args=layer.extra_args,
+                    secret_mounts=secret_mounts,
+                )
+            )
+            layers.append(poetry_layer)
         elif isinstance(layer, Commands):
             commands_layer = image_definition_pb2.Layer(
                 commands=image_definition_pb2.Commands(
@@ -330,4 +365,9 @@ def _get_build_secrets_from_image(image: Image) -> Optional[typing.List[Secret]]
                 else:
                     raise ValueError(f"Unsupported secret_mount type: {type(secret_mount)}")
 
+    image_registry_secret = image._image_registry_secret
+    if image_registry_secret:
+        secrets.append(
+            Secret(key=image_registry_secret.key, group=image_registry_secret.group, mount=DEFAULT_SECRET_DIR)
+        )
     return secrets

flyte/_internal/imagebuild/utils.py

@@ -1,8 +1,12 @@
 import shutil
 from pathlib import Path
+from typing import List, Optional
 
+from flyte._image import DockerIgnore, Image
+from flyte._logging import logger
 
-def copy_files_to_context(src: Path, context_path: Path) -> Path:
+
+def copy_files_to_context(src: Path, context_path: Path, ignore_patterns: list[str] = []) -> Path:
     """
     This helper function ensures that absolute paths that users specify are converted correctly to a path in the
     context directory. Doing this prevents collisions while ensuring files are available in the context.
@@ -23,8 +27,52 @@ def copy_files_to_context(src: Path, context_path: Path) -> Path:
         dst_path = context_path / src
     dst_path.parent.mkdir(parents=True, exist_ok=True)
     if src.is_dir():
-        # TODO: Add support dockerignore
-        shutil.copytree(src, dst_path, dirs_exist_ok=True, ignore=shutil.ignore_patterns(".idea", ".venv"))
+        default_ignore_patterns = [".idea", ".venv"]
+        ignore_patterns = list(set(ignore_patterns + default_ignore_patterns))
+        shutil.copytree(src, dst_path, dirs_exist_ok=True, ignore=shutil.ignore_patterns(*ignore_patterns))
     else:
         shutil.copy(src, dst_path)
     return dst_path
+
+
+def get_and_list_dockerignore(image: Image) -> List[str]:
+    """
+    Get and parse dockerignore patterns from .dockerignore file.
+
+    This function first looks for a DockerIgnore layer in the image's layers. If found, it uses
+    the path specified in that layer. If no DockerIgnore layer is found, it falls back to looking
+    for a .dockerignore file in the root_path directory.
+
+    :param image: The Image object
+    """
+    from flyte._initialize import _get_init_config
+
+    # Look for DockerIgnore layer in the image layers
+    dockerignore_path: Optional[Path] = None
+    patterns: List[str] = []
+
+    for layer in image._layers:
+        if isinstance(layer, DockerIgnore) and layer.path.strip():
+            dockerignore_path = Path(layer.path)
+    # If DockerIgnore layer not specified, set dockerignore_path under root_path
+    init_config = _get_init_config()
+    root_path = init_config.root_dir if init_config else None
+    if not dockerignore_path and root_path:
+        dockerignore_path = Path(root_path) / ".dockerignore"
+    # Return empty list if no .dockerignore file found
+    if not dockerignore_path or not dockerignore_path.exists() or not dockerignore_path.is_file():
+        logger.info(f".dockerignore file not found at path: {dockerignore_path}")
+        return patterns
+
+    try:
+        with open(dockerignore_path, "r", encoding="utf-8") as f:
+            for line in f:
+                stripped_line = line.strip()
+                # Skip empty lines, whitespace-only lines, and comments
+                if not stripped_line or stripped_line.startswith("#"):
+                    continue
+                patterns.append(stripped_line)
+    except Exception as e:
+        logger.error(f"Failed to read .dockerignore file at {dockerignore_path}: {e}")
+        return []
+    return patterns

flyte/_internal/runtime/convert.py

@@ -8,27 +8,33 @@ from dataclasses import dataclass
 from types import NoneType
 from typing import Any, Dict, List, Tuple, Union, get_args
 
-from flyteidl.core import execution_pb2, interface_pb2, literals_pb2
+from flyteidl2.core import execution_pb2, interface_pb2, literals_pb2
+from flyteidl2.task import common_pb2, task_definition_pb2
 
 import flyte.errors
 import flyte.storage as storage
-from flyte._protos.workflow import common_pb2, run_definition_pb2, task_definition_pb2
+from flyte._context import ctx
 from flyte.models import ActionID, NativeInterface, TaskContext
 from flyte.types import TypeEngine, TypeTransformerFailedError
 
 
 @dataclass(frozen=True)
 class Inputs:
-    proto_inputs: run_definition_pb2.Inputs
+    proto_inputs: common_pb2.Inputs
 
     @classmethod
     def empty(cls) -> "Inputs":
-        return cls(proto_inputs=run_definition_pb2.Inputs())
+        return cls(proto_inputs=common_pb2.Inputs())
+
+    @property
+    def context(self) -> Dict[str, str]:
+        """Get the context as a dictionary."""
+        return {kv.key: kv.value for kv in self.proto_inputs.context}
 
 
 @dataclass(frozen=True)
 class Outputs:
-    proto_outputs: run_definition_pb2.Outputs
+    proto_outputs: common_pb2.Outputs
 
 
 @dataclass
@@ -102,15 +108,30 @@ def is_optional_type(tp) -> bool:
     return NoneType in get_args(tp)  # fastest check
 
 
-async def convert_from_native_to_inputs(interface: NativeInterface, *args, **kwargs) -> Inputs:
+async def convert_from_native_to_inputs(
+    interface: NativeInterface, *args, custom_context: Dict[str, str] | None = None, **kwargs
+) -> Inputs:
     kwargs = interface.convert_to_kwargs(*args, **kwargs)
 
     missing = [key for key in interface.required_inputs() if key not in kwargs]
     if missing:
         raise ValueError(f"Missing required inputs: {', '.join(missing)}")
 
+    # Read custom_context from TaskContext if available (inside task execution)
+    # Otherwise use the passed parameter (for remote run initiation)
+    context_kvs = None
+    tctx = ctx()
+    if tctx and tctx.custom_context:
+        # Inside a task - read from TaskContext
+        context_to_use = tctx.custom_context
+        context_kvs = [literals_pb2.KeyValuePair(key=k, value=v) for k, v in context_to_use.items()]
+    elif custom_context:
+        # Remote run initiation
+        context_kvs = [literals_pb2.KeyValuePair(key=k, value=v) for k, v in custom_context.items()]
+
     if len(interface.inputs) == 0:
-        return Inputs.empty()
+        # Handle context even for empty inputs
+        return Inputs(proto_inputs=common_pb2.Inputs(context=context_kvs))
 
     # fill in defaults if missing
     type_hints: Dict[str, type] = {}
@@ -144,12 +165,12 @@ async def convert_from_native_to_inputs(interface: NativeInterface, *args, **kwa
         for k, v in already_converted_kwargs.items():
             copied_literals[k] = v
         literal_map = literals_pb2.LiteralMap(literals=copied_literals)
+
     # Make sure we the interface, not literal_map or kwargs, because those may have a different order
     return Inputs(
-        proto_inputs=run_definition_pb2.Inputs(
-            literals=[
-                run_definition_pb2.NamedLiteral(name=k, value=literal_map.literals[k]) for k in interface.inputs.keys()
-            ]
+        proto_inputs=common_pb2.Inputs(
+            literals=[common_pb2.NamedLiteral(name=k, value=literal_map.literals[k]) for k in interface.inputs.keys()],
+            context=context_kvs,
         )
     )
 
@@ -191,11 +212,11 @@ async def convert_from_native_to_outputs(o: Any, interface: NativeInterface, tas
     for (output_name, python_type), v in zip(interface.outputs.items(), o):
         try:
             lit = await TypeEngine.to_literal(v, python_type, TypeEngine.to_literal_type(python_type))
-            named.append(run_definition_pb2.NamedLiteral(name=output_name, value=lit))
+            named.append(common_pb2.NamedLiteral(name=output_name, value=lit))
         except TypeTransformerFailedError as e:
             raise flyte.errors.RuntimeDataValidationError(output_name, e, task_name)
 
-    return Outputs(proto_outputs=run_definition_pb2.Outputs(literals=named))
+    return Outputs(proto_outputs=common_pb2.Outputs(literals=named))
 
 
 async def convert_outputs_to_native(interface: NativeInterface, outputs: Outputs) -> Union[Any, Tuple[Any, ...]]:
@@ -222,7 +243,7 @@ def convert_error_to_native(err: execution_pb2.ExecutionError | Exception | Erro
     if isinstance(err, Error):
         err = err.err
 
-    user_code, server_code = _clean_error_code(err.code)
+    user_code, _server_code = _clean_error_code(err.code)
     match err.kind:
         case execution_pb2.ExecutionError.UNKNOWN:
             return flyte.errors.RuntimeUnknownError(code=user_code, message=err.message, worker=err.worker)
@@ -351,7 +372,7 @@ def generate_inputs_repr_for_literal(literal: literals_pb2.Literal) -> bytes:
     return literal.SerializeToString(deterministic=True)
 
 
-def generate_inputs_hash_for_named_literals(inputs: list[run_definition_pb2.NamedLiteral]) -> str:
+def generate_inputs_hash_for_named_literals(inputs: list[common_pb2.NamedLiteral]) -> str:
     """
     Generate a hash for the inputs using the new literal representation approach that respects
     hash values already present in literals. This is used to uniquely identify the inputs for a task
@@ -375,7 +396,7 @@ def generate_inputs_hash_for_named_literals(inputs: list[run_definition_pb2.Name
     return hash_data(combined_bytes)
 
 
-def generate_inputs_hash_from_proto(inputs: run_definition_pb2.Inputs) -> str:
+def generate_inputs_hash_from_proto(inputs: common_pb2.Inputs) -> str:
     """
     Generate a hash for the inputs. This is used to uniquely identify the inputs for a task.
     :param inputs: The inputs to hash.
@@ -404,7 +425,7 @@ def generate_cache_key_hash(
     task_interface: interface_pb2.TypedInterface,
     cache_version: str,
     ignored_input_vars: List[str],
-    proto_inputs: run_definition_pb2.Inputs,
+    proto_inputs: common_pb2.Inputs,
 ) -> str:
     """
     Generate a cache key hash based on the inputs hash, task name, task interface, and cache version.
@@ -420,7 +441,7 @@ def generate_cache_key_hash(
     """
     if ignored_input_vars:
         filtered = [named_lit for named_lit in proto_inputs.literals if named_lit.name not in ignored_input_vars]
-        final = run_definition_pb2.Inputs(literals=filtered)
+        final = common_pb2.Inputs(literals=filtered)
         final_inputs = generate_inputs_hash_from_proto(final)
     else:
         final_inputs = inputs_hash

flyte/_internal/runtime/io.py

@@ -5,10 +5,11 @@ It uses the storage module to handle the actual uploading and downloading of fil
 TODO: Convert to use streaming apis
 """
 
-from flyteidl.core import errors_pb2, execution_pb2
+from flyteidl.core import errors_pb2
+from flyteidl2.core import execution_pb2
+from flyteidl2.task import common_pb2
 
 import flyte.storage as storage
-from flyte._protos.workflow import run_definition_pb2
 from flyte.models import PathRewrite
 
 from .convert import Inputs, Outputs, _clean_error_code
@@ -70,7 +71,7 @@ async def upload_outputs(outputs: Outputs, output_path: str, max_bytes: int = -1
     await storage.put_stream(data_iterable=outputs.proto_outputs.SerializeToString(), to_path=output_uri)
 
 
-async def upload_error(err: execution_pb2.ExecutionError, output_prefix: str):
+async def upload_error(err: execution_pb2.ExecutionError, output_prefix: str) -> str:
     """
     :param err: execution_pb2.ExecutionError
     :param output_prefix: The output prefix of the remote uri.
@@ -87,7 +88,7 @@ async def upload_error(err: execution_pb2.ExecutionError, output_prefix: str):
         )
     )
     error_uri = error_path(output_prefix)
-    await storage.put_stream(data_iterable=error_document.SerializeToString(), to_path=error_uri)
+    return await storage.put_stream(data_iterable=error_document.SerializeToString(), to_path=error_uri)
 
 
 # ------------------------------- DOWNLOAD Methods ------------------------------- #
@@ -98,7 +99,7 @@ async def load_inputs(path: str, max_bytes: int = -1, path_rewrite_config: PathR
     :param path_rewrite_config: If provided, rewrites paths in the input blobs according to the configuration.
     :return: Inputs object
     """
-    lm = run_definition_pb2.Inputs()
+    lm = common_pb2.Inputs()
 
     if max_bytes == -1:
         proto_str = b"".join([c async for c in storage.get_stream(path=path)])
@@ -137,7 +138,7 @@ async def load_outputs(path: str, max_bytes: int = -1) -> Outputs:
                       If -1, reads the entire file.
     :return: Outputs object
     """
-    lm = run_definition_pb2.Outputs()
+    lm = common_pb2.Outputs()
 
     if max_bytes == -1:
         proto_str = b"".join([c async for c in storage.get_stream(path=path)])
@@ -169,7 +170,7 @@ async def load_error(path: str) -> execution_pb2.ExecutionError:
     err.ParseFromString(proto_str)
 
     if err.error is not None:
-        user_code, server_code = _clean_error_code(err.error.code)
+        user_code, _server_code = _clean_error_code(err.error.code)
         return execution_pb2.ExecutionError(
             code=user_code,
             message=err.error.message,

flyte/_internal/runtime/resources_serde.py

@@ -1,8 +1,8 @@
-from typing import List, Optional, Tuple
+from typing import Dict, List, Optional, Tuple
 
-from flyteidl.core import tasks_pb2
+from flyteidl2.core import tasks_pb2
 
-from flyte._resources import CPUBaseType, Resources
+from flyte._resources import CPUBaseType, DeviceClass, Resources
 
 ACCELERATOR_DEVICE_MAP = {
     "A100": "nvidia-tesla-a100",
@@ -24,6 +24,14 @@ ACCELERATOR_DEVICE_MAP = {
     "V6E": "tpu-v6e-slice",
 }
 
+_DeviceClassToProto: Dict[DeviceClass, "tasks_pb2.GPUAccelerator.DeviceClass"] = {
+    "GPU": tasks_pb2.GPUAccelerator.NVIDIA_GPU,
+    "TPU": tasks_pb2.GPUAccelerator.GOOGLE_TPU,
+    "NEURON": tasks_pb2.GPUAccelerator.AMAZON_NEURON,
+    "AMD_GPU": tasks_pb2.GPUAccelerator.AMD_GPU,
+    "HABANA_GAUDI": tasks_pb2.GPUAccelerator.HABANA_GAUDI,
+}
+
 
 def _get_cpu_resource_entry(cpu: CPUBaseType) -> tasks_pb2.Resources.ResourceEntry:
     return tasks_pb2.Resources.ResourceEntry(
@@ -54,11 +62,17 @@ def _get_gpu_extended_resource_entry(resources: Resources) -> Optional[tasks_pb2
     device = resources.get_device()
     if device is None:
         return None
-    if device.device not in ACCELERATOR_DEVICE_MAP:
-        raise ValueError(f"GPU of type {device.device} unknown, cannot map to device name")
+
+    device_class = _DeviceClassToProto.get(device.device_class, tasks_pb2.GPUAccelerator.NVIDIA_GPU)
+    if device.device is None:
+        raise RuntimeError("Device type must be specified for GPU string.")
+    else:
+        device_type = device.device
+    device_type = ACCELERATOR_DEVICE_MAP.get(device_type, device_type)
     return tasks_pb2.GPUAccelerator(
-        device=ACCELERATOR_DEVICE_MAP[device.device],
+        device=device_type,
         partition_size=device.partition if device.partition else None,
+        device_class=device_class,
     )
 
 

flyte/_internal/runtime/reuse.py

@@ -2,7 +2,7 @@ import hashlib
 import typing
 from venv import logger
 
-from flyteidl.core import tasks_pb2
+from flyteidl2.core import tasks_pb2
 
 import flyte.errors
 from flyte import ReusePolicy

flyte/_internal/runtime/task_serde.py

@@ -8,14 +8,14 @@ import typing
 from datetime import timedelta
 from typing import Optional, cast
 
-from flyteidl.core import identifier_pb2, literals_pb2, security_pb2, tasks_pb2
+from flyteidl2.core import identifier_pb2, literals_pb2, security_pb2, tasks_pb2
+from flyteidl2.task import common_pb2, environment_pb2, task_definition_pb2
 from google.protobuf import duration_pb2, wrappers_pb2
 
 import flyte.errors
 from flyte._cache.cache import VersionParameters, cache_from_request
 from flyte._logging import logger
 from flyte._pod import _PRIMARY_CONTAINER_NAME_FIELD, PodTemplate
-from flyte._protos.workflow import common_pb2, environment_pb2, task_definition_pb2
 from flyte._secret import SecretRequest, secrets_from_request
 from flyte._task import AsyncFunctionTaskTemplate, TaskTemplate
 from flyte.models import CodeBundle, SerializationContext
@@ -172,6 +172,7 @@ def get_proto_task(task: TaskTemplate, serialize_context: SerializationContext)
             pod_template_name=(task.pod_template if task.pod_template and isinstance(task.pod_template, str) else None),
             interruptible=task.interruptible,
             generates_deck=wrappers_pb2.BoolValue(value=task.report),
+            debuggable=task.debuggable,
         ),
         interface=transform_native_to_typed_interface(task.native_interface),
         custom=custom if len(custom) > 0 else None,
@@ -208,17 +209,13 @@ def _get_urun_container(
         else None
     )
     resources = get_proto_resources(task_template.resources)
-    # pr: under what conditions should this return None?
+
     if isinstance(task_template.image, str):
         raise flyte.errors.RuntimeSystemError("BadConfig", "Image is not a valid image")
 
-    env_name = ""
-    if task_template.parent_env is not None:
-        task_env = task_template.parent_env()
-        if task_env is not None:
-            env_name = task_env.name
-    else:
-        raise flyte.errors.RuntimeSystemError("BadConfig", "Task template has no parent environment")
+    env_name = task_template.parent_env_name
+    if env_name is None:
+        raise flyte.errors.RuntimeSystemError("BadConfig", f"Task {task_template.name} has no parent environment name")
 
     if not serialize_context.image_cache:
         # This computes the image uri, computing hashes as necessary so can fail if done remotely.

flyte/_internal/runtime/taskrunner.py

@@ -129,6 +129,14 @@ async def convert_and_run(
     in a context tree.
     """
     ctx = internal_ctx()
+
+    # Load inputs first to get context
+    if input_path:
+        inputs = await load_inputs(input_path, path_rewrite_config=raw_data_path.path_rewrite)
+
+    # Extract context from inputs
+    custom_context = inputs.context if inputs else {}
+
     tctx = TaskContext(
         action=action,
         checkpoints=checkpoints,
@@ -142,9 +150,10 @@ async def convert_and_run(
         report=flyte.report.Report(name=action.name),
         mode="remote" if not ctx.data.task_context else ctx.data.task_context.mode,
         interactive_mode=interactive_mode,
+        custom_context=custom_context,
     )
+
     with ctx.replace_task_context(tctx):
-        inputs = await load_inputs(input_path, path_rewrite_config=raw_data_path.path_rewrite) if input_path else inputs
         inputs_kwargs = await convert_inputs_to_native(inputs, task.native_interface)
         out, err = await run_task(tctx=tctx, controller=controller, task=task, inputs=inputs_kwargs)
         if err is not None:

flyte/_internal/runtime/trigger_serde.py

@@ -1,12 +1,12 @@
 import asyncio
 from typing import Union
 
-from flyteidl.core import interface_pb2, literals_pb2
+from flyteidl2.core import interface_pb2, literals_pb2
+from flyteidl2.task import common_pb2, run_pb2, task_definition_pb2
 from google.protobuf import timestamp_pb2, wrappers_pb2
 
 import flyte.types
 from flyte import Cron, FixedRate, Trigger, TriggerTime
-from flyte._protos.workflow import common_pb2, run_definition_pb2, trigger_definition_pb2
 
 
 def _to_schedule(m: Union[Cron, FixedRate], kickoff_arg_name: str | None = None) -> common_pb2.Schedule:
@@ -36,7 +36,7 @@ async def process_default_inputs(
     task_name: str,
     task_inputs: interface_pb2.VariableMap,
     task_default_inputs: list[common_pb2.NamedParameter],
-) -> list[run_definition_pb2.NamedLiteral]:
+) -> list[common_pb2.NamedLiteral]:
     """
     Process default inputs and convert them to NamedLiteral objects.
 
@@ -68,10 +68,10 @@ async def process_default_inputs(
             keys.append(p.name)
             final_literals.append(p.parameter.default)
 
-    literals: list[run_definition_pb2.NamedLiteral] = []
+    literals: list[common_pb2.NamedLiteral] = []
     for k, lit in zip(keys, final_literals):
         literals.append(
-            run_definition_pb2.NamedLiteral(
+            common_pb2.NamedLiteral(
                 name=k,
                 value=lit,
             )
@@ -85,7 +85,7 @@ async def to_task_trigger(
     task_name: str,
     task_inputs: interface_pb2.VariableMap,
     task_default_inputs: list[common_pb2.NamedParameter],
-) -> trigger_definition_pb2.TaskTrigger:
+) -> task_definition_pb2.TaskTrigger:
     """
     Converts a Trigger object to a TaskTrigger protobuf object.
     Args:
@@ -98,15 +98,15 @@ async def to_task_trigger(
     """
     env = None
     if t.env_vars:
-        env = run_definition_pb2.Envs()
+        env = run_pb2.Envs()
         for k, v in t.env_vars.items():
             env.values.append(literals_pb2.KeyValuePair(key=k, value=v))
 
-    labels = run_definition_pb2.Labels(values=t.labels) if t.labels else None
+    labels = run_pb2.Labels(values=t.labels) if t.labels else None
 
-    annotations = run_definition_pb2.Annotations(values=t.annotations) if t.annotations else None
+    annotations = run_pb2.Annotations(values=t.annotations) if t.annotations else None
 
-    run_spec = run_definition_pb2.RunSpec(
+    run_spec = run_pb2.RunSpec(
         overwrite_cache=t.overwrite_cache,
         envs=env,
         interruptible=wrappers_pb2.BoolValue(value=t.interruptible) if t.interruptible is not None else None,
@@ -139,12 +139,12 @@ async def to_task_trigger(
         kickoff_arg_name=kickoff_arg_name,
     )
 
-    return trigger_definition_pb2.TaskTrigger(
+    return task_definition_pb2.TaskTrigger(
         name=t.name,
-        spec=trigger_definition_pb2.TaskTriggerSpec(
+        spec=task_definition_pb2.TaskTriggerSpec(
             active=t.auto_activate,
             run_spec=run_spec,
-            inputs=run_definition_pb2.Inputs(literals=literals),
+            inputs=common_pb2.Inputs(literals=literals),
         ),
         automation_spec=common_pb2.TriggerAutomationSpec(
             type=common_pb2.TriggerAutomationSpec.Type.TYPE_SCHEDULE,

flyte/_internal/runtime/types_serde.py

@@ -1,6 +1,6 @@
 from typing import Dict, Optional, TypeVar
 
-from flyteidl.core import interface_pb2
+from flyteidl2.core import interface_pb2
 
 from flyte.models import NativeInterface
 from flyte.types._type_engine import TypeEngine

flyte/_keyring/file.py

@@ -72,9 +72,9 @@ class SimplePlainTextKeyring(KeyringBackend):
 
     @property
     def file_path(self) -> Path:
-        from flyte._initialize import get_common_config
+        from flyte._initialize import get_init_config
 
-        config_path = get_common_config().source_config_path
+        config_path = get_init_config().source_config_path
         if config_path and str(config_path.parent) == ".flyte":
             # if the config is in a .flyte directory, use that as the path
             return config_path.parent / "keyring.cfg"