Creating Tasks from Python Functions

Flyte SDK transforms standard Python functions, including asynchronous ones, into executable tasks that can run on the Flyte platform. This transformation is primarily handled by the AsyncFunctionTaskTemplate class, which wraps your Python code and provides the necessary logic for execution, serialization, and containerization.

The AsyncFunctionTaskTemplate

At its core, AsyncFunctionTaskTemplate is a specialized implementation of TaskTemplate designed to encapsulate a Python function. It is automatically generated when you decorate a function with the @env.task decorator.

@dataclass(kw_only=True)
class AsyncFunctionTaskTemplate(TaskTemplate[P, R]):
    func: FunctionTypes
    plugin_config: Optional[Any] = None

    # ... other methods ...

The func attribute stores the actual Python function being wrapped. The plugin_config attribute allows for passing plugin-specific configurations, enabling custom task behaviors. For debugging and traceability, the source_file property can return the path to the file where the function is defined:

@property
def source_file(self) -> Optional[str]:
    if hasattr(self.func, "__code__") and self.func.__code__:
        return self.func.__code__.co_filename
    return None

Transforming Functions with @env.task

The most common way to create a Flyte task from a Python function is by using the task decorator provided by a TaskEnvironment instance. This decorator inspects your function and automatically constructs an AsyncFunctionTaskTemplate behind the scenes.

First, define a TaskEnvironment:

import flyte

env = flyte.TaskEnvironment("my_project_env")

Then, apply the @env.task decorator to your Python functions. This works seamlessly for both asynchronous and synchronous functions.

Here's an example of an asynchronous function being converted into a Flyte task, as seen in flyte/io/_file.py:

import pandas as pd
from flyte.io import File

@env.task
async def my_async_task(file: File[pd.DataFrame]):
    async with file.open() as f:
        df = pd.read_csv(f)
    # Further processing with df

For synchronous functions, the process is identical:

@env.task
def my_sync_task(a: int, b: int) -> int:
    return a + b

Internally, the TaskEnvironment.task decorator (found in flyte/_task_environment.py) is responsible for this creation. It checks if a plugin_config is provided to determine if a custom TaskPluginRegistry should be consulted, otherwise, it defaults to AsyncFunctionTaskTemplate:

# Excerpt from flyte/_task_environment.py
def decorator(func: FunctionTypes) -> AsyncFunctionTaskTemplate[P, R]:
    # ... (snip)
    if self.plugin_config is not None:
        from flyte.extend import TaskPluginRegistry

        task_template_class: type[AsyncFunctionTaskTemplate[P, R]] | None = TaskPluginRegistry.find(
            config_type=type(self.plugin_config)
        )
        # ... (snip)
    else:
        task_template_class = AsyncFunctionTaskTemplate[P, R]

    task_template_class = cast(type[AsyncFunctionTaskTemplate[P, R]], task_template_class)
    tmpl = task_template_class(
        func=func,
        name=task_name,
        # ... other parameters ...
    )
    self._tasks[task_name] = tmpl
    return tmpl

Advanced Task Configuration

The @env.task decorator offers a rich set of parameters to configure your task's behavior, which are then passed to the AsyncFunctionTaskTemplate constructor. These include settings for caching, retries, timeouts, and resource allocation.

from datetime import timedelta
from flyte.models import Resources

@env.task(
    cache=True,
    cache_version="1.0",
    retries=3,
    timeout=timedelta(minutes=5),
    resources=Resources(cpu="1", mem="500Mi"),
    # pod_template=my_pod_template, # if you have a custom pod template
)
def configured_task(data: str) -> str:
    # Task logic here
    return data.upper()

Task Execution Flow

AsyncFunctionTaskTemplate defines two key methods for execution: forward and execute.

Local Execution with forward

When you invoke a task locally (e.g., during development or testing), the forward method is called. This method directly executes the wrapped Python function without any Flyte platform overhead. If the function is asynchronous, forward returns the coroutine, allowing the caller to await it.

def forward(self, *args: P.args, **kwargs: P.kwargs) -> Coroutine[Any, Any, R] | R:
    # In local execution, we want to just call the function. Note we're not awaiting anything here.
    # If the function was a coroutine function, the coroutine is returned and the await that the caller has
    # in front of the task invocation will handle the awaiting.
    return self.func(*args, **kwargs)

Remote Execution with execute

The execute method is invoked when the task runs on the Flyte platform. It handles the execution within the Flyte context, managing task-specific data and ensuring proper handling of both synchronous and asynchronous functions. It also includes pre and post hooks for potential future extensions.

async def execute(self, *args: P.args, **kwargs: P.kwargs) -> R:
    ctx = internal_ctx()
    assert ctx.data.task_context is not None, "Function should have already returned if not in a task context"
    ctx_data = await self.pre(*args, **kwargs)
    tctx = ctx.data.task_context.replace(data=ctx_data)
    with ctx.replace_task_context(tctx):
        if iscoroutinefunction(self.func):
            v = await self.func(*args, **kwargs)
        else:
            v = self.func(*args, **kwargs)
        await self.post(v)
    return v

Container Argument Generation

The container_args method within AsyncFunctionTaskTemplate is crucial for preparing the task to run within a containerized environment. It generates the command-line arguments that the Flyte agent uses to execute your task, including paths for inputs, outputs, and metadata.

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}}",
        "{{.prevCheckpointPrefix}}",
        "--run-name",
        "{{.runName}}",
        "--name",
        "{{.actionName}}",
    ]
    # ... (additional arguments for image cache, code bundle, and resolver)
    return args

This method ensures that the container running your task receives all necessary information to locate inputs, store outputs, and resolve task-specific code, especially when dealing with code bundles or custom resolvers.

Extensibility with Task Plugins

AsyncFunctionTaskTemplate serves as a foundational type for extending Flyte's task capabilities. The TaskPluginRegistry (located in flyte/_task_plugins.py) allows users to register custom task templates that might inherit from or wrap AsyncFunctionTaskTemplate. This enables specialized handling for different types of tasks or external systems, driven by the plugin_config provided during task definition.

Gotchas and Best Practices

• Concurrency for Asynchronous Tasks: When using reusable environments, concurrency greater than one is only supported for asynchronous tasks. If you define a synchronous task within a reusable environment, its concurrency must be explicitly set to one. This is enforced by the TaskEnvironment.task decorator, which will raise a ValueError if violated.
• Plugin Registration: If you provide a plugin_config to the @env.task decorator, ensure that a corresponding task plugin is registered with flyte.extend.TaskPluginRegistry. Failure to do so will result in a ValueError during task definition.
• Internal Development Notes: The AsyncFunctionTaskTemplate includes internal TODO comments, such as in the execute method (# TODO We may need to keep this as the bare func execute, and need a pre and post execute some other func.) and a pr comment in container_args regarding versioning (# pr: should this be serialize_context.version or code_bundle.version?). These indicate areas of active development or ongoing design considerations within the Flyte SDK. While not directly impacting current usage, they offer insight into the evolving nature of the codebase.