flyte 0.0.1b3__py3-none-any.whl → 0.2.0a0__py3-none-any.whl

union/_secret.py

@@ -1,61 +0,0 @@
-import pathlib
-import re
-from dataclasses import dataclass
-from typing import List, Optional, Union
-
-
-@dataclass
-class Secret:
-    """
-    Secrets are used to inject sensitive information into tasks. Secrets can be mounted as environment variables or
-    files. The secret key is the name of the secret in the secret store. The group is optional and maybe used with some
-    secret stores to organize secrets. The secret_mount is used to specify how the secret should be mounted. If the
-    secret_mount is set to "env" the secret will be mounted as an environment variable. If the secret_mount is set to
-    "file" the secret will be mounted as a file. The as_env_var is an optional parameter that can be used to specify the
-    name of the environment variable that the secret should be mounted as.
-
-    Example:
-    ```python
-    @task(secrets="MY_SECRET")
-    async def my_task():
-        os.environ["MY_SECRET"]  # This will be set to the value of the secret
-
-    @task(secrets=Secret("MY_SECRET", mount="/path/to/secret"))
-    async def my_task2():
-        async with open("/path/to/secret") as f:
-            secret_value = f.read()
-    ```
-
-    TODO: Add support for secret versioning (some stores) and secret groups (some stores) and mounting as files.
-
-    :param key: The name of the secret in the secret store.
-    :param group: The group of the secret in the secret store.
-    :param mount: Use this to specify the path where the secret should be mounted.
-    :param as_env_var: The name of the environment variable that the secret should be mounted as.
-    """
-
-    key: str
-    group: Optional[str] = None
-    mount: pathlib.Path | None = None
-    as_env_var: Optional[str] = None
-
-    def __post_init__(self):
-        if self.as_env_var is not None:
-            pattern = r"^[A-Z_][A-Z0-9_]*$"
-            if not re.match(pattern, self.as_env_var):
-                raise ValueError(f"Invalid environment variable name: {self.as_env_var}, must match {pattern}")
-
-
-SecretRequest = Union[str, Secret, List[str | Secret]]
-
-
-def secrets_from_request(secrets: SecretRequest) -> List[Secret]:
-    """
-    Converts a secret request into a list of secrets.
-    """
-    if isinstance(secrets, str):
-        return [Secret(key=secrets)]
-    elif isinstance(secrets, Secret):
-        return [secrets]
-    else:
-        return [Secret(key=s) if isinstance(s, str) else s for s in secrets]

union/_task.py

