From 38eec6c2a0bbb6132c550bdf2fed9dbbe16f2f18 Mon Sep 17 00:00:00 2001 From: James Bourbeau Date: Wed, 4 Oct 2023 15:20:16 -0500 Subject: [PATCH] Update plugins to inherit from ``WorkerPlugin`` (#1230) Upstream in `distributed` we're considering enforcing plugins to inherit from their respective base class (e.g. `WorkerPlugin`, `SchedulerPlugin`, `NannyPlugin`) https://github.com/dask/distributed/pull/8149. This PR updates plugins here to inherhit from `WorkerPlugin`. This makes things a little more future-proof and is probably a good thing to do anyways. Authors: - James Bourbeau (https://github.com/jrbourbeau) - Peter Andreas Entschev (https://github.com/pentschev) Approvers: - Peter Andreas Entschev (https://github.com/pentschev) URL: https://github.com/rapidsai/dask-cuda/pull/1230 --- dask_cuda/utils.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/dask_cuda/utils.py b/dask_cuda/utils.py index a155dc593..1e244bb31 100644 --- a/dask_cuda/utils.py +++ b/dask_cuda/utils.py @@ -18,7 +18,7 @@ import distributed # noqa: required for dask.config.get("distributed.comm.ucx") from dask.config import canonical_name from dask.utils import format_bytes, parse_bytes -from distributed import Worker, wait +from distributed import Worker, WorkerPlugin, wait from distributed.comm import parse_address try: @@ -32,7 +32,7 @@ def nvtx_annotate(message=None, color="blue", domain=None): yield -class CPUAffinity: +class CPUAffinity(WorkerPlugin): def __init__(self, cores): self.cores = cores @@ -40,7 +40,7 @@ def setup(self, worker=None): os.sched_setaffinity(0, self.cores) -class RMMSetup: +class RMMSetup(WorkerPlugin): def __init__( self, initial_pool_size, @@ -135,7 +135,7 @@ def setup(self, worker=None): rmm.mr.set_current_device_resource(rmm.mr.TrackingResourceAdaptor(mr)) -class PreImport: +class PreImport(WorkerPlugin): def __init__(self, libraries): if libraries is None: libraries = []