diff --git a/alchemiscale/cli.py b/alchemiscale/cli.py index c87caa5e..db1129b2 100644 --- a/alchemiscale/cli.py +++ b/alchemiscale/cli.py @@ -489,6 +489,29 @@ def v03_to_v04(url, user, password, dbname): click.echo("Migration completed without errors.") +@migrate.command() +@db_params +def v07_to_v08(url, user, password, dbname): + """Perform migration appropriate for transitioning from alchemiscale v0.7 + to v0.8. + + Note that options here can be set by environment variables, as shown on + each option. + """ + from .storage.statestore import get_n4js + from .settings import Neo4jStoreSettings + from .migrations.v07_to_v08 import migrate + + cli_values = url | user | password | dbname + settings = get_settings_from_options(cli_values, Neo4jStoreSettings) + + n4js = get_n4js(settings) + + migrate(n4js) + + click.echo("Migration completed without errors.") + + def _identity_type_string_to_cls(identity_type: str) -> type[CredentialedEntity]: if identity_type == "user": identity_type_cls = CredentialedUserIdentity diff --git a/alchemiscale/compute/api.py b/alchemiscale/compute/api.py index ea9b1e4d..0087bff7 100644 --- a/alchemiscale/compute/api.py +++ b/alchemiscale/compute/api.py @@ -109,6 +109,7 @@ def register_computeservice( *, compute_manager_id: str | None = Body(None, embed=True), hostname: str | None = Body(None, embed=True), + environment: dict | None = Body(None, embed=True), n4js: Neo4jStore = Depends(get_n4js_depends), ): now = datetime.datetime.now(tz=datetime.UTC) @@ -124,6 +125,7 @@ def register_computeservice( failure_times=[], manager_name=manager_name, hostname=hostname, + environment=environment, ) try: diff --git a/alchemiscale/compute/client.py b/alchemiscale/compute/client.py index 1c1edadf..f68c9858 100644 --- a/alchemiscale/compute/client.py +++ b/alchemiscale/compute/client.py @@ -45,10 +45,15 @@ def register( compute_service_id: ComputeServiceID, compute_manager_id: ComputeManagerID | None = None, hostname: str | None = None, + environment: dict | None = None, ): res = self._post_resource( f"/computeservice/{compute_service_id}/register", - {"compute_manager_id": compute_manager_id, "hostname": hostname}, + { + "compute_manager_id": compute_manager_id, + "hostname": hostname, + "environment": environment, + }, ) return ComputeServiceID(res) diff --git a/alchemiscale/compute/environment.py b/alchemiscale/compute/environment.py new file mode 100644 index 00000000..736c68a5 --- /dev/null +++ b/alchemiscale/compute/environment.py @@ -0,0 +1,101 @@ +""" +:mod:`alchemiscale.compute.environment` --- compute environment capture +======================================================================= + +Best-effort capture of the software environment a compute service executes +`Task`\\ s in, for durable execution provenance (issue #106). + +We try a sequence of package managers --- ``micromamba``, ``mamba``, ``conda``, +then ``pip`` --- and take the first that yields a usable package listing. Some +information about the execution environment is better than none; a service whose +environment cannot be introspected simply records no environment. + +The capture is done once per compute service (the environment is fixed for the +service's lifetime) and copied into durable provenance server-side, deduplicated +so that identical environments across services/claims are stored once. +""" + +import datetime +import json +import shutil +import subprocess + +# (tool, argv) pairs tried in order; the first that returns a parseable package +# listing wins. conda-family tools and pip both support a JSON listing. +_CAPTURE_COMMANDS: list[tuple[str, list[str]]] = [ + ("micromamba", ["micromamba", "list", "--json"]), + ("mamba", ["mamba", "list", "--json"]), + ("conda", ["conda", "list", "--json"]), + ("pip", ["pip", "list", "--format=json"]), +] + + +def _parse_packages(payload: str) -> dict[str, str]: + """Parse a conda/pip ``--json`` listing into a ``{name: version}`` map. + + Both conda-family ``list --json`` and ``pip list --format=json`` emit a JSON + array of objects carrying ``name`` and ``version`` keys. + """ + data = json.loads(payload) + if not isinstance(data, list): + raise ValueError("unexpected package listing shape") + packages = {} + for entry in data: + name = entry.get("name") + version = entry.get("version") + if name is not None and version is not None: + packages[str(name)] = str(version) + if not packages: + raise ValueError("no packages parsed from listing") + return packages + + +def capture_environment( + timeout: float = 60.0, + commands: list[tuple[str, list[str]]] | None = None, +) -> dict | None: + """Capture the current software environment, best-effort. + + Tries ``micromamba``/``mamba``/``conda``/``pip`` in order and returns the + first successful listing as:: + + {"tool": "conda", "packages": {name: version, ...}, "captured_at": ""} + + Returns ``None`` if no tool is available or none produces a usable listing + (a missing tool, a non-zero exit, a timeout, or unparseable output all cause + a fall-through to the next tool). Never raises. + + Parameters + ---------- + timeout + Per-command timeout, in seconds. + commands + Override the (tool, argv) sequence; primarily for testing. + """ + for tool, argv in commands if commands is not None else _CAPTURE_COMMANDS: + if shutil.which(argv[0]) is None: + continue + try: + proc = subprocess.run( + argv, + capture_output=True, + text=True, + timeout=timeout, + check=False, + ) + except (OSError, subprocess.SubprocessError): + continue + if proc.returncode != 0 or not proc.stdout: + continue + try: + packages = _parse_packages(proc.stdout) + except (json.JSONDecodeError, ValueError): + continue + + return { + "tool": tool, + "packages": packages, + "captured_at": datetime.datetime.now(tz=datetime.UTC).isoformat(), + } + + return None diff --git a/alchemiscale/compute/service.py b/alchemiscale/compute/service.py index 479463dd..f287c56f 100644 --- a/alchemiscale/compute/service.py +++ b/alchemiscale/compute/service.py @@ -21,6 +21,7 @@ from .client import AlchemiscaleComputeClient from .execute import execute_DAG from .capture import SynchronousExecutionHooks +from .environment import capture_environment from .settings import ComputeServiceSettings from ..storage.models import ComputeServiceID from ..models import Scope, ScopedKey @@ -85,6 +86,12 @@ def __init__(self, settings: ComputeServiceSettings): # service creates; falls back to the OS hostname when not set self.hostname = self.settings.hostname or socket.gethostname() + # capture the software environment once (fixed for the service's + # lifetime); best-effort, recorded in durable execution provenance + self.environment = ( + capture_environment() if self.settings.capture_environment else None + ) + # shared between the main loop and the heartbeat thread; both wake # on a single ``stop()`` (which calls ``int_sleep.interrupt()``). # If you split these into separate functors, be sure to interrupt @@ -122,6 +129,7 @@ def _register(self): self.compute_service_id, self.settings.compute_manager_id, hostname=self.hostname, + environment=self.environment, ) def _deregister(self): diff --git a/alchemiscale/compute/settings.py b/alchemiscale/compute/settings.py index 5105b3de..47c16eca 100644 --- a/alchemiscale/compute/settings.py +++ b/alchemiscale/compute/settings.py @@ -35,6 +35,16 @@ class Config: "`socket.gethostname()`." ), ) + capture_environment: bool = Field( + True, + description=( + "If True, capture the software environment (package versions) this " + "service executes in --- via `micromamba`/`mamba`/`conda`/`pip` --- " + "once at startup, and record it in durable execution provenance. " + "Best-effort: if no package manager is available, no environment is " + "recorded." + ), + ) compute_manager_id: str | None = Field( None, description=( diff --git a/alchemiscale/migrations/v07_to_v08.py b/alchemiscale/migrations/v07_to_v08.py new file mode 100644 index 00000000..a6851e97 --- /dev/null +++ b/alchemiscale/migrations/v07_to_v08.py @@ -0,0 +1,28 @@ +""" +:mod:`alchemiscale.migrations.v07_to_v08` --- migration for v0.7 to v0.8 +======================================================================== + +""" + +from ..storage.statestore import Neo4jStore + + +def migrate(n4js: Neo4jStore): + """Migrate state store from alchemiscale v0.7 to v0.8. + + Adds the uniqueness constraint on the new ``ComputeEnvironment`` node label + that backs the deduplicated storage of compute-service execution + environments (issue #106): environments are content-addressed by ``hash``, + and the constraint is load-bearing --- it makes the ``MERGE`` on + ``ComputeEnvironment.hash`` (at compute-service registration) correct under + concurrency and fast, and brings an existing deployment's constraint set in + line with what ``Neo4jStore.check`` expects. + + Idempotent (``CREATE CONSTRAINT ... IF NOT EXISTS``); no data migration is + required. Pre-existing ``Task``\\ s simply have no recorded environment. + """ + + n4js.execute_query(""" + CREATE CONSTRAINT compute_environment_hash IF NOT EXISTS + FOR (n:ComputeEnvironment) REQUIRE n.hash IS UNIQUE + """) diff --git a/alchemiscale/storage/models.py b/alchemiscale/storage/models.py index 94b8e38b..14846bb2 100644 --- a/alchemiscale/storage/models.py +++ b/alchemiscale/storage/models.py @@ -9,6 +9,7 @@ import datetime from enum import Enum, StrEnum from uuid import uuid4, UUID +import json import re import hashlib @@ -83,6 +84,7 @@ class ComputeServiceRegistration(BaseModel): failure_times: list[datetime.datetime] = [] manager_name: str | None = None hostname: str | None = None + environment: dict | None = None model_config = ConfigDict(arbitrary_types_allowed=True) @@ -113,6 +115,78 @@ def from_dict(cls, dct): return cls(**dct_) +class ComputeEnvironment(BaseModel): + """A software environment a compute service executes `Task`\\ s in. + + Content-addressed by `hash` (a digest of the capturing tool plus the + ``{package: version}`` map), so that identical environments across services + and claims are stored as a single node and referenced from many + `TaskProvenance` attempts. It outlives the `ComputeServiceRegistration`\\ s + that reference it, so an attempt's environment survives the service's + teardown. + + Attributes + ---------- + hash + Content digest identifying this environment. + tool + The package manager that produced the listing (``micromamba``, + ``mamba``, ``conda``, or ``pip``). + packages + Mapping of package name to version. + captured_at + When the environment was captured on the compute service. + """ + + hash: str + tool: str + packages: dict[str, str] + captured_at: datetime.datetime | None = None + + model_config = ConfigDict(arbitrary_types_allowed=True) + + @staticmethod + def content_hash(tool: str, packages: dict[str, str]) -> str: + """Deterministic digest of ``(tool, packages)`` for deduplication.""" + canonical = json.dumps({"tool": tool, "packages": packages}, sort_keys=True) + return hashlib.sha256(canonical.encode()).hexdigest() + + @classmethod + def from_capture(cls, environment: dict) -> "ComputeEnvironment": + """Build from a `capture_environment` result + (``{"tool", "packages", "captured_at"}``).""" + tool = environment["tool"] + packages = {str(k): str(v) for k, v in environment["packages"].items()} + captured_at = environment.get("captured_at") + if isinstance(captured_at, str): + captured_at = datetime.datetime.fromisoformat(captured_at) + return cls( + hash=cls.content_hash(tool, packages), + tool=tool, + packages=packages, + captured_at=captured_at, + ) + + def to_capture_dict(self) -> dict: + """Render as the client-facing environment mapping.""" + return { + "tool": self.tool, + "packages": self.packages, + "captured_at": _iso(self.captured_at), + } + + @classmethod + def from_node(cls, node) -> "ComputeEnvironment": + """Build from a raw ``ComputeEnvironment`` Neo4j node (``packages`` is + stored as a JSON string property).""" + return cls( + hash=node["hash"], + tool=node["tool"], + packages=json.loads(node["packages"]), + captured_at=_coerce_datetime(node.get("captured_at")), + ) + + class ComputeManagerInstruction(StrEnum): OK = "OK" SKIP = "SKIP" @@ -843,6 +917,7 @@ class TaskAttempt(BaseModel): units_completed: int | None = None units_total: int | None = None protocoldagresultref: ScopedKey | None = None + environment: dict | None = None model_config = ConfigDict(arbitrary_types_allowed=True) @@ -861,6 +936,7 @@ def to_dict(self): if self.protocoldagresultref is not None else None ), + "environment": self.environment, } @classmethod @@ -881,6 +957,7 @@ def from_dict(cls, d): if d.get("protocoldagresultref") is not None else None ), + environment=d.get("environment"), ) diff --git a/alchemiscale/storage/statestore.py b/alchemiscale/storage/statestore.py index 11d6a5c6..6d53db5d 100644 --- a/alchemiscale/storage/statestore.py +++ b/alchemiscale/storage/statestore.py @@ -33,6 +33,7 @@ from stratocaster.base import Strategy from .models import ( + ComputeEnvironment, ComputeServiceID, ComputeServiceRegistration, ComputeManagerRegistration, @@ -179,6 +180,7 @@ def _status_write( // create CLAIMS relationship with given compute service MATCH (csreg:ComputeServiceRegistration {{identifier: $compute_service_id}}) + OPTIONAL MATCH (csreg)-[:HAS_ENVIRONMENT]->(ce:ComputeEnvironment) CREATE (t)<-[cl:CLAIMS {{claimed: datetime($datetimestr)}}]-(csreg) // create an immutable TaskProvenance record for this execution attempt, @@ -195,6 +197,12 @@ def _status_write( }}) CREATE (tp)-[:PROVENANCE_OF]->(t) + // link the attempt to the service's (deduplicated) execution environment, + // if one was captured; RAN_IN survives the registration's deletion + FOREACH (_ IN CASE WHEN ce IS NULL THEN [] ELSE [1] END | + CREATE (tp)-[:RAN_IN]->(ce) + ) + {_status_write('t', TaskStatusEnum.running.value, time_param='datetimestr')} RETURN t @@ -219,6 +227,10 @@ class Neo4jStore(AlchemiscaleStateStore): "name": "compute_service_registration_identifier", "property": "identifier", }, + "ComputeEnvironment": { + "name": "compute_environment_hash", + "property": "hash", + }, } def __init__(self, settings: Neo4jStoreSettings): @@ -1530,9 +1542,12 @@ def register_computeservice( """ - node = Node( - "ComputeServiceRegistration", **compute_service_registration.to_dict() - ) + reg_dict = compute_service_registration.to_dict() + # the captured environment is a nested map, not a valid node property; + # it is stored on a deduplicated ComputeEnvironment node below + environment = reg_dict.pop("environment", None) + + node = Node("ComputeServiceRegistration", **reg_dict) with self.transaction() as tx: create_subgraph(tx, Subgraph() | node) @@ -1553,6 +1568,33 @@ def register_computeservice( if not len(list(results)): raise ValueError("Could not find ComputeManagerRegistration") + # deduplicate and link the compute environment, if one was captured. + # The ComputeEnvironment node is content-addressed by its hash and + # shared across services/claims; it is deliberately not deleted with + # the registration, so an attempt's environment survives teardown. + if environment: + ce = ComputeEnvironment.from_capture(environment) + tx.run( + """ + MERGE (ce:ComputeEnvironment {hash: $hash}) + ON CREATE SET ce.tool = $tool, + ce.packages = $packages, + ce.captured_at = $captured_at + WITH ce + MATCH (csr:ComputeServiceRegistration {identifier: $identifier}) + MERGE (csr)-[:HAS_ENVIRONMENT]->(ce) + """, + hash=ce.hash, + tool=ce.tool, + packages=json.dumps(ce.packages), + captured_at=( + ce.captured_at.isoformat() + if ce.captured_at is not None + else None + ), + identifier=str(compute_service_registration.identifier), + ) + return compute_service_registration.identifier def deregister_computeservice(self, compute_service_id: ComputeServiceID): @@ -3692,10 +3734,18 @@ def add_protocol_dag_result_ref_tracebacks( merge_subgraph(tx, subgraph, "GufeTokenizable", "_scoped_key") @staticmethod - def _task_provenance_node_to_attempt(tp, pdrr_sk) -> TaskAttempt: - """Build a `TaskAttempt` record from a `TaskProvenance` node and the - `ScopedKey` string of its produced `ProtocolDAGResultRef` (or `None`).""" + def _task_provenance_node_to_attempt( + tp, pdrr_sk, environment_node=None + ) -> TaskAttempt: + """Build a `TaskAttempt` record from a `TaskProvenance` node, the + `ScopedKey` string of its produced `ProtocolDAGResultRef` (or `None`), + and the linked `ComputeEnvironment` node (or `None`).""" outcome = tp.get("outcome") + environment = ( + ComputeEnvironment.from_node(environment_node).to_capture_dict() + if environment_node is not None + else None + ) return TaskAttempt( compute_service_id=tp["compute_service_id"], hostname=tp.get("hostname"), @@ -3708,6 +3758,7 @@ def _task_provenance_node_to_attempt(tp, pdrr_sk) -> TaskAttempt: protocoldagresultref=( ScopedKey.from_str(pdrr_sk) if pdrr_sk is not None else None ), + environment=environment, ) def get_task_history( @@ -3723,7 +3774,8 @@ def get_task_history( q = """ MATCH (t:Task {_scoped_key: $task})<-[:PROVENANCE_OF]-(tp:TaskProvenance) OPTIONAL MATCH (tp)-[:PROVENANCE_OF]->(pdrr:ProtocolDAGResultRef) - RETURN tp, pdrr._scoped_key AS pdrr_sk + OPTIONAL MATCH (tp)-[:RAN_IN]->(ce:ComputeEnvironment) + RETURN tp, pdrr._scoped_key AS pdrr_sk, ce ORDER BY tp.datetime_claimed DESC """ if limit is not None: @@ -3738,7 +3790,7 @@ def get_task_history( for record in tx.run(q, **params): attempts.append( self._task_provenance_node_to_attempt( - record["tp"], record["pdrr_sk"] + record["tp"], record["pdrr_sk"], record["ce"] ) ) return attempts @@ -3762,6 +3814,7 @@ def get_tasks_details(self, tasks: list[ScopedKey]) -> list[TaskDetails | None]: } OPTIONAL MATCH (latest_tp)-[:PROVENANCE_OF]->(latest_pdrr:ProtocolDAGResultRef) + OPTIONAL MATCH (latest_tp)-[:RAN_IN]->(latest_ce:ComputeEnvironment) OPTIONAL MATCH (t)<-[cl:CLAIMS]-(csreg:ComputeServiceRegistration) OPTIONAL MATCH (t)<-[:PROVENANCE_OF]-(claim_tp:TaskProvenance {compute_service_id: csreg.identifier}) WHERE claim_tp.datetime_end IS NULL @@ -3771,6 +3824,7 @@ def get_tasks_details(self, tasks: list[ScopedKey]) -> list[TaskDetails | None]: num_claims, latest_tp, latest_pdrr._scoped_key AS latest_pdrr_sk, + latest_ce, cl.claimed AS claimed, csreg.identifier AS csid, csreg.hostname AS cs_hostname, @@ -3798,7 +3852,9 @@ def get_tasks_details(self, tasks: list[ScopedKey]) -> list[TaskDetails | None]: most_recent_attempt = None if record["latest_tp"] is not None: most_recent_attempt = self._task_provenance_node_to_attempt( - record["latest_tp"], record["latest_pdrr_sk"] + record["latest_tp"], + record["latest_pdrr_sk"], + record["latest_ce"], ) by_task[record["task_sk"]] = TaskDetails( diff --git a/alchemiscale/tests/integration/storage/test_statestore_introspection.py b/alchemiscale/tests/integration/storage/test_statestore_introspection.py index b10f5e6a..f0da60df 100644 --- a/alchemiscale/tests/integration/storage/test_statestore_introspection.py +++ b/alchemiscale/tests/integration/storage/test_statestore_introspection.py @@ -33,8 +33,10 @@ def _register( compute_service_id: ComputeServiceID, hostname: str | None = "host-a", manager_name: str | None = None, + environment: dict | None = None, ) -> ComputeServiceID: - """Register a compute service carrying a ``hostname`` (and optional manager).""" + """Register a compute service carrying a ``hostname`` (and optional manager + and captured environment).""" now = datetime.datetime.now(tz=datetime.UTC) registration = ComputeServiceRegistration( identifier=compute_service_id, @@ -43,6 +45,7 @@ def _register( failure_times=[], hostname=hostname, manager_name=manager_name, + environment=environment, ) return n4js.register_computeservice(registration) @@ -75,6 +78,7 @@ def _claimed_task( compute_service_id: ComputeServiceID, hostname: str | None = "host-a", manager_name: str | None = None, + environment: dict | None = None, ): """Assemble a network, create+action a single Task, and claim it. @@ -85,7 +89,7 @@ def _claimed_task( transformation_sk = n4js.get_scoped_key(transformation, scope_test) task_sk = n4js.create_task(transformation_sk) n4js.action_tasks([task_sk], taskhub_sk) - _register(n4js, compute_service_id, hostname, manager_name) + _register(n4js, compute_service_id, hostname, manager_name, environment) claimed = n4js.claim_taskhub_tasks(taskhub_sk, compute_service_id) assert claimed[0] == task_sk return task_sk, taskhub_sk @@ -671,3 +675,107 @@ def running_tasks(scope, count, name): # a scope with no org cannot be leveled with pytest.raises(ValueError): n4js.get_scope_compute_share(Scope()) + + # --- compute environment (issue #106 comment) ------------------------- + + ENV = { + "tool": "conda", + "packages": {"gufe": "1.10.0", "python": "3.11.9", "openmm": "8.1.1"}, + "captured_at": "2026-07-20T00:00:00+00:00", + } + + def _count_environment_nodes(self, n4js) -> int: + return n4js.execute_query( + "MATCH (ce:ComputeEnvironment) RETURN count(ce) AS n" + ).records[0]["n"] + + def test_environment_surfaced_on_task_history( + self, n4js, network_tyk2, transformation, scope_test + ): + csid = ComputeServiceID.new_from_name("env.history") + task_sk, _ = self._claimed_task( + n4js, + network_tyk2, + transformation, + scope_test, + csid, + environment=self.ENV, + ) + + # one ComputeEnvironment node created, linked to the attempt + assert self._count_environment_nodes(n4js) == 1 + + attempt = n4js.get_task_history(task_sk)[0] + assert attempt.environment is not None + assert attempt.environment["tool"] == "conda" + assert attempt.environment["packages"] == self.ENV["packages"] + + def test_environment_absent_when_not_captured( + self, n4js, network_tyk2, transformation, scope_test + ): + csid = ComputeServiceID.new_from_name("env.none") + task_sk, _ = self._claimed_task( + n4js, network_tyk2, transformation, scope_test, csid, environment=None + ) + assert self._count_environment_nodes(n4js) == 0 + assert n4js.get_task_history(task_sk)[0].environment is None + + def test_environment_deduplicated_across_services(self, n4js): + # the ComputeEnvironment node is created at registration and content- + # addressed, so two services with the SAME environment share one node + for i in range(2): + _register( + n4js, + ComputeServiceID.new_from_name(f"env.dedup.{i}"), + hostname=f"h{i}", + environment=self.ENV, + ) + assert self._count_environment_nodes(n4js) == 1 + + # a service with a DIFFERENT environment gets its own node + other_env = { + **self.ENV, + "packages": {**self.ENV["packages"], "openmm": "8.2.0"}, + } + _register( + n4js, + ComputeServiceID.new_from_name("env.dedup.other"), + environment=other_env, + ) + assert self._count_environment_nodes(n4js) == 2 + + def test_environment_survives_registration_expiry( + self, n4js, network_tyk2, transformation, scope_test + ): + csid = ComputeServiceID.new_from_name("env.expire") + task_sk, _ = self._claimed_task( + n4js, + network_tyk2, + transformation, + scope_test, + csid, + environment=self.ENV, + ) + + # expire the registration (deletes it and its HAS_ENVIRONMENT edge) + n4js.execute_query( + """ + MATCH (csreg:ComputeServiceRegistration {identifier: $csid}) + SET csreg.heartbeat = datetime($past) + """, + csid=str(csid), + past=( + datetime.datetime.now(tz=datetime.UTC) - timedelta(hours=1) + ).isoformat(), + ) + n4js.expire_registrations( + datetime.datetime.now(tz=datetime.UTC) - timedelta(minutes=1) + ) + + # the ComputeEnvironment node and the attempt's RAN_IN link survive, so + # the attempt's environment is still reported + assert self._count_environment_nodes(n4js) == 1 + attempt = n4js.get_task_history(task_sk)[0] + assert attempt.outcome == TaskOutcomeEnum.expired + assert attempt.environment is not None + assert attempt.environment["packages"] == self.ENV["packages"] diff --git a/alchemiscale/tests/unit/compute/test_environment.py b/alchemiscale/tests/unit/compute/test_environment.py new file mode 100644 index 00000000..604781c4 --- /dev/null +++ b/alchemiscale/tests/unit/compute/test_environment.py @@ -0,0 +1,133 @@ +"""Unit tests for best-effort compute-environment capture +(:mod:`alchemiscale.compute.environment`).""" + +import json +import subprocess + +import pytest + +from alchemiscale.compute import environment as envmod +from alchemiscale.compute.environment import capture_environment, _parse_packages + + +def _fake_run_factory(outputs: dict[str, tuple[int, str]]): + """Build a fake ``subprocess.run`` returning per-tool ``(returncode, stdout)``.""" + + def _fake_run(argv, **kwargs): + tool = argv[0] + returncode, stdout = outputs.get(tool, (1, "")) + return subprocess.CompletedProcess( + args=argv, returncode=returncode, stdout=stdout, stderr="" + ) + + return _fake_run + + +def _install(monkeypatch, present: set[str], outputs: dict[str, tuple[int, str]]): + monkeypatch.setattr( + envmod.shutil, "which", lambda name: name if name in present else None + ) + monkeypatch.setattr(envmod.subprocess, "run", _fake_run_factory(outputs)) + + +CONDA_JSON = json.dumps( + [ + {"name": "gufe", "version": "1.10.0", "channel": "conda-forge"}, + {"name": "python", "version": "3.11.9", "channel": "conda-forge"}, + ] +) +PIP_JSON = json.dumps( + [{"name": "gufe", "version": "1.10.0"}, {"name": "pip", "version": "24.0"}] +) + + +class TestParsePackages: + def test_conda_shape(self): + assert _parse_packages(CONDA_JSON) == {"gufe": "1.10.0", "python": "3.11.9"} + + def test_pip_shape(self): + assert _parse_packages(PIP_JSON) == {"gufe": "1.10.0", "pip": "24.0"} + + def test_empty_list_raises(self): + with pytest.raises(ValueError): + _parse_packages("[]") + + def test_non_list_raises(self): + with pytest.raises(ValueError): + _parse_packages('{"not": "a list"}') + + def test_bad_json_raises(self): + with pytest.raises(json.JSONDecodeError): + _parse_packages("not json") + + +class TestCaptureEnvironment: + def test_first_tool_wins(self, monkeypatch): + # micromamba present and successful -> used, later tools untried + _install( + monkeypatch, + present={"micromamba", "conda", "pip"}, + outputs={"micromamba": (0, CONDA_JSON)}, + ) + env = capture_environment() + assert env["tool"] == "micromamba" + assert env["packages"] == {"gufe": "1.10.0", "python": "3.11.9"} + assert env["captured_at"] + + def test_falls_through_missing_tools(self, monkeypatch): + # micromamba/mamba/conda absent -> pip used + _install(monkeypatch, present={"pip"}, outputs={"pip": (0, PIP_JSON)}) + env = capture_environment() + assert env["tool"] == "pip" + assert env["packages"] == {"gufe": "1.10.0", "pip": "24.0"} + + def test_falls_through_nonzero_exit(self, monkeypatch): + # conda present but errors -> falls through to pip + _install( + monkeypatch, + present={"conda", "pip"}, + outputs={"conda": (1, ""), "pip": (0, PIP_JSON)}, + ) + assert capture_environment()["tool"] == "pip" + + def test_falls_through_unparseable_output(self, monkeypatch): + # conda present, returns garbage -> falls through to pip + _install( + monkeypatch, + present={"conda", "pip"}, + outputs={"conda": (0, "not json"), "pip": (0, PIP_JSON)}, + ) + assert capture_environment()["tool"] == "pip" + + def test_falls_through_empty_listing(self, monkeypatch): + # conda present, returns empty list (no packages) -> falls through + _install( + monkeypatch, + present={"conda", "pip"}, + outputs={"conda": (0, "[]"), "pip": (0, PIP_JSON)}, + ) + assert capture_environment()["tool"] == "pip" + + def test_no_tools_returns_none(self, monkeypatch): + _install(monkeypatch, present=set(), outputs={}) + assert capture_environment() is None + + def test_all_tools_fail_returns_none(self, monkeypatch): + _install( + monkeypatch, + present={"micromamba", "mamba", "conda", "pip"}, + outputs={t: (1, "") for t in ("micromamba", "mamba", "conda", "pip")}, + ) + assert capture_environment() is None + + def test_subprocess_error_is_swallowed(self, monkeypatch): + monkeypatch.setattr(envmod.shutil, "which", lambda name: name) + + def _boom(argv, **kwargs): + if argv[0] == "pip": + return subprocess.CompletedProcess(argv, 0, PIP_JSON, "") + raise subprocess.TimeoutExpired(argv, 1) + + monkeypatch.setattr(envmod.subprocess, "run", _boom) + # the conda-family tools time out; pip succeeds -> never raises + assert capture_environment()["tool"] == "pip" diff --git a/alchemiscale/tests/unit/test_introspection_records.py b/alchemiscale/tests/unit/test_introspection_records.py index db0ccbf1..69021d3e 100644 --- a/alchemiscale/tests/unit/test_introspection_records.py +++ b/alchemiscale/tests/unit/test_introspection_records.py @@ -5,12 +5,14 @@ """ import datetime +import json import pytest from gufe.tokenization import GufeKey from alchemiscale.models import Scope, ScopedKey from alchemiscale.storage.models import ( + ComputeEnvironment, ComputeServiceID, ProtocolDAGResultRec, ProtocolUnitResultRec, @@ -102,6 +104,7 @@ class TestTaskAttempt: ], ) def test_roundtrip(self, outcome, pdrr): + env = {"tool": "conda", "packages": {"gufe": "1.10.0"}, "captured_at": None} ta = TaskAttempt( compute_service_id=str(CSID), hostname="h", @@ -112,11 +115,13 @@ def test_roundtrip(self, outcome, pdrr): units_completed=1 if outcome else None, units_total=4 if outcome else None, protocoldagresultref=pdrr, + environment=env if outcome is not None else None, ) ta2 = TaskAttempt.from_dict(ta.to_dict()) assert ta2.compute_service_id == str(CSID) assert ta2.outcome is outcome assert ta2.protocoldagresultref == pdrr + assert ta2.environment == (env if outcome is not None else None) assert ta2.datetime_claimed == NOW @@ -306,3 +311,52 @@ def test_gufe_roundtrip(self): assert purr2.has_logs is True assert purr2.start_time == NOW assert str(purr2.key) == str(key1) + + +class TestComputeEnvironment: + CAP = { + "tool": "conda", + "packages": {"gufe": "1.10.0", "python": "3.11.9"}, + "captured_at": "2026-07-20T00:00:00+00:00", + } + + def test_content_hash_order_independent(self): + h1 = ComputeEnvironment.content_hash("conda", {"a": "1", "b": "2"}) + h2 = ComputeEnvironment.content_hash("conda", {"b": "2", "a": "1"}) + assert h1 == h2 + + def test_content_hash_distinguishes_tool_and_versions(self): + base = ComputeEnvironment.content_hash("conda", {"a": "1"}) + assert base != ComputeEnvironment.content_hash("pip", {"a": "1"}) + assert base != ComputeEnvironment.content_hash("conda", {"a": "2"}) + + def test_from_capture(self): + ce = ComputeEnvironment.from_capture(self.CAP) + assert ce.tool == "conda" + assert ce.packages == self.CAP["packages"] + assert ce.hash == ComputeEnvironment.content_hash("conda", self.CAP["packages"]) + assert ce.captured_at == datetime.datetime.fromisoformat( + self.CAP["captured_at"] + ) + + def test_to_capture_dict_roundtrip(self): + ce = ComputeEnvironment.from_capture(self.CAP) + assert ce.to_capture_dict() == self.CAP + + def test_from_node(self): + ce = ComputeEnvironment.from_capture(self.CAP) + + class FakeNode(dict): + def get(self, k, d=None): + return super().get(k, d) + + node = FakeNode( + hash=ce.hash, + tool="conda", + packages=json.dumps(ce.packages), + captured_at=self.CAP["captured_at"], + ) + ce2 = ComputeEnvironment.from_node(node) + assert ce2.hash == ce.hash + assert ce2.packages == ce.packages + assert ce2.tool == "conda" diff --git a/docs/compute.rst b/docs/compute.rst index b24fab18..2de4ba5b 100644 --- a/docs/compute.rst +++ b/docs/compute.rst @@ -204,6 +204,11 @@ All of them have sensible defaults, so you only need to set them to change the d This is the ``hostname`` surfaced through :py:meth:`~alchemiscale.interface.client.AlchemiscaleClient.get_task_history` and :py:meth:`~alchemiscale.interface.client.AlchemiscaleClient.get_tasks_details`. If unset (``null``), the service uses ``socket.gethostname()``. +``capture_environment`` + If ``true`` (the default), the service captures its software environment (package versions) once at startup, trying ``micromamba``, ``mamba``, ``conda``, then ``pip`` and taking the first that succeeds. + The environment each ``Task`` execution attempt ran in is recorded in durable provenance (deduplicated across services) and surfaced on the :py:class:`~alchemiscale.storage.models.TaskAttempt` records returned by :py:meth:`~alchemiscale.interface.client.AlchemiscaleClient.get_task_history`. + Best-effort: if no package manager is available, no environment is recorded. + ``capture_streams`` If ``true`` (the default), each :external+gufe:py:class:`~gufe.protocols.protocolunit.ProtocolUnit`\'s :external+gufe:py:class:`~gufe.protocols.protocolunit.Context` is constructed with per-attempt stdout/stderr directories, so ``gufe``'s native per-unit stream-capture mechanism archives whatever the :external+gufe:py:class:`~gufe.protocols.protocol.Protocol` directs into them. This is *protocol opt-in*: the compute service only provides the capture directories, and each :external+gufe:py:class:`~gufe.protocols.protocol.Protocol` chooses what, if anything, to write there. diff --git a/docs/operations.rst b/docs/operations.rst index 7c18b2b6..7129e90b 100644 --- a/docs/operations.rst +++ b/docs/operations.rst @@ -171,3 +171,18 @@ Migrate schema from ``alchemiscale`` 0.3 to 0.4 4. Shut down the ``neo4j`` service (``Ctrl+C`` of running instance in step 2), then bring up the full set of services:: USER_ID=$(id -u) GROUP_ID=$(id -g) docker-compose up -d + + +Migrate schema from ``alchemiscale`` 0.7 to 0.8 +----------------------------------------------- +``alchemiscale`` 0.8 introduces durable Task execution provenance, including a +record of the software environment each compute service runs in. +This requires a lightweight, idempotent schema migration that adds a +``neo4j`` uniqueness constraint for the new ``ComputeEnvironment`` node label; +no data migration is required. + +Perform the schema migration against your running deployment:: + + docker run --rm -it --network alchemiscale-server_db -e NEO4J_URL=bolt://neo4j:7687 -e NEO4J_USER= -e NEO4J_PASS= \ + ghcr.io/openforcefield/alchemiscale-server:v0.8.0 \ + database migrate v07-to-v08 diff --git a/docs/user_guide/introspection.rst b/docs/user_guide/introspection.rst index 68723530..2c1ecac8 100644 --- a/docs/user_guide/introspection.rst +++ b/docs/user_guide/introspection.rst @@ -34,6 +34,7 @@ Each :py:class:`~alchemiscale.storage.models.TaskAttempt` records: * ``outcome`` — one of ``complete``, ``error``, ``expired`` (the compute service lost its registration before producing a result), or ``released`` (you forced the :py:class:`~alchemiscale.storage.models.Task` to another status before it finished) * ``units_completed`` and ``units_total`` — how far the attempt progressed through its :external+gufe:py:class:`~gufe.protocols.protocolunit.ProtocolUnit`\s * ``protocoldagresultref`` — the :py:class:`~alchemiscale.models.ScopedKey` of the :external+gufe:py:class:`~gufe.protocols.protocoldag.ProtocolDAGResult` the attempt produced, where one exists (``expired`` and ``released`` attempts have none) +* ``environment`` — the software environment the attempt ran in, as ``{"tool": ..., "packages": {name: version, ...}, "captured_at": ...}``, captured by the compute service at startup (``None`` if the service did not capture one; see :ref:`compute`) You can limit the history to the most recent attempts with the ``limit`` keyword argument:: diff --git a/news/issue-106.rst b/news/issue-106.rst index 22f560ff..2928ec61 100644 --- a/news/issue-106.rst +++ b/news/issue-106.rst @@ -2,6 +2,7 @@ * Durable per-attempt execution provenance: each ``Task`` execution attempt is now recorded as a ``TaskProvenance`` record, capturing details such as the compute service that claimed it and when. * Compute services now register with a ``hostname``, recorded alongside their execution provenance. +* Compute services capture their software environment (package versions, via ``micromamba``/``mamba``/``conda``/``pip``) at startup; each ``Task`` execution attempt records the environment it ran in (deduplicated across services), surfaced in ``AlchemiscaleClient.get_task_history``. * ``AlchemiscaleClient.get_task_history`` returns the full per-attempt history of a ``Task``, and ``AlchemiscaleClient.get_tasks_details`` returns detailed per-``Task`` information. **Changed:**