flyte 0.1.0__py3-none-any.whl → 0.2.0a0__py3-none-any.whl

flyte/_resources.py

@@ -0,0 +1,226 @@
+from dataclasses import dataclass
+from typing import Literal, Optional, Tuple, Union, get_args
+
+import rich.repr
+
+GPUType = Literal["T4", "A100", "A100 80G", "H100", "L4", "L40s"]
+GPUQuantity = Literal[1, 2, 3, 4, 5, 6, 7, 8]
+A100Parts = Literal["1g.5gb", "2g.10gb", "3g.20gb", "4g.20gb", "7g.40gb"]
+"""
+Partitions for NVIDIA A100 GPU.
+"""
+
+A100_80GBParts = Literal["1g.10gb", "2g.20gb", "3g.40gb", "4g.40gb", "7g.80gb"]
+"""
+Partitions for NVIDIA A100 80GB GPU.
+"""
+
+TPUType = Literal["V5P", "V6E"]
+V5EParts = Literal["1x1", "2x2", "2x4", "4x4", "4x8", "8x8", "8x16", "16x16"]
+
+V5PParts = Literal[
+    "2x2x1", "2x2x2", "2x4x4", "4x4x4", "4x4x8", "4x8x8", "8x8x8", "8x8x16", "8x16x16", "16x16x16", "16x16x24"
+]
+"""
+Slices for Google Cloud TPU v5p.
+"""
+
+V6EParts = Literal["1x1", "2x2", "2x4", "4x4", "4x8", "8x8", "8x16", "16x16"]
+"""
+Slices for Google Cloud TPU v6e.
+"""
+
+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",
+    "A100:1",
+    "A100:2",
+    "A100:3",
+    "A100:4",
+    "A100:5",
+    "A100:6",
+    "A100:7",
+    "A100:8",
+    "A100 80G:1",
+    "A100 80G:2",
+    "A100 80G:3",
+    "A100 80G:4",
+    "A100 80G:5",
+    "A100 80G:6",
+    "A100 80G:7",
+    "A100 80G:8",
+    "H100:1",
+    "H100:2",
+    "H100:3",
+    "H100:4",
+    "H100:5",
+    "H100:6",
+    "H100:7",
+    "H100:8",
+]
+
+
+@rich.repr.auto
+@dataclass(frozen=True, slots=True)
+class Device:
+    """
+     Represents a device type, its quantity and partition if applicable.
+    :param device: The type of device (e.g., "T4", "A100").
+    :param quantity: The number of devices of this type.
+    :param partition: The partition of the device (e.g., "1g.5gb", "2g.10gb" for gpus) or ("1x1", ... for tpus).
+    """
+
+    quantity: int
+    device: str | None = None
+    partition: str | None = None
+
+    def __post_init__(self):
+        if self.quantity < 1:
+            raise ValueError("GPU quantity must be at least 1")
+
+
+def GPU(device: GPUType, quantity: GPUQuantity, partition: A100Parts | A100_80GBParts | None = None) -> Device:
+    """
+    Create a GPU device instance.
+    :param device: The type of GPU (e.g., "T4", "A100").
+    :param quantity: The number of GPUs of this type.
+    :param partition: The partition of the GPU (e.g., "1g.5gb", "2g.10gb" for gpus) or ("1x1", ... for tpus).
+    :return: Device instance.
+    """
+    if quantity < 1:
+        raise ValueError("GPU quantity must be at least 1")
+    if device not in get_args(GPUType):
+        raise ValueError(f"Invalid GPU type: {device}. Must be one of {get_args(GPUType)}")
+    if partition is not None and device == "A100":
+        if partition not in get_args(A100Parts):
+            raise ValueError(f"Invalid partition for A100: {partition}. Must be one of {get_args(A100Parts)}")
+    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)
+
+
+def TPU(device: TPUType, partition: V5PParts | V6EParts | None = None):
+    """
+    Create a TPU device instance.
+    :param device: Device type (e.g., "V5P", "V6E").
+    :param partition: Partition of the TPU (e.g., "1x1", "2x2", ...).
+    :return: Device instance.
+    """
+    if device not in get_args(TPUType):
+        raise ValueError(f"Invalid TPU type: {device}. Must be one of {get_args(TPUType)}")
+    if partition is not None and device == "V5P":
+        if partition not in get_args(V5PParts):
+            raise ValueError(f"Invalid partition for V5P: {partition}. Must be one of {get_args(V5PParts)}")
+    elif partition is not None and device == "V6E":
+        if partition not in get_args(V6EParts):
+            raise ValueError(f"Invalid partition for V6E: {partition}. Must be one of {get_args(V6EParts)}")
+    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)
+
+
+CPUBaseType = int | float | str
+
+
+@dataclass
+class Resources:
+    """
+    Resources such as CPU, Memory, and GPU that can be allocated to a task.
+
+    Example:
+    - Single CPU, 1GiB of memory, and 1 T4 GPU:
+    ```python
+    @task(resources=Resources(cpu=1, memory="1GiB", gpu="T4:1"))
+    def my_task() -> int:
+        return 42
+    ```
+    - 1CPU with limit upto 2CPU, 2GiB of memory, and 8 A100 GPUs and 10GiB of disk:
+    ```python
+    @task(resources=Resources(cpu=(1, 2), memory="2GiB", gpu="A100:8", disk="10GiB"))
+    def my_task() -> int:
+        return 42
+    ```
+
+    :param cpu: The amount of CPU to allocate to the task. This can be a string, int, float, list of ints or strings,
+        or a tuple of two ints or strings.
+    :param memory: The amount of memory to allocate to the task. This can be a string, int, float, list of ints or
+        strings, or a tuple of two ints or strings.
+    :param gpu: The amount of GPU to allocate to the task. This can be an Accelerators enum, an int, or None.
+    :param disk: The amount of disk to allocate to the task. This is a string of the form "10GiB".
+    """
+
+    cpu: Union[CPUBaseType, Tuple[CPUBaseType, CPUBaseType], None] = None
+    memory: Union[str, Tuple[str, str], None] = None
+    gpu: Union[Accelerators, int, Device, None] = None
+    disk: Union[str, None] = None
+    shm: Union[str, Literal["auto"], None] = None
+
+    def __post_init__(self):
+        if isinstance(self.cpu, tuple):
+            if len(self.cpu) != 2:
+                raise ValueError("cpu tuple must have exactly two elements")
+        if isinstance(self.memory, tuple):
+            if len(self.memory) != 2:
+                raise ValueError("memory tuple must have exactly two elements")
+        if isinstance(self.cpu, (int, float)):
+            if self.cpu < 0:
+                raise ValueError("cpu must be greater than or equal to 0")
+        if self.gpu is not None:
+            if isinstance(self.gpu, int):
+                if self.gpu < 0:
+                    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}")
+
+    def get_device(self) -> Optional[Device]:
+        """
+        Get the accelerator string for the task.
+
+        :return: If GPUs are requested, return a tuple of the device name, and potentially a partition string.
+                 Default cloud provider labels typically use the following values: `1g.5gb`, `2g.10gb`, etc.
+        """
+        if self.gpu is None:
+            return None
+        if isinstance(self.gpu, int):
+            return Device(quantity=self.gpu)
+        if isinstance(self.gpu, str):
+            device, portion = self.gpu.split(":")
+            return Device(device=device, quantity=int(portion))
+        return self.gpu
+
+    def get_shared_memory(self) -> Optional[str]:
+        """
+        Get the shared memory string for the task.
+
+        :return: The shared memory string.
+        """
+        if self.shm is None:
+            return None
+        if self.shm == "auto":
+            return ""
+        return self.shm

