Strategist service initial implementation, with help from Claude Code. (#421)

* Strategist service initial implementation, with help from Claude Code.

* Strategist refinements, with Claude's help

* Working through remaining bits of Strategist._execute_strategy

* Added initial Strategy guide to docs

* Completed StrategistService with API points.

* Working on Strategist tests

* Added connection closing to Neo4jStore, StrategistService

* Strategist tests now all passing

* Working on remaining test failures in interface API and AlchemiscaleClient.

* Fixed API validation error test

* Black!

* Adjustments to api tests for strategies, added unit tests for StrategyState

* datetime module to class fixes

* Add stratocaster to conda envs

* Remove outdated import

* Changes from @ianmkenney review

* Black!

* More fixes from @ianmkenney review

* Refactored long StrategistService._execute_strategy method to be more readable

* Black!

* Fix and simplify strategist example config; adjust CLI accordingly

* DummyStrategy now terminates after enough results accumulate

* Black!

* More test fixes from @ianmkenney review

* Strategy guide incremental changes

* Update docs/strategy_guide.rst

Co-authored-by: Ian Kenney <ianmichaelkenney@gmail.com>

* Reorganized user guide to allow expansion with Strategy guide

* Doc fixes, plus rewrite of Strategy user guide doc

* Fix sphinx warnings

* More doc formatting fixes

* Added top matter to user guide index

* Added news item

---------

Co-authored-by: Ian Kenney <ianmichaelkenney@gmail.com>

Author: dotsdl

alchemiscale/base/client.py

@@ -24,10 +24,7 @@
 from gufe.tokenization import GufeTokenizable, JSON_HANDLER
 
 from ..models import ScopedKey
-
-
-def json_to_gufe(jsondata):
-    return GufeTokenizable.from_dict(json.loads(jsondata, cls=JSON_HANDLER.decoder))
+from ..compression import json_to_gufe
 
 
 class AlchemiscaleBaseClientError(Exception):

alchemiscale/cli.py

@@ -516,6 +516,40 @@ def scope(func):
     return scope(func)
 
 
+@cli.command(help="Start the strategist service.")
+@click.option(
+    "--config-file",
+    type=click.File("r"),
+    help="YAML configuration file for the strategist service.",
+    required=True,
+)
+def strategist(config_file):
+    """Start the strategist service for executing strategies on networks."""
+    from alchemiscale.models import Scope
+    from alchemiscale.strategist.service import StrategistService
+    from alchemiscale.strategist.settings import StrategistSettings
+
+    params = yaml.safe_load(config_file)
+
+    if "scopes" in params:
+        params["scopes"] = [Scope.from_str(scope) for scope in params["scopes"]]
+
+    service = StrategistService(StrategistSettings(**params))
+
+    # add signal handling
+    for signame in {"SIGHUP", "SIGINT", "SIGTERM"}:
+
+        def stop(*args, **kwargs):
+            service.stop()
+
+        signal.signal(getattr(signal, signame), stop)
+
+    try:
+        service.start()
+    except KeyboardInterrupt:
+        pass
+
+
 @cli.group(help="Subcommands for managing identities")
 def identity(): ...
 

alchemiscale/compression.py

@@ -3,6 +3,10 @@
 import zstandard as zstd
 
 
+def json_to_gufe(jsondata):
+    return GufeTokenizable.from_dict(json.loads(jsondata, cls=JSON_HANDLER.decoder))
+
+
 def compress_keyed_chain_zstd(keyed_chain: list[tuple[str, dict]]) -> bytes:
     """Compress a keyed chain using zstandard compression.
 

alchemiscale/interface/api.py

@@ -12,6 +12,7 @@
 
 import json
 from gufe.tokenization import JSON_HANDLER, KeyedChain
+from pydantic import ValidationError
 
 from ..base.api import (
     GufeJSONResponse,
@@ -30,7 +31,7 @@
 from ..settings import get_base_api_settings
 from ..storage.statestore import Neo4jStore
 from ..storage.objectstore import S3ObjectStore
-from ..storage.models import TaskStatusEnum
+from ..storage.models import TaskStatusEnum, StrategyState
 from ..models import Scope, ScopedKey
 from ..security.models import TokenData, CredentialedUserIdentity
 
@@ -393,10 +394,6 @@ def get_chemicalsystem(
 ### compute
 
 
-@router.post("/networks/{scoped_key}/strategy")
-def set_strategy(scoped_key: str, *, strategy: dict = Body(...), scope: Scope): ...
-
-
 @router.post("/transformations/{transformation_scoped_key}/tasks")
 def create_tasks(
     transformation_scoped_key,
@@ -1180,6 +1177,148 @@ def get_task_failures(
     return [str(sk) for sk in n4js.get_task_failures(sk)]
 
 
+### strategies
+
+
+@router.post("/networks/{network_scoped_key}/strategy")
+async def set_network_strategy(
+    network_scoped_key,
+    *,
+    request: Request,
+    n4js: Neo4jStore = Depends(get_n4js_depends),
+    token: TokenData = Depends(get_token_data_depends),
+):
+    """Set a Strategy for the given AlchemicalNetwork.
+
+    Expected request body:
+    {
+        "strategy": {...},  // GUFE strategy object, or null to remove
+        "max_tasks_per_transformation": 3,
+        "task_scaling": "exponential",
+        "mode": "partial",
+        "sleep_interval": 3600
+    }
+    """
+    sk = ScopedKey.from_str(network_scoped_key)
+    validate_scopes(sk.scope, token)
+
+    # Handle request body with custom JSON decoder for GUFE objects
+    body = await request.body()
+    body_ = json.loads(body.decode("utf-8"), cls=JSON_HANDLER.decoder)
+
+    try:
+        strategy_keyed_chain = body_.pop("strategy")
+
+        # Convert KeyedChain to GufeTokenizable if strategy is provided
+        if strategy_keyed_chain is not None:
+            strategy_kc = KeyedChain(strategy_keyed_chain)
+            strategy = strategy_kc.to_gufe()
+        else:
+            strategy = None
+    except Exception as e:
+        raise HTTPException(
+            status_code=http_status.HTTP_422_UNPROCESSABLE_ENTITY,
+            detail=str(e),
+        )
+
+    if strategy is not None:
+        # Create strategy state from body parameters
+        try:
+            strategy_state = StrategyState(**body_)
+        except ValidationError as e:
+            raise HTTPException(
+                status_code=http_status.HTTP_422_UNPROCESSABLE_ENTITY,
+                detail=str(e),
+            )
+
+        try:
+            strategy_sk = n4js.set_network_strategy(sk, strategy, strategy_state)
+        except ValueError:
+            raise HTTPException(
+                status_code=http_status.HTTP_400_BAD_REQUEST,
+                detail=str(e),
+            )
+
+        return str(strategy_sk) if strategy_sk is not None else None
+    else:
+        # Remove strategy
+        n4js.set_network_strategy(sk, None)
+        return None
+
+
+@router.get("/networks/{network_scoped_key}/strategy")
+def get_network_strategy(
+    network_scoped_key: str,
+    *,
+    n4js: Neo4jStore = Depends(get_n4js_depends),
+    token: TokenData = Depends(get_token_data_depends),
+):
+    """Get the Strategy for the given AlchemicalNetwork."""
+    sk = ScopedKey.from_str(network_scoped_key)
+    validate_scopes(sk.scope, token)
+
+    strategy = n4js.get_network_strategy(sk)
+    return GufeJSONResponse(strategy) if strategy is not None else None
+
+
+@router.get("/networks/{network_scoped_key}/strategy/state")
+def get_network_strategy_state(
+    network_scoped_key: str,
+    *,
+    n4js: Neo4jStore = Depends(get_n4js_depends),
+    token: TokenData = Depends(get_token_data_depends),
+):
+    """Get the StrategyState for the given AlchemicalNetwork."""
+    sk = ScopedKey.from_str(network_scoped_key)
+    validate_scopes(sk.scope, token)
+
+    strategy_state = n4js.get_network_strategy_state(sk)
+
+    return strategy_state.to_dict() if strategy_state is not None else None
+
+
+@router.get("/networks/{network_scoped_key}/strategy/status")
+def get_network_strategy_status(
+    network_scoped_key: str,
+    *,
+    n4js: Neo4jStore = Depends(get_n4js_depends),
+    token: TokenData = Depends(get_token_data_depends),
+):
+    """Get the status of the Strategy for the given AlchemicalNetwork."""
+    sk = ScopedKey.from_str(network_scoped_key)
+    validate_scopes(sk.scope, token)
+
+    strategy_state = n4js.get_network_strategy_state(sk)
+
+    return strategy_state.status.value if strategy_state is not None else None
+
+
+@router.post("/networks/{network_scoped_key}/strategy/awake")
+def set_network_strategy_awake(
+    network_scoped_key: str,
+    *,
+    n4js: Neo4jStore = Depends(get_n4js_depends),
+    token: TokenData = Depends(get_token_data_depends),
+):
+    """Set the Strategy status to 'awake' for the given AlchemicalNetwork."""
+    sk = ScopedKey.from_str(network_scoped_key)
+    validate_scopes(sk.scope, token)
+
+    strategy_state = n4js.get_network_strategy_state(sk)
+
+    if strategy_state is None:
+        return
+
+    # Update strategy state to awake and clear error info
+    strategy_state.status = "awake"
+    strategy_state.exception = None
+    strategy_state.traceback = None
+
+    updated = n4js.update_strategy_state(sk, strategy_state)
+
+    return str(updated) if updated is not None else None
+
+
 ### add router
 
 app.include_router(router)

alchemiscale/interface/client.py

@@ -21,16 +21,16 @@
 from ..base.client import (
     AlchemiscaleBaseClient,
     AlchemiscaleBaseClientError,
-    json_to_gufe,
     use_session,
 )
-from ..compression import decompress_gufe_zstd, compress_keyed_chain_zstd
+from ..compression import decompress_gufe_zstd, compress_keyed_chain_zstd, json_to_gufe
 from ..models import Scope, ScopedKey
 from ..storage.models import (
     TaskStatusEnum,
     NetworkStateEnum,
+    StrategyState,
 )
-from ..strategies import Strategy
+from stratocaster.base import Strategy
 from ..validators import validate_network_nonself
 
 from warnings import warn
@@ -690,14 +690,131 @@ def _get_chemicalsystem():
 
     ### compute
 
-    def set_strategy(self, network: ScopedKey, strategy: Strategy):
-        """Set the Strategy for evaluating the given AlchemicalNetwork.
+    #### strategies
 
-        The Strategy will be applied to create and action tasks for the
-        Transformations in the AlchemicalNetwork without user interaction.
+    def set_network_strategy(
+        self,
+        network: ScopedKey,
+        strategy: GufeTokenizable | None,
+        max_tasks_per_transformation: int = 3,
+        task_scaling: str = "exponential",
+        mode: str = "partial",
+        sleep_interval: int = 3600,
+    ) -> ScopedKey | None:
+        """Set a Strategy for the given AlchemicalNetwork.
+
+        Parameters
+        ----------
+        network
+            ScopedKey of the AlchemicalNetwork.
+        strategy
+            Strategy object (GufeTokenizable) or None to remove strategy.
+        max_tasks_per_transformation.
+            Maximum number of actioned Tasks allowed on a Transformation at once.
+        task_scaling
+            How to translate weights into Task counts: "linear" or "exponential".
+        mode
+            Strategy mode: "full", "partial", or "disabled".
+        sleep_interval
+            Wait time between iterations of the Strategy in seconds.
+
+        Returns
+        -------
+        ScopedKey | None
+            ScopedKey of the Strategy that was set, or ``None`` if strategy was
+            removed.
+        """
+        if strategy is not None:
+            # Convert strategy to keyed chain for serialization
+            strategy_keyed_chain = KeyedChain.gufe_to_keyed_chain_rep(strategy)
+            data = {
+                "strategy": strategy_keyed_chain,
+                "max_tasks_per_transformation": max_tasks_per_transformation,
+                "task_scaling": task_scaling,
+                "mode": mode,
+                "sleep_interval": sleep_interval,
+            }
+        else:
+            # Remove strategy
+            data = {"strategy": None}
+
+        result = self._post_resource(f"/networks/{network}/strategy", data)
+        return ScopedKey.from_str(result) if result else None
+
+    def get_network_strategy(self, network: ScopedKey) -> Strategy | None:
+        """Get the Strategy for the given AlchemicalNetwork.
+
+        Parameters
+        ----------
+        network
+            ScopedKey of the AlchemicalNetwork.
+
+        Returns
+        -------
+        strategy
+            Strategy object for this AlchemicalNetwork; ``None`` if no Strategy
+            is set.
+
+        """
+        keyed_chain = self._get_resource(f"/networks/{network}/strategy")
+        return GufeTokenizable.from_keyed_chain(keyed_chain) if keyed_chain else None
+
+    def get_network_strategy_state(self, network: ScopedKey) -> StrategyState | None:
+        """Get the StrategyState for the given AlchemicalNetwork.
+
+        Parameters
+        ----------
+        network
+            ScopedKey of the AlchemicalNetwork.
+
+        Returns
+        -------
+        strategy_state
+            Strategy state with execution metadata; ``None`` if no Strategy is
+            set.
 
         """
-        raise NotImplementedError
+        state_dict = self._get_resource(f"/networks/{network}/strategy/state")
+        return StrategyState.from_dict(state_dict) if state_dict else None
+
+    def get_network_strategy_status(self, network: ScopedKey) -> str | None:
+        """Get the status of the Strategy for the given AlchemicalNetwork.
+
+        Parameters
+        ----------
+        network
+            ScopedKey of the AlchemicalNetwork.
+
+        Returns
+        -------
+        status
+            Strategy status: "awake", "dormant", or "error"; ``None`` if no
+            Strategy is set.
+
+        """
+        return self._get_resource(f"/networks/{network}/strategy/status")
+
+    def set_network_strategy_awake(self, network: ScopedKey) -> ScopedKey | None:
+        """Set the Strategy status to 'awake' for the given AlchemicalNetwork.
+
+        This resets a dormant or errored strategy to active status.
+
+        Parameters
+        ----------
+        network
+            ScopedKey of the AlchemicalNetwork to set Strategy 'awake' for.
+
+        Returns
+        -------
+        ScopedKey | None
+            ScopedKey of the AlchemicalNetwork if Strategy status set to
+            'awake'; ``None`` otherwise.
+
+        """
+        result = self._post_resource(f"/networks/{network}/strategy/awake", {})
+        return ScopedKey.from_str(result) if result is not None else None
+
+    #### tasks
 
     def create_tasks(
         self,