From 87f1435131535e560d03c20cfb932be85e4a285d Mon Sep 17 00:00:00 2001 From: Maxim Lippeveld Date: Wed, 22 May 2024 09:29:49 +0000 Subject: [PATCH] Assign cellpose resource to only one localcluster worker --- src/scip/utils/util.py | 18 ++++++++++++------ 1 file changed, 12 insertions(+), 6 deletions(-) diff --git a/src/scip/utils/util.py b/src/scip/utils/util.py index b6800ba..3ecdcac 100644 --- a/src/scip/utils/util.py +++ b/src/scip/utils/util.py @@ -69,12 +69,18 @@ def __enter__(self): self.client = Client(n_workers=1, threads_per_worker=1) elif self.mode == "local": from dask.distributed import LocalCluster - with dask.config.set({"distributed.worker.resources.cellpose": 1}): - self.port = self.port or "8787" - self.cluster = LocalCluster( - n_workers=self.n_workers, threads_per_worker=self.threads_per_process, - processes=True, dashboard_address=f":{self.port}" - ) + + self.port = self.port or "8787" + self.cluster = LocalCluster( + n_workers=self.n_workers, threads_per_worker=self.threads_per_process, + processes=True, dashboard_address=f":{self.port}" + ) + + if resources:= self.cluster.workers[0].resources: + resources["cellpose"] = 1 + else: + self.cluster.workers[0].resources = dict(cellpose=1) + self.client = Client(self.cluster) elif self.mode == "jobqueue": from dask_jobqueue import PBSCluster