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`) dask/distributed#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: #1230

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,15 +32,15 @@ def nvtx_annotate(message=None, color="blue", domain=None):
         yield
 
 
-class CPUAffinity:
+class CPUAffinity(WorkerPlugin):
     def __init__(self, cores):
         self.cores = cores
 
     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 = []