From 9cb9ce5069d9973d135fab2f6af4c37fba8f5550 Mon Sep 17 00:00:00 2001 From: David Dotson Date: Fri, 29 Nov 2024 19:32:00 -0700 Subject: [PATCH] `SynchronousComputeService` now properly claims tasks with protocols filter We weren't actually using the `protocols` setting for the `SynchronousComputeService`, so `Task` claims weren't being filtered if this was set to anything but `None`. --- alchemiscale/compute/service.py | 33 ++++++--------------------------- 1 file changed, 6 insertions(+), 27 deletions(-) diff --git a/alchemiscale/compute/service.py b/alchemiscale/compute/service.py index 2955555d..31b4a922 100644 --- a/alchemiscale/compute/service.py +++ b/alchemiscale/compute/service.py @@ -157,32 +157,6 @@ def heartbeat(self): self.beat() time.sleep(self.heartbeat_interval) - def claim_tasks( - self, count=1, protocols: Optional[List[str]] = None - ) -> List[Optional[ScopedKey]]: - """Get a Task to execute from compute API. - - Returns `None` if no Task was available matching service configuration. - - Parameters - ---------- - count - The maximum number of Tasks to claim. - protocols - Protocol names to restrict Task claiming to. `None` means no restriction. - Regex patterns are allowed. - - """ - - tasks = self.client.claim_tasks( - scopes=self.scopes, - compute_service_id=self.compute_service_id, - count=count, - protocols=protocols, - ) - - return tasks - def task_to_protocoldag( self, task: ScopedKey ) -> Tuple[ProtocolDAG, Transformation, Optional[ProtocolDAGResult]]: @@ -306,7 +280,12 @@ def cycle(self, max_tasks: Optional[int] = None, max_time: Optional[int] = None) # claim tasks from the compute API self.logger.info("Claiming tasks") - tasks: List[ScopedKey] = self.claim_tasks(self.claim_limit) + tasks: List[ScopedKey] = self.client.claim_tasks( + scopes=self.scopes, + compute_service_id=self.compute_service_id, + count=self.claim_limit, + protocols=self.settings.protocols, + ) self.logger.info("Claimed %d tasks", len([t for t in tasks if t is not None])) # if no tasks claimed, sleep