Repository: dask/dask
Pull request #4792 (2019-05-10 20:13:22)

Add document for how Dask works with GPUs
This could use review both from people who know about GPUs (@pentschev @quasiben) and from people who are unassociated with NVIDIA (to keep things honest) (@jcrist @martindurant)

I'm getting more questions about how Dask works with GPUs these days, and I'd like to have a place where I can start to consolidate information.  A page like this seems ideal to me, but I'm happy to move this elsewhere if desired.  I'm slightly concerned about the fast-moving nature of the work here.  I intentionally didn't point to more experimental projects like dask-cuda or dask-cuml or UCX.

- [ ] Tests added / passed
- [ ] Passes `flake8 dask`



Patch:
diff --git a/dask/utils.py b/dask/utils.py
index da1a3ffdb..6cacae859 100644
--- a/dask/utils.py
+++ b/dask/utils.py
@@ -1,5 +1,6 @@
 from __future__ import absolute_import, division, print_function
 
+from datetime import timedelta
 import functools
 import inspect
 import os
@@ -10,7 +11,7 @@ import re
 from errno import ENOENT
 from contextlib import contextmanager
 from importlib import import_module
-from numbers import Integral
+from numbers import Integral, Number
 from threading import Lock
 import uuid
 from weakref import WeakValueDictionary
@@ -1091,6 +1092,119 @@ byte_sizes.update({k[0]: v for k, v in byte_sizes.items() if k and 'i' not in k}
 byte_sizes.update({k[:-1]: v for k, v in byte_sizes.items() if k and 'i' in k})
 
 
+def format_time(n):
+    """ format integers as time
+
+    >>> format_time(1)
+    '1.00 s'
+    >>> format_time(0.001234)
+    '1.23 ms'
+    >>> format_time(0.00012345)
+    '123.45 us'
+    >>> format_time(123.456)
+    '123.46 s'
+    """
+    if n >= 1:
+        return "%.2f s" % n
+    if n >= 1e-3:
+        return "%.2f ms" % (n * 1e3)
+    return "%.2f us" % (n * 1e6)
+
+
+def format_bytes(n):
+    """ Format bytes as text
+
+    >>> format_bytes(1)
+    '1 B'
+    >>> format_bytes(1234)
+    '1.23 kB'
+    >>> format_bytes(12345678)
+    '12.35 MB'
+    >>> format_bytes(1234567890)
+    '1.23 GB'
+    >>> format_bytes(1234567890000)
+    '1.23 TB'
+    >>> format_bytes(1234567890000000)
+    '1.23 PB'
+    """
+    if n > 1e15:
+        return "%0.2f PB" % (n / 1e15)
+    if n > 1e12:
+        return "%0.2f TB" % (n / 1e12)
+    if n > 1e9:
+        return "%0.2f GB" % (n / 1e9)
+    if n > 1e6:
+        return "%0.2f MB" % (n / 1e6)
+    if n > 1e3:
+        return "%0.2f kB" % (n / 1000)
+    return "%d B" % n
+
+
+timedelta_sizes = {
+    "s": 1,
+    "ms": 1e-3,
+    "us": 1e-6,
+    "ns": 1e-9,
+    "m": 60,
+    "h": 3600,
+    "d": 3600 * 24,
+}
+
+tds2 = {
+    "second": 1,
+    "minute": 60,
+    "hour": 60 * 60,
+    "day": 60 * 60 * 24,
+    "millisecond": 1e-3,
+    "microsecond": 1e-6,
+    "nanosecond": 1e-9,
+}
+tds2.update({k + "s": v for k, v in tds2.items()})
+timedelta_sizes.update(tds2)
+timedelta_sizes.update({k.upper(): v for k, v in timedelta_sizes.items()})
+
+
+def parse_timedelta(s, default="seconds"):
+    """ Parse timedelta string to number of seconds
+
+    Examples
+    --------
+    >>> parse_timedelta('3s')
+    3
+    >>> parse_timedelta('3.5 seconds')
+    3.5
+    >>> parse_timedelta('300ms')
+    0.3
+    >>> parse_timedelta(timedelta(seconds=3))  # also supports timedeltas
+    3
+    """
+    if isinstance(s, timedelta):
+        s = s.total_seconds()
+        return int(s) if int(s) == s else s
+    if isinstance(s, Number):
+        s = str(s)
+    s = s.replace(" ", "")
+    if not s[0].isdigit():
+        s = "1" + s
+
+    for i in range(len(s) - 1, -1, -1):
+        if not s[i].isalpha():
+            break
+    index = i + 1
+
+    prefix = s[:index]
+    suffix = s[index:] or default
+
+    n = float(prefix)
+
+    multiplier = timedelta_sizes[suffix.lower()]
+
+    result = n * multiplier
+    if int(result) == result:
+        result = int(result)
+    return result
+
+
 def has_keyword(func, keyword):
     try:
         if PY3:
diff --git a/docs/requirements-docs.txt b/docs/requirements-docs.txt
index 7d43c8c37..3070f13b9 100644
--- a/docs/requirements-docs.txt
+++ b/docs/requirements-docs.txt
@@ -1,6 +1,7 @@
 numpydoc==0.8
 sphinx
 dask_sphinx_theme>=1.1.0
+sphinx-click
 toolz
 cloudpickle
 pandas>=0.19.0
diff --git a/docs/source/conf.py b/docs/source/conf.py
index 3508e784b..7ef37d354 100644
--- a/docs/source/conf.py
+++ b/docs/source/conf.py
@@ -27,7 +27,8 @@ from shutil import copyfile
 # Add any Sphinx extension module names here, as strings. They can be extensions
 # coming with Sphinx (named 'sphinx.ext.*') or your custom ones.
 extensions = ['sphinx.ext.autodoc', 'sphinx.ext.mathjax', 'sphinx.ext.intersphinx',
-              'sphinx.ext.autosummary', 'sphinx.ext.extlinks', 'numpydoc']
+              'sphinx.ext.autosummary', 'sphinx.ext.extlinks', 'numpydoc',
+              'sphinx_click.ext']
 
 numpydoc_show_class_members = False
 
