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

flyte/_map.py

@@ -1,13 +1,13 @@
 import asyncio
 import functools
 import logging
-from typing import Any, AsyncGenerator, AsyncIterator, Generic, Iterable, Iterator, List, Union, cast
+from typing import Any, AsyncGenerator, AsyncIterator, Generic, Iterable, Iterator, List, Union, cast, overload
 
 from flyte.syncify import syncify
 
 from ._group import group
 from ._logging import logger
-from ._task import P, R, TaskTemplate
+from ._task import AsyncFunctionTaskTemplate, F, P, R
 
 
 class MapAsyncIterator(Generic[P, R]):
@@ -15,7 +15,7 @@ class MapAsyncIterator(Generic[P, R]):
 
     def __init__(
         self,
-        func: TaskTemplate[P, R] | functools.partial[R],
+        func: AsyncFunctionTaskTemplate[P, R, F] | functools.partial[R],
         args: tuple,
         name: str,
         concurrency: int,
@@ -78,7 +78,7 @@ class MapAsyncIterator(Generic[P, R]):
 
         if isinstance(self.func, functools.partial):
             # Handle partial functions by merging bound args/kwargs with mapped args
-            base_func = cast(TaskTemplate, self.func.func)
+            base_func = cast(AsyncFunctionTaskTemplate, self.func.func)
             bound_args = self.func.args
             bound_kwargs = self.func.keywords or {}
 
@@ -144,7 +144,7 @@ class _Mapper(Generic[P, R]):
         :param func: partial function to validate
         :raises TypeError: if the partial function is not valid for mapping
         """
-        f = cast(TaskTemplate, func.func)
+        f = cast(AsyncFunctionTaskTemplate, func.func)
         inputs = f.native_interface.inputs
         params = list(inputs.keys())
         total_params = len(params)
@@ -172,9 +172,28 @@ class _Mapper(Generic[P, R]):
                         f"in partial function {f.name}."
                     )
 
+    @overload
     def __call__(
         self,
-        func: TaskTemplate[P, R] | functools.partial[R],
+        func: AsyncFunctionTaskTemplate[P, R, F] | functools.partial[R],
+        *args: Iterable[Any],
+        group_name: str | None = None,
+        concurrency: int = 0,
+    ) -> Iterator[R]: ...
+
+    @overload
+    def __call__(
+        self,
+        func: AsyncFunctionTaskTemplate[P, R, F] | functools.partial[R],
+        *args: Iterable[Any],
+        group_name: str | None = None,
+        concurrency: int = 0,
+        return_exceptions: bool = True,
+    ) -> Iterator[R]: ...
+
+    def __call__(
+        self,
+        func: AsyncFunctionTaskTemplate[P, R, F] | functools.partial[R],
         *args: Iterable[Any],
         group_name: str | None = None,
         concurrency: int = 0,
@@ -194,10 +213,10 @@ class _Mapper(Generic[P, R]):
             return
 
         if isinstance(func, functools.partial):
-            f = cast(TaskTemplate, func.func)
+            f = cast(AsyncFunctionTaskTemplate, func.func)
             self.validate_partial(func)
         else:
-            f = cast(TaskTemplate, func)
+            f = cast(AsyncFunctionTaskTemplate, func)
 
         name = self._get_name(f.name, group_name)
         logger.debug(f"Blocking Map for {name}")
@@ -234,7 +253,7 @@ class _Mapper(Generic[P, R]):
 
     async def aio(
         self,
-        func: TaskTemplate[P, R] | functools.partial[R],
+        func: AsyncFunctionTaskTemplate[P, R, F] | functools.partial[R],
         *args: Iterable[Any],
         group_name: str | None = None,
         concurrency: int = 0,
@@ -244,10 +263,10 @@ class _Mapper(Generic[P, R]):
             return
 
         if isinstance(func, functools.partial):
-            f = cast(TaskTemplate, func.func)
+            f = cast(AsyncFunctionTaskTemplate, func.func)
             self.validate_partial(func)
         else:
-            f = cast(TaskTemplate, func)
+            f = cast(AsyncFunctionTaskTemplate, func)
 
         name = self._get_name(f.name, group_name)
         with group(name):
@@ -277,7 +296,7 @@ class _Mapper(Generic[P, R]):
 
 @syncify
 async def _map(
-    func: TaskTemplate[P, R] | functools.partial[R],
+    func: AsyncFunctionTaskTemplate[P, R, F] | functools.partial[R],
     *args: Iterable[Any],
     name: str = "map",
     concurrency: int = 0,
@@ -290,4 +309,37 @@ async def _map(
         yield result
 
 
-map: _Mapper = _Mapper()
+@overload
+def map(
+    func: AsyncFunctionTaskTemplate[P, R, F] | functools.partial[R],
+    *args: Iterable[Any],
+    group_name: str | None = None,
+    concurrency: int = 0,
+) -> Iterator[R]: ...
+
+
+@overload
+def map(
+    func: AsyncFunctionTaskTemplate[P, R, F] | functools.partial[R],
+    *args: Iterable[Any],
+    group_name: str | None = None,
+    concurrency: int = 0,
+    return_exceptions: bool = True,
+) -> Iterator[Union[R, Exception]]: ...
+
+
+def map(
+    func: AsyncFunctionTaskTemplate[P, R, F] | functools.partial[R],
+    *args: Iterable[Any],
+    group_name: str | None = None,
+    concurrency: int = 0,
+    return_exceptions: bool = True,
+) -> Iterator[Union[R, Exception]]:
+    map: _Mapper = _Mapper()
+    return map(
+        func,
+        *args,
+        group_name=group_name,
+        concurrency=concurrency,
+        return_exceptions=return_exceptions,
+    )

flyte/_pod.py

@@ -2,7 +2,7 @@ from dataclasses import dataclass, field
 from typing import TYPE_CHECKING, Dict, Optional
 
 if TYPE_CHECKING:
-    from flyteidl.core.tasks_pb2 import K8sPod
+    from flyteidl2.core.tasks_pb2 import K8sPod
     from kubernetes.client import V1PodSpec
 
 
@@ -20,7 +20,7 @@ class PodTemplate(object):
     annotations: Optional[Dict[str, str]] = None
 
     def to_k8s_pod(self) -> "K8sPod":
-        from flyteidl.core.tasks_pb2 import K8sObjectMetadata, K8sPod
+        from flyteidl2.core.tasks_pb2 import K8sObjectMetadata, K8sPod
         from kubernetes.client import ApiClient
 
         return K8sPod(

flyte/_resources.py

@@ -1,5 +1,6 @@
+import typing
 from dataclasses import dataclass, fields
-from typing import TYPE_CHECKING, Literal, Optional, Tuple, Union, get_args
+from typing import TYPE_CHECKING, Dict, Literal, Optional, Tuple, Union, get_args
 
 import rich.repr
 
@@ -10,7 +11,7 @@ if TYPE_CHECKING:
 
 PRIMARY_CONTAINER_DEFAULT_NAME = "primary"
 
-GPUType = Literal["T4", "A100", "A100 80G", "H100", "L4", "L40s"]
+GPUType = Literal["A10", "A10G", "A100", "A100 80G", "B200", "H100", "L4", "L40s", "T4", "V100", "RTX PRO 6000"]
 GPUQuantity = Literal[1, 2, 3, 4, 5, 6, 7, 8]
 A100Parts = Literal["1g.5gb", "2g.10gb", "3g.20gb", "4g.20gb", "7g.40gb"]
 """
@@ -37,31 +38,32 @@ V6EParts = Literal["1x1", "2x2", "2x4", "4x4", "4x8", "8x8", "8x16", "16x16"]
 Slices for Google Cloud TPU v6e.
 """
 
+NeuronType = Literal["Inf1", "Inf2", "Trn1", "Trn1n", "Trn2", "Trn2u"]
+
+AMD_GPUType = Literal["MI100", "MI210", "MI250", "MI250X", "MI300A", "MI300X", "MI325X", "MI350X", "MI355X"]
+
+HABANA_GAUDIType = Literal["Gaudi1"]
+
 Accelerators = Literal[
-    "T4:1",
-    "T4:2",
-    "T4:3",
-    "T4:4",
-    "T4:5",
-    "T4:6",
-    "T4:7",
-    "T4:8",
-    "L4:1",
-    "L4:2",
-    "L4:3",
-    "L4:4",
-    "L4:5",
-    "L4:6",
-    "L4:7",
-    "L4:8",
-    "L40s:1",
-    "L40s:2",
-    "L40s:3",
-    "L40s:4",
-    "L40s:5",
-    "L40s:6",
-    "L40s:7",
-    "L40s:8",
+    # A10
+    "A10:1",
+    "A10:2",
+    "A10:3",
+    "A10:4",
+    "A10:5",
+    "A10:6",
+    "A10:7",
+    "A10:8",
+    # A10G
+    "A10G:1",
+    "A10G:2",
+    "A10G:3",
+    "A10G:4",
+    "A10G:5",
+    "A10G:6",
+    "A10G:7",
+    "A10G:8",
+    # A100
     "A100:1",
     "A100:2",
     "A100:3",
@@ -70,6 +72,7 @@ Accelerators = Literal[
     "A100:6",
     "A100:7",
     "A100:8",
+    # A100 80G
     "A100 80G:1",
     "A100 80G:2",
     "A100 80G:3",
@@ -78,6 +81,16 @@ Accelerators = Literal[
     "A100 80G:6",
     "A100 80G:7",
     "A100 80G:8",
+    # B200
+    "B200:1",
+    "B200:2",
+    "B200:3",
+    "B200:4",
+    "B200:5",
+    "B200:6",
+    "B200:7",
+    "B200:8",
+    # H100
     "H100:1",
     "H100:2",
     "H100:3",
@@ -86,8 +99,97 @@ Accelerators = Literal[
     "H100:6",
     "H100:7",
     "H100:8",
+    # H200
+    "H200:1",
+    "H200:2",
+    "H200:3",
+    "H200:4",
+    "H200:5",
+    "H200:6",
+    "H200:7",
+    "H200:8",
+    # L4
+    "L4:1",
+    "L4:2",
+    "L4:3",
+    "L4:4",
+    "L4:5",
+    "L4:6",
+    "L4:7",
+    "L4:8",
+    # L40s
+    "L40s:1",
+    "L40s:2",
+    "L40s:3",
+    "L40s:4",
+    "L40s:5",
+    "L40s:6",
+    "L40s:7",
+    "L40s:8",
+    # V100
+    "V100:1",
+    "V100:2",
+    "V100:3",
+    "V100:4",
+    "V100:5",
+    "V100:6",
+    "V100:7",
+    "V100:8",
+    # RTX 6000
+    "RTX PRO 6000:1",
+    # T4
+    "T4:1",
+    "T4:2",
+    "T4:3",
+    "T4:4",
+    "T4:5",
+    "T4:6",
+    "T4:7",
+    "T4:8",
+    # Trn1
+    "Trn1:1",
+    # Trn1n
+    "Trn1n:1",
+    # Trn2
+    "Trn2:1",
+    # Trn2u
+    "Trn2u:1",
+    # Inf1
+    "Inf1:1",
+    # Inf2
+    "Inf2:1",
+    # MI100
+    "MI100:1",
+    # MI210
+    "MI210:1",
+    # MI250
+    "MI250:1",
+    # MI250X
+    "MI250X:1",
+    # MI300A
+    "MI300A:1",
+    # MI300X
+    "MI300X:1",
+    # MI325X
+    "MI325X:1",
+    # MI350X
+    "MI350X:1",
+    # MI355X
+    "MI355X:1",
+    # Habana Gaudi
+    "Gaudi1:1",
 ]
 
+DeviceClass = Literal["GPU", "TPU", "NEURON", "AMD_GPU", "HABANA_GAUDI"]
+
+_DeviceClassType: Dict[typing.Any, str] = {
+    GPUType: "GPU",
+    TPUType: "TPU",
+    NeuronType: "NEURON",
+    AMD_GPUType: "AMD_GPU",
+    HABANA_GAUDIType: "HABANA_GAUDI",
+}
+
 
 @rich.repr.auto
 @dataclass(frozen=True, slots=True)
@@ -100,6 +202,7 @@ class Device:
     """
 
     quantity: int
+    device_class: DeviceClass
     device: str | None = None
     partition: str | None = None
 
@@ -126,7 +229,7 @@ def GPU(device: GPUType, quantity: GPUQuantity, partition: A100Parts | A100_80GB
     elif partition is not None and device == "A100 80G":
         if partition not in get_args(A100_80GBParts):
             raise ValueError(f"Invalid partition for A100 80G: {partition}. Must be one of {get_args(A100_80GBParts)}")
-    return Device(device=device, quantity=quantity, partition=partition)
+    return Device(device=device, quantity=quantity, partition=partition, device_class="GPU")
 
 
 def TPU(device: TPUType, partition: V5PParts | V6EParts | None = None):
@@ -147,7 +250,42 @@ def TPU(device: TPUType, partition: V5PParts | V6EParts | None = None):
     elif partition is not None and device == "V5E":
         if partition not in get_args(V5EParts):
             raise ValueError(f"Invalid partition for V5E: {partition}. Must be one of {get_args(V5EParts)}")
-    return Device(1, device, partition)
+    return Device(1, "TPU", device, partition)
+
+
+def Neuron(device: NeuronType) -> Device:
+    """
+    Create a Neuron device instance.
+    :param device: Device type (e.g., "Inf1", "Inf2", "Trn1", "Trn1n", "Trn2", "Trn2u").
+    :param quantity: The number of Neuron devices of this type.
+    :return: Device instance.
+    """
+    if device not in get_args(NeuronType):
+        raise ValueError(f"Invalid Neuron type: {device}. Must be one of {get_args(NeuronType)}")
+    return Device(device=device, quantity=1, device_class="NEURON")
+
+
+def AMD_GPU(device: AMD_GPUType) -> Device:
+    """
+    Create an AMD GPU device instance.
+    :param device: Device type (e.g., "MI100", "MI210", "MI250", "MI250X", "MI300A", "MI300X", "MI325X", "MI350X",
+                   "MI355X").
+    :return: Device instance.
+    """
+    if device not in get_args(AMD_GPUType):
+        raise ValueError(f"Invalid AMD GPU type: {device}. Must be one of {get_args(AMD_GPUType)}")
+    return Device(device=device, quantity=1, device_class="AMD_GPU")
+
+
+def HABANA_GAUDI(device: HABANA_GAUDIType) -> Device:
+    """
+    Create a Habana Gaudi device instance.
+    :param device: Device type (e.g., "DL1").
+    :return: Device instance.
+    """
+    if device not in get_args(HABANA_GAUDIType):
+        raise ValueError(f"Invalid Habana Gaudi type: {device}. Must be one of {get_args(HABANA_GAUDIType)}")
+    return Device(device=device, quantity=1, device_class="HABANA_GAUDI")
 
 
 CPUBaseType = int | float | str
@@ -202,7 +340,7 @@ class Resources:
                     raise ValueError("gpu must be greater than or equal to 0")
             elif isinstance(self.gpu, str):
                 if self.gpu not in get_args(Accelerators):
-                    raise ValueError(f"gpu must be one of {Accelerators}")
+                    raise ValueError(f"gpu must be one of {Accelerators}, got {self.gpu}")
 
     def get_device(self) -> Optional[Device]:
         """
@@ -214,10 +352,16 @@ class Resources:
         if self.gpu is None:
             return None
         if isinstance(self.gpu, int):
-            return Device(quantity=self.gpu)
+            return Device(quantity=self.gpu, device_class="GPU")
         if isinstance(self.gpu, str):
             device, portion = self.gpu.split(":")
-            return Device(device=device, quantity=int(portion))
+            for cls, cls_name in _DeviceClassType.items():
+                if device in get_args(cls):
+                    device_class = cls_name
+                    break
+            else:
+                raise ValueError(f"Invalid device type: {device}. Must be one of {list(_DeviceClassType.keys())}")
+            return Device(device=device, device_class=device_class, quantity=int(portion))  # type: ignore
         return self.gpu
 
     def get_shared_memory(self) -> Optional[str]:

flyte/_run.py

@@ -12,13 +12,13 @@ from flyte._environment import Environment
 from flyte._initialize import (
     _get_init_config,
     get_client,
-    get_common_config,
+    get_init_config,
     get_storage,
     requires_initialization,
     requires_storage,
 )
 from flyte._logging import logger
-from flyte._task import P, R, TaskTemplate
+from flyte._task import F, P, R, TaskTemplate
 from flyte.models import (
     ActionID,
     Checkpoints,
@@ -94,6 +94,7 @@ class _Runner:
         log_level: int | None = None,
         disable_run_cache: bool = False,
         queue: Optional[str] = None,
+        custom_context: Dict[str, str] | None = None,
     ):
         from flyte._tools import ipython_check
 
@@ -124,11 +125,15 @@ class _Runner:
         self._log_level = log_level
         self._disable_run_cache = disable_run_cache
         self._queue = queue
+        self._custom_context = custom_context or {}
 
     @requires_initialization
-    async def _run_remote(self, obj: TaskTemplate[P, R] | LazyEntity, *args: P.args, **kwargs: P.kwargs) -> Run:
+    async def _run_remote(self, obj: TaskTemplate[P, R, F] | LazyEntity, *args: P.args, **kwargs: P.kwargs) -> Run:
         import grpc
-        from flyteidl.core import literals_pb2
+        from flyteidl2.common import identifier_pb2
+        from flyteidl2.core import literals_pb2
+        from flyteidl2.task import run_pb2
+        from flyteidl2.workflow import run_definition_pb2, run_service_pb2
         from google.protobuf import wrappers_pb2
 
         from flyte.remote import Run
@@ -138,21 +143,21 @@ class _Runner:
         from ._deploy import build_images
         from ._internal.runtime.convert import convert_from_native_to_inputs
         from ._internal.runtime.task_serde import translate_task_to_wire
-        from ._protos.common import identifier_pb2
-        from ._protos.workflow import run_definition_pb2, run_service_pb2
 
-        cfg = get_common_config()
+        cfg = get_init_config()
         project = self._project or cfg.project
         domain = self._domain or cfg.domain
 
         if isinstance(obj, LazyEntity):
             task = await obj.fetch.aio()
             task_spec = task.pb2.spec
-            inputs = await convert_from_native_to_inputs(task.interface, *args, **kwargs)
+            inputs = await convert_from_native_to_inputs(
+                task.interface, *args, custom_context=self._custom_context, **kwargs
+            )
             version = task.pb2.task_id.version
             code_bundle = None
         else:
-            task = cast(TaskTemplate[P, R], obj)
+            task = cast(TaskTemplate[P, R, F], obj)
             if obj.parent_env is None:
                 raise ValueError("Task is not attached to an environment. Please attach the task to an environment")
 
@@ -204,7 +209,9 @@ class _Runner:
                 root_dir=cfg.root_dir,
             )
             task_spec = translate_task_to_wire(obj, s_ctx)
-            inputs = await convert_from_native_to_inputs(obj.native_interface, *args, **kwargs)
+            inputs = await convert_from_native_to_inputs(
+                obj.native_interface, *args, custom_context=self._custom_context, **kwargs
+            )
 
         env = self._env_vars or {}
         if env.get("LOG_LEVEL") is None:
@@ -254,9 +261,9 @@ class _Runner:
                     raise ValueError(f"Environment variable {k} must be a string, got {type(v)}")
                 kv_pairs.append(literals_pb2.KeyValuePair(key=k, value=v))
 
-            env_kv = run_definition_pb2.Envs(values=kv_pairs)
-            annotations = run_definition_pb2.Annotations(values=self._annotations)
-            labels = run_definition_pb2.Labels(values=self._labels)
+            env_kv = run_pb2.Envs(values=kv_pairs)
+            annotations = run_pb2.Annotations(values=self._annotations)
+            labels = run_pb2.Labels(values=self._labels)
 
             try:
                 resp = await get_client().run_service.CreateRun(
@@ -265,7 +272,7 @@ class _Runner:
                         project_id=project_id,
                         task_spec=task_spec,
                         inputs=inputs.proto_inputs,
-                        run_spec=run_definition_pb2.RunSpec(
+                        run_spec=run_pb2.RunSpec(
                             overwrite_cache=self._overwrite_cache,
                             interruptible=wrappers_pb2.BoolValue(value=self._interruptible)
                             if self._interruptible is not None
@@ -318,7 +325,7 @@ class _Runner:
 
     @requires_storage
     @requires_initialization
-    async def _run_hybrid(self, obj: TaskTemplate[P, R], *args: P.args, **kwargs: P.kwargs) -> R:
+    async def _run_hybrid(self, obj: TaskTemplate[P, R, F], *args: P.args, **kwargs: P.kwargs) -> R:
         """
         Run a task in hybrid mode. This means that the parent action will be run locally, but the child actions will be
         run in the cluster remotely. This is currently only used for testing,
@@ -333,7 +340,7 @@ class _Runner:
         from ._internal import create_controller
         from ._internal.runtime.taskrunner import run_task
 
-        cfg = get_common_config()
+        cfg = get_init_config()
 
         if obj.parent_env is None:
             raise ValueError("Task is not attached to an environment. Please attach the task to an environment.")
@@ -411,6 +418,7 @@ class _Runner:
                 compiled_image_cache=image_cache,
                 run_base_dir=run_base_dir,
                 report=flyte.report.Report(name=action.name),
+                custom_context=self._custom_context,
             )
             async with ctx.replace_task_context(tctx):
                 return await run_task(tctx=tctx, controller=controller, task=obj, inputs=inputs)
@@ -420,10 +428,11 @@ class _Runner:
             raise err
         return outputs
 
-    async def _run_local(self, obj: TaskTemplate[P, R], *args: P.args, **kwargs: P.kwargs) -> Run:
+    async def _run_local(self, obj: TaskTemplate[P, R, F], *args: P.args, **kwargs: P.kwargs) -> Run:
+        from flyteidl2.common import identifier_pb2
+
         from flyte._internal.controllers import create_controller
         from flyte._internal.controllers._local_controller import LocalController
-        from flyte._protos.common import identifier_pb2
         from flyte.remote import Run
         from flyte.report import Report
 
@@ -461,7 +470,9 @@ class _Runner:
             compiled_image_cache=None,
             report=Report(name=action.name),
             mode="local",
+            custom_context=self._custom_context,
         )
