flyte 0.2.0b1__py3-none-any.whl → 2.0.0b46__py3-none-any.whl

flyte/_metrics.py

@@ -0,0 +1,59 @@
+"""
+Timer utilities for emitting timing metrics via structured logging.
+"""
+
+import time
+from typing import Any, Dict, Optional
+
+from flyte._logging import logger
+
+
+class Stopwatch:
+    """
+    Simple stopwatch for timing code blocks.
+    Emits timing metrics via structured logging when stopped.
+
+    Example:
+        sw = Stopwatch("download_inputs")
+        sw.start()
+        # code to time
+        sw.stop()
+
+    :param metric_name: Name of the metric to emit
+    :param extra_fields: Additional fields to include in the log record
+    """
+
+    def __init__(self, metric_name: str, extra_fields: Optional[Dict[str, Any]] = None):
+        self.metric_name = metric_name
+        self.extra_fields = extra_fields
+        self._start_time: Optional[float] = None
+
+    def start(self):
+        """Start the stopwatch."""
+        self._start_time = time.perf_counter()
+
+    def stop(self):
+        """Stop the stopwatch and emit the timing metric."""
+        if self._start_time is None:
+            raise RuntimeError(f"Stopwatch '{self.metric_name}' was never started")
+        duration = time.perf_counter() - self._start_time
+        _emit_metric(self.metric_name, duration, self.extra_fields)
+
+
+def _emit_metric(metric_name: str, duration: float, extra_fields: Optional[Dict[str, Any]] = None):
+    """
+    Emit a timing metric via structured logging.
+
+    :param metric_name: Name of the metric (may be hierarchical with dots)
+    :param duration: Duration in seconds
+    :param extra_fields: Additional fields to include in the log record
+    """
+    extra = {
+        "metric_type": "timer",
+        "metric_name": metric_name,
+        "duration_seconds": duration,
+    }
+    if extra_fields:
+        extra.update(extra_fields)
+
+    logger.info(f"Stopwatch: {metric_name} completed in {duration:.4f}s", extra=extra)

flyte/_module.py

@@ -0,0 +1,74 @@
+import inspect
+import os
+import pathlib
+import sys
+from types import ModuleType
+from typing import Tuple
+
+
+def extract_obj_module(obj: object, /, source_dir: pathlib.Path | None = None) -> Tuple[str, ModuleType]:
+    """
+    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__
+        elif "site-packages" in str(file_path) or "dist-packages" in str(file_path):
+            raise ValueError("Object from a library")
+        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, entity_module

flyte/_pod.py

@@ -0,0 +1,30 @@
+from dataclasses import dataclass, field
+from typing import TYPE_CHECKING, Dict, Optional
+
+if TYPE_CHECKING:
+    from flyteidl2.core.tasks_pb2 import K8sPod
+    from kubernetes.client import V1PodSpec
+
+
+_PRIMARY_CONTAINER_NAME_FIELD = "primary_container_name"
+_PRIMARY_CONTAINER_DEFAULT_NAME = "primary"
+
+
+@dataclass(init=True, repr=True, eq=True, frozen=False)
+class PodTemplate(object):
+    """Custom PodTemplate specification for a Task."""
+
+    pod_spec: Optional["V1PodSpec"] = field(default_factory=lambda: V1PodSpec())
+    primary_container_name: str = _PRIMARY_CONTAINER_DEFAULT_NAME
+    labels: Optional[Dict[str, str]] = None
+    annotations: Optional[Dict[str, str]] = None
+
+    def to_k8s_pod(self) -> "K8sPod":
+        from flyteidl2.core.tasks_pb2 import K8sObjectMetadata, K8sPod
+        from kubernetes.client import ApiClient
+
+        return K8sPod(
+            metadata=K8sObjectMetadata(labels=self.labels, annotations=self.annotations),
+            pod_spec=ApiClient().sanitize_for_serialization(self.pod_spec),
+            primary_container_name=self.primary_container_name,
+        )

flyte/_resources.py

@@ -1,9 +1,19 @@
-from dataclasses import dataclass
-from typing import Literal, Optional, Tuple, Union, get_args
+import typing
+from dataclasses import dataclass, fields
+from typing import TYPE_CHECKING, Dict, Literal, Optional, Tuple, Union, get_args
 
 import rich.repr
 
-GPUType = Literal["T4", "A100", "A100 80G", "H100", "L4", "L40s"]
+from flyte._pod import _PRIMARY_CONTAINER_DEFAULT_NAME
+
+if TYPE_CHECKING:
+    from kubernetes.client import V1PodSpec
+
+PRIMARY_CONTAINER_DEFAULT_NAME = "primary"
+
+GPUType = Literal[
+    "A10", "A10G", "A100", "A100 80G", "B200", "H100", "H200", "L4", "L40s", "T4", "V100", "RTX PRO 6000", "GB10"
+]
 GPUQuantity = Literal[1, 2, 3, 4, 5, 6, 7, 8]
 A100Parts = Literal["1g.5gb", "2g.10gb", "3g.20gb", "4g.20gb", "7g.40gb"]
 """
