Skip to content

Commit

Permalink
Put claim_tasks method back in SynchronousComputeService to simplify …
Browse files Browse the repository at this point in the history
…testing
  • Loading branch information
dotsdl committed Dec 3, 2024
1 parent 9cb9ce5 commit d29965e
Showing 1 changed file with 14 additions and 6 deletions.
20 changes: 14 additions & 6 deletions alchemiscale/compute/service.py
Original file line number Diff line number Diff line change
Expand Up @@ -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]]:
Expand Down Expand Up @@ -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
Expand Down

0 comments on commit d29965e

Please sign in to comment.