diff --git a/docs/source/gpu.rst b/docs/source/gpu.rst
new file mode 100644
index 000000000..18d69161f
--- /dev/null
+++ b/docs/source/gpu.rst
@@ -0,0 +1,166 @@
+GPUs
+====
+
+Dask works with GPUs in a few ways.
+
+
+Custom Computations
+-------------------
+
+Many people use Dask alongside GPU-accelerated libraries like PyTorch and
+TensorFlow to manage workloads across several machines.  They typically use
+Dask's custom APIs, notably :doc:`Delayed <delayed>` and :doc:`Futures
+<futures>`.
+
+Dask doesn't need to know that these functions use GPUs.  It just runs Python
+functions.  Whether or not those Python functions use a GPU is orthogonal to
+Dask.  It will work regardless.
+
+As a worked example, you may want to view this talk:
+
+.. raw:: html
+
+    <iframe width="560"
+            height="315"
+            src="https://developer.download.nvidia.com/video/gputechconf/gtc/2019/video/S9198/s9198-dask-and-v100s-for-fast-distributed-batch-scoring-of-computer-vision-workloads.mp4"
+            frameborder="0"
+            allow="autoplay; encrypted-media"
+            allowfullscreen>
+    </iframe>
+
+
+
+High Level Collections
+----------------------
+
+Dask can also help to scale out large array and dataframe computations by
+combining the Dask Array and DataFrame collections with a GPU-accelerated
+array or dataframe library.
+
+Recall that :doc:`Dask Array <array>` creates a large array out of many NumPy
+arrays and :doc:`Dask DataFrame <dataframe>` creates a large dataframe out of
+many Pandas dataframes.  We can use these same systems with GPUs if we swap out
+the NumPy/Pandas components with GPU-accelerated versions of those same
+libraries, as long as the GPU accelerated version looks enough like
+NumPy/Pandas in order to interoperate with Dask.
+
+Fortunately, libraries that mimic NumPy, Pandas, and Scikit-Learn on the GPU do
+exist.
+
+
+DataFrames
+~~~~~~~~~~
+
+The `RAPIDS <https://rapids.ai>`_ libraries provide a GPU accelerated
+Pandas-like library,
+`cuDF <https://rapidsai.github.io/projects/cudf/en/latest/>`_,
+which interoperates well and is tested against Dask DataFrame.
+
+If you have cudf installed then you should be able to convert a Pandas-backed
+Dask DataFrame to a cuDF-backed Dask DataFrame as follows:
+
+.. code-block:: python
+
+   import cudf
+
+   df = df.map_partitions(cudf.from_pandas)  # convert pandas partitions into cudf partitions
+
+However, cuDF does not support the entire Pandas interface, and so a variety of
+Dask DataFrame operations will not function properly. Check the 
+`cudf API Reference <https://rapidsai.github.io/projects/cudf/en/latest/api.html>`_
+for currently supported interface.
+
+
+Arrays
+~~~~~~
+
+.. note:: Dask's integration with CuPy relies on features recently added to
+   NumPy and CuPy, particularly in version ``numpy>=1.17`` and ``cupy>=6``
+
+`Chainer's CuPy <https://cupy.chainer.org/>`_ library provides a GPU
+accelerated NumPy-like library that interoperates nicely with Dask Array.
+
+If you have CuPy installed then you should be able to convert a NumPy-backed
+Dask Array into a CuPy backed Dask Array as follows:
+
+.. code-block:: python
+
+   import cupy
+
+   x = x.map_blocks(cupy.asarray)
+
+CuPy is fairly mature and adheres closely to the NumPy API.  However, small
+differences do exist and these can cause Dask Array operations to function
+improperly. Check the
+`CuPy Reference Manual <https://docs-cupy.chainer.org/en/stable/reference/index.html>`_
+for API compatibility.
+
+
+Scikit-Learn
+~~~~~~~~~~~~
+
+There are a variety of GPU accelerated machine learning libraries that follow
+the Scikit-Learn Estimator API of fit, transform, and predict.  These can
+generally be used within `Dask-ML's <https://ml.dask.org>`_ meta estimators,
+such as `hyper parameter optimization <https://ml.dask.org/hyper-parameter-search.html>`_.
+
+Some of these include:
+
+-  `Skorch <https://skorch.readthedocs.io/>`_
+-  `cuML <https://rapidsai.github.io/projects/cuml/en/latest/>`_
+-  `LightGBM <https://github.com/Microsoft/LightGBM>`_
+-  `XGBoost <https://xgboost.readthedocs.io/en/latest/>`_
+-  `Thunder SVM <https://github.com/Xtra-Computing/thundersvm>`_
+-  `Thunder GBM <https://github.com/Xtra-Computing/thundergbm>`_
+
+
+Setup
+-----
+
+From the examples above we can see that the user experience of using Dask with
+GPU-backed libraries isn't very different from using it with CPU-backed
+libraries.  However, there are some changes you might consider making when
+setting up your cluster.
+
+Restricting Work
+~~~~~~~~~~~~~~~~
+
+By default Dask allows as many tasks as you have CPU cores to run concurrently.
+However if your tasks primarily use a GPU then you probably want far fewer
+tasks running at once.  There are a few ways to limit parallelism here:
+
+-   Limit the number of threads explicitly on your workers using the
+    ``--nthreads`` keyword in the CLI or the ``ncores=`` keyword the
+    Cluster constructor.
+-   Use `worker resources <https://distributed.dask.org/en/latest/resources.html>`_ and tag certain
+    tasks as GPU tasks so that the scheduler will limit them, while leaving the
+    rest of your CPU cores for other work
+
+Specifying GPUs per Machine
+~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+Some configurations may have many GPU devices per node.  Dask is often used to
+balance and coordinate work between these devices.
+
+In these situations it is common to start one Dask worker per device, and use
+the CUDA environment varible ``CUDA_VISIBLE_DEVICES`` to pin each worker to
+prefer one device.
+
+.. code-block:: bash
+
+   # If we have four GPUs on one machine
+   CUDA_VISIBLE_DEVICES=0 dask-worker ...
+   CUDA_VISIBLE_DEVICES=1 dask-worker ...
+   CUDA_VISIBLE_DEVICES=2 dask-worker ...
+   CUDA_VISIBLE_DEVICES=3 dask-worker ...
+
+The `Dask CUDA <https://github.com/rapidsai/dask-cuda>`_ project contains some
+convenience CLI and Python utilities to automate this process.
+
+Work in Progress
+----------------
+
+GPU computing is a quickly moving field today and as a result the information
+in this page is likely to go out of date quickly.  We encourage interested
+readers to check out `Dask's Blog <https://blog.dask.org>`_ which has more
+timely updates on ongoing work.
diff --git a/docs/source/index.rst b/docs/source/index.rst
index 4d1bf335c..6016c6f26 100644
--- a/docs/source/index.rst
+++ b/docs/source/index.rst
@@ -252,6 +252,7 @@ often a better choice. If you are a *core developer*, then you should start here
 * :doc:`caching`
 * :doc:`bytes`
 * :doc:`remote-data-services`
