Thanks a lot to Jacob for the original answer – it helped me a lot to get GPU support working with Dask Gateway quickly.

After following this setup, I wanted to go one step further and make it possible for users to choose between CPU and GPU modes directly from the Dask Gateway client (via the UI or programmatically), and also specify how many GPUs each worker should use.

I didn’t find an example of this in the documentation or online, so I ended up solving it like this by extending the cluster_options section in the Gateway configuration:

gateway:
  extraConfig:
    clusteroptions: |
        from dask_gateway_server.options import Options, Integer, Float, String, Select

        def option_handler(options):
            config = {
                "worker_cores": options.worker_cores,
                "worker_memory": f"{options.worker_memory}G",
                "image": options.image,
            }

            if options.worker_type == "gpu":
                config["worker_cmd"] = ["dask-cuda-worker"]
                config["worker_extra_container_config"] = {
                    "resources": {
                        "limits": {
                            "nvidia.com/gpu": options.gpu_count
                        }
                    }
                }
                config["scheduler_extra_container_config"] = {
                    "resources": {
                        "limits": {
                            "nvidia.com/gpu": 1
                        }
                    }
                }
            else:
                config["worker_cmd"] = ["dask-worker"]

            return config

        c.Backend.cluster_options = Options(
            Select("worker_type", ["cpu", "gpu"], default="cpu", label="Worker Type"),

            Integer("worker_cores", 2, min=1, max=24, label="Worker Cores (CPU only)"),
            Float("worker_memory", 4, min=1, max=64, label="Worker Memory (GiB)"),

            Integer("gpu_count", 1, min=1, max=4, label="GPUs per worker (GPU only)"),

            String("image", default="daskgateway/dask-gateway:latest", label="Image"),
            handler=option_handler,
        )