Add optional user to cluster options handler (dask#329)

The `handler` function provided to the `Options` object for
`cluster_options` can now take an optional `user` argument. This can be
used to provide user-specific configuration, without requiring stashing
the `user` in a closure from `cluster_options`.

Author: jcrist

dask-gateway-server/dask_gateway_server/backends/base.py

@@ -96,7 +96,7 @@ async def process_cluster_options(self, user, request):
         try:
             cluster_options = await self.get_cluster_options(user)
             requested_options = cluster_options.parse_options(request)
-            overrides = cluster_options.get_configuration(requested_options)
+            overrides = cluster_options.get_configuration(requested_options, user)
             config = self.cluster_config_class(parent=self, **overrides)
         except asyncio.CancelledError:
             raise

dask-gateway-server/dask_gateway_server/options.py

@@ -1,4 +1,5 @@
 import copy
+import inspect
 import textwrap
 from collections import OrderedDict
 from collections.abc import Sequence
@@ -17,10 +18,12 @@ class Options(object):
     *fields : Field
         Zero or more configurable fields.
     handler : callable, optional
-        A callable with the signature ``handler(options)``, where ``options``
-        is the validated dict of user options. Should return a dict of
-        configuration overrides to forward to the cluster manager. If not
-        provided, the default will return the options unchanged.
+        A callable with the signature ``handler(options)`` or
+        ``handler(options, user)``, where ``options`` is the validated dict of
+        user options, and ``user`` is a ``User`` model for that user. Should
+        return a dict of configuration overrides to forward to the cluster
+        manager. If not provided, the default will return the options
+        unchanged.
 
     Example
     -------
@@ -59,6 +62,24 @@ def __init__(self, *fields, handler=None):
         self.fields = fields
         self.handler = handler
 
+    @property
+    def handler(self):
+        return self._handler
+
+    @handler.setter
+    def handler(self, handler):
+        if handler is None:
+            self._handler = None
+        else:
+            sig = inspect.signature(handler)
+            if len(sig.parameters) == 1 and not any(
+                a.kind == inspect.Parameter.VAR_POSITIONAL
+                for a in sig.parameters.values()
+            ):
+                self._handler = lambda options, user: handler(options)
+            else:
+                self._handler = handler
+
     def get_specification(self):
         return [f.json_spec() for f in self.fields]
 
@@ -81,11 +102,11 @@ def transform_options(self, options):
             for f in self.fields
         }
 
-    def get_configuration(self, options):
+    def get_configuration(self, options, user):
         options = self.transform_options(options)
         if self.handler is None:
             return options
-        return self.handler(FrozenAttrDict(options))
+        return self.handler(FrozenAttrDict(options), user)
 
 
 _field_doc_template = """\

docs/source/cluster-options.rst

@@ -69,7 +69,10 @@ A :class:`dask_gateway_server.options.Options` object takes two arguments:
 
 - ``handler``: An optional handler function for translating the values set by
   those options into configuration values to set on the corresponding
-  :ref:`ClusterConfig <cluster-config>`.
+  :ref:`ClusterConfig <cluster-config>`. Should have the signature
+  ``handler(options)`` or ``handler(options, user)``, where ``options`` is the
+  validated dict of user options, and ``user`` is a ``User`` model for that
+  user.
 
 ``Field`` objects provide typed specifications for a user facing option. There
 are several different ``Field`` classes available, each representing a
@@ -173,7 +176,7 @@ Different Options per User Group
 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
 
 Cluster options may be configured to differ based on the user by providing a
-function for :data:`c.Backend.cluster_options`. This function recieves a
+function for :data:`c.Backend.cluster_options`. This function receives a
 :class:`dask_gateway_server.models.User` object and should return a
 :class:`dask_gateway_server.options.Options` object. It may optionally be an
 ``async`` function.
@@ -196,18 +199,46 @@ member of the "power-users" group.
     def generate_options(user):
         if "power-users" in user.groups:
             options = Options(
-                Integer("worker_cores", default=1, min=1, max=4, label="Worker Cores"),
-                Float("worker_memory", default=1, min=1, max=8, label="Worker Memory (GiB)"),
+                Integer("worker_cores", default=1, min=1, max=8, label="Worker Cores"),
+                Float("worker_memory", default=1, min=1, max=16, label="Worker Memory (GiB)"),
                 handler=options_handler,
-                )
+            )
         else:
             options = Options(
-                Integer("worker_cores", default=1, min=1, max=8, label="Worker Cores"),
-                Float("worker_memory", default=1, min=1, max=16, label="Worker Memory (GiB)"),
+                Integer("worker_cores", default=1, min=1, max=4, label="Worker Cores"),
+                Float("worker_memory", default=1, min=1, max=8, label="Worker Memory (GiB)"),
                 handler=options_handler,
-                )
+            )
 
     c.Backend.cluster_options = generate_options
 
+User-specific Configuration
+^^^^^^^^^^^^^^^^^^^^^^^^^^^
+
+Since the ``handler`` function can optionally take in the ``User`` object, you
+can use this to add user-specific configuration. Note that you don't have to
+expose any configuration options to make use of this, the options handler is
+called regardless.
+
+Here we configure the worker cores and memory based on the user's groups:
+
+.. code-block:: python
+
+    from dask_gateway_server.options import Options
+
+    def options_handler(options, user):
+        if "power-users" in user.groups:
+            return {
+                "worker_cores": 8,
+                "worker_memory": "16 G"
+            }
+        else:
+            return {
+                "worker_cores": 4,
+                "worker_memory": "8 G"
+            }
+
+    c.Backend.cluster_options = Options(handler=options_handler)
+
 
 .. _ipywidgets: https://ipywidgets.readthedocs.io/en/latest/

tests/test_options.py

@@ -4,6 +4,7 @@
 import dask_gateway.options as client_options
 import dask_gateway_server.options as server_options
 from dask_gateway_server.options import FrozenAttrDict
+from dask_gateway_server.models import User
 
 
 def test_string():
@@ -381,19 +382,29 @@ def test_server_options_parse_options(server_opts):
 def test_server_options_get_configuration(server_opts):
     options = {"integer_field": 2}
     sol = {"integer_field2": 2, "float_field": 2.5, "select_field": 5}
+    user = User("alice")
 
     # Default handler
-    config = server_opts.get_configuration(options)
+    config = server_opts.get_configuration(options, user)
     assert config == sol
 
-    # Custom handler
+    # Custom handler, no user
     def handler(options):
         assert isinstance(options, FrozenAttrDict)
         assert dict(options) == sol
         return {"a": 1}
 
     server_opts.handler = handler
-    assert server_opts.get_configuration(options) == {"a": 1}
+    assert server_opts.get_configuration(options, user) == {"a": 1}
+
+    # Custom handler, with user
+    def handler(options, user):
+        assert isinstance(options, FrozenAttrDict)
+        assert dict(options) == sol
+        return {"a": 1, "username": user.name}
+
+    server_opts.handler = handler
+    assert server_opts.get_configuration(options, user) == {"a": 1, "username": "alice"}
 
 
 @pytest.fixture