Capture compute execution environment in Task provenance (#106)

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

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:

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)
 

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": "<iso>"}
+
+    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

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):

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=(

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
+        """)

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"),
         )