Skip to content

Dask

This example showcases how to run Dask within a Hera submitted Argo workflow.

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
from hera.workflows import (
    Steps,
    Workflow,
    script,
)


@script(image="ghcr.io/dask/dask:latest")
def dask_computation(namespace: str = "default", n_workers: int = 1) -> None:
    import subprocess

    # this is required for otherwise the dask distributed and kubernetes clients packages are not included by default
    # ideally, you'd have a package in your organization that takes care of the following cluster details :)
    subprocess.run(
        ["pip", "install", "dask-kubernetes", "dask[distributed]"], stdout=subprocess.PIPE, universal_newlines=True
    )

    import dask.array as da
    from dask.distributed import Client
    from dask_kubernetes.classic import KubeCluster, make_pod_spec

    cluster = KubeCluster(
        pod_template=make_pod_spec(
            image="ghcr.io/dask/dask:latest",
            memory_limit="4G",
            memory_request="2G",
            cpu_limit=1,
            cpu_request=1,
        ),
        namespace=namespace,
        n_workers=n_workers,
    )

    # once the `Client` is initialized all dask calls are actually implicitly performed against it
    client = Client(cluster)
    array = da.ones((1000, 1000, 1000))
    print("Array mean = {array_mean}, expected = 1.0".format(array_mean=array.mean().compute()))
    client.close()


with Workflow(generate_name="dask-", entrypoint="s") as w:
    with Steps(name="s"):
        dask_computation()

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
apiVersion: argoproj.io/v1alpha1
kind: Workflow
metadata:
  generateName: dask-
spec:
  entrypoint: s
  templates:
  - name: s
    steps:
    - - name: dask-computation
        template: dask-computation
  - inputs:
      parameters:
      - default: default
        name: namespace
      - default: '1'
        name: n_workers
    name: dask-computation
    script:
      command:
      - python
      image: ghcr.io/dask/dask:latest
      source: 'import os

        import sys

        sys.path.append(os.getcwd())

        import json

        try: n_workers = json.loads(r''''''{{inputs.parameters.n_workers}}'''''')

        except: n_workers = r''''''{{inputs.parameters.n_workers}}''''''

        try: namespace = json.loads(r''''''{{inputs.parameters.namespace}}'''''')

        except: namespace = r''''''{{inputs.parameters.namespace}}''''''


        import subprocess

        subprocess.run([''pip'', ''install'', ''dask-kubernetes'', ''dask[distributed]''],
        stdout=subprocess.PIPE, universal_newlines=True)

        import dask.array as da

        from dask.distributed import Client

        from dask_kubernetes.classic import KubeCluster, make_pod_spec

        cluster = KubeCluster(pod_template=make_pod_spec(image=''ghcr.io/dask/dask:latest'',
        memory_limit=''4G'', memory_request=''2G'', cpu_limit=1, cpu_request=1), namespace=namespace,
        n_workers=n_workers)

        client = Client(cluster)

        array = da.ones((1000, 1000, 1000))

        print(''Array mean = {array_mean}, expected = 1.0''.format(array_mean=array.mean().compute()))

        client.close()'

Comments