+* :doc:`gpu`
 * :doc:`cite`
 * :doc:`funding`
 * :doc:`logos`
@@ -270,6 +271,7 @@ often a better choice. If you are a *core developer*, then you should start here
    caching.rst
    bytes.rst
    remote-data-services.rst
+   gpu.rst
    cite.rst
    funding.rst
    logos.rst
diff --git a/docs/source/setup/cli.rst b/docs/source/setup/cli.rst
index 50f44d70f..8e5892676 100644
--- a/docs/source/setup/cli.rst
+++ b/docs/source/setup/cli.rst
@@ -88,87 +88,25 @@ API
 
 .. warning::
 
-   These may be out-dated.  We recommend referring to the ``--help`` text of your
-   installed version.
-
-dask-scheduler
-~~~~~~~~~~~~~~
-
-.. code-block:: bash
-
-   $ dask-scheduler --help
-   Usage: dask-scheduler [OPTIONS]
-
-   Options:
-     --host TEXT             URI, IP or hostname of this server
-     --port INTEGER          Serving port
-     --interface TEXT        Preferred network interface like 'eth0' or 'ib0'
-     --tls-ca-file PATH      CA cert(s) file for TLS (in PEM format)
-     --tls-cert PATH         certificate file for TLS (in PEM format)
-     --tls-key PATH          private key file for TLS (in PEM format)
-     --bokeh-port INTEGER    Bokeh port for visual diagnostics
-     --bokeh / --no-bokeh    Launch Bokeh Web UI  [default: True]
-     --show / --no-show      Show web UI
-     --bokeh-whitelist TEXT  IP addresses to whitelist for bokeh
-     --bokeh-prefix TEXT     Prefix for the bokeh app
-     --use-xheaders BOOLEAN  User xheaders in bokeh app for ssl termination in
-                             header  [default: False]
-     --pid-file TEXT         File to write the process PID
-     --scheduler-file TEXT   File to write connection information. This may be a
-                             good way to share connection information if your
-                             cluster is on a shared network file system
-     --local-directory TEXT  Directory to place scheduler files
-     --preload TEXT          Module that should be loaded by each worker process
-                             like "foo.bar" or "/path/to/foo.py"
-     --help                  Show this message and exit
-
-
-dask-worker
-~~~~~~~~~~~
-
-.. code-block:: bash
-
-   $ dask-worker --help
-   Usage: dask-worker [OPTIONS] [SCHEDULER]
-
-   Options:
-     --tls-ca-file PATH            CA cert(s) file for TLS (in PEM format)
-     --tls-cert PATH               certificate file for TLS (in PEM format)
-     --tls-key PATH                private key file for TLS (in PEM format)
-     --worker-port INTEGER         Serving computation port, defaults to random
-     --nanny-port INTEGER          Serving nanny port, defaults to random
-     --bokeh-port INTEGER          Bokeh port, defaults to 8789
-     --bokeh / --no-bokeh          Launch Bokeh Web UI  [default: True]
-     --listen-address TEXT         The address to which the worker binds.
-                                   Example: tcp://0.0.0.0:9000
-     --contact-address TEXT        The address the worker advertises to the
-                                   scheduler for communication with it and other
-                                   workers. Example: tcp://127.0.0.1:9000
-     --host TEXT                   Serving host. Should be an ip address that is
-                                   visible to the scheduler and other workers.
-                                   See --listen-address and --contact-address if
-                                   you need different listen and contact
-                                   addresses. See --interface
-     --interface TEXT              Network interface like 'eth0' or 'ib0'
-     --nthreads INTEGER            Number of threads per process
-     --nprocs INTEGER              Number of worker processes.  Defaults to one
-     --name TEXT                   A unique name for this worker like 'worker-1'
-     --memory-limit TEXT           Bytes of memory that the worker can use. This
-                                   can be an integer (bytes), float (fraction of
-                                   total system memory), string (like 5GB or
-                                   5000M), 'auto', or zero for no memory
-                                   management
-     --reconnect / --no-reconnect  Reconnect to scheduler if disconnected
-     --nanny / --no-nanny          Start workers in nanny process for management
-     --pid-file TEXT               File to write the process PID
-     --local-directory TEXT        Directory to place worker files
-     --resources TEXT              Resources for task constraints like "GPU=2
-                                   MEM=10e9"
-     --scheduler-file TEXT         Filename to JSON encoded scheduler information. 
-                                   Use with dask-scheduler --scheduler-file
-     --death-timeout FLOAT         Seconds to wait for a scheduler before closing
-     --bokeh-prefix TEXT           Prefix for the bokeh app
-     --preload TEXT                Module that should be loaded by each worker
-                                   process like "foo.bar" or "/path/to/foo.py"
-     --help                        Show this message and exit
+   The command line documentation here may differ depending on your installed
+   version. We recommend referring to the output of ``<command> --help``.
 