@@ -1,354 +0,0 @@
-from __future__ import annotations
-
-import inspect
-import weakref
-from dataclasses import dataclass, field
-from functools import cached_property
-from typing import TYPE_CHECKING, Any, Callable, Dict, Generic, List, Literal, Optional, ParamSpec, TypeVar, Union
-
-from flyteidl.core.tasks_pb2 import DataLoadingConfig
-
-from ._cache import Cache, CacheRequest
-from ._context import internal_ctx
-from ._datastructures import NativeInterface, SerializationContext
-from ._doc import Documentation
-from ._image import Image
-from ._resources import Resources
-from ._retry import RetryStrategy
-from ._reusable_environment import ReusePolicy
-from ._secret import SecretRequest
-from ._timeout import TimeoutType
-from .errors import RuntimeUserError
-
-if TYPE_CHECKING:
-    from kubernetes.client import V1PodTemplate
-
-    from ._task_environment import TaskEnvironment
-
-P = ParamSpec("P")  # capture the function's parameters
-R = TypeVar("R")  # return type
-
-
-@dataclass(kw_only=True)
-class TaskTemplate(Generic[P, R]):
-    """
-    Task template is a template for a task that can be executed. It defines various parameters for the task, which
-    can be defined statically at the time of task definition or dynamically at the time of task invocation using
-    the override method.
-
-    Example usage:
-    ```python
-    @task(name="my_task", image="my_image", resources=Resources(cpu="1", memory="1Gi"))
-    def my_task():
-        pass
-    ```
-
-    :param name: Optional The name of the task (defaults to the function name)
-    :param task_type: Router type for the task, this is used to determine how the task will be executed.
-     This is usually set to match with th execution plugin.
-    :param image: Optional The image to use for the task, if set to "auto" will use the default image for the python
-    version with union installed
-    :param resources: Optional The resources to use for the task
-    :param cache: Optional The cache policy for the task, defaults to auto, which will cache the results of the task.
-    :param interruptable: Optional The interruptable policy for the task, defaults to False, which means the task
-     will not be scheduled on interruptable nodes. If set to True, the task will be scheduled on interruptable nodes,
-     and the code should handle interruptions and resumptions.
-    :param retries: Optional The number of retries for the task, defaults to 0, which means no retries.
-    :param reusable: Optional The reusability policy for the task, defaults to None, which means the task environment
-    will not be reused across task invocations.
-    :param docs: Optional The documentation for the task, if not provided the function docstring will be used.
-    :param env: Optional The environment variables to set for the task.
-    :param secrets: Optional The secrets that will be injected into the task at runtime.
-    :param timeout: Optional The timeout for the task.
-    """
-
-    name: str
-    task_type: str = "python"
-    task_type_version: int = 0
-    image: Union[str, Image, Literal["auto"]] = "auto"
-    resources: Optional[Resources] = None
-    cache: CacheRequest = "auto"
-    interruptable: bool = False
-    retries: Union[int, RetryStrategy] = 0
-    reusable: Union[ReusePolicy, Literal["auto"], None] = (None,)
-    docs: Optional[Documentation] = None
-    env: Optional[Dict[str, str]] = None
-    secrets: Optional[SecretRequest] = None
-    timeout: Optional[TimeoutType] = None
-    primary_container_name: str = "primary"
-    pod_template: Optional[Union[str, V1PodTemplate]] = None
-    interface: Optional[NativeInterface] = None
-
-    parent_env: Optional[weakref.ReferenceType[TaskEnvironment]] = None
-    local: bool = field(default=False, init=False)
-    ref: bool = field(default=False, init=False, repr=False, compare=False)
-
-    def __post_init__(self):
-        # If pod_template is set to a pod, verify
-        if self.pod_template is not None and not isinstance(self.pod_template, str):
-            try:
-                from kubernetes.client import V1PodTemplate  # noqa: F401
-            except ImportError as e:
-                raise ImportError(
-                    "kubernetes is not installed, please install kubernetes package to use pod_template"
-                ) from e
-
-        # Auto set the image based on the image request
-        if self.image == "auto":
-            self.image = Image.auto()
-        elif isinstance(self.image, str):
-            self.image = Image.from_prebuilt(str(self.image))
-
-        # Auto set cache based on the cache request
-        if isinstance(self.cache, str):
-            match self.cache:
-                case "auto":
-                    self.cache = Cache(behavior="auto")
-                case "override":
-                    self.cache = Cache(behavior="override")
-                case "disable":
-                    self.cache = Cache(behavior="disable")
-
-        # if retries is set to int, convert to RetryStrategy
-        if isinstance(self.retries, int):
-            self.retries = RetryStrategy(count=self.retries)
-
-    def __getstate__(self):
-        """
-        This method is called when the object is pickled. We need to remove the parent_env reference
-        to avoid circular references.
-        """
-        state = self.__dict__.copy()
-        state.pop("parent_env", None)
-        return state
-
-    def __setstate__(self, state):
-        """
-        This method is called when the object is unpickled. We need to set the parent_env reference
-        to the environment that created the task.
-        """
-        self.__dict__.update(state)
-        self.parent_env = None
-
-    def clone(self) -> TaskTemplate:
-        import copy
-
-        return copy.deepcopy(self)
-
-    async def pre(self, *args, **kwargs) -> Dict[str, Any]:
-        """
-        This is the preexecute function that will be
-        called before the task is executed
-        """
-        return {}
-
-    async def execute(self, *args, **kwargs) -> Any:
-        """
-        This is the pure python function that will be executed when the task is called.
-        """
-        raise NotImplementedError
-
-    async def post(self, return_vals: Any) -> Any:
-        """
-        This is the postexecute function that will be
-        called after the task is executed
-        """
-        return return_vals
-
-    # ---- Extension points ----
-    def config(self, sctx: SerializationContext) -> Dict[str, str]:
-        """
-        Returns additional configuration for the task. This is a set of key-value pairs that can be used to
-        configure the task execution environment at runtime. This is usually used by plugins.
-        """
-        return {}
-
-    def custom_config(self, sctx: SerializationContext) -> Dict[str, str]:
-        """
-        Returns additional configuration for the task. This is a set of key-value pairs that can be used to
-        configure the task execution environment at runtime. This is usually used by plugins.
-        """
-        return {}
-
-    def data_loading_config(self, sctx: SerializationContext) -> DataLoadingConfig:
-        """
-        This configuration allows executing raw containers in Flyte using the Flyte CoPilot system
-        Flyte CoPilot, eliminates the needs of sdk inside the container. Any inputs required by the users container
-        are side-loaded in the input_path
-        Any outputs generated by the user container - within output_path are automatically uploaded
-        """
-
-    def container_args(self, sctx: SerializationContext) -> List[str]:
-        """
-        Returns the container args for the task. This is a set of key-value pairs that can be used to
-        configure the task execution environment at runtime. This is usually used by plugins.
-        """
-        return []
-
-    def sql(self, sctx: SerializationContext) -> Optional[str]:
-        """
-        Returns the SQL for the task. This is a set of key-value pairs that can be used to
-        configure the task execution environment at runtime. This is usually used by plugins.
-        """
-        return None
-
-    # ---- Extension points ----
-
-    @property
-    def native_interface(self) -> NativeInterface:
-        return self.interface
-
-    async def __call__(self, *args: P.args, **kwargs: P.kwargs) -> R:
-        """
-        This is the entrypoint for an async function task at runtime. It will be called during an execution.
-        Please do not override this method, if you simply want to modify the execution behavior, override the
-        execute method.
-
-        # TODO lets provide one hook to implement, _pre, _execute and _post. We do not want actual execute to
-        """
-        try:
-            ctx = internal_ctx()
-            if ctx.data.task_context:
-                # If we are in a task context, that implies we are executing a Run.
-                # In this scenario, we should submit the task to the controller.
-                # We will also check if we are not initialized, It is not expected to be not initialized
-                from ._internal.controllers import get_controller
-
-                controller = await get_controller()
-                if controller:
-                    return await controller.submit(self, *args, **kwargs)
-            if inspect.iscoroutinefunction(self.execute):
-                return await self.execute(*args, **kwargs)
-            return self.execute(*args, **kwargs)
-        except RuntimeUserError:
-            raise
-        except Exception as e:
-            raise RuntimeUserError(type(e).__name__, str(e)) from e
-
-    def override(
-        self,
-        *,
-        local: Optional[bool] = None,
-        ref: Optional[bool] = None,
-        resources: Optional[Resources] = None,
-        cache: CacheRequest = "auto",
-        retries: Union[int, RetryStrategy] = 0,
-        timeout: Optional[TimeoutType] = None,
-        reusable: Union[ReusePolicy, Literal["auto"], None] = None,
-        env: Optional[Dict[str, str]] = None,
-        secrets: Optional[SecretRequest] = None,
-        **kwargs: Any,
-    ) -> TaskTemplate:
-        """
-        Override various parameters of the task template. This allows for dynamic configuration of the task
-        when it is called, such as changing the image, resources, cache policy, etc.
-        """
-        c = self.clone()
-        c.resources = resources or self.resources
-        c.cache = cache or self.cache
-        c.retries = retries or self.retries
-        c.timeout = timeout or self.timeout
-        c.reusable = reusable or self.reusable
-        c.env = env or self.env
-        c.secrets = secrets or self.secrets
-        c.local = local or self.local
-        c.ref = ref or self.ref
-        for k, v in kwargs.items():
-            if k == "name":
-                raise ValueError("Name cannot be overridden")
-            if k == "image":
-                raise ValueError("Image cannot be overridden")
-            if k == "docs":
-                raise ValueError("Docs cannot be overridden")
-            if k == "interface":
-                raise ValueError("Interface cannot be overridden")
-            if hasattr(c, k):
-                setattr(c, k, v)
-        return c
-
-
-@dataclass(kw_only=True)
-class AsyncFunctionTaskTemplate(TaskTemplate[P, R]):
-    """
-    A task template that wraps an asynchronous functions. This is automatically created when an asynchronous function
-    is decorated with the task decorator.
-    """
-
-    func: Callable[P, R]
-
-    @cached_property
-    def native_interface(self) -> NativeInterface:
-        return NativeInterface.from_callable(self.func)
-
-    async def execute(self, *args: P.args, **kwargs: P.kwargs) -> R:
-        """
-        This is the execute method that will be called when the task is invoked. It will call the actual function.
-        # TODO We may need to keep this as the bare func execute, and need a pre and post execute some other func.
-        """
-        ctx = internal_ctx()
-        ctx_data = await self.pre(*args, **kwargs)
-        if ctx.data.task_context is not None:
-            tctx = ctx.data.task_context.replace(data=ctx_data)
-            with ctx.replace_task_context(tctx):
-                v = await self.func(*args, **kwargs)
-                await self.post(v)
-        else:
-            v = await self.func(*args, **kwargs)
-            await self.post(v)
-        return v
-
-    def container_args(self, serialize_context: SerializationContext) -> List[str]:
-        args = [
-            "urun",
-            "--org",  # todo remove this
-            "testorg",  # todo remove this
-            "--inputs",
-            serialize_context.input_path,
-            "--outputs-path",
-            serialize_context.output_path,
-            "--version",
-            serialize_context.version,  # pr: should this be serialize_context.version or code_bundle.version?
-            "--raw-data-path",
-            "{{.rawOutputDataPrefix}}",
-            "--checkpoint-path",
-            "{{.checkpointOutputPrefix}}",
-            "--prev-checkpoint",
-            "{{.prevCheckpointPrefix}}",
-            "--run-name",
-            "{{.runName}}",
-            "--name",
-            "{{.actionName}}",
-        ]
-        # Add on all the known images
-        if serialize_context.image_cache and serialize_context.image_cache.serialized_form:
-            args = [*args, "--image-cache", serialize_context.image_cache.serialized_form]
-        else:
-            if serialize_context.image_cache:
-                args = [*args, "--image-cache", serialize_context.image_cache.to_transport]
-
-        if serialize_context.code_bundle:
-            if serialize_context.code_bundle.tgz:
-                args = [*args, *["--tgz", f"{serialize_context.code_bundle.tgz}"]]
-            elif serialize_context.code_bundle.pkl:
-                args = [*args, *["--pkl", f"{serialize_context.code_bundle.pkl}"]]
-            args = [*args, *["--dest", f"{serialize_context.code_bundle.destination or '.'}"]]
-
-        if not serialize_context.code_bundle or not serialize_context.code_bundle.pkl:
-            # If we do not have a code bundle, or if we have one, but it is not a pkl, we need to add the resolver
-
-            from union._internal.resolvers.default import DefaultTaskResolver
-
-            _task_resolver = DefaultTaskResolver()
-            args = [
-                *args,
-                *[
-                    "--resolver",
-                    # pr: how do we determine which task resolver to use?
-                    _task_resolver.import_path,
-                    *_task_resolver.loader_args(task=self),
-                ],
-            ]
-
-        assert all(isinstance(item, str) for item in args), f"All args should be strings, non string item = {args}"
-
-        return args