+
         with ctx.replace_task_context(tctx):
             # make the local version always runs on a different thread, returns a wrapped future.
             if obj._call_as_synchronous:
@@ -473,7 +484,7 @@ class _Runner:
 
         class _LocalRun(Run):
             def __init__(self, outputs: Tuple[Any, ...] | Any):
-                from flyte._protos.workflow import run_definition_pb2
+                from flyteidl2.workflow import run_definition_pb2
 
                 self._outputs = outputs
                 super().__init__(
@@ -506,7 +517,7 @@ class _Runner:
     @syncify
     async def run(
         self,
-        task: TaskTemplate[P, Union[R, Run]] | LazyEntity,
+        task: TaskTemplate[P, Union[R, Run], F] | LazyEntity,
         *args: P.args,
         **kwargs: P.kwargs,
     ) -> Union[R, Run]:
@@ -574,6 +585,7 @@ def with_runcontext(
     log_level: int | None = None,
     disable_run_cache: bool = False,
     queue: Optional[str] = None,
+    custom_context: Dict[str, str] | None = None,
 ) -> _Runner:
     """
     Launch a new run with the given parameters as the context.
@@ -618,6 +630,9 @@ def with_runcontext(
         set using `flyte.init()`
     :param disable_run_cache: Optional If true, the run cache will be disabled. This is useful for testing purposes.
     :param queue: Optional The queue to use for the run. This is used to specify the cluster to use for the run.
+    :param custom_context: Optional global input context to pass to the task. This will be available via
+        get_custom_context() within the task and will automatically propagate to sub-tasks.
+        Acts as base/default values that can be overridden by context managers in the code.
 
     :return: runner
     """
@@ -646,11 +661,12 @@ def with_runcontext(
         log_level=log_level,
         disable_run_cache=disable_run_cache,
         queue=queue,
+        custom_context=custom_context,
     )
 
 
 @syncify
-async def run(task: TaskTemplate[P, R], *args: P.args, **kwargs: P.kwargs) -> Union[R, Run]:
+async def run(task: TaskTemplate[P, R, F], *args: P.args, **kwargs: P.kwargs) -> Union[R, Run]:
     """
     Run a task with the given parameters
     :param task: task to run