From 61557324712394becc7101ef68d893009e5797d8 Mon Sep 17 00:00:00 2001 From: Peter Andreas Entschev Date: Mon, 23 Oct 2023 07:27:17 -0700 Subject: [PATCH] Forbid specifying `log_spilling=True` with custom `worker_class` --- dask_cuda/local_cuda_cluster.py | 13 ++++++++----- 1 file changed, 8 insertions(+), 5 deletions(-) diff --git a/dask_cuda/local_cuda_cluster.py b/dask_cuda/local_cuda_cluster.py index e84dce0c7..57fe85439 100644 --- a/dask_cuda/local_cuda_cluster.py +++ b/dask_cuda/local_cuda_cluster.py @@ -2,7 +2,6 @@ import logging import os import warnings -from functools import partial import dask from distributed import LocalCluster, Nanny, Worker @@ -332,10 +331,14 @@ def __init__( enable_rdmacm=enable_rdmacm, ) - worker_class = partial( - LoggedNanny if log_spilling is True else Nanny, - worker_class=worker_class, - ) + if worker_class is not None and log_spilling is True: + raise ValueError( + "Cannot enable `log_spilling` when `worker_class` is specified. If " + "logging is needed, ensure `worker_class` is a subclass of " + "`distributed.local_cuda_cluster.LoggedNanny` or a subclass of " + "`distributed.local_cuda_cluster.LoggedWorker`, and specify " + "`log_spilling=False`." + ) self.pre_import = pre_import