union/_task_environment.py

@@ -1,186 +0,0 @@
-from __future__ import annotations
-
-import asyncio
-import weakref
-from dataclasses import dataclass, field, replace
-from datetime import timedelta
-from functools import wraps
-from typing import TYPE_CHECKING, Callable, Dict, List, Literal, Optional, ParamSpec, TypeVar, Union
-
-import rich.repr
-
-from ._cache import CacheRequest
-from ._doc import Documentation
-from ._environment import Environment
-from ._image import Image
-from ._resources import Resources
-from ._retry import RetryStrategy
-from ._reusable_environment import ReusePolicy
-from ._secret import SecretRequest
-from ._task import AsyncFunctionTaskTemplate, TaskTemplate
-
-if TYPE_CHECKING:
-    from kubernetes.client import V1PodTemplate
-
-P = ParamSpec("P")  # capture the function's parameters
-R = TypeVar("R")  # return type
-
-
-@rich.repr.auto
-@dataclass(init=True, repr=True)
-class TaskEnvironment(Environment):
-    """
-    Environment class to define a new environment for a set of tasks.
-
-    Example usage:
-    ```python
-    env = union.TaskEnvironment(name="my_env", image="my_image", resources=Resources(cpu="1", memory="1Gi"))
-
-    @env.task
-    async def my_task():
-        pass
-    ```
-
-    :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 env_dep_hints: 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.
-    :param cache: Cache policy for the environment.
-    :param reusable: Reuse policy for the environment, if set, a python process may be reused for multiple tasks.
-    """
-
-    cache: Union[CacheRequest] = "auto"
-    reusable: Union[ReusePolicy, Literal["auto"], None] = None
-    # TODO Shall we make this union of string or env? This way we can lookup the env by module/file:name
-    # TODO also we could add list of files that are used by this environment
-
-    _tasks: Dict[str, TaskTemplate] = field(default_factory=dict, init=False)
-
-    def clone_with(
-        self,
-        name: str,
-        image: Optional[Union[str, Image, Literal["auto"]]] = None,
-        resources: Optional[Resources] = None,
-        cache: Union[CacheRequest, None] = None,
-        env: Optional[Dict[str, str]] = None,
-        reusable: Union[ReusePolicy, None] = None,
-        secrets: Optional[SecretRequest] = None,
-        env_dep_hints: Optional[List[TaskEnvironment]] = None,
-    ) -> TaskEnvironment:
-        """
-        Clone the environment with new settings.
-        """
-        if image is None:
-            image = self.image
-        else:
-            image = "auto"
-        return replace(
-            self,
-            cache=cache,
-            reusable=reusable,
-            name=name,
-            image=image,
-            resources=resources,
-            env=env,
-            secrets=secrets,
-            env_dep_hints=env_dep_hints,
-        )
-
-    def _task(
-        self,
-        _func=None,
-        *,
-        name: Optional[str] = None,
-        cache: Union[CacheRequest] | None = None,
-        retries: Union[int, RetryStrategy] = 0,
-        timeout: Union[timedelta, int] = 0,
-        docs: Optional[Documentation] = None,
-        secrets: Optional[SecretRequest] = None,
-        pod_template: Optional[Union[str, "V1PodTemplate"]] = None,
-    ) -> Union[AsyncFunctionTaskTemplate, Callable[P, R]]:
-        """
-        :param name: Optional The name of the task (defaults to the function name)
-        :param cache: Optional The cache policy for the task, defaults to auto, which will cache the results of the
-        task.
-        :param retries: Optional The number of retries for the task, defaults to 0, which means no retries.
-        :param docs: Optional The documentation for the task, if not provided the function docstring will be used.
-        :param secrets: Optional The secrets that will be injected into the task at runtime.
-        :param timeout: Optional The timeout for the task.
-        """
-        if self.reusable is not None:
-            if pod_template is not None:
-                raise ValueError("Cannot set pod_template when environment is reusable.")
-
-        def decorator(func: Callable[P, R]) -> AsyncFunctionTaskTemplate:
-            task_name = name or func.__name__
-
-            @wraps(func)
-            async def wrapper(*args: P.args, **kwargs: P.kwargs) -> R:
-                return await func(*args, **kwargs)
-
-            if not asyncio.iscoroutinefunction(func):
-                raise TypeError(
-                    f"Function {func.__name__} is not a coroutine function. Use @env.task decorator for async tasks."
-                    f"You can simply mark your function as async def {func.__name__} to make it a coroutine function, "
-                    f"it is ok to write sync code in async functions, but not the other way around."
-                )
-            tmpl = AsyncFunctionTaskTemplate(
-                func=wrapper,
-                name=task_name,
-                image=self.image,
-                resources=self.resources,
-                cache=cache or self.cache,
-                retries=retries,
-                timeout=timeout,
-                reusable=self.reusable,
-                docs=docs,
-                env=self.env,
-                secrets=secrets or self.secrets,
-                pod_template=pod_template or self.pod_template,
-                parent_env=weakref.ref(self),
-            )
-            self._tasks[task_name] = tmpl
-            return tmpl
-
-        if _func is None:
-            return decorator
-        return decorator(_func)
-
-    @property
-    def task(self) -> Union[AsyncFunctionTaskTemplate, Callable[P, R]]:
-        """
-        Decorator to create a new task with the environment settings.
-        The task will be executed in its own container with the specified image, resources, and environment variables,
-        unless reusePolicy is set, in which case the same container will be reused for all tasks with the same
-        environment settings.
-
-        :param name: Optional The name of the task (defaults to the function name)
-        :param cache: Optional The cache policy for the task, defaults to auto, which will cache the results of the
-         task.
-        :param retries: Optional The number of retries for the task, defaults to 0, which means no retries.
-        :param docs: Optional The documentation for the task, if not provided the function docstring will be used.
-        :param secrets: Optional The secrets that will be injected into the task at runtime.
-        :param timeout: Optional The timeout for the task.
-
-        :return: New Task instance or Task decorator
-        """
-        return self._task
-
-    @property
-    def tasks(self) -> Dict[str, TaskTemplate]:
-        """
-        Get all tasks defined in the environment.
-        """
-        return self._tasks
-
-    def add_task(self, task: TaskTemplate) -> TaskTemplate:
-        """
-        Add a task to the environment.
-        """
-        if task.name in self._tasks:
-            raise ValueError(f"Task {task.name} already exists in the environment. Task names should be unique.")
-        self._tasks[task.name] = task
-        return task

