Merge pull request #218 from openforcefield/alchemiscale-fah

Changes needed to support execution via `alchemiscale-fah`

alchemiscale/cli.py

@@ -362,6 +362,7 @@ def get_settings_override():
 def synchronous(config_file):
  from alchemiscale.models import Scope
  from alchemiscale.compute.service import SynchronousComputeService
+ from alchemiscale.compute.settings import ComputeServiceSettings
 
  params = yaml.safe_load(config_file)
 
@@ -373,7 +374,7 @@ def synchronous(config_file):
  Scope.from_str(scope) for scope in params_init["scopes"]
  ]
 
- service = SynchronousComputeService(**params_init)
+ service = SynchronousComputeService(ComputeServiceSettings(**params_init))
 
  # add signal handling
  for signame in {"SIGHUP", "SIGINT", "SIGTERM"}:

alchemiscale/compute/api.py

@@ -8,6 +8,7 @@
 import os
 import json
 from datetime import datetime, timedelta
+import random
 
 from fastapi import FastAPI, APIRouter, Body, Depends
 from fastapi.middleware.gzip import GZipMiddleware
@@ -23,6 +24,7 @@
  get_cred_entity,
  validate_scopes,
  validate_scopes_query,
+ minimize_scope_space,
  _check_store_connectivity,
  gufe_to_json,
  GzipRoute,
@@ -177,6 +179,7 @@ def claim_taskhub_tasks(
  *,
  compute_service_id: str = Body(),
  count: int = Body(),
+ protocols: Optional[List[str]] = Body(None, embed=True),
  n4js: Neo4jStore = Depends(get_n4js_depends),
  token: TokenData = Depends(get_token_data_depends),
 ):
@@ -187,13 +190,91 @@ def claim_taskhub_tasks(
  taskhub=taskhub_scoped_key,
  compute_service_id=ComputeServiceID(compute_service_id),
  count=count,
+ protocols=protocols,
  )
 
  return [str(t) if t is not None else None for t in tasks]
 
 
