Skip to main content

Dask and Ray Distributed Execution

Running Dask Workloads​

To execute distributed Dask workloads as Flyte tasks, you define your task using the flyteplugins.dask.task.DaskTask plugin. This plugin requires a plugin_config of type flyteplugins.dask.task.Dask, which allows you to specify the configuration for your Dask cluster, including the scheduler and worker settings.

Here's a basic example of a Dask task that uses a default scheduler and a worker group with four workers:

import asyncio
import typing

from dask.distributed import Client
from flytekit import task, workflow
from flyteplugins.dask.task import Dask, DaskTask, Scheduler, WorkerGroup

dask_config = Dask()

@task(task_config=DaskTask(plugin_config=dask_config))
async def hello_dask_nested(n: int = 3) -> typing.List[int]:
    print("running dask task")
    client = Client()
    futures = client.map(lambda x: x + 1, range(n))
    res = client.gather(futures)
    return res

@workflow
def dask_workflow(n: int = 3) -> typing.List[int]:
    return hello_dask_nested(n=n)

The flyteplugins.dask.task.Dask object provides two primary configuration components: scheduler and workers.

  • The scheduler parameter accepts a flyteplugins.dask.task.Scheduler object. This allows you to configure the Dask scheduler pod, including specifying a custom image and allocating resources (CPU and memory).
  • The workers parameter accepts a flyteplugins.dask.task.WorkerGroup object. Here, you can define the number_of_workers, their image, and resources.

For instance, to specify custom images and resource requests for both the scheduler and workers, you would configure Dask as follows:

from flytekit.models.resources import Resources
from flyteplugins.dask.task import Dask, Scheduler, WorkerGroup

dask_config_with_resources = Dask()

Dask Task Considerations​

  • Image Requirements: Any custom images used for Dask scheduler and worker pods (specified via Scheduler.image and WorkerGroup.image) must have dask[distributed] installed. Additionally, the Python environment within these images should be consistent with the environment of the job runner pod.
  • Code Bundle Distribution: The flyteplugins.dask.task.DaskTask automatically handles the distribution of your code bundle to the Dask scheduler and worker pods. This is managed by DownloadCodeBundleSchedulerPlugin and DownloadCodeBundleWorkerPlugin, which register with the Dask client. This mechanism is active only when the task is executing within a Flyte cluster (flyte.ctx().is_in_cluster() is true) and a code bundle is present.

Running Ray Workloads​

To execute distributed Ray jobs as Flyte tasks, you use the flyteplugins.ray.task.RayFunctionTask plugin. This plugin is configured via a plugin_config of type flyteplugins.ray.task.RayJobConfig, which defines the Ray cluster's head and worker node settings.

Here's an example of a Ray task that uses a single worker group with two replicas and a custom runtime environment:

import asyncio
import typing

import ray
from flytekit import task, workflow
from flyteplugins.ray.task import HeadNodeConfig, RayFunctionTask, RayJobConfig, WorkerNodeConfig

@ray.remote
def f(x):
    return x * x

ray_config = RayJobConfig()

@task(task_config=RayFunctionTask(plugin_config=ray_config))
async def hello_ray_nested(n: int = 3) -> typing.List[int]:
    print("running ray task")
    futures = [f.remote(i) for i in range(n)]
    res = ray.get(futures)
    return res

@workflow
def ray_workflow(n: int = 3) -> typing.List[int]:
    return hello_ray_nested(n=n)

The flyteplugins.ray.task.RayJobConfig object allows for detailed configuration of the Ray cluster:

  • head_node_config: An optional flyteplugins.ray.task.HeadNodeConfig object to configure the Ray head node. You can specify ray_start_params (parameters passed to ray start), a pod_template, and resource requests and limits.
  • worker_node_config: A list of flyteplugins.ray.task.WorkerNodeConfig objects, each defining a worker group. For each worker group, you can set group_name, replicas (number of instances), min_replicas, max_replicas (for autoscaling), ray_start_params, a pod_template, and resource requests and limits.
  • enable_autoscaling: A boolean to enable or disable autoscaling for the Ray cluster.
  • runtime_env: A dictionary to specify the Ray runtime environment, such as Python package dependencies. This is deprecated in KubeRay versions 1.1.0 and later, where runtime_env_yaml should be used instead.
  • address: An optional string specifying the address of the Ray cluster. If not provided, Ray might default to local execution or require specific environment setup for cluster execution.
  • shutdown_after_job_finishes: A boolean indicating whether the Ray cluster should shut down immediately after the job completes.
  • ttl_seconds_after_finished: An optional integer specifying the time-to-live for the Ray cluster after the job finishes, allowing for post-job inspection or cleanup.

For example, to configure a Ray job with specific resource limits for the head node and a worker group with autoscaling enabled:

from flytekit.models.resources import Resources
from flyteplugins.ray.task import HeadNodeConfig, RayJobConfig, WorkerNodeConfig

ray_config_with_resources = RayJobConfig()

head_node_config_instance = HeadNodeConfig()
head_node_config_instance.limits = Resources(cpu="2", memory="4Gi")
ray_config_with_resources.head_node_config = head_node_config_instance

worker_node_config_instance = WorkerNodeConfig()
worker_node_config_instance.group_name = "high-cpu-workers"
worker_node_config_instance.replicas = 1
worker_node_config_instance.min_replicas = 1
worker_node_config_instance.max_replicas = 5
worker_node_config_instance.limits = Resources(cpu="4", memory="8Gi")
ray_config_with_resources.worker_node_config = [worker_node_config_instance]

ray_config_with_resources.enable_autoscaling = True

Ray Task Considerations​

  • Runtime Environment: The runtime_env parameter in RayJobConfig is deprecated in KubeRay versions 1.1.0 and later. For newer KubeRay deployments, use runtime_env_yaml for specifying the Ray runtime environment. The RayFunctionTask handles the conversion of runtime_env to runtime_env_yaml internally if runtime_env is provided.
  • Ray Initialization: The RayFunctionTask automatically initializes Ray within the task's pre method. If flyte.ctx().is_in_cluster() is true, it configures the runtime_env with the current working directory and excludes specific files to ensure proper code distribution within the cluster. If Ray is already initialized, it will not re-initialize it.