Running Dask on the cluster with MLRun

Running Dask on the cluster with MLRun#

The Dask framework enables you to parallelize your Python code and run it as a distributed process on an Iguazio cluster and dramatically accelerate the performance. In this notebook you'll learn how to create a Dask cluster and then an MLRun function running as a Dask client. It also demonstrates how to run parallelize custom algorithm using Dask Delayed option.

For more information on Dask over Kubernetes: https://kubernetes.dask.org/en/latest/.

Note

Dask is currently in Tech Preview status.

In this section

Set up the environment
Create and start the Dask cluster
Creating a function that runs over Dask
Test the function over Dask
Track the progress in the UI

Set up the environment#

# set mlrun api path and artifact path for logging
import mlrun

project = mlrun.get_or_create_project("dask-demo", "./")

> 2024-10-02 08:12:17,022 [info] Created and saved project: {"context":"./","from_template":null,"name":"dask-demo","overwrite":false,"save":true}
> 2024-10-02 08:12:17,026 [info] Project created successfully: {"project_name":"dask-demo","stored_in_db":true}

Create and start the Dask cluster#

Dask functions can be local (local workers), or remote (use containers in the cluster). In the case of remote you can specify the number of replicas (optional) or leave blank for auto-scale.
Use the set() to define the Dask cluster and set the desired configuration of that clustered function.

If the Dask workers need to access the shared file system, apply a shared volume mount (e.g. via v3io mount).

The Dask function spec has several unique attributes (in addition to the standard job attributes):

.remote — bool, use local or clustered dask
.replicas — number of desired replicas, keep 0 for auto-scale
.min_replicas, .max_replicas — set replicas range for auto-scale
.scheduler_timeout — cluster is killed after timeout (inactivity), default is '60 minutes', should be at least 5 minutes to avoid transient issues
.nthreads — number of worker threads

If you want to access the Dask dashboard or scheduler from remote you need to use NodePort service type (set .service_type to 'NodePort'), and the external IP need to be specified in the MLRun configuration (mlconf.remote_host). This is set automatically if you are running on an Iguazio cluster.

Specify the kind (dask) and the container image:

# create an mlrun function that will init the dask cluster
dask_cluster_name = "dask-cluster"
dask_cluster = project.set_function(
    name=dask_cluster_name, kind="dask", image="mlrun/mlrun"
)
dask_cluster.apply(mlrun.mount_v3io())

<mlrun.runtimes.daskjob.DaskCluster at 0x7f965e52b610>

# set range for # of replicas with replicas and max_replicas
dask_cluster.spec.min_replicas = 1
dask_cluster.spec.max_replicas = 4

# set the use of dask remote cluster (distributed)
dask_cluster.spec.remote = True
dask_cluster.spec.service_type = "NodePort"

# set dask memory and cpu limits
dask_cluster.with_worker_requests(mem="2G", cpu="2")

Initialize the Dask Cluster#

When you request the dask cluster client attribute, it verifies that the cluster is up and running:

import warnings

# Ignore mismatched versions warning logged by dask
warnings.simplefilter("ignore")

# init dask client and use the scheduler address as param in the following cell
dask_cluster.client

> 2024-10-02 08:13:56,759 [info] Trying dask client at: tcp://mlrun-dask-cluster-05bbaed5-3.default-tenant:8786
> 2024-10-02 08:13:56,833 [info] Using remote dask scheduler (mlrun-dask-cluster-05bbaed5-3) at: tcp://mlrun-dask-cluster-05bbaed5-3.default-tenant:8786

dashboard link: default-tenant.app.vmdev17.lab.iguazeng.com:31898

Client

Client-4548a53d-8096-11ef-a9bc-6e177956313e

Connection method: Direct
Dashboard: http://mlrun-dask-cluster-05bbaed5-3.default-tenant:8787/status

Scheduler Info

Scheduler

Scheduler-0a2c3809-fca1-4535-bf7f-1457230f1857

Comm: tcp://10.233.92.190:8786	Workers: 0
Dashboard: http://10.233.92.190:8787/status	Total threads: 0
Started: 1 minute ago	Total memory: 0 B

Workers

Creating a function that runs over Dask#

# mlrun: start-code

Import mlrun and dask. Nuclio is only used to convert the code into an MLRun function.

from dask.distributed import Client
from dask import dataframe as dd

import mlrun

Python function code#

This simple function reads a .csv file using dask dataframe. It runs the groupby and describe functions on the dataset, and stores the results as a dataset artifact.

def test_dask(
    context, dataset: mlrun.DataItem, client=None, dask_function: str = None
) -> None:
    # setup dask client from the MLRun dask cluster function
    if dask_function:
        client = mlrun.import_function(dask_function).client
    elif not client:
        client = Client()

    # load the dataitem as dask dataframe (dd)
    df = dataset.as_df(df_module=dd)

    # run describe (get statistics for the dataframe) with dask
    df_describe = df.describe().compute()

    # run groupby and count using dask
    df_grpby = df.groupby("VendorID").count().compute()

    context.log_dataset("describe", df=df_grpby, format="csv", index=True)
    return

# mlrun: end-code

Test the function over Dask#

Load sample data#

DATA_URL = "/User/examples/ytrip.csv"

!mkdir -p /User/examples/
!curl -L "https://s3.wasabisys.com/iguazio/data/Taxi/yellow_tripdata_2019-01_subset.csv" > {DATA_URL}

  % Total    % Received % Xferd  Average Speed   Time    Time     Time  Current
                                 Dload  Upload   Total   Spent    Left  Speed
 54 84.9M   54 46.4M    0     0  1653k      0  0:00:52  0:00:28  0:00:24 1680k5M    0     0  1655k      0  0:00:52  0:00:20  0:00:32 1529k

Convert the code to MLRun function#

Use set_function to convert the code to MLRun function and specify the configuration for the Dask process (e.g. replicas, memory etc.).
Note that the resource configurations are per worker.

# mlrun transforms the code above (up to nuclio: end-code cell) into serverless function
# which runs in k8s pods
fn = project.set_function(name="test-dask", kind="job", handler="test_dask").apply(
    mlrun.mount_v3io()
)

Run the function#

# function URI is db://<project>/<name>
dask_uri = f"db://{project.name}/{dask_cluster_name}"

r = fn.run(
    handler=test_dask,
    inputs={"dataset": DATA_URL},
    params={"dask_function": dask_uri},
    auto_build=True,
)

Track the progress in the UI#

You can view the progress and detailed information in the MLRun UI by clicking on the uid above.
To track the dask progress: in the Dask UI click the "dashboard link" above the "client" section.