+@router.post("/claim")
+def claim_tasks(
+ scopes: List[Scope] = Body(),
+ compute_service_id: str = Body(),
+ count: int = Body(),
+ protocols: Optional[List[str]] = Body(None, embed=True),
+ n4js: Neo4jStore = Depends(get_n4js_depends),
+ token: TokenData = Depends(get_token_data_depends),
+):
+ # intersect query scopes with accessible scopes in the token
+ scopes_reduced = minimize_scope_space(scopes)
+ query_scopes = []
+ for scope in scopes_reduced:
+ query_scopes.extend(validate_scopes_query(scope, token))
+
+ taskhubs = dict()
+ # query each scope for available taskhubs
+ # loop might be more removable in the future with a Union like operator on scopes
+ for single_query_scope in set(query_scopes):
+ taskhubs.update(n4js.query_taskhubs(scope=single_query_scope, return_gufe=True))
+
+ # list of tasks to return
+ tasks = []
+
+ if len(taskhubs) == 0:
+ return []
+
+ # claim tasks from taskhubs based on weight; keep going till we hit our
+ # total desired task count, or we run out of taskhubs to draw from
+ while len(tasks) < count and len(taskhubs) > 0:
+ weights = [th.weight for th in taskhubs.values()]
+
+ if sum(weights) == 0:
+ break
+
+ # based on weights, choose taskhub to draw from
+ taskhub: ScopedKey = random.choices(list(taskhubs.keys()), weights=weights)[0]
+
+ # claim tasks from the taskhub
+ claimed_tasks = n4js.claim_taskhub_tasks(
+ taskhub,
+ compute_service_id=ComputeServiceID(compute_service_id),
+ count=(count - len(tasks)),
+ protocols=protocols,
+ )
+
+ # gather up claimed tasks, if present
+ for t in claimed_tasks:
+ if t is not None:
+ tasks.append(t)
+
+ # remove this taskhub from the options available; repeat
+ taskhubs.pop(taskhub)
+
+ return [str(t) for t in tasks] + [None] * (count - len(tasks))
+
+
 @router.get("/tasks/{task_scoped_key}/transformation")
 def get_task_transformation(
+ task_scoped_key,
+ *,
+ n4js: Neo4jStore = Depends(get_n4js_depends),
+ token: TokenData = Depends(get_token_data_depends),
+):
+ sk = ScopedKey.from_str(task_scoped_key)
+ validate_scopes(sk.scope, token)
+
+ transformation: ScopedKey
+
+ transformation, _ = n4js.get_task_transformation(
+ task=task_scoped_key,
+ return_gufe=False,
+ )
+
+ return str(transformation)
+
+
+@router.get("/tasks/{task_scoped_key}/transformation/gufe")
+def retrieve_task_transformation(
  task_scoped_key,
  *,
  n4js: Neo4jStore = Depends(get_n4js_depends),

alchemiscale/compute/client.py

@@ -35,15 +35,17 @@ class AlchemiscaleComputeClient(AlchemiscaleBaseClient):
  _exception = AlchemiscaleComputeClientError
 
  def register(self, compute_service_id: ComputeServiceID):
- res = self._post_resource(f"computeservice/{compute_service_id}/register", {})
+ res = self._post_resource(f"/computeservice/{compute_service_id}/register", {})
  return ComputeServiceID(res)
 
  def deregister(self, compute_service_id: ComputeServiceID):
- res = self._post_resource(f"computeservice/{compute_service_id}/deregister", {})
+ res = self._post_resource(
+ f"/computeservice/{compute_service_id}/deregister", {}
+ )
  return ComputeServiceID(res)
 
  def heartbeat(self, compute_service_id: ComputeServiceID):
- res = self._post_resource(f"computeservice/{compute_service_id}/heartbeat", {})
+ res = self._post_resource(f"/computeservice/{compute_service_id}/heartbeat", {})
  return ComputeServiceID(res)
 
  def list_scopes(self) -> List[Scope]:
@@ -71,19 +73,48 @@ def query_taskhubs(
  return taskhubs
 
  def claim_taskhub_tasks(
- self, taskhub: ScopedKey, compute_service_id: ComputeServiceID, count: int = 1
+ self,
+ taskhub: ScopedKey,
+ compute_service_id: ComputeServiceID,
+ count: int = 1,
+ protocols: Optional[List[str]] = None,
  ) -> Task:
  """Claim a `Task` from the specified `TaskHub`"""
- data = dict(compute_service_id=str(compute_service_id), count=count)
- tasks = self._post_resource(f"taskhubs/{taskhub}/claim", data)
+ data = dict(
+ compute_service_id=str(compute_service_id), count=count, protocols=protocols
+ )
+ tasks = self._post_resource(f"/taskhubs/{taskhub}/claim", data)
+
+ return [ScopedKey.from_str(t) if t is not None else None for t in tasks]
+
+ def claim_tasks(
+ self,
+ scopes: List[Scope],
+ compute_service_id: ComputeServiceID,
+ count: int = 1,
+ protocols: Optional[List[str]] = None,
+ ):
+ """Claim Tasks from TaskHubs within a list of Scopes."""
+ data = dict(
+ scopes=[scope.dict() for scope in scopes],
+ compute_service_id=str(compute_service_id),
+ count=count,
+ protocols=protocols,
+ )
+ tasks = self._post_resource("/claim", data)
 
  return [ScopedKey.from_str(t) if t is not None else None for t in tasks]
 
- def get_task_transformation(
+ def get_task_transformation(self, task: ScopedKey) -> ScopedKey:
+ """Get the Transformation associated with the given Task."""
+ transformation = self._get_resource(f"/tasks/{task}/transformation")
+ return ScopedKey.from_str(transformation)
+
+ def retrieve_task_transformation(
  self, task: ScopedKey
  ) -> Tuple[Transformation, Optional[ProtocolDAGResult]]:
  transformation, protocoldagresult = self._get_resource(
- f"tasks/{task}/transformation"
+ f"/tasks/{task}/transformation/gufe"
  )
 
  return (
@@ -104,6 +135,6 @@ def set_task_result(
  compute_service_id=str(compute_service_id),
  )
 
- pdr_sk = self._post_resource(f"tasks/{task}/results", data)
+ pdr_sk = self._post_resource(f"/tasks/{task}/results", data)
 
  return ScopedKey.from_dict(pdr_sk)