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

flyte/_map.py

@@ -1,17 +1,26 @@
 import asyncio
-from typing import Any, AsyncGenerator, AsyncIterator, Generic, Iterable, Iterator, List, Union, cast
+import functools
+import logging
+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]):
     """AsyncIterator implementation for the map function results"""
 
-    def __init__(self, func: TaskTemplate[P, R], args: tuple, name: str, concurrency: int, return_exceptions: bool):
+    def __init__(
+        self,
+        func: AsyncFunctionTaskTemplate[P, R, F] | functools.partial[R],
+        args: tuple,
+        name: str,
+        concurrency: int,
+        return_exceptions: bool,
+    ):
         self.func = func
         self.args = args
         self.name = name
@@ -49,13 +58,16 @@ class MapAsyncIterator(Generic[P, R]):
             return result
         except Exception as e:
             self._exception_count += 1
-            logger.debug(f"Task {self._current_index - 1} failed with exception: {e}")
+            logger.debug(
+                f"Task {self._current_index - 1} failed with exception: {e}, return_exceptions={self.return_exceptions}"
+            )
             if self.return_exceptions:
                 return e
             else:
                 # Cancel remaining tasks
                 for remaining_task in self._tasks[self._current_index + 1 :]:
                     remaining_task.cancel()
+                logger.warning("Exception raising is `ON`, raising exception and cancelling remaining tasks")
                 raise e
 
     async def _initialize(self):
@@ -64,10 +76,26 @@ class MapAsyncIterator(Generic[P, R]):
         tasks = []
         task_count = 0
 
-        for arg_tuple in zip(*self.args):
-            task = asyncio.create_task(self.func.aio(*arg_tuple))
-            tasks.append(task)
-            task_count += 1
+        if isinstance(self.func, functools.partial):
+            # Handle partial functions by merging bound args/kwargs with mapped args
+            base_func = cast(AsyncFunctionTaskTemplate, self.func.func)
+            bound_args = self.func.args
+            bound_kwargs = self.func.keywords or {}
+
+            for arg_tuple in zip(*self.args):
+                # Merge bound positional args with mapped args
+                merged_args = bound_args + arg_tuple
+                if logger.isEnabledFor(logging.DEBUG):
+                    logger.debug(f"Running {base_func.name} with args: {merged_args} and kwargs: {bound_kwargs}")
+                task = asyncio.create_task(base_func.aio(*merged_args, **bound_kwargs))
+                tasks.append(task)
+                task_count += 1
+        else:
+            # Handle regular TaskTemplate functions
+            for arg_tuple in zip(*self.args):
+                task = asyncio.create_task(self.func.aio(*arg_tuple))
+                tasks.append(task)
+                task_count += 1
 
         if task_count == 0:
             logger.info(f"Group '{self.name}' has no tasks to process")
@@ -107,9 +135,65 @@ class _Mapper(Generic[P, R]):
         """Get the name of the group, defaulting to 'map' if not provided."""
         return f"{task_name}_{group_name or 'map'}"
 