@@ -15,6 +25,11 @@ A100_80GBParts = Literal["1g.10gb", "2g.20gb", "3g.40gb", "4g.40gb", "7g.80gb"]
 Partitions for NVIDIA A100 80GB GPU.
 """
 
+H200Parts = Literal["1g.18gb", "1g.35gb", "2g.35gb", "3g.71gb", "4g.71gb", "7g.141gb"]
+"""
+Partitions for NVIDIA H200 GPU (141GB HBM3e).
+"""
+
 TPUType = Literal["V5P", "V6E"]
 V5EParts = Literal["1x1", "2x2", "2x4", "4x4", "4x8", "8x8", "8x16", "16x16"]
 
@@ -30,31 +45,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",
@@ -63,6 +79,7 @@ Accelerators = Literal[
     "A100:6",
     "A100:7",
     "A100:8",
+    # A100 80G
     "A100 80G:1",
     "A100 80G:2",
     "A100 80G:3",
@@ -71,6 +88,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",
@@ -79,8 +106,137 @@ 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",
+    # GB10
+    "GB10: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)
@@ -93,6 +249,7 @@ class Device:
     """
 
     quantity: int
+    device_class: DeviceClass
     device: str | None = None
     partition: str | None = None
 
@@ -101,7 +258,9 @@ class Device:
             raise ValueError("GPU quantity must be at least 1")
 
 
-def GPU(device: GPUType, quantity: GPUQuantity, partition: A100Parts | A100_80GBParts | None = None) -> Device:
+def GPU(
+    device: GPUType, quantity: GPUQuantity, partition: A100Parts | A100_80GBParts | H200Parts | None = None
+) -> Device:
     """
     Create a GPU device instance.
     :param device: The type of GPU (e.g., "T4", "A100").
@@ -119,7 +278,10 @@ 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)
+    elif partition is not None and device == "H200":
+        if partition not in get_args(H200Parts):
+            raise ValueError(f"Invalid partition for H200: {partition}. Must be one of {get_args(H200Parts)}")
+    return Device(device=device, quantity=quantity, partition=partition, device_class="GPU")
 
 
 def TPU(device: TPUType, partition: V5PParts | V6EParts | None = None):
@@ -140,7 +302,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., "Gaudi1").
+    :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
@@ -171,6 +368,8 @@ class Resources:
         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".
+    :param shm: The amount of shared memory to allocate to the task. This is a string of the form "10GiB" or "auto".
+        If "auto", then the shared memory will be set to max amount of shared memory available on the node.
     """
 
     cpu: Union[CPUBaseType, Tuple[CPUBaseType, CPUBaseType], None] = None
@@ -195,7 +394,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]:
         """
@@ -207,10 +406,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]:
@@ -224,3 +429,61 @@ class Resources:
         if self.shm == "auto":
             return ""
         return self.shm
+
+
+def _check_resource_is_singular(resource: Resources):
+    """
+    Raise a value error if the resource has a tuple.
+    """
+    for field in fields(resource):
+        value = getattr(resource, field.name)
+        if isinstance(value, (tuple, list)):
+            raise ValueError(f"{value} can not be a list or tuple")
+    return resource
+
+
+def pod_spec_from_resources(
+    primary_container_name: str = _PRIMARY_CONTAINER_DEFAULT_NAME,
+    requests: Optional[Resources] = None,
+    limits: Optional[Resources] = None,
+    k8s_gpu_resource_key: str = "nvidia.com/gpu",
+) -> "V1PodSpec":
+    from kubernetes.client import V1Container, V1PodSpec, V1ResourceRequirements
+
+    def _construct_k8s_pods_resources(resources: Optional[Resources], k8s_gpu_resource_key: str):
+        if resources is None:
+            return None
+
+        resources_map = {
+            "cpu": "cpu",
+            "memory": "memory",
+            "gpu": k8s_gpu_resource_key,
+            "ephemeral_storage": "ephemeral-storage",
+        }
+
+        k8s_pod_resources = {}
+
+        _check_resource_is_singular(resources)
+        for resource in fields(resources):
+            resource_value = getattr(resources, resource.name)
+            if resource_value is not None:
+                k8s_pod_resources[resources_map[resource.name]] = resource_value
+
+        return k8s_pod_resources
+
+    requests = _construct_k8s_pods_resources(resources=requests, k8s_gpu_resource_key=k8s_gpu_resource_key)
+    limits = _construct_k8s_pods_resources(resources=limits, k8s_gpu_resource_key=k8s_gpu_resource_key)
+    requests = requests or limits
+    limits = limits or requests
+
+    return V1PodSpec(
+        containers=[
+            V1Container(
+                name=primary_container_name,
+                resources=V1ResourceRequirements(
+                    requests=requests,
+                    limits=limits,
+                ),
+            )
+        ]
+    )

flyte/_retry.py

@@ -1,6 +1,4 @@
 from dataclasses import dataclass
-from datetime import timedelta
-from typing import Union
 
 
 @dataclass
@@ -17,16 +15,12 @@ class RetryStrategy:
     ```
     - 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))
+    @task(retries=RetryStrategy(count=5))
     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