ClusteredTaskEnvironment

Package: flyte.clustered

A TaskEnvironment that emits a Kubernetes JobSet for distributed multi-node training.

Inherits all fields from TaskEnvironment (name, image, resources, env_vars, secrets, pod_template, queue, cache, reusable). The fields below are specific to clustered execution.

Restarts and retries:

A clustered task can start over at three levels; only the second uses the task’s retries.

1. JobSet restarts (failure_policy): when a worker fails, every pod of the JobSet is recreated in place, up to failure_policy.max_restarts times (node maintenance is free with restart_on_host_maintenance). flyte.ctx().restart_attempt counts them.
2. User retries (retries= on env.task): a new attempt with a new JobSet. Each uses one of the task’s retries and waits RetryStrategy.backoff first. Causes: the task failed after its JobSet restarts were used up, or the gang was preempted by higher-priority work after it started. If the task can be preempted, set retries > 0, otherwise the first preemption fails it. Limit run time with timeout=flyte.Timeout(max_runtime=...).
3. System retries: the same attempt with a new JobSet, after an infrastructure problem such as the queue being stopped or nodes failing. They do not use retries.

The failure message names the cause and, for user retries, how many retries are left.

Waiting and queues:

The task waits until every worker is ready. Bound the wait with timeout=flyte.Timeout(max_queued_time=..., deadline=...) on env.task: max_queued_time covers each attempt until all workers are ready, deadline the whole run across attempts and retries.

With Kueue on the cluster, choose the queue with a pod-template label:

env = ClusteredTaskEnvironment(
    ...,
    pod_template=flyte.PodTemplate(labels={"kueue.x-k8s.io/queue-name": "training"}),
)

Parameters

class ClusteredTaskEnvironment(
    name: str,
    depends_on: List[Environment] = <factory>,
    pod_template: Optional[Union[str, PodTemplate]] = None,
    description: Optional[str] = None,
    secrets: Optional[SecretRequest] = None,
    env_vars: Optional[Dict[str, str]] = None,
    resources: Optional[Resources] = None,
    interruptible: bool = False,
    image: Union[str, Image, Literal['auto'], None] = 'auto',
    include: Tuple[str, ...] = <factory>,
    service_account: Optional[str] = None,
    cache: CacheRequest = 'disable',
    reusable: ReusePolicy | None = None,
    plugin_config: Optional[Any] = None,
    queue: Optional[str] = None,
    replicas: int,
    nproc_per_node: int,
    runtime: Runtime = <factory>,
    interconnect: Literal['tcp'] = 'tcp',
    failure_policy: ClusterFailurePolicy = <factory>,
    ttl_seconds_after_finished: Optional[int] = None,
)

Properties

Methods

add_dependency()

def add_dependency(
    *env: Environment,
)

Add one or more environment dependencies so they are deployed together.

When you deploy this environment, any environments added via add_dependency will also be deployed. This is an alternative to passing depends_on=[...] at construction time, useful when the dependency is defined after the environment is created.

Duplicate dependencies are silently ignored. An environment cannot depend on itself.

clone_with()

def clone_with(
    name: str,
    image: Optional[Union[str, Image, Literal['auto']]] = None,
    resources: Optional[Resources] = None,
    env_vars: Optional[Dict[str, str]] = None,
    secrets: Optional[SecretRequest] = None,
    depends_on: Optional[List[Environment]] = None,
    description: Optional[str] = None,
    interruptible: Optional[bool] = None,
    include: Optional[Tuple[str, ...]] = None,
    service_account: Optional[str] = None,
    **kwargs: Any,
) -> TaskEnvironment

Create a new TaskEnvironment that shares most settings with this one but differs in name and selected overrides.

Use clone_with when you need several environments that share a common base configuration (image, resources, secrets, etc.) but vary in one or two settings, avoiding repetition.

gpu_env = flyte.TaskEnvironment(
    name="gpu_env",
    image=my_image,
    resources=flyte.Resources(gpu="A100:1", memory="16Gi"),
)

# Same image and resources, different name and cache policy
gpu_env_cached = gpu_env.clone_with("gpu_env_cached", cache="auto")

Any parameter not explicitly passed inherits the value from the original environment.

from_task()

def from_task(
    name: str,
    *tasks: TaskTemplate,
    depends_on: Optional[List['Environment']] = None,
) -> TaskEnvironment

Create a TaskEnvironment from a list of tasks. All tasks should have the same image or no Image defined. Similarity of Image is determined by the python reference, not by value.

If images are different, an error is raised. If no image is defined, the image is set to “auto”.

For any other tasks that need to be use these tasks, the returned environment can be used in the depends_on attribute of the other TaskEnvironment.

Returns

The created TaskEnvironment.

Raises

task()

def task(
    short_name: Optional[str] = None,
    prefix: Optional[str] = None,
    cache: CacheRequest | None = None,
    retries: Union[int, RetryStrategy] = 0,
    timeout: TimeoutType = 0,
    docs: Optional[Documentation] = None,
    pod_template: Optional[Union[str, PodTemplate]] = None,
    report: bool = False,
    interruptible: bool | None = None,
    max_inline_io_bytes: int = 10485760,
    queue: Optional[str] = None,
    service_account: Optional[str] = None,
    triggers: Tuple[Trigger, ...] | Trigger = (),
    links: Tuple[Link, ...] | Link = (),
    task_resolver: Any | None = None,
    entrypoint: bool = False,
    produces_artifacts: bool = False,
) -> Callable[[F], AsyncFunctionTaskTemplate[P, R, F]] | AsyncFunctionTaskTemplate[P, R, F]

Decorate a function to be a task.

Returns: A TaskTemplate that can be used to deploy the task.

to_custom_dict()

def to_custom_dict()

Serialize this environment to the dict shape expected by ClusteredTaskSpec proto.