+.. click:: distributed.cli.dask_scheduler:main
+   :prog: dask-scheduler
+   :show-nested:
+
+.. click:: distributed.cli.dask_worker:main
+   :prog: dask-worker
+   :show-nested:
+
+.. click:: distributed.cli.dask_ssh:main
+   :prog: dask-ssh
+   :show-nested:
+
+.. click:: distributed.cli.dask_submit:main
+   :prog: dask-submit
+   :show-nested:
+
+.. click:: distributed.cli.dask_remote:main
+   :prog: dask-remote
+   :show-nested:


Test patch:
diff --git a/dask/tests/test_utils.py b/dask/tests/test_utils.py
index cfa94ace7..eb747781d 100644
--- a/dask/tests/test_utils.py
+++ b/dask/tests/test_utils.py
@@ -1,3 +1,4 @@
+import datetime
 import functools
 import operator
 import pickle
@@ -10,7 +11,8 @@ from dask.utils import (takes_multiple_arguments, Dispatch, random_state_data,
                         memory_repr, methodcaller, M, skip_doctest,
                         SerializableLock, funcname, ndeepmap, ensure_dict,
                         extra_titles, asciitable, itemgetter, partial_by_order,
-                        has_keyword, derived_from)
+                        has_keyword, derived_from, parse_timedelta,
+                        parse_bytes)
 from dask.utils_test import inc
 from dask.highlevelgraph import HighLevelGraph
 
