Added worker resources from config (#4456)

Author: TomAugspurger

distributed/distributed-schema.yaml

@@ -305,6 +305,14 @@ properties:
               Whether or not to run consistency checks during execution.
               This is typically only used for debugging.
 
+          resources:
+            type: object
+            description: |
+              A dictionary specifying resources for workers.
+
+              See https://distributed.dask.org/en/latest/resources.html for more information.
+            properties: {}
+
           lifetime:
             type: object
             description: |

distributed/distributed.yaml

@@ -70,6 +70,7 @@ distributed:
     preload-argv: []        # See https://docs.dask.org/en/latest/setup/custom-startup.html
     daemon: True
     validate: False         # Check worker state at every step for debugging
+    resources: {}           # Key: value pairs specifying worker resources.
     lifetime:
       duration: null        # Time after which to gracefully shutdown the worker
       stagger: 0 seconds    # Random amount by which to stagger lifetimes

distributed/tests/test_resources.py

@@ -392,3 +392,21 @@ def g(dask_worker):
     logs = client.run(g)
     assert logs[a["address"]]
     assert not logs[b["address"]]
+
+
+@gen_cluster(config={"distributed.worker.resources.my_resources": 1}, client=True)
+async def test_resources_from_config(c, s, a, b):
+    info = c.scheduler_info()
+    for worker in [a, b]:
+        assert info["workers"][worker.address]["resources"] == {"my_resources": 1}
+
+
+@gen_cluster(
+    worker_kwargs=dict(resources={"my_resources": 10}),
+    config={"distributed.worker.resources.my_resources": 1},
+    client=True,
+)
+async def test_resources_from_python_override_config(c, s, a, b):
+    info = c.scheduler_info()
+    for worker in [a, b]:
+        assert info["workers"][worker.address]["resources"] == {"my_resources": 10}

distributed/worker.py

@@ -519,6 +519,9 @@ def __init__(
             nthreads = ncores
 
         self.nthreads = nthreads or CPU_COUNT
+        if resources is None:
+            resources = dask.config.get("distributed.worker.resources", None)
+
         self.total_resources = resources or {}
         self.available_resources = (resources or {}).copy()
         self.death_timeout = parse_timedelta(death_timeout)

docs/source/resources.rst

@@ -47,6 +47,49 @@ When we submit tasks to the cluster we specify constraints per task
    processed = [client.submit(process, d, resources={'GPU': 1}) for d in data]
    final = client.submit(aggregate, processed, resources={'MEMORY': 70e9})
 
+Specifying Resources
+--------------------
+
+Resources can be specifed in several ways. The easiest option will depend on exactly
+how your cluster is being created.
+
+**From the command line**
+
+Resources can be provided when starting the worker process, as shown above:
+
+.. code-block:: console
+
+   dask-worker scheduler:8786 --resources "GPU=2"
+
+The keys are used as the resource name and the values are parsed into a numeric value.
+
+**From Dask's configuration system**
+
+Alternatively, resources can be specified using Dask's
+`configuration system <https://docs.dask.org/en/latest/configuration.html>`_.
+
+.. code-block:: python
+
+   from distributed import LocalCluster
+
+   with dask.config.set({"distributed.worker.resources.GPU": 2}):
+       cluster = LocalCluster()
+
+The configuration will need to be set in the process that's spawning the actual worker.
+This might be easiest to achieve by specifying resources as an environment variable
+(shown in the next section).
+
+**From environment variables**
+
+Like any other Dask config value, resources can be specified as environment variables
+before starting the process. Using Bash syntax
+
+.. code-block:: console
+
+   $ DASK_DISTRIBUTED__WORKER__RESOURCES__GPU=2 dask-worker
+   ...
+
+This might be the easiest solution if you aren't able to pass options to the :class:`distributed.Worker` class.
 
 Resources are applied separately to each worker process
 -------------------------------------------------------