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

flyte/_environment.py

@@ -0,0 +1,84 @@
+from __future__ import annotations
+
+import re
+from dataclasses import dataclass, field
+from typing import TYPE_CHECKING, Any, Dict, List, Literal, Optional, Union
+
+import rich.repr
+
+from flyte._secret import SecretRequest
+
+from ._image import Image
+from ._resources import Resources
+
+if TYPE_CHECKING:
+    from kubernetes.client import V1PodTemplate
+
+
+def is_snake_or_kebab_with_numbers(s: str) -> bool:
+    return re.fullmatch(r"^[a-z0-9]+([_-][a-z0-9]+)*$", s) is not None
+
+
+@rich.repr.auto
+@dataclass(init=True, repr=True)
+class Environment:
+    """
+    :param name: Name of the environment
+    :param image: Docker image to use for the environment. If set to "auto", will use the default image.
+    :param resources: Resources to allocate for the environment.
+    :param env: Environment variables to set for the environment.
+    :param secrets: Secrets to inject into the environment.
+    :param depends_on: Environment dependencies to hint, so when you deploy the environment, the dependencies are
+        also deployed. This is useful when you have a set of environments that depend on each other.
+    """
+
+    name: str
+    depends_on: List[Environment] = field(default_factory=list)
+    pod_template: Optional[Union[str, "V1PodTemplate"]] = None
+    description: Optional[str] = None
+    secrets: Optional[SecretRequest] = None
+    env: Optional[Dict[str, str]] = None
+    resources: Optional[Resources] = None
+    image: Union[str, Image, Literal["auto"]] = "auto"
+
+    def __post_init__(self):
+        if not is_snake_or_kebab_with_numbers(self.name):
+            raise ValueError(f"Environment name '{self.name}' must be in snake_case or kebab-case format.")
+
+    def add_dependency(self, *env: Environment):
+        """
+        Add a dependency to the environment.
+        """
+        self.depends_on.extend(env)
+
+    def clone_with(
+        self,
+        name: str,
+        image: Optional[Union[str, Image, Literal["auto"]]] = None,
+        resources: Optional[Resources] = None,
+        env: Optional[Dict[str, str]] = None,
+        secrets: Optional[SecretRequest] = None,
+        depends_on: Optional[List[Environment]] = None,
+        **kwargs: Any,
+    ) -> Environment:
+        raise NotImplementedError
+
+    def _get_kwargs(self) -> Dict[str, Any]:
+        """
+        Get the keyword arguments for the environment.
+        """
+        kwargs: Dict[str, Any] = {
+            "depends_on": self.depends_on,
+            "image": self.image,
+        }
+        if self.resources is not None:
+            kwargs["resources"] = self.resources
+        if self.secrets is not None:
+            kwargs["secrets"] = self.secrets
+        if self.env is not None:
+            kwargs["env"] = self.env
+        if self.pod_template is not None:
+            kwargs["pod_template"] = self.pod_template
+        if self.description is not None:
+            kwargs["description"] = self.description
+        return kwargs

flyte/_excepthook.py

@@ -0,0 +1,37 @@
+import logging
+import sys
+import traceback
+
+from flyte._logging import logger
+
+# Save the original excepthook so we can call it later
+original_excepthook = sys.excepthook
+
+# Filters: exclude frames where filename or function name contains these substrings
+EXCLUDED_MODULE_SUBSTRINGS = ["_internal", "syncify"]
+EXCLUDED_FILE_SUBSTRINGS = ["syncify", "_code_bundle"]
+
+
+def should_include_frame(frame: traceback.FrameSummary) -> bool:
+    return not (
+        any(sub in frame.name for sub in EXCLUDED_MODULE_SUBSTRINGS)
+        or any(sub in frame.filename for sub in EXCLUDED_FILE_SUBSTRINGS)
+    )
+
+
+def custom_excepthook(exc_type, exc_value, exc_tb):
+    """
+    Custom exception hook to filter and format tracebacks.
+    If the logger's level is DEBUG, it uses the original excepthook.
+    """
+
+    if logger.getEffectiveLevel() <= logging.DEBUG:
+        original_excepthook(exc_type, exc_value, exc_tb)
+    else:
+        # Extract and filter traceback
+        tb_list = traceback.extract_tb(exc_tb)
+        filtered_tb = [frame for frame in tb_list if should_include_frame(frame)]
+        # Print the filtered version (custom format)
+        print("Filtered traceback (most recent call last):")
+        print("".join(traceback.format_list(filtered_tb)))
+        print(f"{exc_type.__name__}: {exc_value}\n")

flyte/_group.py

@@ -0,0 +1,32 @@
+from contextlib import contextmanager
+
+from ._context import internal_ctx
+from .models import GroupData
+
+
+@contextmanager
+def group(name: str):
+    """
+    Create a new group with the given name. The method is intended to be used as a context manager.
+
+    Example:
+    ```python
+    @task
+    async def my_task():
+        ...
+        with group("my_group"):
+            t1(x,y)  # tasks in this block will be grouped under "my_group"
+        ...
+    ```
+
+    :param name: The name of the group
+    """
+    ctx = internal_ctx()
+    if ctx.data.task_context is None:
+        yield
+        return
+    tctx = ctx.data.task_context
+    new_tctx = tctx.replace(group_data=GroupData(name))
+    with ctx.replace_task_context(new_tctx):
+        yield
+        # Exit the context and restore the previous context

flyte/_hash.py

@@ -0,0 +1,23 @@
+from typing import Callable, Generic, TypeVar
+
+T = TypeVar("T")
+
+
+class HashOnReferenceMixin(object):
+    def __hash__(self):
+        return hash(id(self))
+
+
+class HashMethod(Generic[T]):
+    """
+    Flyte-specific object used to wrap the hash function for a specific type
+    """
+
+    def __init__(self, function: Callable[[T], str]):
+        self._function = function
+
+    def calculate(self, obj: T) -> str:
+        """
+        Calculate hash for `obj`.
+        """
+        return self._function(obj)