Running Dask on the cluster with mlrun

The dask frameworks enables users to parallelize their python code and run it as a distributed process on Iguazio cluster and dramatically accelerate their 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/

Set up the environment

# set mlrun api path and arrtifact path for logging
import mlrun
project_name = "dask-demo"
mlrun.set_environment(project=project_name, artifact_path = './')
('dask-demo', '/User/dask')

Create and Start Dask Cluster

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

if the dask workers need to access the shared file system we apply a shared volume mount (e.g. via v3io mount).

Dask function spec have 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 will be killed after timeout (inactivity), default is ‘60 minutes’

  • .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 mlrun configuration (mlconf.remote_host), this will be set automatically if you are running on an Iguazio cluster.

We specify the kind (dask) and the container image

# create an mlrun function which will init the dask cluster
dask_cluster_name = "dask-cluster"
dask_cluster = mlrun.new_function(dask_cluster_name, kind='dask', image='mlrun/ml-models')
dask_cluster.apply(mlrun.mount_v3io())
<mlrun.runtimes.daskjob.DaskCluster at 0x7f7dbe4166d0>
# 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_requests(mem='2G', cpu='2')

Initialize the Dask Cluster

When we request the dask cluster client attribute it will verify the cluster is up and running

# init dask client and use the scheduler address as param in the following cell
dask_cluster.client
> 2021-01-24 23:48:54,057 [info] trying dask client at: tcp://mlrun-dask-cluster-b3c6e737-3.default-tenant:8786
> 2021-01-24 23:48:54,067 [info] using remote dask scheduler (mlrun-dask-cluster-b3c6e737-3) at: tcp://mlrun-dask-cluster-b3c6e737-3.default-tenant:8786
/User/.pythonlibs/jupyter/lib/python3.7/site-packages/distributed/client.py:1129: VersionMismatchWarning: Mismatched versions found

+---------+--------+-----------+---------+
| Package | client | scheduler | workers |
+---------+--------+-----------+---------+
| blosc   | 1.7.0  | 1.10.2    | 1.10.2  |
| lz4     | 3.1.0  | 3.1.3     | 3.1.3   |
| msgpack | 1.0.0  | 1.0.2     | 1.0.2   |
| numpy   | 1.19.2 | 1.18.1    | 1.18.1  |
| toolz   | 0.11.1 | 0.10.0    | 0.10.0  |
| tornado | 6.0.4  | 6.0.3     | 6.0.3   |
+---------+--------+-----------+---------+
Notes: 
-  msgpack: Variation is ok, as long as everything is above 0.6
  warnings.warn(version_module.VersionMismatchWarning(msg[0]["warning"]))

Client

Cluster

  • Workers: 1
  • Cores: 1
  • Memory: 4.12 GB

Creating A Function Which Run Over Dask

# mlrun: start-code

Import mlrun and dask. nuclio is used just to convert the code into an mlrun function

import mlrun 
%nuclio config kind = "job"
%nuclio config spec.image = "mlrun/ml-models"
%nuclio: setting kind to 'job'
%nuclio: setting spec.image to 'mlrun/ml-models'
from dask.distributed import Client
from dask import delayed
from dask import dataframe as dd

import warnings
import numpy as np
import os
import mlrun

warnings.filterwarnings("ignore")

python function code

This simple function reads a csv file using dask dataframe and run group by and describe function on the dataset and store 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 Our 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
100 84.9M  100 84.9M    0     0  17.3M      0  0:00:04  0:00:04 --:--:-- 19.1M

Convert the code to MLRun function

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

# mlrun will transform the code above (up to nuclio: end-code cell) into serverless function 
# which will run in k8s pods
fn = mlrun.code_to_function("test_dask",  kind='job', handler="test_dask").apply(mlrun.mount_v3io())

Run the function

When running the function you would see a link as part of the result. click on this link takes you to the dask monitoring dashboard

# 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})
> 2021-01-24 23:49:37,858 [info] starting run test-dask-test_dask uid=6410ec27b63e4a12b025696fcabc2dc9 DB=http://mlrun-api:8080
> 2021-01-24 23:49:38,069 [info] Job is running in the background, pod: test-dask-test-dask-rmgkn
> 2021-01-24 23:49:41,647 [warning] Unable to parse server or client version. Assuming compatible: {'server_version': 'unstable', 'client_version': 'unstable'}
> 2021-01-24 23:49:42,112 [info] using in-cluster config.
> 2021-01-24 23:49:42,113 [info] trying dask client at: tcp://mlrun-dask-cluster-b3c6e737-3.default-tenant:8786
> 2021-01-24 23:49:42,134 [info] using remote dask scheduler (mlrun-dask-cluster-b3c6e737-3) at: tcp://mlrun-dask-cluster-b3c6e737-3.default-tenant:8786
remote dashboard: default-tenant.app.yh57.iguazio-cd0.com:30433
> 2021-01-24 23:49:48,334 [info] run executed, status=completed
final state: completed
project uid iter start state name labels inputs parameters results artifacts
dask-demo 0 Jan 24 23:49:42 completed test-dask-test_dask
v3io_user=admin
kind=job
owner=admin
host=test-dask-test-dask-rmgkn
dataset
dask_function=db://dask-demo/dask-cluster
describe
to track results use .show() or .logs() or in CLI: 
!mlrun get run 6410ec27b63e4a12b025696fcabc2dc9 --project dask-demo , !mlrun logs 6410ec27b63e4a12b025696fcabc2dc9 --project dask-demo
> 2021-01-24 23:49:50,284 [info] run executed, status=completed

Track the progress in the UI

Users can view the progress and detailed information in the mlrun UI by clicking on the uid above.
Also, to track the dask progress in the dask UI click on the “dashboard link” above the “client” section