union/_timeout.py

@@ -1,47 +0,0 @@
-from dataclasses import dataclass
-from datetime import timedelta
-
-
-@dataclass
-class Timeout:
-    """
-    Timeout class to define a timeout for a task.
-    The task timeout can be set to a maximum runtime and a maximum queued time.
-    Maximum runtime is the maximum time the task can run for (in one attempt).
-    Maximum queued time is the maximum time the task can stay in the queue before it starts executing.
-
-    Example usage:
-    ```python
-    timeout = Timeout(max_runtime=timedelta(minutes=5), max_queued_time=timedelta(minutes=10))
-    @env.task(timeout=timeout)
-    async def my_task():
-        pass
-    ```
-    :param max_runtime: timedelta or int - Maximum runtime for the task. If specified int, it will be converted to
-    timedelta as seconds.
-    :param max_queued_time: optional, timedelta or int - Maximum queued time for the task. If specified int,
-    it will be converted to timedelta as seconds. Defaults to None.
-
-    """
-
-    max_runtime: timedelta | int
-    max_queued_time: timedelta | int | None = None
-
-
-TimeoutType = Timeout | int | timedelta
-
-
-def timeout_from_request(timeout: TimeoutType) -> Timeout:
-    """
-    Converts a timeout request into a Timeout object.
-    """
-    if isinstance(timeout, Timeout):
-        return timeout
-    else:
-        if isinstance(timeout, int):
-            timeout = timedelta(seconds=timeout)
-        elif isinstance(timeout, timedelta):
-            pass
-        else:
-            raise ValueError("Timeout must be an instance of Timeout, int, or timedelta.")
-        return Timeout(max_runtime=timeout)

union/_tools.py

@@ -1,27 +0,0 @@
-import os
-
-
-def ipython_check() -> bool:
-    """
-    Check if interface is launching from iPython (not colab)
-    :return is_ipython (bool): True or False
-    """
-    is_ipython = False
-    try:  # Check if running interactively using ipython.
-        from IPython import get_ipython
-
-        if get_ipython() is not None:
-            is_ipython = True
-    except (ImportError, NameError):
-        pass
-    return is_ipython
-
-
-def is_in_cluster() -> bool:
-    """
-    Check if the task is running in a cluster
-    :return is_in_cluster (bool): True or False
-    """
-    if os.getenv("_UN_CLS"):
-        return True
-    return False

union/_utils/__init__.py

@@ -1,11 +0,0 @@
-"""
-Internal utility functions.
-
-Except for logging, modules in this package should not depend on any other part of the repo.
-"""
-
-from .file_handling import filehash_update, update_hasher_for_source
-from .lazy_module import lazy_module
-from .uv_script_parser import parse_uv_script_file
-
-__all__ = ["filehash_update", "lazy_module", "parse_uv_script_file", "update_hasher_for_source"]