+    @staticmethod
+    def validate_partial(func: functools.partial[R]):
+        """
+        This method validates that the provided partial function is valid for mapping, i.e. only the one argument
+        is left for mapping and the rest are provided as keywords or args.
+
+        :param func: partial function to validate
+        :raises TypeError: if the partial function is not valid for mapping
+        """
+        f = cast(AsyncFunctionTaskTemplate, func.func)
+        inputs = f.native_interface.inputs
+        params = list(inputs.keys())
+        total_params = len(params)
+        provided_args = len(func.args)
+        provided_kwargs = len(func.keywords or {})
+
+        # Calculate how many parameters are left unspecified
+        unspecified_count = total_params - provided_args - provided_kwargs
+
+        # Exactly one parameter should be left for mapping
+        if unspecified_count != 1:
+            raise TypeError(
+                f"Partial function must leave exactly one parameter unspecified for mapping. "
+                f"Found {unspecified_count} unspecified parameters in {f.name}, "
+                f"params: {inputs.keys()}"
+            )
+
+        # Validate that no parameter is both in args and keywords
+        if func.keywords:
+            param_names = list(inputs.keys())
+            for i, arg_name in enumerate(param_names[: provided_args + 1]):
+                if arg_name in func.keywords:
+                    raise TypeError(
+                        f"Parameter '{arg_name}' is provided both as positional argument and keyword argument "
+                        f"in partial function {f.name}."
+                    )
+
+    @overload
+    def __call__(
+        self,
+        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: TaskTemplate[P, R],
+        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 __call__(
+        self,
+        func: AsyncFunctionTaskTemplate[P, R, F] | functools.partial[R],
         *args: Iterable[Any],
         group_name: str | None = None,
         concurrency: int = 0,
@@ -128,7 +212,13 @@ class _Mapper(Generic[P, R]):
         if not args:
             return
 
-        name = self._get_name(func.name, group_name)
+        if isinstance(func, functools.partial):
+            f = cast(AsyncFunctionTaskTemplate, func.func)
+            self.validate_partial(func)
+        else:
+            f = cast(AsyncFunctionTaskTemplate, func)
+
+        name = self._get_name(f.name, group_name)
         logger.debug(f"Blocking Map for {name}")
         with group(name):
             import flyte
@@ -154,7 +244,7 @@ class _Mapper(Generic[P, R]):
                     *args,
                     name=name,
                     concurrency=concurrency,
-                    return_exceptions=True,
+                    return_exceptions=return_exceptions,
                 ),
             ):
                 logger.debug(f"Mapped {x}, task {i}")
@@ -163,7 +253,7 @@ class _Mapper(Generic[P, R]):
 
     async def aio(
         self,
-        func: TaskTemplate[P, R],
+        func: AsyncFunctionTaskTemplate[P, R, F] | functools.partial[R],
         *args: Iterable[Any],
         group_name: str | None = None,
         concurrency: int = 0,
@@ -171,7 +261,14 @@ class _Mapper(Generic[P, R]):
     ) -> AsyncGenerator[Union[R, Exception], None]:
         if not args:
             return
-        name = self._get_name(func.name, group_name)
+
+        if isinstance(func, functools.partial):
+            f = cast(AsyncFunctionTaskTemplate, func.func)
+            self.validate_partial(func)
+        else:
+            f = cast(AsyncFunctionTaskTemplate, func)
+
+        name = self._get_name(f.name, group_name)
         with group(name):
             import flyte
 