flyte/_retry.py

@@ -0,0 +1,32 @@
+from dataclasses import dataclass
+from datetime import timedelta
+from typing import Union
+
+
+@dataclass
+class RetryStrategy:
+    """
+    Retry strategy for the task or task environment. Retry strategy is optional or can be a simple number of retries.
+
+    Example:
+    - This will retry the task 5 times.
+    ```
+    @task(retries=5)
+    def my_task():
+        pass
+    ```
+    - This will retry the task 5 times with a maximum backoff of 10 seconds and a backoff factor of 2.
+    ```
+    @task(retries=RetryStrategy(count=5, max_backoff=10, backoff=2))
+    def my_task():
+        pass
+    ```
+
+    :param count: The number of retries.
+    :param backoff: The maximum backoff time for retries. This can be a float or a timedelta.
+    :param backoff: The backoff exponential factor. This can be an integer or a float.
+    """
+
+    count: int
+    backoff: Union[float, timedelta, None] = None
+    backoff_factor: Union[int, float, None] = None

flyte/_reusable_environment.py

@@ -0,0 +1,25 @@
+from dataclasses import dataclass
+from datetime import timedelta
+from typing import Optional, Tuple, Union
+
+
+@dataclass
+class ReusePolicy:
+    """
+    ReusePolicy can be used to configure a task to reuse the environment. This is useful when the environment creation
+    is expensive and the runtime of the task is short. The environment will be reused for the next invocation of the
+    task, even the python process maybe be reused by subsequent task invocations. A good mental model is to think of
+    the environment as a container that is reused for multiple tasks, more like a long-running service.
+
+    Caution: It is important to note that the environment is shared, so managing memory and resources is important.
+
+    :param replicas: Either a single int representing number of replicas or a tuple of two ints representing
+     the min and max
+    :param idle_ttl: The maximum idle duration for an environment replica, specified as either seconds (int) or a
+        timedelta. If not set, the environment's global default will be used.
+        When a replica remains idle — meaning no tasks are running — for this duration, it will be automatically
+        terminated.
+    """
+
+    replicas: Union[int, Tuple[int, int]] = 1
+    idle_ttl: Optional[Union[int, timedelta]] = None