flyte 0.0.1b0__py3-none-any.whl

flyte/_secret.py

@@ -0,0 +1,61 @@
+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]

flyte/_task.py

@@ -0,0 +1,367 @@
+from __future__ import annotations
+
+import weakref
+from dataclasses import dataclass, field, replace
+from functools import cached_property
+from typing import (
+    TYPE_CHECKING,
+    Any,
+    Awaitable,
+    Callable,
+    Coroutine,
+    Dict,
+    Generic,
+    List,
+    Literal,
+    Optional,
+    ParamSpec,
+    TypeVar,
+    Union,
+)
+
+from flyteidl.core.tasks_pb2 import DataLoadingConfig
+
+from flyte.errors import RuntimeSystemError, RuntimeUserError
+
+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
+
+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 flyte 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
+    interface: NativeInterface
+    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
+    report: bool = False
+
+    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
+
+    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) -> Coroutine[Any, Any, 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.is_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)
+            return await self.execute(*args, **kwargs)
+        except RuntimeSystemError:
+            raise
+        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.
+        """
+        resources = resources or self.resources
+        cache = cache or self.cache
+        retries = retries or self.retries
+        timeout = timeout or self.timeout
+        reusable = reusable or self.reusable
+        env = env or self.env
+        secrets = secrets or self.secrets
+        local = local or self.local
+        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")
+        return replace(
+            self,
+            resources=resources,
+            cache=cache,
+            retries=retries,
+            timeout=timeout,
+            reusable=reusable,
+            env=env,
+            secrets=secrets,
+        )
+
+
+@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, Awaitable[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 = [
+            "a0",
+            "--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 flyte._internal.resolvers.default import DefaultTaskResolver
+
+            _task_resolver = DefaultTaskResolver()
+            args = [
+                *args,
+                *[
+                    "--resolver",
+                    _task_resolver.import_path,
+                    *_task_resolver.loader_args(task=self, root_dir=serialize_context.root_dir),
+                ],
+            ]
+
+        assert all(isinstance(item, str) for item in args), f"All args should be strings, non string item = {args}"
+
+        return args

flyte/_task_environment.py

@@ -0,0 +1,200 @@
+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, Awaitable, Callable, Dict, List, Literal, Optional, ParamSpec, TypeVar, Union, cast
+
+import rich.repr
+
+from ._cache import CacheRequest
+from ._datastructures import NativeInterface
+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 = flyte.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[Environment]] = None,
+    ) -> TaskEnvironment:
+        """
+        Clone the environment with new settings.
+        """
+        if image is None:
+            image = self.image
+        else:
+            image = "auto"
+        return replace(
+            self,
+            cache=cache if cache else "auto",
+            reusable=reusable,
+            name=name,
+            image=image,
+            resources=resources,
+            env=env,
+            secrets=secrets,
+            env_dep_hints=env_dep_hints if env_dep_hints else [],
+        )
+
+    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,
+        report: bool = False,
+    ) -> 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.
+        :param pod_template: Optional The pod template for the task, if not provided the default pod template will be
+        used.
+        :param report: Optional Whether to generate the html report for the task, defaults to False.
+        """
+        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, Awaitable[R]]) -> AsyncFunctionTaskTemplate[P, R]:
+            task_name = name or func.__name__
+            task_name = self.name + "." + task_name
+            if len(task_name) > 30:
+                # delete this if we can remove the restriction that task names must be <= 30 characters
+                task_name = task_name[-30:]
+
+            @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),
+                interface=NativeInterface.from_callable(func),
+                report=report,
+            )
+            self._tasks[task_name] = tmpl
+            return tmpl
+
+        if _func is None:
+            return cast(AsyncFunctionTaskTemplate, decorator)
+        return cast(AsyncFunctionTaskTemplate, decorator(_func))
+
+    @property
+    def task(self) -> Callable:
+        """
+        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.
+        :param pod_template: Optional The pod template for the task, if not provided the default pod template will be
+         used.
+        :param report: Optional Whether to generate the html report for the task, defaults to False.
+
+        :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

flyte/_timeout.py

@@ -0,0 +1,47 @@
+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)