@@ -199,7 +296,7 @@ class _Mapper(Generic[P, R]):
 
 @syncify
 async def _map(
-    func: TaskTemplate[P, R],
+    func: AsyncFunctionTaskTemplate[P, R, F] | functools.partial[R],
     *args: Iterable[Any],
     name: str = "map",
     concurrency: int = 0,

flyte/_module.py

@@ -0,0 +1,70 @@
+import inspect
+import os
+import pathlib
+import sys
+
+
+def extract_obj_module(obj: object, /, source_dir: pathlib.Path) -> str:
+    """
+    Extract the module from the given object. If source_dir is provided, the module will be relative to the source_dir.
+
+    Args:
+        obj: The object to extract the module from.
+        source_dir: The source directory to use for relative paths.
+
+    Returns:
+        The module name as a string.
+    """
+    if source_dir is None:
+        raise ValueError("extract_obj_module: source_dir cannot be None - specify root-dir")
+    # Get the module containing the object
+    entity_module = inspect.getmodule(obj)
+    if entity_module is None:
+        obj_name = getattr(obj, "__name__", str(obj))
+        raise ValueError(f"Object {obj_name} has no module.")
+
+    fp = entity_module.__file__
+    if fp is None:
+        obj_name = getattr(obj, "__name__", str(obj))
+        raise ValueError(f"Object {obj_name} has no module.")
+
+    file_path = pathlib.Path(fp)
+    try:
+        # Get the relative path to the current directory
+        # Will raise ValueError if the file is not in the source directory
+        relative_path = file_path.relative_to(str(pathlib.Path(source_dir).absolute()))
+
+        if relative_path == pathlib.Path("_internal/resolvers"):
+            entity_module_name = entity_module.__name__
+        else:
+            # Replace file separators with dots and remove the '.py' extension
+            dotted_path = os.path.splitext(str(relative_path))[0].replace(os.sep, ".")
+            entity_module_name = dotted_path
+    except ValueError:
+        # If source_dir is not provided or file is not in source_dir, fallback to module name
+        # File is not relative to source_dir - check if it's an installed package
+        file_path_str = str(file_path)
+        if "site-packages" in file_path_str or "dist-packages" in file_path_str:
+            # It's an installed package - use the module's __name__ directly
+            # This will be importable via importlib.import_module()
+            entity_module_name = entity_module.__name__
+        else:
+            # File is not in source_dir and not in site-packages - re-raise the error
+            obj_name = getattr(obj, "__name__", str(obj))
+            raise ValueError(
+                f"Object {obj_name} module file {file_path} is not relative to "
+                f"source directory {source_dir} and is not an installed package."
+            )
+
+    if entity_module_name == "__main__":
+        """
+        This case is for the case in which the object is run from the main module.
+        """
+        fp = sys.modules["__main__"].__file__
+        if fp is None:
+            obj_name = getattr(obj, "__name__", str(obj))
+            raise ValueError(f"Object {obj_name} has no module.")
+        main_path = pathlib.Path(fp)
+        entity_module_name = main_path.stem
+
+    return entity_module_name

flyte/_pod.py

@@ -2,8 +2,8 @@ from dataclasses import dataclass, field
 from typing import TYPE_CHECKING, Dict, Optional
 
 if TYPE_CHECKING:
-    from flyteidl.core.tasks_pb2 import K8sPod
-    from kubernetes.client import ApiClient, V1PodSpec
+    from flyteidl2.core.tasks_pb2 import K8sPod
+    from kubernetes.client import V1PodSpec
 
 
 _PRIMARY_CONTAINER_NAME_FIELD = "primary_container_name"
@@ -20,7 +20,8 @@ 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(
             metadata=K8sObjectMetadata(labels=self.labels, annotations=self.annotations),

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,135 @@ 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",
+    "Trn1:4",
+    "Trn1:8",
+    "Trn1:16",
+    # Trn1n
+    "Trn1n:1",
+    "Trn1n:4",
+    "Trn1n:8",
+    "Trn1n:16",
+    # Trn2
+    "Trn2:1",
+    "Trn2:4",
+    "Trn2:8",
+    "Trn2:16",
+    # Trn2u
+    "Trn2u:1",
+    "Trn2u:4",
+    "Trn2u:8",
+    "Trn2u:16",
+    # Inf1
+    "Inf1:1",
+    "Inf1:2",
+    "Inf1:3",
+    "Inf1:4",
+    "Inf1:5",
+    "Inf1:6",
+    "Inf1:7",
+    "Inf1:8",
+    "Inf1:9",
+    "Inf1:10",
+    "Inf1:11",
+    "Inf1:12",
+    "Inf1:13",
+    "Inf1:14",
+    "Inf1:15",
+    "Inf1:16",
+    # Inf2
+    "Inf2:1",
+    "Inf2:2",
+    "Inf2:3",
+    "Inf2:4",
+    "Inf2:5",
+    "Inf2:6",
+    "Inf2:7",
+    "Inf2:8",
+    "Inf2:9",
+    "Inf2:10",
+    "Inf2:11",
+    "Inf2:12",
+    # 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 +240,7 @@ class Device:
     """
 
     quantity: int
+    device_class: DeviceClass
     device: str | None = None
     partition: str | None = None
 
@@ -126,7 +267,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 +288,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 +378,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 +390,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]: