Creating Your First Task

In flyte-sdk, you define units of work by decorating Python functions with the @task decorator. This process transforms a standard function into an AsyncFunctionTaskTemplate, which encapsulates the function logic along with execution metadata like resource requirements, retry strategies, and caching policies.

Prerequisites

To follow this guide, ensure you have flyte-sdk installed and can import the core components:

import flyte
from flyte import Resources, RetryStrategy

Step 1: Initialize a Task Environment

Before defining tasks, you must create a TaskEnvironment. This environment acts as a container and configuration provider for your tasks, defining default settings like the Docker image and hardware resources.

# Define an environment for your data processing tasks
data_env = flyte.TaskEnvironment(
    name="data_processing",
    image="ghcr.io/flyteorg/flytekit:py3.11-latest",
    resources=Resources(cpu="2", memory="4Gi")
)

The TaskEnvironment class (found in _task_environment.py) manages the lifecycle and registration of tasks within its scope.

Step 2: Define a Synchronous Task

To create a task, apply the .task decorator from your environment instance to a Python function. For simple computations, a standard synchronous function is sufficient.

@data_env.task(cache="auto", retries=3)
def square(n: int) -> int:
    return n * n

When you apply @data_env.task, flyte-sdk creates an AsyncFunctionTaskTemplate (defined in _task.py). Even though the function is synchronous, the SDK sets _call_as_synchronous = True internally to ensure it is executed correctly within the Flyte engine while still allowing it to be called normally in local scripts.

Step 3: Define an Asynchronous Task

flyte-sdk natively supports async functions. These are particularly useful for I/O-bound operations or when you need to manage concurrency within a task.

import asyncio

@data_env.task(timeout=60)
async def fetch_data(url: str) -> str:
    # Simulate an async network request
    await asyncio.sleep(1)
    return f"Data from {url}"

For asynchronous functions, the AsyncFunctionTaskTemplate uses await self.func(*args, **kwargs) during its execute method.

Note: If you use a reusable environment (where a Python process is kept alive for multiple tasks), flyte-sdk requires tasks to be async if you want to run them with a concurrency greater than 1. This is enforced in TaskEnvironment.task to prevent blocking the shared process.

Step 4: Overriding Task Configuration

You can customize individual tasks by passing arguments to the decorator. These arguments override the defaults set in the TaskEnvironment.

@data_env.task(
    short_name="heavy_task",
    resources=Resources(cpu="4", memory="16Gi"),
    retries=RetryStrategy(count=5)
)
def memory_intensive_op(data: list[int]) -> int:
    return sum(data)

The TaskTemplate base class handles these overrides, ensuring that specific requirements like pod_template or env_vars are correctly associated with the task template before it is registered in the environment's _tasks dictionary.

Complete Example

Here is how your task definitions look when combined in a single module:

import flyte
from flyte import Resources
import asyncio

# 1. Setup Environment
env = flyte.TaskEnvironment(name="tutorial_env")

# 2. Define Sync Task
@env.task
def add(a: int, b: int) -> int:
    return a + b

# 3. Define Async Task
@env.task
async def process(x: int) -> int:
    result = await asyncio.to_thread(lambda: x * 10)
    return result

# Local execution works like standard Python
if __name__ == "__main__":
    print(f"Sync result: {add(5, 10)}")
    print(f"Async result: {asyncio.run(process(5))}")

When you run this code locally, the AsyncFunctionTaskTemplate.forward method is invoked, which simply calls your original function, allowing for easy testing and debugging outside of a Flyte cluster.