diff --git a/plugins/coreml/package.json b/plugins/coreml/package.json index c7970ea56..0909423af 100644 --- a/plugins/coreml/package.json +++ b/plugins/coreml/package.json @@ -41,6 +41,7 @@ ], "labels": { "require": [ + "compute", "darwin" ] } diff --git a/plugins/coreml/src/coreml/__init__.py b/plugins/coreml/src/coreml/__init__.py index 45ed35d4f..f7de7bd54 100644 --- a/plugins/coreml/src/coreml/__init__.py +++ b/plugins/coreml/src/coreml/__init__.py @@ -147,6 +147,9 @@ class CoreMLPlugin( if not self.forked: asyncio.ensure_future(self.prepareRecognitionModels(), loop=self.loop) + def getClusterLabels(self): + return {"labels": {"require": ["compute", "darwin"]}} + async def prepareRecognitionModels(self): try: devices = [ diff --git a/plugins/coreml/src/coreml/face_recognition.py b/plugins/coreml/src/coreml/face_recognition.py index f93f95326..0b4113e44 100644 --- a/plugins/coreml/src/coreml/face_recognition.py +++ b/plugins/coreml/src/coreml/face_recognition.py @@ -34,6 +34,9 @@ class CoreMLFaceRecognition(FaceRecognizeDetection): self.detectExecutor = concurrent.futures.ThreadPoolExecutor(1, "detect-face") self.recogExecutor = concurrent.futures.ThreadPoolExecutor(1, "recog-face") + def getClusterLabels(self): + return {"labels": {"require": ["compute", "darwin"]}} + def downloadModel(self, model: str): model_version = "v7" mlmodel = "model" diff --git a/plugins/coreml/src/coreml/text_recognition.py b/plugins/coreml/src/coreml/text_recognition.py index 52f47ade2..2a29926c9 100644 --- a/plugins/coreml/src/coreml/text_recognition.py +++ b/plugins/coreml/src/coreml/text_recognition.py @@ -19,6 +19,9 @@ class CoreMLTextRecognition(TextRecognition): self.detectExecutor = concurrent.futures.ThreadPoolExecutor(1, "detect-text") self.recogExecutor = concurrent.futures.ThreadPoolExecutor(1, "recog-text") + def getClusterLabels(self): + return {"labels": {"require": ["compute", "darwin"]}} + def downloadModel(self, model: str): model_version = "v8" mlmodel = "model" diff --git a/plugins/openvino/src/predict/__init__.py b/plugins/openvino/src/predict/__init__.py index 99e7a9a5f..396ee4c08 100644 --- a/plugins/openvino/src/predict/__init__.py +++ b/plugins/openvino/src/predict/__init__.py @@ -310,6 +310,9 @@ class PredictPlugin(DetectPlugin, scrypted_sdk.ClusterForkInterface): finally: data.close() + def getClusterLabels(self): + return {"labels": {"require": ["compute"]}} + async def forkInterfaceInternal(self, options: dict): if self.plugin: return await self.plugin.forkInterfaceInternal(options) @@ -323,9 +326,7 @@ class PredictPlugin(DetectPlugin, scrypted_sdk.ClusterForkInterface): forked = self.forks.get(clusterWorkerId, None) if not forked: - forked = scrypted_sdk.fork( - {"labels": {"require": ["compute"]}, **(options or {})} - ) + forked = scrypted_sdk.fork({**self.getClusterLabels(), **(options or {})}) def clusterWorkerExit(result): print("cluster worker exit", clusterWorkerId) @@ -380,7 +381,8 @@ class PredictPlugin(DetectPlugin, scrypted_sdk.ClusterForkInterface): {"clusterWorkerId": clusterWorkerId} ) except: - traceback.print_exc() + # traceback.print_exc() + pass asyncio.ensure_future(startClusterWorker(), loop=self.loop)