From d29965e55cb3316a604f463478bc68aa4d2c8041 Mon Sep 17 00:00:00 2001 From: David Dotson Date: Mon, 2 Dec 2024 20:26:49 -0700 Subject: [PATCH] Put claim_tasks method back in SynchronousComputeService to simplify testing --- alchemiscale/compute/service.py | 20 ++++++++++++++------ 1 file changed, 14 insertions(+), 6 deletions(-) diff --git a/alchemiscale/compute/service.py b/alchemiscale/compute/service.py index 31b4a922..d9332e57 100644 --- a/alchemiscale/compute/service.py +++ b/alchemiscale/compute/service.py @@ -157,6 +157,19 @@ def heartbeat(self): self.beat() time.sleep(self.heartbeat_interval) + def claim_tasks(self) -> List[Optional[ScopedKey]]: + """Get a Task to execute from compute API. + + Returns `None` if no Task was available matching service configuration. + + """ + return self.client.claim_tasks( + scopes=self.scopes, + compute_service_id=self.compute_service_id, + count=self.claim_limit, + protocols=self.settings.protocols, + ) + def task_to_protocoldag( self, task: ScopedKey ) -> Tuple[ProtocolDAG, Transformation, Optional[ProtocolDAGResult]]: @@ -280,12 +293,7 @@ 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.client.claim_tasks( - scopes=self.scopes, - compute_service_id=self.compute_service_id, - count=self.claim_limit, - protocols=self.settings.protocols, - ) + tasks: List[ScopedKey] = self.claim_tasks() self.logger.info("Claimed %d tasks", len([t for t in tasks if t is not None])) # if no tasks claimed, sleep