@@ -427,3 +429,42 @@ def test_derived_from_dask_dataframe():
     [axis_arg] = [line for line in dd.DataFrame.dropna.__doc__.split('\n') if 'axis :' in line]
     assert "not supported" in axis_arg.lower()
     assert "dask" in axis_arg.lower()
+
+
+def test_parse_bytes():
+    assert parse_bytes("100") == 100
+    assert parse_bytes("100 MB") == 100000000
+    assert parse_bytes("100M") == 100000000
+    assert parse_bytes("5kB") == 5000
+    assert parse_bytes("5.4 kB") == 5400
+    assert parse_bytes("1kiB") == 1024
+    assert parse_bytes("1Mi") == 2 ** 20
+    assert parse_bytes("1e6") == 1000000
+    assert parse_bytes("1e6 kB") == 1000000000
+    assert parse_bytes("MB") == 1000000
+
+
+def test_parse_timedelta():
+    for text, value in [
+        ("1s", 1),
+        ("100ms", 0.1),
+        ("5S", 5),
+        ("5.5s", 5.5),
+        ("5.5 s", 5.5),
+        ("1 second", 1),
+        ("3.3 seconds", 3.3),
+        ("3.3 milliseconds", 0.0033),
+        ("3500 us", 0.0035),
+        ("1 ns", 1e-9),
+        ("2m", 120),
+        ("2 minutes", 120),
+        (datetime.timedelta(seconds=2), 2),
+        (datetime.timedelta(milliseconds=100), 0.1),
+    ]:
+        result = parse_timedelta(text)
+        assert abs(result - value) < 1e-14
+
+    assert parse_timedelta("1ms", default="seconds") == 0.001
+    assert parse_timedelta("1", default="seconds") == 1
+    assert parse_timedelta("1", default="ms") == 0.001
+    assert parse_timedelta(1, default="ms") == 0.001