From 4dbbbecbbe4097b6464a6ad2a9467e18fa87bf83 Mon Sep 17 00:00:00 2001 From: David Dotson Date: Thu, 9 Jul 2026 21:35:06 -0600 Subject: [PATCH 01/18] Add Task execution and failure introspection (v0.8.0) Implements the v0.8.0 introspection design (milestone 9): durable execution provenance, live progress reporting, failure introspection (tracebacks, per-unit logs, stdout/stderr), and compute share reporting. Addresses #106, #211, #347, #195, #415, #295, #349, #389. - Durable provenance: an immutable TaskProvenance record per execution attempt (created at claim, finalized at result/expiry/deregistration/ release, race-safe against late results), hostname on compute service registrations, and Task.datetime_status_changed / Task.reason indicators with centralized status writes across every mutation site. - Live progress: an alchemiscale-owned execute_DAG mirroring gufe's semantics plus unit-attempt hooks (guarded by a gufe-equivalence suite), driving event-driven, fire-and-forget progress reporting. - Failure introspection: fast bulk tracebacks; per-unit log capture over the gufekey logger namespace plus gufe-native stdout/stderr capture; ProtocolUnitResultRef nodes and artifact retrieval; and graceful ProtocolDAG creation-failure handling (Task -> error with reason rather than killing the compute service). - Compute share: get_scope_compute_share. - Client methods, API routes, ComputeServiceSettings, a v04_to_v05 index migration, Sphinx docs, news fragments, and unit + integration tests. Co-Authored-By: Claude Opus 4.8 (1M context) --- alchemiscale/cli.py | 23 + alchemiscale/compute/api.py | 147 ++- alchemiscale/compute/capture.py | 144 +++ alchemiscale/compute/client.py | 88 +- alchemiscale/compute/execute.py | 257 +++++ alchemiscale/compute/service.py | 137 ++- alchemiscale/compute/settings.py | 48 + alchemiscale/interface/api.py | 274 +++++- alchemiscale/interface/client.py | 299 +++++- alchemiscale/migrations/v04_to_v05.py | 52 + alchemiscale/storage/models.py | 517 +++++++++- alchemiscale/storage/objectstore.py | 88 ++ alchemiscale/storage/statestore.py | 921 ++++++++++++++++-- .../storage/test_statestore_introspection.py | 662 +++++++++++++ alchemiscale/tests/unit/compute/__init__.py | 0 .../tests/unit/compute/test_capture.py | 331 +++++++ .../unit/compute/test_execute_equivalence.py | 759 +++++++++++++++ docs/compute.rst | 52 + docs/user_guide/handling_errors.rst | 106 ++ docs/user_guide/index.rst | 1 + docs/user_guide/introspection.rst | 109 +++ news/issue-106.rst | 9 + news/issue-195.rst | 3 + news/issue-211.rst | 3 + news/issue-295.rst | 4 + news/issue-347.rst | 3 + news/issue-349.rst | 3 + news/issue-389.rst | 3 + news/issue-415.rst | 7 + 29 files changed, 4964 insertions(+), 86 deletions(-) create mode 100644 alchemiscale/compute/capture.py create mode 100644 alchemiscale/compute/execute.py create mode 100644 alchemiscale/migrations/v04_to_v05.py create mode 100644 alchemiscale/tests/integration/storage/test_statestore_introspection.py create mode 100644 alchemiscale/tests/unit/compute/__init__.py create mode 100644 alchemiscale/tests/unit/compute/test_capture.py create mode 100644 alchemiscale/tests/unit/compute/test_execute_equivalence.py create mode 100644 docs/user_guide/introspection.rst create mode 100644 news/issue-106.rst create mode 100644 news/issue-195.rst create mode 100644 news/issue-211.rst create mode 100644 news/issue-295.rst create mode 100644 news/issue-347.rst create mode 100644 news/issue-349.rst create mode 100644 news/issue-389.rst create mode 100644 news/issue-415.rst diff --git a/alchemiscale/cli.py b/alchemiscale/cli.py index c87caa5e..d49fdd7d 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 v04_to_v05(url, user, password, dbname): + """Perform migration appropriate for transitioning from alchemiscale v0.4 + to v0.5. + + 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.v04_to_v05 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 f5fd5d95..51247346 100644 --- a/alchemiscale/compute/api.py +++ b/alchemiscale/compute/api.py @@ -5,6 +5,7 @@ """ import json +import os import datetime from datetime import timedelta import random @@ -12,7 +13,7 @@ from fastapi import FastAPI, APIRouter, Body, Depends, HTTPException, Request from fastapi import status as http_status from fastapi.middleware.gzip import GZipMiddleware -from gufe.tokenization import JSON_HANDLER +from gufe.tokenization import JSON_HANDLER, GufeKey from gufe.protocols import ProtocolDAGResult from ..base.api import ( @@ -107,6 +108,7 @@ def register_computeservice( compute_service_id, *, compute_manager_id: str | None = Body(None, embed=True), + hostname: str | None = Body(None, embed=True), n4js: Neo4jStore = Depends(get_n4js_depends), ): now = datetime.datetime.now(tz=datetime.UTC) @@ -121,6 +123,7 @@ def register_computeservice( heartbeat=now, failure_times=[], manager_name=manager_name, + hostname=hostname, ) try: @@ -400,11 +403,38 @@ async def set_task_result( creator=compute_service_id, ) - # push the reference to the state store + # push the reference to the state store; this also finalizes the open + # TaskProvenance attempt for this (task, compute_service_id) pair with the + # appropriate outcome and links it to the new ProtocolDAGResultRef result_sk: ScopedKey = n4js.set_task_result( - task=task_sk, protocoldagresultref=protocoldagresultref + task=task_sk, + protocoldagresultref=protocoldagresultref, + compute_service_id=ComputeServiceID(compute_service_id), ) + # derive one ProtocolUnitResultRef per unit result, and extract any embedded + # stdout/stderr from the (already-deserialized) ProtocolDAGResult into + # per-unit, retrieval-optimized artifacts --- flipping has_stdout/has_stderr. + # Streams need no new compute-facing routes: they ride inside the PDR blob. + refs_map = n4js.add_protocol_unit_result_refs(protocoldagresultref, result_sk, pdr) + if protocoldagresultref.location: + base_location = os.path.dirname(protocoldagresultref.location) + for unit_result in pdr.protocol_unit_results: + purr_sk = refs_map.get(unit_result.key) + if purr_sk is None: + continue + unit_location = os.path.join(base_location, "units", str(unit_result.key)) + if unit_result.stdout: + s3os.push_protocol_unit_result_streams( + unit_location, "stdout", unit_result.stdout + ) + n4js.set_protocol_unit_result_ref_artifacts(purr_sk, has_stdout=True) + if unit_result.stderr: + s3os.push_protocol_unit_result_streams( + unit_location, "stderr", unit_result.stderr + ) + n4js.set_protocol_unit_result_ref_artifacts(purr_sk, has_stderr=True) + # if success, set task complete, remove from all hubs # otherwise, set as errored, leave in hubs if protocoldagresultref.ok: @@ -413,6 +443,8 @@ async def set_task_result( n4js.add_protocol_dag_result_ref_tracebacks( pdr.protocol_unit_failures, result_sk ) + # provenance already finalized by set_task_result above, so no + # compute_service_id passed here (avoids a redundant finalization) n4js.set_task_error(tasks=[task_sk]) # report that the compute service experienced a failure @@ -423,6 +455,115 @@ async def set_task_result( return result_sk +@router.post("/tasks/{task_scoped_key}/error", response_model=str) +async def set_task_error( + task_scoped_key, + *, + reason: str = Body(..., embed=True), + compute_service_id: str = Body(..., embed=True), + n4js: Neo4jStore = Depends(get_n4js_depends), + token: TokenData = Depends(get_token_data_depends), +): + """Set a Task to `error` with a `reason`, for `ProtocolDAG` creation failures. + + No `ProtocolDAGResult` exists in this case (creation failed before any unit + ran), which is exactly why `reason` lives on the `Task`. The open + `TaskProvenance` attempt for this ``(task, compute_service_id)`` pair is + finalized with `outcome = error`. + + `ProtocolDAG` creation failures deliberately do not participate in restart + policies: `resolve_task_restarts` matches against `Tracebacks` nodes only, + and none exist here, so the Task's restarts are cancelled rather than + retried --- these tend to be systematic problems with the `Transformation` + itself, which auto-retrying would only mask. + """ + task_sk = ScopedKey.from_str(task_scoped_key) + validate_scopes(task_sk.scope, token) + + n4js.set_task_error( + tasks=[task_sk], + reason=reason, + compute_service_id=ComputeServiceID(compute_service_id), + ) + + # no Tracebacks node exists, so this cancels (does not renew) the Task's + # restart patterns + n4js.resolve_task_restarts(task_scoped_keys=[task_sk]) + + return str(task_sk) + + +@router.post("/computeservice/{compute_service_id}/progress") +def update_task_progress( + compute_service_id, + *, + progress: dict[str, dict[str, int]] = Body(..., embed=True), + n4js: Neo4jStore = Depends(get_n4js_depends), +): + """Record live progress counts for a service's claimed Tasks. + + Body maps `Task` ScopedKey strings to + ``{"units_completed": int, "units_total": int}`` --- one batched request per + push event, regardless of claim count. Like a heartbeat, this route never + rejects: updates for Tasks the service no longer claims are silently dropped + server-side (the claim expired mid-flight). + """ + progress_ = { + task_sk: (counts["units_completed"], counts["units_total"]) + for task_sk, counts in progress.items() + } + n4js.update_task_progress(ComputeServiceID(compute_service_id), progress_) + return None + + +@router.post( + "/tasks/{task_scoped_key}/results/{protocoldagresultref_scoped_key}/units/{unit_result_key}/artifacts/logs" +) +def set_unit_result_logs( + task_scoped_key, + protocoldagresultref_scoped_key, + unit_result_key, + *, + logs: str = Body(..., embed=True), + n4js: Neo4jStore = Depends(get_n4js_depends), + s3os: S3ObjectStore = Depends(get_s3os_depends), + token: TokenData = Depends(get_token_data_depends), +): + """Upload captured log text for a single unit result, flipping `has_logs`. + + Ordered after streams and unit-ref creation (which happen in + `set_task_result`), so a service dying mid-upload leaves consistent state: + refs and streams exist, missing logs are simply flagged absent. + """ + task_sk = ScopedKey.from_str(task_scoped_key) + pdrr_sk = ScopedKey.from_str(protocoldagresultref_scoped_key) + # authorize BOTH the Task and the ProtocolDAGResultRef scopes, and require + # the ref to actually be a result of the Task --- otherwise a caller + # credentialed for one scope could flip flags / overwrite artifacts on a + # ref in another scope by pairing it with an in-scope Task path parameter + validate_scopes(task_sk.scope, token) + validate_scopes(pdrr_sk.scope, token) + + purr_sk = n4js.get_protocol_unit_result_ref_scoped_key( + pdrr_sk, GufeKey(unit_result_key), task=task_sk + ) + if purr_sk is None: + raise HTTPException( + status_code=http_status.HTTP_404_NOT_FOUND, + detail=( + f"No ProtocolUnitResultRef for unit result '{unit_result_key}' " + f"under '{protocoldagresultref_scoped_key}' for task " + f"'{task_scoped_key}'" + ), + ) + + protocolunitresultref = n4js.get_gufe(purr_sk) + s3os.push_protocol_unit_result_logs(protocolunitresultref.location, logs) + n4js.set_protocol_unit_result_ref_artifacts(purr_sk, has_logs=True) + + return str(purr_sk) + + def process_compute_manager_id_string( compute_manager_id_string: str, ) -> ComputeManagerID: diff --git a/alchemiscale/compute/capture.py b/alchemiscale/compute/capture.py new file mode 100644 index 00000000..0917c897 --- /dev/null +++ b/alchemiscale/compute/capture.py @@ -0,0 +1,144 @@ +""" +:mod:`alchemiscale.compute.capture` --- per-unit log capture and progress hooks +================================================================================ + +Execution hooks for :class:`~alchemiscale.compute.service.SynchronousComputeService` +that plug into the alchemiscale executor (:mod:`alchemiscale.compute.execute`) to: + +- capture log records emitted through ``gufe``'s sanctioned logging channel + (:attr:`gufe.tokenization.GufeTokenizable.logger`, the ``gufekey.{module}.{qualname}`` + namespace), scoped to a single unit attempt, so attribution is exact per + `ProtocolUnitResult`/`ProtocolUnitFailure`; +- push live progress counts, fire-and-forget, at unit boundaries. + +We deliberately capture *only* the ``gufekey`` namespace --- what protocols emit +through ``ProtocolUnit.logger`` --- not third-party library loggers (OpenMM, +openff-toolkit, RDKit, ...), whose volume is unbounded and uncurated. A protocol +wanting library logs kept routes them through a ``Context.stdout``/``stderr`` +`FileHandler` instead (the stream channel). +""" + +import logging +import time + +from gufe.protocols.protocolunit import ProtocolUnit, ProtocolUnitResult + +from ..models import ScopedKey +from ..storage.models import ComputeServiceID +from .execute import ExecutionHooks + +# the logger namespace every GufeTokenizable.logger writes to; hierarchy +# propagation delivers all descendants to a handler attached here +GUFEKEY_LOGGER_NAME = "gufekey" + + +class GufeKeyLogHandler(logging.Handler): + """A `logging.Handler` that accumulates formatted, timestamped log lines. + + Attached to the ``gufekey`` logger for the duration of a single unit + attempt. Each record carries ``record.gufekey`` (the emitting unit's gufe + key), which the formatter includes for labeling and sanity-checking. + """ + + def __init__(self, level: int | str = logging.NOTSET): + super().__init__(level=level) + # `defaults` covers records that reach the `gufekey` logger without the + # `GufeTokenizable.logger` adapter's `record.gufekey` stamp (the + # sanctioned channel always sets it, but we must not raise on a record + # that doesn't). + formatter = logging.Formatter( + "[%(asctime)s] [%(gufekey)s] [%(levelname)s] %(message)s", + defaults={"gufekey": "-"}, + ) + formatter.converter = time.gmtime # UTC timestamps + self.setFormatter(formatter) + self.lines: list[str] = [] + + def emit(self, record: logging.LogRecord) -> None: + try: + self.lines.append(self.format(record)) + except Exception: # pragma: no cover - defensive, mirrors logging + self.handleError(record) + + def text(self) -> str: + return "\n".join(self.lines) + + +class SynchronousExecutionHooks(ExecutionHooks): + """Execution hooks binding one `Task`'s DAG execution to log capture and + progress reporting. + + Parameters + ---------- + task + The `Task` whose DAG is being executed (for the progress payload). + compute_service_id + The claiming compute service (for the progress payload). + progress_callback + A fire-and-forget callable ``(task, units_completed, units_total)`` that + pushes progress; it must not raise or block (the service wraps its own + transport in a swallow-and-log). + gufekey_loglevel + The level to which the ``gufekey`` logger is set so protocol ``INFO`` + logs reach the capture handler (the logger otherwise inherits the root + ``WARNING``). Applied once, when the hooks are constructed. + log_cap_bytes + Per-unit-result cap on captured log text; the tail (where errors live) + is kept. + """ + + def __init__( + self, + *, + task: ScopedKey, + compute_service_id: ComputeServiceID, + progress_callback, + capture_logs: bool = True, + gufekey_loglevel: int | str = logging.INFO, + log_cap_bytes: int = 1_048_576, + ): + self.task = task + self.compute_service_id = compute_service_id + self.progress_callback = progress_callback + self.capture_logs = capture_logs + self.log_cap_bytes = log_cap_bytes + + # captured log text per unit-result gufe key, ready for upload + self.unit_logs: dict[str, str] = {} + + self._logger = logging.getLogger(GUFEKEY_LOGGER_NAME) + # ensure protocol logs at gufekey_loglevel reach the handler; only mutate + # the (process-global) logger level when capture is actually enabled + if self.capture_logs: + self._logger.setLevel(gufekey_loglevel) + self._handler: GufeKeyLogHandler | None = None + + def on_unit_attempt_start(self, unit: ProtocolUnit, attempt: int) -> None: + if not self.capture_logs: + return + # open a fresh capture handler scoped to this single unit attempt + self._handler = GufeKeyLogHandler() + self._logger.addHandler(self._handler) + + def on_unit_attempt_end( + self, + unit: ProtocolUnit, + attempt: int, + result: ProtocolUnitResult | None, + ) -> None: + if self._handler is None: + return + self._logger.removeHandler(self._handler) + if result is not None: + text = self._handler.text() + if text: + data = text.encode("utf-8") + if len(data) > self.log_cap_bytes: + # keep the tail + data = data[-self.log_cap_bytes :] + text = data.decode("utf-8", errors="replace") + self.unit_logs[str(result.key)] = text + self._handler = None + + def on_progress(self, units_completed: int, units_total: int) -> None: + self.progress_callback(self.task, units_completed, units_total) diff --git a/alchemiscale/compute/client.py b/alchemiscale/compute/client.py index 532b329f..1c1edadf 100644 --- a/alchemiscale/compute/client.py +++ b/alchemiscale/compute/client.py @@ -5,10 +5,15 @@ """ +import json +from urllib.parse import urljoin + +import requests import zstandard as zstd from gufe import Transformation from gufe.protocols import ProtocolDAGResult +from gufe.tokenization import JSON_HANDLER from ..base.client import ( AlchemiscaleBaseClient, @@ -39,10 +44,11 @@ def register( self, compute_service_id: ComputeServiceID, compute_manager_id: ComputeManagerID | None = None, + hostname: str | None = None, ): res = self._post_resource( f"/computeservice/{compute_service_id}/register", - {"compute_manager_id": compute_manager_id}, + {"compute_manager_id": compute_manager_id, "hostname": hostname}, ) return ComputeServiceID(res) @@ -165,6 +171,86 @@ def set_task_result( return ScopedKey.from_dict(pdr_sk) + def set_task_error( + self, + task: ScopedKey, + reason: str, + compute_service_id: ComputeServiceID | None = None, + ) -> ScopedKey: + """Set a `Task` to `error` with a human-readable `reason`. + + Used by the compute service when `ProtocolDAG` creation fails, where + there is no `ProtocolDAGResult` to submit. The `reason` (a traceback) + is stored on `Task.reason` and the open `TaskProvenance` attempt is + finalized with `outcome = error`. + """ + data = dict(reason=reason, compute_service_id=str(compute_service_id)) + task_sk = self._post_resource(f"/tasks/{task}/error", data) + return ScopedKey.from_str(task_sk) + + def update_task_progress( + self, + compute_service_id: ComputeServiceID, + progress: dict[str, dict[str, int]], + timeout: float = 5.0, + ) -> None: + """Push live progress counts for this service's claimed Tasks. + + `progress` maps `Task` ScopedKey strings to + ``{"units_completed": int, "units_total": int}``. One batched request + per push event. + + This is **fire-and-forget**: a single attempt with a short timeout and + NO retry/backoff (the usual retry machinery would stall the DAG between + units on a flaky API). Progress is best-effort telemetry, refreshed at + the next unit boundary. Errors propagate to the caller, which is + expected to log-and-continue. + + This deliberately does **not** fetch or refresh a JWT: `_get_token` + issues its request with ``timeout=None`` and would block the execution + thread indefinitely if the token endpoint hangs. A token is virtually + always already present by the time a DAG executes (registration/claim + obtained one); if it is missing or stale, the push is simply skipped and + the retrying transport (result push, next claim) refreshes it before the + next boundary. + """ + if self._jwtoken is None: + # no token yet and we won't block to get one; skip this push + return + + url = urljoin(self.api_url, f"/computeservice/{compute_service_id}/progress") + jsondata = json.dumps(progress, cls=JSON_HANDLER.encoder) + + resp = requests.post( + url, + data=jsondata, + headers=self._headers, + timeout=timeout, + verify=self.verify, + ) + if not 200 <= resp.status_code < 300: + raise self._exception( + f"Status Code {resp.status_code} : {resp.reason}", + status_code=resp.status_code, + ) + + def set_task_result_unit_logs( + self, + task: ScopedKey, + protocoldagresultref: ScopedKey, + unit_result_key: str, + logs: str, + ) -> None: + """Upload captured log text for a single unit result. + + Uses the normal (retrying) transport --- this happens after DAG + execution, not between units, so a retry cannot stall execution. + """ + self._post_resource( + f"/tasks/{task}/results/{protocoldagresultref}/units/{unit_result_key}/artifacts/logs", + {"logs": logs}, + ) + class AlchemiscaleComputeManagerClientError(AlchemiscaleBaseClientError): ... diff --git a/alchemiscale/compute/execute.py b/alchemiscale/compute/execute.py new file mode 100644 index 00000000..ee8786d1 --- /dev/null +++ b/alchemiscale/compute/execute.py @@ -0,0 +1,257 @@ +""" +:mod:`alchemiscale.compute.execute` --- alchemiscale-owned DAG executor +======================================================================= + +An alchemiscale-maintained executor for :class:`~gufe.protocols.ProtocolDAG`\\ s, +used by :class:`~alchemiscale.compute.service.SynchronousComputeService`. + +It mirrors the execution semantics of :func:`gufe.protocols.protocoldag.execute_DAG` +--- topological iteration, per-attempt :class:`~gufe.protocols.protocolunit.Context` +construction (including stream directories), retry behavior, all-attempts result +accumulation, halt on persistent failure, ``KeyboardInterrupt``/``ExecutionInterrupt`` +pass-through, and ``gufe``-compatible unit-result caching --- while adding +**unit-attempt start/end hooks**. Those hooks are the single seam that exact +progress reporting and per-attempt log capture plug into. + +This is a deliberate direction choice: ``SynchronousComputeService`` is the +reference implementation *for alchemiscale compute services*, and it should +support every behavior we want alchemiscale services to have. The fork risk --- +``gufe``'s executor keeps evolving --- is mitigated by a behavioral-equivalence +test suite (``tests/.../test_execute_equivalence.py``) that runs identical DAGs +through both implementations and asserts equivalent ``ProtocolDAGResult``\\ s, +exercised on every ``gufe`` upgrade. + +The ``gufe`` version this mirrors is pinned in ``pyproject.toml``. ``gufe``'s +private input-mapping and cache-validation helpers (``_pu_to_pur``, +``_get_valid_unit_results``) are reused directly so that resume/cache semantics +cannot silently diverge; if a future ``gufe`` removes or changes them, the +equivalence suite fails loudly. +""" + +import shutil +import warnings +from json import JSONDecodeError +from pathlib import Path + +from gufe.protocols.protocoldag import ( + ProtocolDAG, + ProtocolDAGResult, + _get_valid_unit_results, + _pu_to_pur, +) +from gufe.protocols.protocolunit import ( + Context, + ProtocolUnit, + ProtocolUnitResult, +) +from gufe.tokenization import GufeKey + + +class ExecutionHooks: + """Hooks invoked by :func:`execute_DAG` at DAG and unit-attempt boundaries. + + Subclass and override the methods of interest; every default is a no-op, so + a bare ``ExecutionHooks()`` reproduces plain ``gufe`` execution semantics. + + The hooks fire in the execution thread and must be cheap and non-blocking: + progress pushes are fire-and-forget, and log-capture open/close must not + stall DAG execution. + """ + + def on_dag_start(self, protocoldag: ProtocolDAG, units_total: int) -> None: + """Called once, before any unit executes.""" + + def on_unit_attempt_start(self, unit: ProtocolUnit, attempt: int) -> None: + """Called immediately before a single unit-attempt executes. + + The log-capture handler for this attempt should be opened here, so that + records emitted during the attempt are attributed to its result. + """ + + def on_unit_attempt_end( + self, + unit: ProtocolUnit, + attempt: int, + result: ProtocolUnitResult | None, + ) -> None: + """Called after a single unit-attempt executes (or is interrupted). + + Always called exactly once per ``on_unit_attempt_start``, even if the + attempt raised (in which case ``result`` is ``None``) --- so the + log-capture handler is guaranteed to be closed. When ``result`` is not + ``None``, captured logs should be associated with ``result.key`` for + later upload. + """ + + def on_progress(self, units_completed: int, units_total: int) -> None: + """Called with the running count of distinct successfully completed + units against the DAG's total unit count. + + Fired once at DAG start with ``(0, units_total)`` --- so the denominator + is visible immediately --- and again after each unit successfully + completes. + """ + + +def execute_DAG( + protocoldag: ProtocolDAG, + *, + shared_basedir: Path, + scratch_basedir: Path, + cache_basedir: Path | None = None, + stderr_basedir: Path | None = None, + stdout_basedir: Path | None = None, + keep_shared: bool = False, + keep_scratch: bool = False, + keep_cache: bool = False, + raise_error: bool = True, + n_retries: int = 0, + hooks: ExecutionHooks | None = None, +) -> ProtocolDAGResult: + """Locally execute a full :class:`ProtocolDAG` in serial and in-process. + + A behavioral mirror of :func:`gufe.protocols.protocoldag.execute_DAG` with + added unit-attempt hooks. All keyword parameters have identical meaning to + the ``gufe`` function; see its docstring. The only addition is ``hooks``. + + Parameters + ---------- + hooks + :class:`ExecutionHooks` invoked at DAG start and at each unit-attempt + start/end, plus progress updates. Defaults to a no-op set, in which + case execution is semantically identical to ``gufe``'s. + + Raises + ------ + ProtocolDAGExecutionError + If the ``ProtocolDAG`` cannot be executed due to an invalid cache state. + """ + if n_retries < 0: + raise ValueError("Must give positive number of retries") + + if hooks is None: + hooks = ExecutionHooks() + + # `protocol_units` is in DAG-dependency order + units_total = len(protocoldag.protocol_units) + hooks.on_dag_start(protocoldag, units_total) + # publish the denominator immediately (0 of N complete) + hooks.on_progress(0, units_total) + + # load any cached unit results (disabled by SynchronousComputeService, but + # implemented for gufe equivalence) + all_cached_results: list[ProtocolUnitResult] = [] + if cache_basedir is not None: + dag_unitresults_cache = cache_basedir / f"{str(protocoldag.key)}-results_cache" + dag_unitresults_cache.mkdir(exist_ok=True, parents=True) + + for file in dag_unitresults_cache.rglob("*.json"): + try: + unit_result = ProtocolUnitResult.from_json(file) + except JSONDecodeError as e: + warnings.warn(f"Unable to read file, skipping {file}: {e}") + else: + all_cached_results.append(unit_result) + + # handle results & optionally caching + results: dict[GufeKey, ProtocolUnitResult] = _get_valid_unit_results( + protocoldag, all_cached_results + ) + all_results = [] # successes AND failures + shared_paths = [] + for unit in protocoldag.protocol_units: + # If we already have results (from cache), skip execution + if unit.key in results: + all_results.append(results[unit.key]) + continue + + # translate each `ProtocolUnit` in input into corresponding `ProtocolUnitResult` + inputs = _pu_to_pur(unit.inputs, results) + + attempt = 0 + while attempt <= n_retries: + shared = shared_basedir / f"shared_{str(unit.key)}_attempt_{attempt}" + shared_paths.append(shared) + shared.mkdir(exist_ok=True) + + scratch = scratch_basedir / f"scratch_{str(unit.key)}_attempt_{attempt}" + scratch.mkdir(exist_ok=True) + + stderr = None + if stderr_basedir: + stderr = stderr_basedir / f"stderr_{str(unit.key)}_attempt_{attempt}" + stderr.mkdir(exist_ok=True) + + stdout = None + if stdout_basedir: + stdout = stdout_basedir / f"stdout_{str(unit.key)}_attempt_{attempt}" + stdout.mkdir(exist_ok=True) + + context = Context( + shared=shared, scratch=scratch, stderr=stderr, stdout=stdout + ) + + # execute this unit-attempt, guaranteeing the end hook fires exactly + # once (so log capture is always closed). KeyboardInterrupt and + # gufe ExecutionInterrupt derive from BaseException and propagate. + # Contract: hook methods must not raise --- a hook exception in this + # `finally` would replace an in-flight interrupt/error being + # propagated. The supplied `ExecutionHooks` implementations honor + # this (they only mutate in-memory state and logging handlers). + hooks.on_unit_attempt_start(unit, attempt) + result = None + try: + result = unit.execute( + context=context, raise_error=raise_error, **inputs + ) + all_results.append(result) + finally: + hooks.on_unit_attempt_end(unit, attempt, result) + + # clean up outputs + if stderr: + shutil.rmtree(stderr) + if stdout: + shutil.rmtree(stdout) + + if not keep_scratch: + shutil.rmtree(scratch) + + if result.ok(): + # attach result to this `ProtocolUnit` + results[unit.key] = result + + # Serialize results if requested + if cache_basedir is not None: + result.to_json( + dag_unitresults_cache / f"{str(unit.key)}_unitresults.json" + ) + + # progress: one more distinct unit successfully completed. + # NOTE: `results` includes any cache-resumed entries, so under + # caching (disabled by SynchronousComputeService) the count + # would include pre-completed units. Harmless while caching is + # off; revisit when #180-style resume lands. + hooks.on_progress(len(results), units_total) + break + attempt += 1 + + if not result.ok(): + # persistent failure halts DAG execution; downstream units yield no + # results at all, freezing progress at the point of failure + break + + if not keep_shared: + for shared_path in shared_paths: + shutil.rmtree(shared_path) + + if not keep_cache and cache_basedir is not None: + shutil.rmtree(dag_unitresults_cache) + + return ProtocolDAGResult( + name=protocoldag.name, + protocol_units=protocoldag.protocol_units, + protocol_unit_results=all_results, + transformation_key=protocoldag.transformation_key, + extends_key=protocoldag.extends_key, + ) diff --git a/alchemiscale/compute/service.py b/alchemiscale/compute/service.py index 29b9b96a..479463dd 100644 --- a/alchemiscale/compute/service.py +++ b/alchemiscale/compute/service.py @@ -6,17 +6,21 @@ from contextlib import contextmanager import gc +import socket import time import logging +import traceback from uuid import uuid4 import threading from pathlib import Path import shutil from gufe import Transformation -from gufe.protocols.protocoldag import execute_DAG, ProtocolDAG, ProtocolDAGResult +from gufe.protocols.protocoldag import ProtocolDAG, ProtocolDAGResult from .client import AlchemiscaleComputeClient +from .execute import execute_DAG +from .capture import SynchronousExecutionHooks from .settings import ComputeServiceSettings from ..storage.models import ComputeServiceID from ..models import Scope, ScopedKey @@ -77,6 +81,10 @@ def __init__(self, settings: ComputeServiceSettings): self.compute_service_id = ComputeServiceID.new_from_name(self.name) + # hostname is copied onto the registration and every TaskProvenance this + # service creates; falls back to the OS hostname when not set + self.hostname = self.settings.hostname or socket.gethostname() + # 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 @@ -110,7 +118,11 @@ def __init__(self, settings: ComputeServiceSettings): def _register(self): """Register this compute service with the compute API.""" - self.client.register(self.compute_service_id, self.settings.compute_manager_id) + self.client.register( + self.compute_service_id, + self.settings.compute_manager_id, + hostname=self.hostname, + ) def _deregister(self): """Deregister this compute service with the compute API.""" @@ -202,15 +214,52 @@ def push_result( return sk - def execute(self, task: ScopedKey) -> ScopedKey: + def execute(self, task: ScopedKey) -> ScopedKey | None: """Executes given Task. - Returns ScopedKey of ProtocolDAGResultRef following push to database. + Returns ScopedKey of ProtocolDAGResultRef following push to database, + or ``None`` if `ProtocolDAG` creation failed and the Task was errored. """ # obtain a ProtocolDAG from the task self.logger.info("Creating ProtocolDAG from '%s'...", task) - protocoldag, transformation, extends = self.task_to_protocoldag(task) + try: + protocoldag, transformation, extends = self.task_to_protocoldag(task) + except Exception: + # `ProtocolDAG` creation failed --- typically a systematic problem + # with the `Transformation` itself, not a random execution-environment + # failure. Record it as a Task `error` with the traceback as its + # `reason` and continue the cycle with the next Task, rather than + # letting the exception propagate out of the service loop (which + # would deregister the service and silently bounce its Tasks back to + # `waiting`). We catch `Exception` only: `KeyboardInterrupt` and + # `gufe` `ExecutionInterrupt` derive from `BaseException` and are + # deliberately allowed to propagate. + tb = traceback.format_exc() + self.logger.error( + "Failed to create ProtocolDAG from '%s'; setting Task to error:\n%s", + task, + tb, + ) + # best-effort: an unreachable/old server (the /error route 404s) must + # not propagate out of the loop and tear the service down --- that is + # the exact pre-#195 failure this handler exists to prevent. Swallow + # and continue to the next Task, as heartbeats already do. The Task + # stays `running` until its registration expires, which the server + # then handles. + try: + self.client.set_task_error( + task, reason=tb, compute_service_id=self.compute_service_id + ) + except Exception: + self.logger.warning( + "Failed to report ProtocolDAG creation error for '%s'; " + "continuing", + task, + exc_info=True, + ) + return None + self.logger.info( "Created '%s' from '%s' performing '%s'", protocoldag, @@ -226,15 +275,38 @@ def execute(self, task: ScopedKey) -> ScopedKey: scratch = self.scratch_basedir / str(protocoldag.key) scratch.mkdir() + # per-attempt stdout/stderr archiving (gufe's native mechanism); the + # executor creates per-attempt subdirs under these base dirs + if self.settings.capture_streams: + stdout_basedir = scratch / "_stdout" + stdout_basedir.mkdir() + stderr_basedir = scratch / "_stderr" + stderr_basedir.mkdir() + else: + stdout_basedir = stderr_basedir = None + + # hooks: per-unit log capture + event-driven, fire-and-forget progress + hooks = SynchronousExecutionHooks( + task=task, + compute_service_id=self.compute_service_id, + progress_callback=self._push_progress, + capture_logs=self.settings.capture_logs, + gufekey_loglevel=self.settings.gufekey_loglevel, + log_cap_bytes=self.settings.log_cap_bytes, + ) + self.logger.info("Executing '%s'...", protocoldag) try: protocoldagresult = execute_DAG( protocoldag, shared_basedir=shared, scratch_basedir=scratch, + stdout_basedir=stdout_basedir, + stderr_basedir=stderr_basedir, keep_scratch=self.keep_scratch, raise_error=False, n_retries=self.settings.n_retries, + hooks=hooks, ) finally: if not self.keep_shared: @@ -259,8 +331,63 @@ def execute(self, task: ScopedKey) -> ScopedKey: result_sk = self.push_result(task, protocoldagresult) self.logger.info("Pushed result `%s'", protocoldagresult) + # upload any captured per-unit logs; best-effort, so an old server + # (which 404s the artifact route) or a transient failure never fails the + # Task --- the unit refs simply keep `has_logs` false + if self.settings.capture_logs and hooks.unit_logs: + for unit_result_key, logtext in hooks.unit_logs.items(): + try: + self.client.set_task_result_unit_logs( + task, result_sk, unit_result_key, logtext + ) + except Exception: + # a failure here is almost always systemic for this push + # (e.g. an old server 404s the artifact route, or the API is + # down) --- retrying every remaining unit would burn the + # client's whole retry budget per unit. Log once and stop + # uploading logs for this Task; the unit refs simply keep + # `has_logs` false. + self.logger.warning( + "Failed to upload logs for unit result '%s' of task '%s'; " + "skipping remaining log uploads for this Task", + unit_result_key, + task, + exc_info=True, + ) + break + return result_sk + def _push_progress( + self, task: ScopedKey, units_completed: int, units_total: int + ) -> None: + """Fire-and-forget progress push for a running Task. + + Runs in the execution thread at unit boundaries, so failures are + swallowed and logged rather than retried --- best-effort telemetry must + never stall the DAG. Progress plays no liveness role; that remains the + heartbeat's job. + """ + try: + self.client.update_task_progress( + self.compute_service_id, + { + str(task): { + "units_completed": units_completed, + "units_total": units_total, + } + }, + timeout=self.settings.progress_push_timeout, + ) + except Exception: + self.logger.debug( + "Progress push failed for task '%s' (%d/%d); continuing", + task, + units_completed, + units_total, + exc_info=True, + ) + def _check_max_tasks(self, max_tasks): if max_tasks is not None: if self._tasks_counter >= max_tasks: diff --git a/alchemiscale/compute/settings.py b/alchemiscale/compute/settings.py index e5269f90..5105b3de 100644 --- a/alchemiscale/compute/settings.py +++ b/alchemiscale/compute/settings.py @@ -27,6 +27,14 @@ class Config: "resources, e.g. different hosts or HPC clusters." ), ) + hostname: str | None = Field( + None, + description=( + "Hostname to record on this compute service's registration and copy " + "onto every `TaskProvenance` it creates. If `None`, the service uses " + "`socket.gethostname()`." + ), + ) compute_manager_id: str | None = Field( None, description=( @@ -52,6 +60,46 @@ class Config: 3, description="Number of times to attempt a given Task on failure.", ) + capture_streams: bool = Field( + True, + description=( + "If True, construct each `ProtocolUnit`'s `Context` with per-attempt " + "stdout/stderr directories so `gufe`'s native stream-capture " + "mechanism archives whatever the protocol directs into them." + ), + ) + capture_logs: bool = Field( + True, + description=( + "If True, capture log records emitted through `gufe`'s `gufekey` " + "logger namespace (`ProtocolUnit.logger`) per unit result and upload " + "them alongside results." + ), + ) + gufekey_loglevel: str = Field( + "INFO", + description=( + "Level to set the `gufekey` logger to for per-unit log capture; the " + "logger otherwise inherits the root logger's level (typically " + "WARNING), which would drop protocol INFO logs." + ), + ) + log_cap_bytes: int = Field( + 1048576, + description=( + "Per-unit-result cap in bytes on captured log text; the tail (where " + "errors live) is kept. Default 1 MiB." + ), + ) + progress_push_timeout: float = Field( + 5.0, + description=( + "Timeout in seconds for a single fire-and-forget progress push. " + "Progress is best-effort telemetry: pushes are not retried, and " + "failures are logged and swallowed so a flaky API never stalls the " + "DAG between units." + ), + ) sleep_interval: int = Field( 30, description="Time in seconds to sleep if no Tasks claimed from compute API." ) diff --git a/alchemiscale/interface/api.py b/alchemiscale/interface/api.py index 341069e9..07963bfa 100644 --- a/alchemiscale/interface/api.py +++ b/alchemiscale/interface/api.py @@ -4,6 +4,7 @@ """ +import re from collections import Counter from fastapi import FastAPI, APIRouter, Body, Depends, HTTPException, Request @@ -565,6 +566,19 @@ def get_scope_status( return dict(status_counts) +@router.get("/scopes/{scope}/compute-share") +def get_scope_compute_share( + scope, + *, + n4js: Neo4jStore = Depends(get_n4js_depends), + token: TokenData = Depends(get_token_data_depends), +) -> float: + scope_obj = Scope.from_str(scope) + validate_scopes(scope_obj, token) + + return n4js.get_scope_compute_share(scope_obj) + + @router.get("/networks/{network_scoped_key}/status") def get_network_status( network_scoped_key, @@ -885,6 +899,7 @@ def tasks_status_set( *, tasks: list[ScopedKey] = Body(), status: str = Body(), + reason: str | None = Body(None), n4js: Neo4jStore = Depends(get_n4js_depends), token: TokenData = Depends(get_token_data_depends), ) -> list[str | None]: @@ -907,7 +922,7 @@ def tasks_status_set( except HTTPException: valid_tasks.append(None) - tasks_updated = n4js.set_task_status(valid_tasks, status) + tasks_updated = n4js.set_task_status(valid_tasks, status, reason=reason) return [str(t) if t is not None else None for t in tasks_updated] @@ -916,6 +931,7 @@ def tasks_status_set( def set_task_status( task_scoped_key, status: str = Body(), + reason: str | None = Body(None), n4js: Neo4jStore = Depends(get_n4js_depends), token: TokenData = Depends(get_token_data_depends), ): @@ -931,7 +947,7 @@ def set_task_status( ) task_sk = ScopedKey.from_str(task_scoped_key) validate_scopes(task_sk.scope, token) - tasks_statused = n4js.set_task_status([task_sk], status) + tasks_statused = n4js.set_task_status([task_sk], status, reason=reason) return [str(t) if t is not None else None for t in tasks_statused][0] @@ -1176,6 +1192,260 @@ def get_task_failures( return [str(sk) for sk in n4js.get_task_failures(sk)] +### task introspection + + +@router.get("/tasks/{task_scoped_key}/history") +def get_task_history( + task_scoped_key, + *, + limit: int | None = None, + 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) + + return [attempt.to_dict() for attempt in n4js.get_task_history(sk, limit)] + + +@router.post("/bulk/tasks/details") +def get_tasks_details( + *, + tasks: list[str] = Body(embed=True), + n4js: Neo4jStore = Depends(get_n4js_depends), + token: TokenData = Depends(get_token_data_depends), +): + task_sks = [ScopedKey.from_str(task) for task in tasks] + + for task_sk in task_sks: + validate_scopes(task_sk.scope, token) + + details = n4js.get_tasks_details(task_sks) + + return [detail.to_dict() if detail is not None else None for detail in details] + + +@router.get("/tasks/{task_scoped_key}/tracebacks") +def get_task_tracebacks( + task_scoped_key, + *, + limit: int | None = None, + 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) + + return [tb.to_dict() for tb in n4js.get_task_tracebacks(sk, limit)] + + +### result and artifact retrieval (section 3.4) + + +_TIMESTAMP_RE = re.compile(r"^\s*\[([^\]]+)\]") + + +def _unit_label(rec) -> str: + """Human-readable label for a `ProtocolUnitResultRec`.""" + return f"{rec.name or str(rec.source_key)} ({rec.scoped_key})" + + +@router.get("/tasks/{task_scoped_key}/resultrecs") +def get_task_result_recs( + task_scoped_key, + *, + ok: bool | None = None, + 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) + + return [rec.to_dict() for rec in n4js.get_task_result_recs(sk, ok)] + + +@router.get("/protocoldagresultrefs/{protocoldagresultref_scoped_key}/unitresultrecs") +def get_result_unit_recs( + protocoldagresultref_scoped_key, + *, + n4js: Neo4jStore = Depends(get_n4js_depends), + token: TokenData = Depends(get_token_data_depends), +): + pdrr_sk = ScopedKey.from_str(protocoldagresultref_scoped_key) + validate_scopes(pdrr_sk.scope, token) + + return [rec.to_dict() for rec in n4js.get_result_unit_recs(pdrr_sk)] + + +@router.get("/protocolunitresultrefs/{protocolunitresultref_scoped_key}/logs") +def get_result_unit_logs( + protocolunitresultref_scoped_key, + *, + n4js: Neo4jStore = Depends(get_n4js_depends), + s3os: S3ObjectStore = Depends(get_s3os_depends), + token: TokenData = Depends(get_token_data_depends), +): + purr_sk = ScopedKey.from_str(protocolunitresultref_scoped_key) + validate_scopes(purr_sk.scope, token) + + purr = n4js.get_gufe(purr_sk) + if not purr.has_logs: + return None + + return s3os.pull_protocol_unit_result_logs(purr.location) + + +@router.get("/protocolunitresultrefs/{protocolunitresultref_scoped_key}/stdout") +def get_result_unit_stdout( + protocolunitresultref_scoped_key, + *, + n4js: Neo4jStore = Depends(get_n4js_depends), + s3os: S3ObjectStore = Depends(get_s3os_depends), + token: TokenData = Depends(get_token_data_depends), +): + purr_sk = ScopedKey.from_str(protocolunitresultref_scoped_key) + validate_scopes(purr_sk.scope, token) + + purr = n4js.get_gufe(purr_sk) + if not purr.has_stdout: + return None + + return s3os.pull_protocol_unit_result_streams(purr.location, "stdout") + + +@router.get("/protocolunitresultrefs/{protocolunitresultref_scoped_key}/stderr") +def get_result_unit_stderr( + protocolunitresultref_scoped_key, + *, + n4js: Neo4jStore = Depends(get_n4js_depends), + s3os: S3ObjectStore = Depends(get_s3os_depends), + token: TokenData = Depends(get_token_data_depends), +): + purr_sk = ScopedKey.from_str(protocolunitresultref_scoped_key) + validate_scopes(purr_sk.scope, token) + + purr = n4js.get_gufe(purr_sk) + if not purr.has_stderr: + return None + + return s3os.pull_protocol_unit_result_streams(purr.location, "stderr") + + +@router.get("/protocoldagresultrefs/{protocoldagresultref_scoped_key}/logs") +def get_result_logs( + protocoldagresultref_scoped_key, + *, + order: str = "unit", + n4js: Neo4jStore = Depends(get_n4js_depends), + s3os: S3ObjectStore = Depends(get_s3os_depends), + token: TokenData = Depends(get_token_data_depends), +) -> str: + if order not in ("unit", "time"): + raise HTTPException( + status_code=http_status.HTTP_400_BAD_REQUEST, + detail=f"`order` takes 'unit' or 'time', not '{order}'", + ) + + pdrr_sk = ScopedKey.from_str(protocoldagresultref_scoped_key) + validate_scopes(pdrr_sk.scope, token) + + unit_recs = n4js.get_result_unit_recs(pdrr_sk) + + if order == "unit": + sections = [] + for rec in unit_recs: + if not rec.has_logs: + continue + purr = n4js.get_gufe(rec.scoped_key) + logs = s3os.pull_protocol_unit_result_logs(purr.location) + sections.append(f"=== unit {_unit_label(rec)} ===\n{logs}") + return "\n".join(sections) + + # order == "time": interleave all units' log lines by their leading + # [timestamp] prefix, each labeled with its unit + entries = [] # (sort_key, order_index, line) + idx = 0 + for rec in unit_recs: + if not rec.has_logs: + continue + purr = n4js.get_gufe(rec.scoped_key) + logs = s3os.pull_protocol_unit_result_logs(purr.location) + label = _unit_label(rec) + for line in logs.splitlines(): + m = _TIMESTAMP_RE.match(line) + sort_key = (0, m.group(1)) if m is not None else (1, "") + entries.append((sort_key, idx, f"[{label}] {line}")) + idx += 1 + + if not entries: + return "" + + # stable sort by parsed timestamp; lines without a timestamp sort last + # while preserving their original relative order + entries.sort(key=lambda e: (e[0], e[1])) + return "\n".join(line for _, _, line in entries) + + +def _render_task_stream(task_scoped_key, stream, n4js, s3os, token) -> str: + has_attr = "has_stdout" if stream == "stdout" else "has_stderr" + + sk = ScopedKey.from_str(task_scoped_key) + validate_scopes(sk.scope, token) + + sections = [] + for pdrr_rec in n4js.get_task_result_recs(sk): + for unit_rec in n4js.get_result_unit_recs(pdrr_rec.scoped_key): + if not getattr(unit_rec, has_attr): + continue + purr = n4js.get_gufe(unit_rec.scoped_key) + files = s3os.pull_protocol_unit_result_streams(purr.location, stream) + for filename, text in files.items(): + sections.append( + f"=== result {pdrr_rec.scoped_key} :: " + f"unit {_unit_label(unit_rec)} :: {filename} ===\n{text}" + ) + + return "\n".join(sections) + + +@router.get("/tasks/{task_scoped_key}/stdout") +def get_task_stdout( + task_scoped_key, + *, + n4js: Neo4jStore = Depends(get_n4js_depends), + s3os: S3ObjectStore = Depends(get_s3os_depends), + token: TokenData = Depends(get_token_data_depends), +) -> str: + return _render_task_stream(task_scoped_key, "stdout", n4js, s3os, token) + + +@router.get("/tasks/{task_scoped_key}/stderr") +def get_task_stderr( + task_scoped_key, + *, + n4js: Neo4jStore = Depends(get_n4js_depends), + s3os: S3ObjectStore = Depends(get_s3os_depends), + token: TokenData = Depends(get_token_data_depends), +) -> str: + return _render_task_stream(task_scoped_key, "stderr", n4js, s3os, token) + + +@router.post("/bulk/tasks/progress") +def get_tasks_progress( + *, + tasks: list[str] = Body(embed=True), + n4js: Neo4jStore = Depends(get_n4js_depends), + token: TokenData = Depends(get_token_data_depends), +): + task_sks = [ScopedKey.from_str(task) for task in tasks] + + for task_sk in task_sks: + validate_scopes(task_sk.scope, token) + + return n4js.get_tasks_progress(task_sks) + + ### strategies diff --git a/alchemiscale/interface/client.py b/alchemiscale/interface/client.py index f4fb7e0d..dbcae1ee 100644 --- a/alchemiscale/interface/client.py +++ b/alchemiscale/interface/client.py @@ -34,6 +34,11 @@ TaskStatusEnum, NetworkStateEnum, StrategyState, + TaskAttempt, + TaskDetails, + TaskTracebacks, + ProtocolDAGResultRec, + ProtocolUnitResultRec, ) from stratocaster.base import Strategy from ..utils import pdr_from_bytes @@ -1364,10 +1369,15 @@ async def async_request(self): return self._run_async(async_request(self)) async def _set_task_status( - self, tasks: list[ScopedKey], status: TaskStatusEnum + self, + tasks: list[ScopedKey], + status: TaskStatusEnum, + reason: str | None = None, ) -> list[ScopedKey | None]: """Set the statuses for many Tasks""" - data = dict(tasks=[t.to_dict() for t in tasks], status=status.value) + data = dict( + tasks=[t.to_dict() for t in tasks], status=status.value, reason=reason + ) tasks_updated = await self._post_resource_async( "/bulk/tasks/status/set", data=data ) @@ -1381,6 +1391,7 @@ def set_tasks_status( tasks: list[ScopedKey], status: TaskStatusEnum | str, batch_size: int = 1000, + reason: str | None = None, ) -> list[ScopedKey | None]: """Set the status of one or multiple Tasks. @@ -1394,6 +1405,9 @@ def set_tasks_status( status The status to set the Tasks to. Can be one of 'waiting', 'invalid', or 'deleted'. + reason + An optional human-readable reason for the status change. This is + only recorded when setting the status to 'invalid' or 'deleted'. Returns ------- @@ -1406,7 +1420,7 @@ def set_tasks_status( status = TaskStatusEnum(status) return self._batched_attribute_setter( - tasks, self._set_task_status, (status,), batch_size + tasks, self._set_task_status, (status, reason), batch_size ) async def _get_task_status(self, tasks: list[ScopedKey]) -> list[TaskStatusEnum]: @@ -2150,6 +2164,285 @@ def get_task_failures( return pdrs + def get_task_history( + self, task: ScopedKey, limit: int | None = None + ) -> list[TaskAttempt]: + """Get the execution history of a `Task`. + + Parameters + ---------- + task + The `ScopedKey` of the `Task` to retrieve the history for. + limit + If given, return at most this many of the most recent attempts. + + Returns + ------- + list[TaskAttempt] + A list of `TaskAttempt`s, one per execution attempt of the `Task`, + most recent first. + """ + params = dict(limit=limit) + attempts = self._get_resource(f"/tasks/{task}/history", params=params) + return [TaskAttempt.from_dict(attempt) for attempt in attempts] + + def get_tasks_details(self, tasks: list[ScopedKey]) -> list[TaskDetails | None]: + """Get summary details for multiple `Task`s. + + Parameters + ---------- + tasks + The `ScopedKey`s of the `Task`s to retrieve details for. + + Returns + ------- + list[TaskDetails | None] + A list of `TaskDetails`, in the same order as given in `tasks`. If + a given `Task` doesn't exist, ``None`` will be returned in its + place. + """ + data = dict(tasks=[str(task) for task in tasks]) + details = self._post_resource("/bulk/tasks/details", data=data) + return [ + TaskDetails.from_dict(detail) if detail is not None else None + for detail in details + ] + + def get_task_tracebacks( + self, task: ScopedKey, limit: int | None = None + ) -> list[TaskTracebacks]: + """Get the tracebacks from failed `ProtocolDAGResult`s of a `Task`. + + Parameters + ---------- + task + The `ScopedKey` of the `Task` to retrieve tracebacks for. + limit + If given, return tracebacks for at most this many of the most + recent failed `ProtocolDAGResult`s. + + Returns + ------- + list[TaskTracebacks] + A list of `TaskTracebacks`, one per failed `ProtocolDAGResult` of + the `Task`, most recent first. + """ + params = dict(limit=limit) + tracebacks = self._get_resource(f"/tasks/{task}/tracebacks", params=params) + return [TaskTracebacks.from_dict(tb) for tb in tracebacks] + + def get_scope_compute_share(self, scope: Scope) -> float: + """Get this identity's fractional compute share within the given `Scope`. + + The share is computed server-side as this `Scope`'s aggregate fraction + relative to its sibling `Scope`s; only the aggregate fraction is + returned. The identity must hold the given `Scope`. + + Parameters + ---------- + scope + The `Scope` to retrieve the compute share for. + + Returns + ------- + float + The fractional compute share for the given `Scope`. + """ + return self._get_resource(f"/scopes/{scope}/compute-share") + + @staticmethod + def _as_scoped_key(obj: ScopedKey | Any) -> ScopedKey: + """Coerce a `ScopedKey` or a `*Rec` carrying one to a `ScopedKey`.""" + return getattr(obj, "scoped_key", obj) + + def get_task_result_recs( + self, task: ScopedKey, ok: bool | None = None + ) -> list[ProtocolDAGResultRec]: + """Get records describing the `ProtocolDAGResult`s of a `Task`. + + Parameters + ---------- + task + The `ScopedKey` of the `Task` to retrieve result records for. + ok + If ``True``, return only records for successful results; if + ``False``, only failures; if ``None`` (default), all of them. + + Returns + ------- + list[ProtocolDAGResultRec] + A list of `ProtocolDAGResultRec`s, one per `ProtocolDAGResult` of + the `Task`, most recent first. + """ + params = dict(ok=ok) + recs = self._get_resource(f"/tasks/{task}/resultrecs", params=params) + return [ProtocolDAGResultRec.from_dict(rec) for rec in recs] + + def get_result_unit_recs( + self, pdrr: ScopedKey | ProtocolDAGResultRec + ) -> list[ProtocolUnitResultRec]: + """Get records describing the `ProtocolUnitResult`s of a `ProtocolDAGResult`. + + Parameters + ---------- + pdrr + The `ScopedKey` of the `ProtocolDAGResultRef` (or the + `ProtocolDAGResultRec` describing it) to retrieve unit records for. + + Returns + ------- + list[ProtocolUnitResultRec] + A list of `ProtocolUnitResultRec`s, one per `ProtocolUnitResult`, + in dependency order. + """ + pdrr_sk = self._as_scoped_key(pdrr) + recs = self._get_resource(f"/protocoldagresultrefs/{pdrr_sk}/unitresultrecs") + return [ProtocolUnitResultRec.from_dict(rec) for rec in recs] + + def get_result_unit_logs( + self, purr: ScopedKey | ProtocolUnitResultRec + ) -> str | None: + """Get the captured logs for a single `ProtocolUnitResult`. + + Parameters + ---------- + purr + The `ScopedKey` of the `ProtocolUnitResultRef` (or the + `ProtocolUnitResultRec` describing it) to retrieve logs for. + + Returns + ------- + str | None + The captured log text, or ``None`` if no logs were captured. + """ + purr_sk = self._as_scoped_key(purr) + return self._get_resource(f"/protocolunitresultrefs/{purr_sk}/logs") + + def get_result_unit_stdout( + self, purr: ScopedKey | ProtocolUnitResultRec + ) -> dict[str, str] | None: + """Get the captured stdout for a single `ProtocolUnitResult`. + + Parameters + ---------- + purr + The `ScopedKey` of the `ProtocolUnitResultRef` (or the + `ProtocolUnitResultRec` describing it) to retrieve stdout for. + + Returns + ------- + dict[str, str] | None + A mapping of filename to captured stdout text, or ``None`` if no + stdout was captured. + """ + purr_sk = self._as_scoped_key(purr) + return self._get_resource(f"/protocolunitresultrefs/{purr_sk}/stdout") + + def get_result_unit_stderr( + self, purr: ScopedKey | ProtocolUnitResultRec + ) -> dict[str, str] | None: + """Get the captured stderr for a single `ProtocolUnitResult`. + + Parameters + ---------- + purr + The `ScopedKey` of the `ProtocolUnitResultRef` (or the + `ProtocolUnitResultRec` describing it) to retrieve stderr for. + + Returns + ------- + dict[str, str] | None + A mapping of filename to captured stderr text, or ``None`` if no + stderr was captured. + """ + purr_sk = self._as_scoped_key(purr) + return self._get_resource(f"/protocolunitresultrefs/{purr_sk}/stderr") + + def get_result_logs( + self, pdrr: ScopedKey | ProtocolDAGResultRec, order: str = "unit" + ) -> str: + """Get a human-readable rendering of all unit logs of a `ProtocolDAGResult`. + + Parameters + ---------- + pdrr + The `ScopedKey` of the `ProtocolDAGResultRef` (or the + `ProtocolDAGResultRec` describing it) to retrieve logs for. + order + How to order the rendered logs. ``'unit'`` (default) groups each + unit's logs under a header; ``'time'`` interleaves all units' log + lines by their leading ``[timestamp]`` prefix. + + Returns + ------- + str + The rendered logs, or ``""`` if none were captured. + """ + pdrr_sk = self._as_scoped_key(pdrr) + params = dict(order=order) + return self._get_resource( + f"/protocoldagresultrefs/{pdrr_sk}/logs", params=params + ) + + def get_task_stdout(self, task: ScopedKey) -> str: + """Get a human-readable rendering of all captured stdout for a `Task`. + + Concatenates stdout across all of the `Task`'s `ProtocolDAGResult`s + (most recent first), with section headers identifying each result, + unit, and filename. + + Parameters + ---------- + task + The `ScopedKey` of the `Task` to retrieve stdout for. + + Returns + ------- + str + The rendered stdout, or ``""`` if none was captured. + """ + return self._get_resource(f"/tasks/{task}/stdout") + + def get_task_stderr(self, task: ScopedKey) -> str: + """Get a human-readable rendering of all captured stderr for a `Task`. + + Concatenates stderr across all of the `Task`'s `ProtocolDAGResult`s + (most recent first), with section headers identifying each result, + unit, and filename. + + Parameters + ---------- + task + The `ScopedKey` of the `Task` to retrieve stderr for. + + Returns + ------- + str + The rendered stderr, or ``""`` if none was captured. + """ + return self._get_resource(f"/tasks/{task}/stderr") + + def get_tasks_progress( + self, tasks: list[ScopedKey] + ) -> list[tuple[int, int] | None]: + """Get execution progress for multiple `Task`s. + + Parameters + ---------- + tasks + The `ScopedKey`s of the `Task`s to retrieve progress for. + + Returns + ------- + list[tuple[int, int] | None] + A list in the same order as `tasks`. Each element is a + ``(units_completed, units_total)`` tuple for a `running` `Task` + reporting progress, or ``None`` otherwise. + """ + data = dict(tasks=[str(task) for task in tasks]) + progress = self._post_resource("/bulk/tasks/progress", data=data) + return [tuple(p) if p is not None else None for p in progress] + def add_task_restart_patterns( self, network_scoped_key: ScopedKey, diff --git a/alchemiscale/migrations/v04_to_v05.py b/alchemiscale/migrations/v04_to_v05.py new file mode 100644 index 00000000..b1ba9fb1 --- /dev/null +++ b/alchemiscale/migrations/v04_to_v05.py @@ -0,0 +1,52 @@ +""" +:mod:`alchemiscale.migrations.v04_to_v05` --- migration for v0.4 to v0.5 +======================================================================== + +""" + +from ..storage.statestore import Neo4jStore + + +def migrate(n4js: Neo4jStore): + """Migrate state store from alchemiscale v0.4 to v0.5. + + Changes: + - adds indexes on the new ``TaskProvenance`` node label to support the + introspection queries introduced in v0.5. ``TaskProvenance`` is a plain + labeled node (not a ``GufeTokenizable``), identified by its properties and + reached from a ``Task`` via the ``PROVENANCE_OF`` relationship: + - ``TaskProvenance.compute_service_id``: provenance records are matched + by the id of the compute service that produced them, both when + finalizing an attempt (``set_task_result``, expiry/deregistration) and + when reading live progress for the current claimant. + - ``TaskProvenance.datetime_claimed``: attempt histories and + most-recent-attempt lookups are ordered by claim time. + + (There is nothing to index for the ``PROVENANCE_OF`` traversal itself: the + ``Task`` endpoint is already covered by the ``GufeTokenizable._scoped_key`` + uniqueness constraint, and the ``PROVENANCE_OF`` relationship carries no + properties.) + + This migration is idempotent (all indexes are created with + ``IF NOT EXISTS``) and requires NO data migration. All new properties are + optional-valued: pre-existing ``Task`` nodes simply have an empty attempt + history and are unaffected. + + """ + + indexes = { + "TaskProvenance_compute_service_id_index": ( + "TaskProvenance", + "compute_service_id", + ), + "TaskProvenance_datetime_claimed_index": ( + "TaskProvenance", + "datetime_claimed", + ), + } + + for name, (label, property_) in indexes.items(): + n4js.execute_query(f""" + CREATE INDEX {name} IF NOT EXISTS + FOR (n:{label}) ON (n.{property_}) + """) diff --git a/alchemiscale/storage/models.py b/alchemiscale/storage/models.py index db9dc695..807629d0 100644 --- a/alchemiscale/storage/models.py +++ b/alchemiscale/storage/models.py @@ -82,6 +82,7 @@ class ComputeServiceRegistration(BaseModel): heartbeat: datetime.datetime failure_times: list[datetime.datetime] = [] manager_name: str | None = None + hostname: str | None = None model_config = ConfigDict(arbitrary_types_allowed=True) @@ -150,14 +151,84 @@ def to_compute_manager_id(self): return ComputeManagerID("-".join([self.name, self.uuid])) +class TaskOutcomeEnum(Enum): + """Terminal outcome of a single execution attempt of a `Task`. + + Attributes + ---------- + complete + The attempt produced a successful `ProtocolDAGResult`. + error + The attempt produced a failed `ProtocolDAGResult`, or errored during + `ProtocolDAG` creation. + expired + The attempt's compute service lost its registration (expiry or + deregistration) before the attempt produced a result. + released + A user forced the claimed `Task` to another status (e.g. `waiting`, + `invalid`, `deleted`) before the attempt produced a result. + """ + + complete = "complete" + error = "error" + expired = "expired" + released = "released" + + class TaskProvenance(BaseModel): - computeserviceid: ComputeServiceID - datetime_start: datetime.datetime - datetime_end: datetime.datetime + """An immutable record of a single execution attempt of a `Task`. + + A `TaskProvenance` node is created at claim time and finalized when the + attempt ends. It survives claim teardown and registration expiry, so that + the history of *who ran what, when* is preserved. The identifying + information (compute service id, hostname, manager name) is copied onto the + record rather than held as a relationship to the (potentially deleted) + `ComputeServiceRegistration`. + + Attributes + ---------- + compute_service_id + The identifier of the compute service that claimed the `Task`. + hostname + The hostname of the compute service, copied from its registration. + manager_name + The name of the compute manager responsible for the compute service, + if any. + datetime_claimed + When the `Task` was claimed for this attempt. + datetime_end + When the attempt was finalized; `None` while the attempt is open. + outcome + The terminal outcome of the attempt; `None` while the attempt is open. + units_completed + The number of distinct `ProtocolUnit`s successfully completed in this + attempt, as of the last progress update. + units_total + The total number of `ProtocolUnit`s in the attempt's `ProtocolDAG`. + """ + + compute_service_id: ComputeServiceID + hostname: str | None = None + manager_name: str | None = None + datetime_claimed: datetime.datetime + datetime_end: datetime.datetime | None = None + outcome: TaskOutcomeEnum | None = None + units_completed: int | None = None + units_total: int | None = None model_config = ConfigDict(arbitrary_types_allowed=True) - # this should include versions of various libraries + def to_dict(self): + dct = self.model_dump() + dct["compute_service_id"] = str(self.compute_service_id) + dct["outcome"] = self.outcome.value if self.outcome is not None else None + return dct + + @classmethod + def from_dict(cls, dct): + dct_ = copy(dct) + dct_["compute_service_id"] = ComputeServiceID(dct_["compute_service_id"]) + return cls(**dct_) class TaskStatusEnum(Enum): @@ -189,6 +260,8 @@ class Task(GufeTokenizable): priority: int claim: str | None datetime_created: datetime.datetime | None + datetime_status_changed: datetime.datetime | None + reason: str | None creator: str | None extends: str | None @@ -198,6 +271,8 @@ def __init__( status: str | TaskStatusEnum = TaskStatusEnum.waiting, priority: int = 10, datetime_created: datetime.datetime | None = None, + datetime_status_changed: datetime.datetime | None = None, + reason: str | None = None, creator: str | None = None, extends: str | None = None, claim: str | None = None, @@ -215,6 +290,8 @@ def __init__( else datetime.datetime.now(tz=datetime.UTC) ) + self.datetime_status_changed = datetime_status_changed + self.reason = reason self.creator = creator self.extends = extends self.claim = claim @@ -228,6 +305,8 @@ def _to_dict(self): "status": self.status.value, "priority": self.priority, "datetime_created": self.datetime_created, + "datetime_status_changed": self.datetime_status_changed, + "reason": self.reason, "creator": self.creator, "extends": self.extends, "claim": self.claim, @@ -552,6 +631,119 @@ def _from_dict(cls, d): return super()._from_dict(d_) +class ProtocolUnitResultRef(ObjectStoreRef): + """A reference to the artifacts of a single `ProtocolUnitResult` or + `ProtocolUnitFailure` within a `ProtocolDAGResult`. + + One `ProtocolUnitResultRef` node is derived per unit result when a + `ProtocolDAGResultRef` is stored, keyed by the unit result's gufe key. The + `has_*` flags record which per-unit artifacts (logs, stdout, stderr) are + present in the object store under `location`. + + Attributes + ---------- + obj_key + The gufe key of the `ProtocolUnitResult`/`ProtocolUnitFailure`. + source_key + The gufe key of the originating `ProtocolUnit`. + name + The name of the unit result, if any. + ok + Whether the unit result is a success (`True`) or failure (`False`). + start_time, end_time + When execution of the unit attempt began and ended. + location + The object store prefix under which this unit result's artifacts live. + has_logs, has_stdout, has_stderr + Whether captured log/stdout/stderr artifacts exist for this unit result. + + Note + ---- + The `has_*` flags and (nothing else) are *mutated in place* via Cypher after + the node is created, as artifacts arrive. The node's `_scoped_key`/`_gufe_key` + are computed once at creation and never recomputed, so lookups stay stable + even though these tokenizable-contributing fields change. This is safe only + because `ProtocolUnitResultRef`s are an internal state-store detail, never + re-tokenized after creation; keep it that way. + """ + + ok: bool + source_key: GufeKey + name: str | None + start_time: datetime.datetime | None + end_time: datetime.datetime | None + has_logs: bool + has_stdout: bool + has_stderr: bool + + def __init__( + self, + *, + location: str | None = None, + obj_key: GufeKey, + source_key: GufeKey, + scope: Scope, + ok: bool, + name: str | None = None, + start_time: datetime.datetime | None = None, + end_time: datetime.datetime | None = None, + has_logs: bool = False, + has_stdout: bool = False, + has_stderr: bool = False, + ): + self.location = location + self.obj_key = GufeKey(obj_key) + self.source_key = GufeKey(source_key) + self.scope = scope + self.ok = ok + self.name = name + self.start_time = start_time + self.end_time = end_time + self.has_logs = has_logs + self.has_stdout = has_stdout + self.has_stderr = has_stderr + + def _to_dict(self): + return { + "location": self.location, + "obj_key": str(self.obj_key), + "source_key": str(self.source_key), + "scope": str(self.scope), + "ok": self.ok, + "name": self.name, + "start_time": ( + self.start_time.isoformat() if self.start_time is not None else None + ), + "end_time": ( + self.end_time.isoformat() if self.end_time is not None else None + ), + "has_logs": self.has_logs, + "has_stdout": self.has_stdout, + "has_stderr": self.has_stderr, + } + + @classmethod + def _from_dict(cls, d): + d_ = copy(d) + d_["scope"] = Scope.from_str(d["scope"]) + d_["source_key"] = GufeKey(d["source_key"]) + d_["start_time"] = ( + datetime.datetime.fromisoformat(d["start_time"]) + if d.get("start_time") is not None + else None + ) + d_["end_time"] = ( + datetime.datetime.fromisoformat(d["end_time"]) + if d.get("end_time") is not None + else None + ) + return cls(**d_) + + @classmethod + def _defaults(cls): + return super()._defaults() + + class StrategyModeEnum(StrEnum): full = "full" partial = "partial" @@ -606,3 +798,320 @@ def to_dict(self): @classmethod def from_dict(cls, d): return cls(**d) + + +def _coerce_datetime(v) -> datetime.datetime | None: + """Coerce a neo4j ``DateTime``, ISO string, or ``datetime`` to ``datetime``.""" + if v is None: + return None + if hasattr(v, "to_native"): + return v.to_native() + if isinstance(v, str): + return datetime.datetime.fromisoformat(v) + return v + + +def _iso(v: datetime.datetime | None) -> str | None: + return v.isoformat() if v is not None else None + + +# --- client-facing API record models -------------------------------------- +# +# These models are the user-facing surface for Task introspection. They are +# deliberately decoupled from the state-store node types (`TaskProvenance`, +# `ProtocolDAGResultRef`, `ProtocolUnitResultRef`): the two families evolve +# independently, joined only by `ScopedKey`s. User-level names are used +# throughout; no storage jargon (`Ref` suffixes, `pdrr`/`purr`) leaks in. + + +class TaskAttempt(BaseModel): + """A single execution attempt of a `Task`, as reported by `get_task_history`. + + Bundles a `TaskProvenance` record's properties with the `ScopedKey` of the + `ProtocolDAGResultRef` the attempt produced (via `PROVENANCE_OF`), where one + exists; `expired`/`released` attempts have none. + """ + + compute_service_id: str + hostname: str | None = None + manager_name: str | None = None + datetime_claimed: datetime.datetime + datetime_end: datetime.datetime | None = None + outcome: TaskOutcomeEnum | None = None + units_completed: int | None = None + units_total: int | None = None + protocoldagresultref: ScopedKey | None = None + + model_config = ConfigDict(arbitrary_types_allowed=True) + + def to_dict(self): + return { + "compute_service_id": self.compute_service_id, + "hostname": self.hostname, + "manager_name": self.manager_name, + "datetime_claimed": _iso(self.datetime_claimed), + "datetime_end": _iso(self.datetime_end), + "outcome": self.outcome.value if self.outcome is not None else None, + "units_completed": self.units_completed, + "units_total": self.units_total, + "protocoldagresultref": ( + str(self.protocoldagresultref) + if self.protocoldagresultref is not None + else None + ), + } + + @classmethod + def from_dict(cls, d): + return cls( + compute_service_id=d["compute_service_id"], + hostname=d.get("hostname"), + manager_name=d.get("manager_name"), + datetime_claimed=_coerce_datetime(d["datetime_claimed"]), + datetime_end=_coerce_datetime(d.get("datetime_end")), + outcome=( + TaskOutcomeEnum(d["outcome"]) if d.get("outcome") is not None else None + ), + units_completed=d.get("units_completed"), + units_total=d.get("units_total"), + protocoldagresultref=( + ScopedKey.from_str(d["protocoldagresultref"]) + if d.get("protocoldagresultref") is not None + else None + ), + ) + + +class TaskClaim(BaseModel): + """The live claim currently held on a `running` `Task`.""" + + compute_service_id: str + hostname: str | None = None + datetime_claimed: datetime.datetime | None = None + units_completed: int | None = None + units_total: int | None = None + + model_config = ConfigDict(arbitrary_types_allowed=True) + + def to_dict(self): + return { + "compute_service_id": self.compute_service_id, + "hostname": self.hostname, + "datetime_claimed": _iso(self.datetime_claimed), + "units_completed": self.units_completed, + "units_total": self.units_total, + } + + @classmethod + def from_dict(cls, d): + return cls( + compute_service_id=d["compute_service_id"], + hostname=d.get("hostname"), + datetime_claimed=_coerce_datetime(d.get("datetime_claimed")), + units_completed=d.get("units_completed"), + units_total=d.get("units_total"), + ) + + +class TaskDetails(BaseModel): + """Bulk indicator summary for a `Task`, as returned by `get_tasks_details`.""" + + task: ScopedKey + status: TaskStatusEnum + datetime_status_changed: datetime.datetime | None = None + reason: str | None = None + num_claims: int = 0 + current_claim: TaskClaim | None = None + most_recent_attempt: TaskAttempt | None = None + + model_config = ConfigDict(arbitrary_types_allowed=True) + + def to_dict(self): + return { + "task": str(self.task), + "status": self.status.value, + "datetime_status_changed": _iso(self.datetime_status_changed), + "reason": self.reason, + "num_claims": self.num_claims, + "current_claim": ( + self.current_claim.to_dict() if self.current_claim is not None else None + ), + "most_recent_attempt": ( + self.most_recent_attempt.to_dict() + if self.most_recent_attempt is not None + else None + ), + } + + @classmethod + def from_dict(cls, d): + return cls( + task=ScopedKey.from_str(d["task"]), + status=TaskStatusEnum(d["status"]), + datetime_status_changed=_coerce_datetime(d.get("datetime_status_changed")), + reason=d.get("reason"), + num_claims=d.get("num_claims", 0), + current_claim=( + TaskClaim.from_dict(d["current_claim"]) + if d.get("current_claim") is not None + else None + ), + most_recent_attempt=( + TaskAttempt.from_dict(d["most_recent_attempt"]) + if d.get("most_recent_attempt") is not None + else None + ), + ) + + +class TaskUnitTraceback(BaseModel): + """A single `ProtocolUnitFailure`'s traceback within a `TaskTracebacks`.""" + + failure_key: GufeKey + source_key: GufeKey + traceback: str + protocolunitresultref: ScopedKey | None = None + + model_config = ConfigDict(arbitrary_types_allowed=True) + + def to_dict(self): + return { + "failure_key": str(self.failure_key), + "source_key": str(self.source_key), + "traceback": self.traceback, + "protocolunitresultref": ( + str(self.protocolunitresultref) + if self.protocolunitresultref is not None + else None + ), + } + + @classmethod + def from_dict(cls, d): + return cls( + failure_key=GufeKey(d["failure_key"]), + source_key=GufeKey(d["source_key"]), + traceback=d["traceback"], + protocolunitresultref=( + ScopedKey.from_str(d["protocolunitresultref"]) + if d.get("protocolunitresultref") is not None + else None + ), + ) + + +class TaskTracebacks(BaseModel): + """Tracebacks for one failed `ProtocolDAGResult` of a `Task`. + + Returned by `get_task_tracebacks`, one record per failed + `ProtocolDAGResultRef`, most recent first. + """ + + protocoldagresultref: ScopedKey + datetime_created: datetime.datetime | None = None + creator: str | None = None + tracebacks: list[TaskUnitTraceback] = [] + + model_config = ConfigDict(arbitrary_types_allowed=True) + + def to_dict(self): + return { + "protocoldagresultref": str(self.protocoldagresultref), + "datetime_created": _iso(self.datetime_created), + "creator": self.creator, + "tracebacks": [tb.to_dict() for tb in self.tracebacks], + } + + @classmethod + def from_dict(cls, d): + return cls( + protocoldagresultref=ScopedKey.from_str(d["protocoldagresultref"]), + datetime_created=_coerce_datetime(d.get("datetime_created")), + creator=d.get("creator"), + tracebacks=[TaskUnitTraceback.from_dict(tb) for tb in d["tracebacks"]], + ) + + +class ProtocolDAGResultRec(BaseModel): + """A record describing one `ProtocolDAGResult` of a `Task`. + + Returned by `get_task_result_recs`. Carries the `ScopedKey` of the + underlying `ProtocolDAGResultRef` as `scoped_key`, which every drill-down + method accepts directly. + """ + + scoped_key: ScopedKey + ok: bool + datetime_created: datetime.datetime | None = None + creator: str | None = None + + model_config = ConfigDict(arbitrary_types_allowed=True) + + def to_dict(self): + return { + "scoped_key": str(self.scoped_key), + "ok": self.ok, + "datetime_created": _iso(self.datetime_created), + "creator": self.creator, + } + + @classmethod + def from_dict(cls, d): + return cls( + scoped_key=ScopedKey.from_str(d["scoped_key"]), + ok=d["ok"], + datetime_created=_coerce_datetime(d.get("datetime_created")), + creator=d.get("creator"), + ) + + +class ProtocolUnitResultRec(BaseModel): + """A record describing one `ProtocolUnitResult` of a `ProtocolDAGResult`. + + Returned by `get_result_unit_recs`. Carries the `ScopedKey` of the + underlying `ProtocolUnitResultRef` as `scoped_key`, plus the `obj_key`/ + `source_key` that let a user correlate it against a deserialized + `ProtocolDAGResult`. + """ + + scoped_key: ScopedKey + obj_key: GufeKey + source_key: GufeKey + name: str | None = None + ok: bool + start_time: datetime.datetime | None = None + end_time: datetime.datetime | None = None + has_logs: bool = False + has_stdout: bool = False + has_stderr: bool = False + + model_config = ConfigDict(arbitrary_types_allowed=True) + + def to_dict(self): + return { + "scoped_key": str(self.scoped_key), + "obj_key": str(self.obj_key), + "source_key": str(self.source_key), + "name": self.name, + "ok": self.ok, + "start_time": _iso(self.start_time), + "end_time": _iso(self.end_time), + "has_logs": self.has_logs, + "has_stdout": self.has_stdout, + "has_stderr": self.has_stderr, + } + + @classmethod + def from_dict(cls, d): + return cls( + scoped_key=ScopedKey.from_str(d["scoped_key"]), + obj_key=GufeKey(d["obj_key"]), + source_key=GufeKey(d["source_key"]), + name=d.get("name"), + ok=d["ok"], + start_time=_coerce_datetime(d.get("start_time")), + end_time=_coerce_datetime(d.get("end_time")), + has_logs=d.get("has_logs", False), + has_stdout=d.get("has_stdout", False), + has_stderr=d.get("has_stderr", False), + ) diff --git a/alchemiscale/storage/objectstore.py b/alchemiscale/storage/objectstore.py index 352b5bee..87c1fb3c 100644 --- a/alchemiscale/storage/objectstore.py +++ b/alchemiscale/storage/objectstore.py @@ -9,6 +9,8 @@ from boto3.session import Session from functools import lru_cache +import zstandard as zstd + from gufe.tokenization import GufeKey from ..models import ScopedKey @@ -18,6 +20,12 @@ # default filename for object store files OBJECT_FILENAME = "obj.json.zst" +# per-unit-result artifact layout, relative to a ProtocolUnitResultRef location +# (``.../{pdr_key}/units/{unit_result_key}``) +LOGS_FILENAME = "logs.txt.zst" +STDOUT_DIRNAME = "stdout" +STDERR_DIRNAME = "stderr" + def get_s3os(settings: S3ObjectStoreSettings) -> "S3ObjectStore": """Convenience function for getting an S3ObjectStore directly from settings.""" @@ -297,3 +305,83 @@ def pull_protocoldagresult( pdr_bytes = self._get_bytes(location) return pdr_bytes + + # --- per-unit-result artifacts (logs, stdout, stderr) ----------------- + # + # These live under the unit result's location prefix + # (`ProtocolUnitResultRef.location`), a retrieval-optimized copy so that log + # retrieval never requires pulling and deserializing the whole + # `ProtocolDAGResult`. Bytes are zstd-compressed, consistent with + # `alchemiscale.compression`. + + def push_protocol_unit_result_logs(self, unit_location: str, logtext: str) -> str: + """Store captured log text for a unit result. Returns the artifact key. + + `logtext` is already truncated to the per-unit cap by the compute + service; here it is simply zstd-compressed and stored. + """ + location = os.path.join(unit_location, LOGS_FILENAME) + compressed = zstd.ZstdCompressor().compress(logtext.encode("utf-8")) + self._store_bytes(location, compressed) + return location + + def pull_protocol_unit_result_logs(self, unit_location: str) -> str: + """Return decompressed log text for a unit result.""" + location = os.path.join(unit_location, LOGS_FILENAME) + compressed = self._get_bytes(location) + return ( + zstd.ZstdDecompressor() + .decompress(compressed) + .decode("utf-8", errors="replace") + ) + + def push_protocol_unit_result_streams( + self, unit_location: str, stream: str, files: dict[str, bytes] + ) -> list[str]: + """Store a unit result's captured stream files (stdout or stderr). + + `files` maps filename -> raw bytes (the `dict[filename, bytes]` gufe + embeds on the `ProtocolUnitResult`). Each is zstd-compressed and stored + under the stream subdirectory, keeping the protocol-given filename. + """ + if stream not in (STDOUT_DIRNAME, STDERR_DIRNAME): + raise ValueError("`stream` must be 'stdout' or 'stderr'") + locations = [] + compressor = zstd.ZstdCompressor() + for filename, data in files.items(): + location = os.path.join(unit_location, stream, f"{filename}.zst") + self._store_bytes(location, compressor.compress(data)) + locations.append(location) + return locations + + def pull_protocol_unit_result_streams( + self, unit_location: str, stream: str + ) -> dict[str, str]: + """Return a unit result's captured stream files, filename -> decoded text. + + Bytes are decoded as UTF-8 with ``errors="replace"``; protocols + overwhelmingly archive text (binary outputs are #180 `ResultFile` + territory). + """ + if stream not in (STDOUT_DIRNAME, STDERR_DIRNAME): + raise ValueError("`stream` must be 'stdout' or 'stderr'") + prefix = os.path.join(unit_location, stream) + "/" + decompressor = zstd.ZstdDecompressor() + out = {} + for obj in self._get_filename_prefix_contents(prefix): + # key includes self.prefix and the full location; recover the + # filename relative to the stream directory, dropping the .zst suffix + key = obj.key + filename = key.rsplit("/", 1)[-1] + if filename.endswith(".zst"): + filename = filename[: -len(".zst")] + data = obj.get()["Body"].read() + out[filename] = decompressor.decompress(data).decode( + "utf-8", errors="replace" + ) + return out + + def _get_filename_prefix_contents(self, prefix: str): + """Iterate S3 objects under a location prefix (excluding ``self.prefix``).""" + filter_prefix = os.path.join(self.prefix, prefix) + return self.resource.Bucket(self.bucket).objects.filter(Prefix=filter_prefix) diff --git a/alchemiscale/storage/statestore.py b/alchemiscale/storage/statestore.py index 343bdcde..b8a89af8 100644 --- a/alchemiscale/storage/statestore.py +++ b/alchemiscale/storage/statestore.py @@ -9,6 +9,7 @@ import datetime from contextlib import contextmanager import json +import os import re from functools import lru_cache, update_wrapper from collections import defaultdict @@ -41,16 +42,28 @@ NetworkMark, NetworkStateEnum, ProtocolDAGResultRef, + ProtocolUnitResultRec, + ProtocolUnitResultRef, + ProtocolDAGResultRec, StrategyState, StrategyModeEnum, StrategyStatusEnum, StrategyTaskScalingEnum, Task, + TaskAttempt, + TaskClaim, + TaskDetails, TaskHub, + TaskOutcomeEnum, + TaskProvenance, TaskRestartPattern, TaskStatusEnum, + TaskTracebacks, + TaskUnitTraceback, Tracebacks, + _coerce_datetime, ) +from gufe.protocols import ProtocolDAGResult, ProtocolUnitResult from ..models import Scope, ScopedKey from .cypher import cypher_or @@ -110,6 +123,52 @@ def _select_tasks_from_taskpool(taskpool: list[tuple[str, float]], count) -> lis return list(np.random.choice(tasks, count, replace=False, p=prob)) +def _status_write( + var: str, + status_value: str, + *, + time_param: str = "statuschange_time", + reason_expr: str = "null", +) -> str: + """Return a Cypher ``SET`` clause writing a ``Task`` status change. + + Every site that mutates ``Task.status`` must also update the + ``datetime_status_changed`` indicator and reset ``reason`` (which describes + the *current* status). Centralizing those three writes here makes a missed + site structurally impossible: any status write routes through this helper. + + Parameters + ---------- + var + The Cypher variable bound to the ``Task`` node being mutated. + status_value + The new status value (a ``TaskStatusEnum`` value string). + time_param + Name of the query parameter carrying the change timestamp as an ISO + string. + reason_expr + Cypher expression for the new ``reason`` value; ``"null"`` clears it + (the default), or a parameter reference such as ``"$reason"``. + """ + # Only refresh `datetime_status_changed` and `reason` when the status + # actually changes: several setters admit their own target status for + # idempotent-return semantics (e.g. `waiting -> waiting`), and a no-op + # re-set must not reset "how long in current status" or wipe a `reason` + # (e.g. a DAG-creation traceback, or a user-forced invalid/deleted reason). + # The CASE expressions read `{var}.status` *before* this SET clause's + # assignments take effect (Cypher evaluates a SET clause's right-hand sides + # against the pre-clause state), so they compare against the old status. + return ( + f"SET {var}.datetime_status_changed = CASE " + f"WHEN {var}.status = '{status_value}' THEN {var}.datetime_status_changed " + f"ELSE datetime(${time_param}) END, " + f"{var}.reason = CASE " + f"WHEN {var}.status = '{status_value}' THEN {var}.reason " + f"ELSE {reason_expr} END, " + f"{var}.status = '{status_value}'" + ) + + CLAIM_QUERY = f""" // only match the task if it doesn't have an existing CLAIMS relationship UNWIND $tasks_list AS task_sk @@ -122,7 +181,21 @@ def _select_tasks_from_taskpool(taskpool: list[tuple[str, float]], count) -> lis MATCH (csreg:ComputeServiceRegistration {{identifier: $compute_service_id}}) CREATE (t)<-[cl:CLAIMS {{claimed: datetime($datetimestr)}}]-(csreg) - SET t.status = '{TaskStatusEnum.running.value}' + // create an immutable TaskProvenance record for this execution attempt, + // copying identifying info off the registration (which may later be + // deleted on expiry/deregistration) + CREATE (tp:TaskProvenance {{ + compute_service_id: $compute_service_id, + hostname: csreg.hostname, + manager_name: csreg.manager_name, + datetime_claimed: datetime($datetimestr), + _org: t._org, + _campaign: t._campaign, + _project: t._project + }}) + CREATE (tp)-[:PROVENANCE_OF]->(t) + + {_status_write('t', TaskStatusEnum.running.value, time_param='datetimestr')} RETURN t """ @@ -1494,11 +1567,19 @@ def deregister_computeservice(self, compute_service_id: ComputeServiceID): """ + now = datetime.datetime.now(tz=datetime.UTC).isoformat() q = f""" MATCH (n:ComputeServiceRegistration {{identifier: $compute_service_id}}) OPTIONAL MATCH (n)-[cl:CLAIMS]->(t:Task {{status: '{TaskStatusEnum.running.value}'}}) - SET t.status = '{TaskStatusEnum.waiting.value}' + {_status_write('t', TaskStatusEnum.waiting.value)} + + // finalize the open provenance attempt of each returned Task as expired + WITH n, t + OPTIONAL MATCH (t)<-[:PROVENANCE_OF]-(tp:TaskProvenance {{compute_service_id: n.identifier}}) + WHERE tp.datetime_end IS NULL + SET tp.outcome = '{TaskOutcomeEnum.expired.value}', + tp.datetime_end = datetime($statuschange_time) WITH n, n.identifier as identifier @@ -1508,7 +1589,11 @@ def deregister_computeservice(self, compute_service_id: ComputeServiceID): """ with self.transaction() as tx: - res = tx.run(q, compute_service_id=str(compute_service_id)) + res = tx.run( + q, + compute_service_id=str(compute_service_id), + statuschange_time=now, + ) identifier = next(res)["identifier"] return ComputeServiceID(identifier) @@ -1530,6 +1615,7 @@ def heartbeat_computeservice( def expire_registrations(self, expire_time: datetime.datetime): """Remove all registrations with last heartbeat prior to the given `expire_time`.""" + now = datetime.datetime.now(tz=datetime.UTC).isoformat() q = f""" MATCH (n:ComputeServiceRegistration) WHERE n.heartbeat < datetime('{expire_time.isoformat()}') @@ -1537,7 +1623,14 @@ def expire_registrations(self, expire_time: datetime.datetime): WITH n OPTIONAL MATCH (n)-[cl:CLAIMS]->(t:Task {{status: '{TaskStatusEnum.running.value}'}}) - SET t.status = '{TaskStatusEnum.waiting.value}' + {_status_write('t', TaskStatusEnum.waiting.value)} + + // finalize the open provenance attempt of each returned Task as expired + WITH n, t + OPTIONAL MATCH (t)<-[:PROVENANCE_OF]-(tp:TaskProvenance {{compute_service_id: n.identifier}}) + WHERE tp.datetime_end IS NULL + SET tp.outcome = '{TaskOutcomeEnum.expired.value}', + tp.datetime_end = datetime($statuschange_time) WITH n, n.identifier as ident @@ -1546,7 +1639,7 @@ def expire_registrations(self, expire_time: datetime.datetime): RETURN ident """ with self.transaction() as tx: - res = tx.run(q) + res = tx.run(q, statuschange_time=now) identities = set() for rec in res: @@ -3444,9 +3537,24 @@ def get_transformation_status(self, transformation: ScopedKey) -> dict[str, int] return counts def set_task_result( - self, task: ScopedKey, protocoldagresultref: ProtocolDAGResultRef + self, + task: ScopedKey, + protocoldagresultref: ProtocolDAGResultRef, + compute_service_id: ComputeServiceID | None = None, ) -> ScopedKey: - """Set a `ProtocolDAGResultRef` pointing to a `ProtocolDAGResult` for the given `Task`.""" + """Set a `ProtocolDAGResultRef` pointing to a `ProtocolDAGResult` for the given `Task`. + + If `compute_service_id` is given, the `TaskProvenance` attempt for this + ``(task, compute_service_id)`` pair is finalized: its `outcome` is set + from `protocoldagresultref.ok`, its `datetime_end` recorded, and a + `PROVENANCE_OF` edge created to the new `ProtocolDAGResultRef`. The + match is scoped to this service's own attempt record — never "whichever + record is open" — so a late result (posted after the service's + registration expired and the `Task` was reclaimed) finalizes its own + attempt without corrupting another service's open record. A record + previously closed as `expired` is overwritten to `complete`/`error`, + since the attempt did in fact finish. + """ if task.qualname != "Task": raise ValueError("`task` ScopedKey does not correspond to a `Task`") @@ -3470,6 +3578,42 @@ def set_task_result( with self.transaction() as tx: merge_subgraph(tx, subgraph, "GufeTokenizable", "_scoped_key") + if compute_service_id is not None: + outcome = ( + TaskOutcomeEnum.complete.value + if protocoldagresultref.ok + else TaskOutcomeEnum.error.value + ) + finalize_q = f""" + MATCH (t:Task {{_scoped_key: $task}}) + MATCH (pdrr:ProtocolDAGResultRef {{_scoped_key: $pdrr}}) + // this service's own attempt: prefer a still-open record; else a + // record prematurely closed as `expired` (which this result + // proves finished after all). A record closed as `released` (a + // user forced the Task off-attempt) or already `complete`/`error` + // is left untouched --- the immutable history stands, and another + // service's open record is never matched (scoped to this csid). + OPTIONAL MATCH (t)<-[:PROVENANCE_OF]-(tp:TaskProvenance {{compute_service_id: $compute_service_id}}) + WHERE tp.datetime_end IS NULL + OR tp.outcome = '{TaskOutcomeEnum.expired.value}' + WITH t, pdrr, tp + ORDER BY (tp.datetime_end IS NULL) DESC, tp.datetime_claimed DESC + WITH pdrr, collect(tp)[0] AS tp + FOREACH (_ IN CASE WHEN tp IS NULL THEN [] ELSE [1] END | + SET tp.outcome = $outcome, + tp.datetime_end = datetime($now) + MERGE (tp)-[:PROVENANCE_OF]->(pdrr) + ) + """ + tx.run( + finalize_q, + task=str(task), + pdrr=str(scoped_key), + compute_service_id=str(compute_service_id), + outcome=outcome, + now=datetime.datetime.now(tz=datetime.UTC).isoformat(), + ) + return scoped_key def get_task_results(self, task: ScopedKey) -> list[ProtocolDAGResultRef]: @@ -3547,8 +3691,561 @@ 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`).""" + outcome = tp.get("outcome") + return TaskAttempt( + compute_service_id=tp["compute_service_id"], + hostname=tp.get("hostname"), + manager_name=tp.get("manager_name"), + datetime_claimed=_coerce_datetime(tp.get("datetime_claimed")), + datetime_end=_coerce_datetime(tp.get("datetime_end")), + outcome=TaskOutcomeEnum(outcome) if outcome is not None else None, + units_completed=tp.get("units_completed"), + units_total=tp.get("units_total"), + protocoldagresultref=( + ScopedKey.from_str(pdrr_sk) if pdrr_sk is not None else None + ), + ) + + def get_task_history( + self, task: ScopedKey, limit: int | None = None + ) -> list[TaskAttempt]: + """Return the execution attempts of a `Task`, most recent first. + + Each `TaskAttempt` bundles a `TaskProvenance` record with the + `ScopedKey` of the `ProtocolDAGResultRef` it produced (where one + exists). If `limit` is given, only the `limit` most recent attempts are + returned. + """ + 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 + ORDER BY tp.datetime_claimed DESC + """ + if limit is not None: + q += "\n LIMIT $limit" + + params = {"task": str(task)} + if limit is not None: + params["limit"] = limit + + attempts = [] + with self.transaction() as tx: + for record in tx.run(q, **params): + attempts.append( + self._task_provenance_node_to_attempt( + record["tp"], record["pdrr_sk"] + ) + ) + return attempts + + def get_tasks_details(self, tasks: list[ScopedKey]) -> list[TaskDetails | None]: + """Return `TaskDetails` for each given `Task`, in input order. + + `None` is returned in place of any `Task` that does not exist. The + `current_claim`'s live progress fields stay `None` until a compute + service reports progress (section 2 of the design). + """ + q = """ + UNWIND $tasks AS task_sk + OPTIONAL MATCH (t:Task {_scoped_key: task_sk}) + + CALL { + WITH t + OPTIONAL MATCH (t)<-[:PROVENANCE_OF]-(tp:TaskProvenance) + WITH tp ORDER BY tp.datetime_claimed DESC + RETURN count(tp) AS num_claims, collect(tp)[0] AS latest_tp + } + + OPTIONAL MATCH (latest_tp)-[:PROVENANCE_OF]->(latest_pdrr:ProtocolDAGResultRef) + 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 + + RETURN task_sk, + t, + num_claims, + latest_tp, + latest_pdrr._scoped_key AS latest_pdrr_sk, + cl.claimed AS claimed, + csreg.identifier AS csid, + csreg.hostname AS cs_hostname, + claim_tp.units_completed AS units_completed, + claim_tp.units_total AS units_total + """ + by_task = {} + with self.transaction() as tx: + for record in tx.run(q, tasks=[str(t) for t in tasks]): + t = record["t"] + if t is None: + by_task[record["task_sk"]] = None + continue + + current_claim = None + if record["csid"] is not None: + current_claim = TaskClaim( + compute_service_id=record["csid"], + hostname=record["cs_hostname"], + datetime_claimed=_coerce_datetime(record["claimed"]), + units_completed=record["units_completed"], + units_total=record["units_total"], + ) + + 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"] + ) + + by_task[record["task_sk"]] = TaskDetails( + task=ScopedKey.from_str(record["task_sk"]), + status=TaskStatusEnum(t["status"]), + datetime_status_changed=_coerce_datetime( + t.get("datetime_status_changed") + ), + reason=t.get("reason"), + num_claims=record["num_claims"], + current_claim=current_claim, + most_recent_attempt=most_recent_attempt, + ) + + return [by_task.get(str(t)) for t in tasks] + + def get_task_tracebacks( + self, task: ScopedKey, limit: int | None = None + ) -> list[TaskTracebacks]: + """Return tracebacks for the failed `ProtocolDAGResult`s of a `Task`. + + One `TaskTracebacks` record per failed `ProtocolDAGResultRef`, most + recent first (by `datetime_created`). Where per-unit + `ProtocolUnitResultRef`s exist (section 3.4), each failure entry carries + the `ScopedKey` of the corresponding unit ref; otherwise it is `None`. + """ + q = """ + MATCH (t:Task {_scoped_key: $task})-[:RESULTS_IN]->(pdrr:ProtocolDAGResultRef {ok: false})<-[:DETAILS]-(tb:Tracebacks) + OPTIONAL MATCH (pdrr)-[:CONTAINS]->(purr:ProtocolUnitResultRef) + WITH pdrr, tb, collect(purr) AS purrs + RETURN pdrr, tb, purrs + ORDER BY pdrr.datetime_created DESC + """ + if limit is not None: + q += "\n LIMIT $limit" + + params = {"task": str(task)} + if limit is not None: + params["limit"] = limit + + records = [] + with self.transaction() as tx: + for record in tx.run(q, **params): + pdrr = record["pdrr"] + tb = record["tb"] + + # map unit-result gufe key -> ProtocolUnitResultRef ScopedKey + purr_by_obj_key = { + purr["obj_key"]: purr["_scoped_key"] for purr in record["purrs"] + } + + tracebacks = tb["tracebacks"] + source_keys = tb["source_keys"] + failure_keys = tb["failure_keys"] + + unit_tracebacks = [] + for traceback, source_key, failure_key in zip( + tracebacks, source_keys, failure_keys + ): + purr_sk = purr_by_obj_key.get(failure_key) + unit_tracebacks.append( + TaskUnitTraceback( + failure_key=GufeKey(failure_key), + source_key=GufeKey(source_key), + traceback=traceback, + protocolunitresultref=( + ScopedKey.from_str(purr_sk) + if purr_sk is not None + else None + ), + ) + ) + + records.append( + TaskTracebacks( + protocoldagresultref=ScopedKey.from_str(pdrr["_scoped_key"]), + datetime_created=_coerce_datetime(pdrr.get("datetime_created")), + creator=pdrr.get("creator"), + tracebacks=unit_tracebacks, + ) + ) + + return records + + def get_scope_compute_share(self, scope: Scope) -> float: + """Return the fraction of currently-`running` `Task`s in `scope` + relative to all `Scope`s at the same level. + + - `Scope('org')` -> the org's running Tasks / all running Tasks; + - `Scope('org', 'campaign')` -> the campaign's / all campaigns in that org; + - `Scope('org', 'campaign', 'project')` -> the project's / all projects + in that org-campaign. + + Only the aggregate fraction is returned; no per-sibling counts are + disclosed. Returns 0.0 when there are no running Tasks in the relevant + population. + """ + if scope.project is not None: + level = "_project" + filters = {"_org": scope.org, "_campaign": scope.campaign} + target = scope.project + elif scope.campaign is not None: + level = "_campaign" + filters = {"_org": scope.org} + target = scope.campaign + elif scope.org is not None: + level = "_org" + filters = {} + target = scope.org + else: + raise ValueError( + "`scope` must specify at least an org to compute a compute share" + ) + + if filters: + filter_string = " {" + ", ".join(f"{k}: ${k}" for k in filters) + "}" + else: + filter_string = "" + + q = f""" + MATCH (t:Task{filter_string}) + WHERE t.status = '{TaskStatusEnum.running.value}' + RETURN t.{level} AS grouping, count(t) AS counts + """ + + with self.transaction() as tx: + res = tx.run(q, **filters) + counts = {rec["grouping"]: rec["counts"] for rec in res} + + total = sum(counts.values()) + if total == 0: + return 0.0 + return counts.get(target, 0) / total + + ## live progress reporting (section 2) + + def update_task_progress( + self, + compute_service_id: ComputeServiceID, + progress: dict[str, tuple[int, int]], + ) -> None: + """Write live progress counts onto the open `TaskProvenance` attempts. + + `progress` maps `Task` ScopedKey strings to + ``(units_completed, units_total)``. An update is dropped for any Task + the sending service no longer holds a `CLAIMS` relationship to (its + claim expired mid-flight) --- there is no live attempt to update. + """ + q = """ + UNWIND $items AS item + MATCH (t:Task {_scoped_key: item.task})<-[:CLAIMS]-(csreg:ComputeServiceRegistration {identifier: $compute_service_id}) + MATCH (t)<-[:PROVENANCE_OF]-(tp:TaskProvenance {compute_service_id: $compute_service_id}) + WHERE tp.datetime_end IS NULL + SET tp.units_completed = item.units_completed, + tp.units_total = item.units_total + """ + items = [ + {"task": task, "units_completed": uc, "units_total": ut} + for task, (uc, ut) in progress.items() + ] + if not items: + return + with self.transaction() as tx: + tx.run(q, items=items, compute_service_id=str(compute_service_id)) + + def get_tasks_progress( + self, tasks: list[ScopedKey] + ) -> list[tuple[int, int] | None]: + """Return ``(units_completed, units_total)`` for each `running` `Task` + with reported progress, `None` otherwise, in input order.""" + q = """ + UNWIND $tasks AS task_sk + OPTIONAL MATCH (t:Task {_scoped_key: task_sk})<-[:CLAIMS]-(csreg:ComputeServiceRegistration) + OPTIONAL MATCH (t)<-[:PROVENANCE_OF]-(tp:TaskProvenance {compute_service_id: csreg.identifier}) + WHERE tp.datetime_end IS NULL AND t.status = '%s' + RETURN task_sk, tp.units_completed AS uc, tp.units_total AS ut + """ % TaskStatusEnum.running.value + + by_task = {} + with self.transaction() as tx: + for rec in tx.run(q, tasks=[str(t) for t in tasks]): + uc, ut = rec["uc"], rec["ut"] + by_task[rec["task_sk"]] = ( + (uc, ut) if uc is not None and ut is not None else None + ) + return [by_task.get(str(t)) for t in tasks] + + ## per-unit result refs and artifacts (section 3.4) + + def add_protocol_unit_result_refs( + self, + protocoldagresultref: ProtocolDAGResultRef, + protocoldagresultref_scoped_key: ScopedKey, + protocoldagresult: ProtocolDAGResult, + ) -> dict[GufeKey, ScopedKey]: + """Derive one `ProtocolUnitResultRef` per `ProtocolUnitResult`/`Failure`. + + Creates a `ProtocolUnitResultRef` node for every unit result in the + `ProtocolDAGResult` (successes and failures, one per *result* --- so a + retried unit yields several), links each to the `ProtocolDAGResultRef` + via `CONTAINS`, and reproduces the execution topology from + `ProtocolDAGResult.result_graph` as `UNIT_DEPENDS_ON` edges. The + dedicated `UNIT_DEPENDS_ON` type is used deliberately (never + `DEPENDS_ON`), so these payload-less topology edges are never fed into + the generic gufe-object reconstruction machinery. + + Artifact-presence flags start `False`; the object store layout for each + unit result's artifacts is recorded on `location`. Returns a mapping of + unit-result gufe key -> `ProtocolUnitResultRef` ScopedKey. + + Idempotent: a `ProtocolDAGResult` gufe key is deterministic, so a + replayed result push matches the same `ProtocolDAGResultRef`. If unit + refs already exist for it, this returns the existing mapping without + re-creating them --- a re-merge would reset `has_logs`/`has_stdout`/ + `has_stderr` (flipped by later, separate requests) back to `False`. + """ + scope = protocoldagresultref_scoped_key.scope + + # short-circuit if unit refs already exist for this ProtocolDAGResultRef + existing = {} + with self.transaction() as tx: + res = tx.run( + """ + MATCH (:ProtocolDAGResultRef {_scoped_key: $pdrr})-[:CONTAINS]->(purr:ProtocolUnitResultRef) + RETURN purr.obj_key AS obj_key, purr._scoped_key AS sk + """, + pdrr=str(protocoldagresultref_scoped_key), + ) + for rec in res: + existing[GufeKey(rec["obj_key"])] = ScopedKey.from_str(rec["sk"]) + if existing: + return existing + + pdrr_node = self._get_node(protocoldagresultref_scoped_key) + + base_location = ( + os.path.dirname(protocoldagresultref.location) + if protocoldagresultref.location + else None + ) + + subgraph = Subgraph() + result_key_to_node = {} + result_key_to_sk: dict[GufeKey, ScopedKey] = {} + + for result in protocoldagresult.protocol_unit_results: + unit_location = ( + os.path.join(base_location, "units", str(result.key)) + if base_location is not None + else None + ) + purr = ProtocolUnitResultRef( + location=unit_location, + obj_key=result.key, + source_key=result.source_key, + scope=scope, + ok=result.ok(), + name=result.name, + start_time=result.start_time, + end_time=result.end_time, + ) + _, purr_node, purr_sk = self._keyed_chain_to_subgraph( + KeyedChain.from_gufe(purr), + scope=scope, + ) + subgraph = subgraph | Relationship.type("CONTAINS")( + pdrr_node, + purr_node, + _org=scope.org, + _campaign=scope.campaign, + _project=scope.project, + ) + result_key_to_node[result.key] = purr_node + result_key_to_sk[result.key] = purr_sk + + # reproduce execution topology (result -> its dependency result) + for node, dependency in protocoldagresult.result_graph.edges(): + na = result_key_to_node.get(node.key) + nb = result_key_to_node.get(dependency.key) + if na is not None and nb is not None: + subgraph = subgraph | Relationship.type("UNIT_DEPENDS_ON")( + na, + nb, + _org=scope.org, + _campaign=scope.campaign, + _project=scope.project, + ) + + with self.transaction() as tx: + merge_subgraph(tx, subgraph, "GufeTokenizable", "_scoped_key") + + return result_key_to_sk + + def get_protocol_unit_result_ref_scoped_key( + self, + protocoldagresultref_scoped_key: ScopedKey, + unit_result_key: GufeKey, + task: ScopedKey | None = None, + ) -> ScopedKey | None: + """Return the `ProtocolUnitResultRef` ScopedKey for a given unit-result + gufe key under a `ProtocolDAGResultRef`, or `None` if absent. + + If `task` is given, the `ProtocolDAGResultRef` must be a result of that + `Task` (`(task)-[:RESULTS_IN]->(pdrr)`); otherwise `None` is returned. + This lets callers refuse a mismatched ``(task, pdrr)`` pair. + """ + if task is not None: + q = """ + MATCH (t:Task {_scoped_key: $task})-[:RESULTS_IN]->(pdrr:ProtocolDAGResultRef {_scoped_key: $pdrr})-[:CONTAINS]->(purr:ProtocolUnitResultRef {obj_key: $obj_key}) + RETURN purr._scoped_key AS sk + """ + else: + q = """ + MATCH (pdrr:ProtocolDAGResultRef {_scoped_key: $pdrr})-[:CONTAINS]->(purr:ProtocolUnitResultRef {obj_key: $obj_key}) + RETURN purr._scoped_key AS sk + """ + with self.transaction() as tx: + res = tx.run( + q, + task=str(task) if task is not None else None, + pdrr=str(protocoldagresultref_scoped_key), + obj_key=str(unit_result_key), + ).to_eager_result() + if not res.records: + return None + return ScopedKey.from_str(res.records[0]["sk"]) + + def set_protocol_unit_result_ref_artifacts( + self, + protocol_unit_result_ref_scoped_key: ScopedKey, + *, + has_logs: bool | None = None, + has_stdout: bool | None = None, + has_stderr: bool | None = None, + ) -> None: + """Flip artifact-presence flags on a `ProtocolUnitResultRef` as artifacts + are stored. Only the flags passed (non-`None`) are written.""" + sets = [] + params = {"purr": str(protocol_unit_result_ref_scoped_key)} + if has_logs is not None: + sets.append("purr.has_logs = $has_logs") + params["has_logs"] = has_logs + if has_stdout is not None: + sets.append("purr.has_stdout = $has_stdout") + params["has_stdout"] = has_stdout + if has_stderr is not None: + sets.append("purr.has_stderr = $has_stderr") + params["has_stderr"] = has_stderr + if not sets: + return + q = f""" + MATCH (purr:ProtocolUnitResultRef {{_scoped_key: $purr}}) + SET {', '.join(sets)} + """ + with self.transaction() as tx: + tx.run(q, **params) + + def get_task_result_recs( + self, task: ScopedKey, ok: bool | None = None + ) -> list[ProtocolDAGResultRec]: + """Return one `ProtocolDAGResultRec` per `ProtocolDAGResult` of the + `Task`, most recent first. `ok` filters to successes/failures.""" + where = "" + params = {"task": str(task)} + if ok is not None: + where = "WHERE pdrr.ok = $ok" + params["ok"] = ok + q = f""" + MATCH (t:Task {{_scoped_key: $task}})-[:RESULTS_IN]->(pdrr:ProtocolDAGResultRef) + {where} + RETURN pdrr + ORDER BY pdrr.datetime_created DESC + """ + recs = [] + with self.transaction() as tx: + for rec in tx.run(q, **params): + pdrr = rec["pdrr"] + recs.append( + ProtocolDAGResultRec( + scoped_key=ScopedKey.from_str(pdrr["_scoped_key"]), + ok=pdrr["ok"], + datetime_created=_coerce_datetime(pdrr.get("datetime_created")), + creator=pdrr.get("creator"), + ) + ) + return recs + + def get_result_unit_recs( + self, protocoldagresultref: ScopedKey + ) -> list[ProtocolUnitResultRec]: + """Return one `ProtocolUnitResultRec` per `ProtocolUnitResult` of the + `ProtocolDAGResult`, in dependency order (via `UNIT_DEPENDS_ON`).""" + q = """ + MATCH (pdrr:ProtocolDAGResultRef {_scoped_key: $pdrr})-[:CONTAINS]->(purr:ProtocolUnitResultRef) + OPTIONAL MATCH (purr)-[:UNIT_DEPENDS_ON]->(dep:ProtocolUnitResultRef)<-[:CONTAINS]-(pdrr) + RETURN purr, collect(dep._scoped_key) AS deps + """ + nodes_by_sk = {} + deps_by_sk = {} + with self.transaction() as tx: + for rec in tx.run(q, pdrr=str(protocoldagresultref)): + purr = rec["purr"] + sk = purr["_scoped_key"] + nodes_by_sk[sk] = purr + deps_by_sk[sk] = [d for d in rec["deps"] if d is not None] + + # topologically order: a unit result's dependencies come before it + g = nx.DiGraph() + g.add_nodes_from(nodes_by_sk) + for sk, deps in deps_by_sk.items(): + for dep in deps: + # edge dependency -> dependent, so topological_sort yields deps first + g.add_edge(dep, sk) + try: + ordered = list(nx.topological_sort(g)) + except nx.NetworkXUnfeasible: + # cyclic (should never happen for a DAG); fall back to start_time + ordered = sorted( + nodes_by_sk, + key=lambda s: nodes_by_sk[s].get("start_time") or "", + ) + + recs = [] + for sk in ordered: + purr = nodes_by_sk[sk] + recs.append( + ProtocolUnitResultRec( + scoped_key=ScopedKey.from_str(sk), + obj_key=GufeKey(purr["obj_key"]), + source_key=GufeKey(purr["source_key"]), + name=purr.get("name"), + ok=purr["ok"], + start_time=_coerce_datetime(purr.get("start_time")), + end_time=_coerce_datetime(purr.get("end_time")), + has_logs=purr.get("has_logs", False), + has_stdout=purr.get("has_stdout", False), + has_stderr=purr.get("has_stderr", False), + ) + ) + return recs + def set_task_status( - self, tasks: list[ScopedKey], status: TaskStatusEnum, raise_error: bool = False + self, + tasks: list[ScopedKey], + status: TaskStatusEnum, + raise_error: bool = False, + reason: str | None = None, ) -> list[ScopedKey | None]: """Set the status of a list of Tasks. @@ -3563,6 +4260,9 @@ def set_task_status( The status to set the Task to. raise_error If `True`, raise a `ValueError` if the status of a given Task cannot be changed. + reason + Optional human-readable reason for the status change; only recorded + for `invalid`/`deleted` transitions (ignored otherwise). Returns ------- @@ -3572,6 +4272,8 @@ def set_task_status( """ method = getattr(self, f"set_task_{status.value}") + if status in (TaskStatusEnum.invalid, TaskStatusEnum.deleted): + return method(tasks, raise_error=raise_error, reason=reason) return method(tasks, raise_error=raise_error) def get_task_status(self, tasks: list[ScopedKey]) -> list[TaskStatusEnum]: @@ -3606,11 +4308,17 @@ def get_task_status(self, tasks: list[ScopedKey]) -> list[TaskStatusEnum]: return statuses def _set_task_status( - self, tasks, q: str, err_msg_func, raise_error + self, tasks, q: str, err_msg_func, raise_error, extra_params: dict | None = None ) -> list[ScopedKey | None]: tasks_statused = [] + params = {"scoped_keys": [str(t) for t in tasks]} + # every status-mutation query writes `datetime_status_changed` via the + # `_status_write` helper, keyed to the `statuschange_time` parameter + params["statuschange_time"] = datetime.datetime.now(tz=datetime.UTC).isoformat() + if extra_params: + params.update(extra_params) with self.transaction() as tx: - res = tx.run(q, scoped_keys=[str(t) for t in tasks]) + res = tx.run(q, **params) for record in res: task_i = record["t"] @@ -3648,7 +4356,16 @@ def set_task_waiting( OPTIONAL MATCH (t_:Task {{_scoped_key: scoped_key}}) WHERE t_.status IN ['{TaskStatusEnum.waiting.value}', '{TaskStatusEnum.running.value}', '{TaskStatusEnum.error.value}'] - SET t_.status = '{TaskStatusEnum.waiting.value}' + {_status_write('t_', TaskStatusEnum.waiting.value)} + + WITH scoped_key, t, t_ + + // if we forced a `running` Task back to `waiting`, its open provenance + // attempt was released before it could produce a result + OPTIONAL MATCH (t_)<-[:PROVENANCE_OF]-(tp:TaskProvenance) + WHERE tp.datetime_end IS NULL + SET tp.outcome = '{TaskOutcomeEnum.released.value}', + tp.datetime_end = datetime($statuschange_time) WITH scoped_key, t, t_ @@ -3682,7 +4399,7 @@ def set_task_running( OPTIONAL MATCH (t_:Task {{_scoped_key: scoped_key}}) WHERE t_.status IN ['{TaskStatusEnum.running.value}', '{TaskStatusEnum.waiting.value}'] - SET t_.status = '{TaskStatusEnum.running.value}' + {_status_write('t_', TaskStatusEnum.running.value)} RETURN scoped_key, t, t_ """ @@ -3709,7 +4426,7 @@ def set_task_complete( OPTIONAL MATCH (t_:Task {{_scoped_key: scoped_key}}) WHERE t_.status IN ['{TaskStatusEnum.complete.value}', '{TaskStatusEnum.running.value}'] - SET t_.status = '{TaskStatusEnum.complete.value}' + {_status_write('t_', TaskStatusEnum.complete.value)} WITH scoped_key, t, t_ @@ -3736,12 +4453,48 @@ def err_msg(t, status): return self._set_task_status(tasks, q, err_msg, raise_error=raise_error) def set_task_error( - self, tasks: list[ScopedKey], raise_error: bool = False + self, + tasks: list[ScopedKey], + raise_error: bool = False, + reason: str | None = None, + compute_service_id: ComputeServiceID | None = None, ) -> list[ScopedKey | None]: """Set the status of a list of Tasks to `error`. Only `running` Tasks can be set to `error`. + Parameters + ---------- + tasks + The Tasks to set to `error`. + raise_error + If `True`, raise a `ValueError` for any Task whose status cannot be + changed. + reason + If given, recorded on `Task.reason`; used by the `ProtocolDAG` + creation-failure path to hand the user the failure traceback. When + `None`, `reason` is cleared. + compute_service_id + If given, the open `TaskProvenance` attempt for this + ``(task, compute_service_id)`` pair is finalized with + `outcome = error`. Used by the creation-failure path, where no + `ProtocolDAGResult` (and hence no `set_task_result`) exists to + finalize provenance. + """ + + reason_expr = "$reason" if reason is not None else "null" + + finalize_provenance = "" + if compute_service_id is not None: + finalize_provenance = f""" + WITH scoped_key, t, t_ + + // no result exists in the creation-failure path, so finalize the open + // provenance attempt for this service directly + OPTIONAL MATCH (t_)<-[:PROVENANCE_OF]-(tp:TaskProvenance {{compute_service_id: $compute_service_id}}) + WHERE tp.datetime_end IS NULL + SET tp.outcome = '{TaskOutcomeEnum.error.value}', + tp.datetime_end = datetime($statuschange_time) """ q = f""" @@ -3752,7 +4505,7 @@ def set_task_error( OPTIONAL MATCH (t_:Task {{_scoped_key: scoped_key}}) WHERE t_.status IN ['{TaskStatusEnum.error.value}', '{TaskStatusEnum.running.value}'] - SET t_.status = '{TaskStatusEnum.error.value}' + {_status_write('t_', TaskStatusEnum.error.value, reason_expr=reason_expr)} WITH scoped_key, t, t_ @@ -3760,96 +4513,118 @@ def set_task_error( // drop CLAIMS relationship OPTIONAL MATCH (t_)<-[cl:CLAIMS]-(csreg:ComputeServiceRegistration) DELETE cl - + {finalize_provenance} RETURN scoped_key, t, t_ """ def err_msg(t, status): return f"Cannot set task {t} with current status: {status} to `error` as it is not currently `running`." - return self._set_task_status(tasks, q, err_msg, raise_error=raise_error) + extra_params = {} + if reason is not None: + extra_params["reason"] = reason + if compute_service_id is not None: + extra_params["compute_service_id"] = str(compute_service_id) + + return self._set_task_status( + tasks, q, err_msg, raise_error=raise_error, extra_params=extra_params + ) def set_task_invalid( - self, tasks: list[ScopedKey], raise_error: bool = False + self, + tasks: list[ScopedKey], + raise_error: bool = False, + reason: str | None = None, ) -> list[ScopedKey | None]: """Set the status of a list of Tasks to `invalid`. Any Task can be set to `invalid`; an `invalid` Task cannot change to any other status. + Parameters + ---------- + reason + If given, recorded on `Task.reason` for the targeted Tasks; + otherwise `reason` is cleared. """ # set the status and delete the ACTIONS relationship # make sure we follow the extends chain and set all tasks to invalid # and remove actions relationships - q = f""" - WITH $scoped_keys AS batch - UNWIND batch AS scoped_key - - OPTIONAL MATCH (t:Task {{_scoped_key: scoped_key}}) - - OPTIONAL MATCH (t_:Task {{_scoped_key: scoped_key}}) - WHERE NOT t_.status IN ['{TaskStatusEnum.deleted.value}'] - SET t_.status = '{TaskStatusEnum.invalid.value}' - - WITH scoped_key, t, t_ - - OPTIONAL MATCH (t_)<-[er:EXTENDS*]-(extends_task:Task) - SET extends_task.status = '{TaskStatusEnum.invalid.value}' - - WITH scoped_key, t, t_, extends_task - - OPTIONAL MATCH (t_)<-[ar:ACTIONS]-(th:TaskHub) - OPTIONAL MATCH (extends_task)<-[ar_e:ACTIONS]-(th:TaskHub) - OPTIONAL MATCH (t_)<-[applies:APPLIES]-(:TaskRestartPattern) - OPTIONAL MATCH (extends_task)<-[applies_e:APPLIES]-(:TaskRestartPattern) - - DELETE ar - DELETE ar_e - DELETE applies - DELETE applies_e - - WITH scoped_key, t, t_ - - // drop CLAIMS relationship if present - OPTIONAL MATCH (t_)<-[cl:CLAIMS]-(csreg:ComputeServiceRegistration) - DELETE cl - - RETURN scoped_key, t, t_ - """ + q = self._invalidate_or_delete_query( + TaskStatusEnum.invalid.value, + excluded_status=TaskStatusEnum.deleted.value, + reason=reason, + ) def err_msg(t, status): return f"Cannot set task {t} with current status: {status} to `invalid` as it is `deleted`." - return self._set_task_status(tasks, q, err_msg, raise_error=raise_error) + extra_params = {"reason": reason} if reason is not None else {} + return self._set_task_status( + tasks, q, err_msg, raise_error=raise_error, extra_params=extra_params + ) def set_task_deleted( - self, tasks: list[ScopedKey], raise_error: bool = False + self, + tasks: list[ScopedKey], + raise_error: bool = False, + reason: str | None = None, ) -> list[ScopedKey | None]: """Set the status of a list of Tasks to `deleted`. Any Task can be set to `deleted`; a `deleted` Task cannot change to any other status. + Parameters + ---------- + reason + If given, recorded on `Task.reason` for the targeted Tasks; + otherwise `reason` is cleared. """ # set the status and delete the ACTIONS relationship # make sure we follow the extends chain and set all tasks to deleted # and remove actions relationships - q = f""" + q = self._invalidate_or_delete_query( + TaskStatusEnum.deleted.value, + excluded_status=TaskStatusEnum.invalid.value, + reason=reason, + ) + + def err_msg(t, status): + return f"Cannot set task {t} with current status: {status} to `deleted` as it is `invalid`." + + extra_params = {"reason": reason} if reason is not None else {} + return self._set_task_status( + tasks, q, err_msg, raise_error=raise_error, extra_params=extra_params + ) + + @staticmethod + def _invalidate_or_delete_query( + status_value: str, *, excluded_status: str, reason: str | None + ) -> str: + """Build the shared Cypher for `set_task_invalid`/`set_task_deleted`. + + Both set a target status on the Task and its `EXTENDS` descendants, + drop ACTIONS/APPLIES/CLAIMS, and release any open provenance attempts + (the user forced a `running` Task off-attempt). + """ + reason_expr = "$reason" if reason is not None else "null" + return f""" WITH $scoped_keys AS batch UNWIND batch AS scoped_key OPTIONAL MATCH (t:Task {{_scoped_key: scoped_key}}) OPTIONAL MATCH (t_:Task {{_scoped_key: scoped_key}}) - WHERE NOT t_.status IN ['{TaskStatusEnum.invalid.value}'] - SET t_.status = '{TaskStatusEnum.deleted.value}' + WHERE NOT t_.status IN ['{excluded_status}'] + {_status_write('t_', status_value, reason_expr=reason_expr)} WITH scoped_key, t, t_ OPTIONAL MATCH (t_)<-[er:EXTENDS*]-(extends_task:Task) - SET extends_task.status = '{TaskStatusEnum.deleted.value}' + {_status_write('extends_task', status_value)} WITH scoped_key, t, t_, extends_task @@ -3865,18 +4640,28 @@ def set_task_deleted( WITH scoped_key, t, t_ + // if a `running` Task (or running descendant) was forced off its + // attempt, that attempt was released before producing a result + OPTIONAL MATCH (released_task:Task)<-[:PROVENANCE_OF]-(tp:TaskProvenance) + WHERE (released_task = t_ OR (t_)<-[:EXTENDS*]-(released_task)) + AND tp.datetime_end IS NULL + SET tp.outcome = '{TaskOutcomeEnum.released.value}', + tp.datetime_end = datetime($statuschange_time) + + WITH scoped_key, t, t_ + // drop CLAIMS relationship if present OPTIONAL MATCH (t_)<-[cl:CLAIMS]-(csreg:ComputeServiceRegistration) DELETE cl + // collapse the row fan-out introduced by the EXTENDS-descendant and + // released-provenance matches, so `_set_task_status` returns exactly one + // entry per input Task (aligned with the input order) + WITH DISTINCT scoped_key, t, t_ + RETURN scoped_key, t, t_ """ - def err_msg(t, status): - return f"Cannot set task {t} with current status: {status} to `deleted` as it is `invalid`." - - return self._set_task_status(tasks, q, err_msg, raise_error=raise_error) - ## task restart policies def add_task_restart_patterns( @@ -4186,17 +4971,17 @@ def resolve_task_restarts(self, task_scoped_keys: Iterable[ScopedKey], *, tx=Non self.cancel_tasks(tasks, taskhub, tx=tx) # any tasks that are still associated with a TaskHub and a TaskRestartPattern must then be okay to switch to waiting - renew_waiting_status_query = """ + renew_waiting_status_query = f""" UNWIND $task_scoped_keys AS task_scoped_key - MATCH (task:Task {status: $error, `_scoped_key`: task_scoped_key})<-[app:APPLIES]-(trp:TaskRestartPattern)-[:ENFORCES]->(taskhub:TaskHub) - SET task.status = $waiting + MATCH (task:Task {{status: $error, `_scoped_key`: task_scoped_key}})<-[app:APPLIES]-(trp:TaskRestartPattern)-[:ENFORCES]->(taskhub:TaskHub) + {_status_write('task', TaskStatusEnum.waiting.value)} """ tx.run( renew_waiting_status_query, task_scoped_keys=list(map(str, task_scoped_keys)), - waiting=TaskStatusEnum.waiting.value, error=TaskStatusEnum.error.value, + statuschange_time=datetime.datetime.now(tz=datetime.UTC).isoformat(), ) ## authentication diff --git a/alchemiscale/tests/integration/storage/test_statestore_introspection.py b/alchemiscale/tests/integration/storage/test_statestore_introspection.py new file mode 100644 index 00000000..4467a0ac --- /dev/null +++ b/alchemiscale/tests/integration/storage/test_statestore_introspection.py @@ -0,0 +1,662 @@ +"""Integration tests for the v0.8.0 Task-introspection state-store surface. + +Covers the ``Neo4jStore`` additions from the introspection work: durable +``TaskProvenance`` records (creation at claim; finalization at result/expiry/ +deregistration/release; the late-result overwrite rules), the +``datetime_status_changed``/``reason`` indicators, per-unit +``ProtocolUnitResultRef``s, live progress, and compute share --- exercised +against a real Neo4j instance via the same harness as ``test_statestore.py``. +""" + +import datetime +from datetime import timedelta + +import pytest +from gufe.protocols import ProtocolUnitFailure + +from alchemiscale.storage.statestore import Neo4jStore +from alchemiscale.storage.models import ( + ComputeServiceID, + ComputeServiceRegistration, + ProtocolDAGResultRef, + TaskAttempt, + TaskDetails, + TaskOutcomeEnum, + TaskStatusEnum, + TaskTracebacks, +) +from alchemiscale.models import Scope, ScopedKey + + +def _register( + n4js: Neo4jStore, + compute_service_id: ComputeServiceID, + hostname: str | None = "host-a", + manager_name: str | None = None, +) -> ComputeServiceID: + """Register a compute service carrying a ``hostname`` (and optional manager).""" + now = datetime.datetime.now(tz=datetime.UTC) + registration = ComputeServiceRegistration( + identifier=compute_service_id, + registered=now, + heartbeat=now, + failure_times=[], + hostname=hostname, + manager_name=manager_name, + ) + return n4js.register_computeservice(registration) + + +def _provenance_nodes(n4js: Neo4jStore, task: ScopedKey) -> list: + """Return the raw ``TaskProvenance`` nodes for a Task, newest claim first. + + ``PROVENANCE_OF`` points from the provenance node to the ``Task``. + """ + q = """ + MATCH (tp:TaskProvenance)-[:PROVENANCE_OF]->(t:Task {_scoped_key: $task}) + RETURN tp + ORDER BY tp.datetime_claimed DESC + """ + return [rec["tp"] for rec in n4js.execute_query(q, task=str(task)).records] + + +class TestStateStoreIntrospection: + + @pytest.fixture + def n4js(self, n4js_fresh): + return n4js_fresh + + def _claimed_task( + self, + n4js: Neo4jStore, + network, + transformation, + scope_test, + compute_service_id: ComputeServiceID, + hostname: str | None = "host-a", + manager_name: str | None = None, + ): + """Assemble a network, create+action a single Task, and claim it. + + Returns ``(task_sk, taskhub_sk)`` with the Task now ``running`` and an + open ``TaskProvenance`` attempt created at claim. + """ + _, taskhub_sk, _ = n4js.assemble_network(network, scope_test) + 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) + claimed = n4js.claim_taskhub_tasks(taskhub_sk, compute_service_id) + assert claimed[0] == task_sk + return task_sk, taskhub_sk + + # --- provenance creation at claim ------------------------------------- + + def test_provenance_created_at_claim( + self, n4js, network_tyk2, transformation, scope_test + ): + csid = ComputeServiceID.new_from_name("prov.claim") + task_sk, _ = self._claimed_task( + n4js, + network_tyk2, + transformation, + scope_test, + csid, + hostname="cluster-node-7", + manager_name=None, + ) + + nodes = _provenance_nodes(n4js, task_sk) + assert len(nodes) == 1 + tp = nodes[0] + assert tp["compute_service_id"] == str(csid) + assert tp["hostname"] == "cluster-node-7" + assert tp.get("manager_name") is None + assert tp.get("datetime_claimed") is not None + # open attempt: not yet finalized + assert tp.get("datetime_end") is None + assert tp.get("outcome") is None + + # the Task itself flipped to running with a status-change timestamp + task_node = n4js.execute_query( + "MATCH (t:Task {_scoped_key: $task}) RETURN t", task=str(task_sk) + ).records[0]["t"] + assert task_node["status"] == TaskStatusEnum.running.value + assert task_node.get("datetime_status_changed") is not None + + # --- finalization: complete / error via set_task_result --------------- + + def test_provenance_finalized_complete( + self, n4js, network_tyk2, transformation, scope_test + ): + csid = ComputeServiceID.new_from_name("prov.complete") + task_sk, _ = self._claimed_task( + n4js, network_tyk2, transformation, scope_test, csid + ) + + pdrr = ProtocolDAGResultRef( + scope=task_sk.scope, obj_key=task_sk.gufe_key, ok=True + ) + pdrr_sk = n4js.set_task_result(task_sk, pdrr, compute_service_id=csid) + + tp = _provenance_nodes(n4js, task_sk)[0] + assert tp["outcome"] == TaskOutcomeEnum.complete.value + assert tp.get("datetime_end") is not None + + # PROVENANCE_OF edge to the produced ProtocolDAGResultRef + linked = n4js.execute_query( + """ + MATCH (tp:TaskProvenance {compute_service_id: $csid})-[:PROVENANCE_OF]->(pdrr:ProtocolDAGResultRef) + RETURN pdrr._scoped_key AS sk + """, + csid=str(csid), + ).records + assert linked[0]["sk"] == str(pdrr_sk) + + def test_provenance_finalized_error( + self, n4js, network_tyk2, transformation, scope_test + ): + csid = ComputeServiceID.new_from_name("prov.error") + task_sk, _ = self._claimed_task( + n4js, network_tyk2, transformation, scope_test, csid + ) + + pdrr = ProtocolDAGResultRef( + scope=task_sk.scope, obj_key=task_sk.gufe_key, ok=False + ) + n4js.set_task_result(task_sk, pdrr, compute_service_id=csid) + + tp = _provenance_nodes(n4js, task_sk)[0] + assert tp["outcome"] == TaskOutcomeEnum.error.value + assert tp.get("datetime_end") is not None + + # --- finalization: expired / released --------------------------------- + + def test_provenance_expired_on_expire_registrations( + self, n4js, network_tyk2, transformation, scope_test + ): + csid = ComputeServiceID.new_from_name("prov.expire") + task_sk, _ = self._claimed_task( + n4js, network_tyk2, transformation, scope_test, csid + ) + + # force the heartbeat into the past so the registration expires + 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) + ) + + tp = _provenance_nodes(n4js, task_sk)[0] + assert tp["outcome"] == TaskOutcomeEnum.expired.value + assert tp.get("datetime_end") is not None + # units_* last-reported values would remain here; the Task returns to waiting + task_node = n4js.execute_query( + "MATCH (t:Task {_scoped_key: $task}) RETURN t", task=str(task_sk) + ).records[0]["t"] + assert task_node["status"] == TaskStatusEnum.waiting.value + + def test_provenance_expired_on_deregister( + self, n4js, network_tyk2, transformation, scope_test + ): + csid = ComputeServiceID.new_from_name("prov.dereg") + task_sk, _ = self._claimed_task( + n4js, network_tyk2, transformation, scope_test, csid + ) + + n4js.deregister_computeservice(csid) + + tp = _provenance_nodes(n4js, task_sk)[0] + assert tp["outcome"] == TaskOutcomeEnum.expired.value + assert tp.get("datetime_end") is not None + + @pytest.mark.parametrize("force_status", ("waiting", "invalid", "deleted")) + def test_provenance_released_on_user_status_change( + self, n4js, network_tyk2, transformation, scope_test, force_status + ): + csid = ComputeServiceID.new_from_name(f"prov.release.{force_status}") + task_sk, _ = self._claimed_task( + n4js, network_tyk2, transformation, scope_test, csid + ) + + n4js.set_task_status([task_sk], TaskStatusEnum(force_status)) + + tp = _provenance_nodes(n4js, task_sk)[0] + assert tp["outcome"] == TaskOutcomeEnum.released.value + assert tp.get("datetime_end") is not None + + # --- late-result race rules (M4) -------------------------------------- + + def test_late_result_overwrites_expired_record( + self, n4js, network_tyk2, transformation, scope_test + ): + csid = ComputeServiceID.new_from_name("prov.late.expired") + task_sk, _ = self._claimed_task( + n4js, network_tyk2, transformation, scope_test, csid + ) + + # registration expires; the open attempt closes as expired + n4js.deregister_computeservice(csid) + tp = _provenance_nodes(n4js, task_sk)[0] + assert tp["outcome"] == TaskOutcomeEnum.expired.value + + # a late result for the SAME service arrives; the attempt did finish, so + # its expired record is overwritten to complete + pdrr = ProtocolDAGResultRef( + scope=task_sk.scope, obj_key=task_sk.gufe_key, ok=True + ) + n4js.set_task_result(task_sk, pdrr, compute_service_id=csid) + + tp = _provenance_nodes(n4js, task_sk)[0] + assert tp["outcome"] == TaskOutcomeEnum.complete.value + + def test_late_result_does_not_overwrite_released_record( + self, n4js, network_tyk2, transformation, scope_test + ): + csid = ComputeServiceID.new_from_name("prov.late.released") + task_sk, _ = self._claimed_task( + n4js, network_tyk2, transformation, scope_test, csid + ) + + # user forces the running Task back to waiting; the attempt is released + n4js.set_task_waiting([task_sk]) + tp = _provenance_nodes(n4js, task_sk)[0] + assert tp["outcome"] == TaskOutcomeEnum.released.value + + # a late result for the same service must NOT resurrect the released + # record (the immutable history that a user ended the attempt stands) + pdrr = ProtocolDAGResultRef( + scope=task_sk.scope, obj_key=task_sk.gufe_key, ok=True + ) + n4js.set_task_result(task_sk, pdrr, compute_service_id=csid) + + tp = _provenance_nodes(n4js, task_sk)[0] + assert tp["outcome"] == TaskOutcomeEnum.released.value + + def test_restart_churn_yields_distinct_records( + self, n4js, network_tyk2, transformation, scope_test + ): + # two attempts by two services: first expires, second completes + csid1 = ComputeServiceID.new_from_name("prov.attempt.one") + task_sk, taskhub_sk = self._claimed_task( + n4js, network_tyk2, transformation, scope_test, csid1 + ) + n4js.deregister_computeservice(csid1) # attempt 1 -> expired, back to waiting + + csid2 = ComputeServiceID.new_from_name("prov.attempt.two") + _register(n4js, csid2, hostname="host-b") + claimed = n4js.claim_taskhub_tasks(taskhub_sk, csid2) + assert claimed[0] == task_sk + pdrr = ProtocolDAGResultRef( + scope=task_sk.scope, obj_key=task_sk.gufe_key, ok=True + ) + n4js.set_task_result(task_sk, pdrr, compute_service_id=csid2) + + nodes = _provenance_nodes(n4js, task_sk) + assert len(nodes) == 2 + outcomes = {n["compute_service_id"]: n["outcome"] for n in nodes} + assert outcomes[str(csid1)] == TaskOutcomeEnum.expired.value + assert outcomes[str(csid2)] == TaskOutcomeEnum.complete.value + + # --- status-change indicator + reason --------------------------------- + + def test_datetime_status_changed_and_reason_on_change( + self, n4js, network_tyk2, transformation, scope_test + ): + _, taskhub_sk, _ = n4js.assemble_network(network_tyk2, scope_test) + transformation_sk = n4js.get_scoped_key(transformation, scope_test) + task_sk = n4js.create_task(transformation_sk) + + # a genuine transition sets the timestamp and the reason + n4js.set_task_invalid([task_sk], reason="operator marked bad input") + node = n4js.execute_query( + "MATCH (t:Task {_scoped_key: $task}) RETURN t", task=str(task_sk) + ).records[0]["t"] + assert node["status"] == TaskStatusEnum.invalid.value + assert node["reason"] == "operator marked bad input" + ts1 = node["datetime_status_changed"] + assert ts1 is not None + + def test_status_write_idempotent_noop_preserves_indicators( + self, n4js, network_tyk2, transformation, scope_test + ): + """A no-op re-set must not reset datetime_status_changed or wipe reason. + + This validates the ``_status_write`` CASE guards against real Neo4j SET + semantics (the guard relies on the CASE reading the pre-clause status). + """ + _, taskhub_sk, _ = n4js.assemble_network(network_tyk2, scope_test) + transformation_sk = n4js.get_scoped_key(transformation, scope_test) + task_sk = n4js.create_task(transformation_sk) + + n4js.set_task_invalid([task_sk], reason="first reason") + node = n4js.execute_query( + "MATCH (t:Task {_scoped_key: $task}) RETURN t", task=str(task_sk) + ).records[0]["t"] + ts1 = node["datetime_status_changed"] + + # re-assert the SAME status with a DIFFERENT reason: no-op transition + n4js.set_task_invalid([task_sk], reason="second reason") + node = n4js.execute_query( + "MATCH (t:Task {_scoped_key: $task}) RETURN t", task=str(task_sk) + ).records[0]["t"] + + assert node["status"] == TaskStatusEnum.invalid.value + # neither the change-timestamp nor the reason was clobbered (compare the + # native datetimes to avoid any neo4j DateTime equality quirk) + assert node["datetime_status_changed"].to_native() == ts1.to_native() + assert node["reason"] == "first reason" + + def test_reason_cleared_on_transition_to_waiting( + self, n4js, network_tyk2, transformation, scope_test + ): + csid = ComputeServiceID.new_from_name("reason.clear") + task_sk, _ = self._claimed_task( + n4js, network_tyk2, transformation, scope_test, csid + ) + # error the running Task with a reason (DAG-creation-failure path) + n4js.set_task_error( + [task_sk], reason="boom during create", compute_service_id=csid + ) + node = n4js.execute_query( + "MATCH (t:Task {_scoped_key: $task}) RETURN t", task=str(task_sk) + ).records[0]["t"] + assert node["reason"] == "boom during create" + + # transition back to waiting clears the reason (describes current status) + n4js.set_task_waiting([task_sk]) + node = n4js.execute_query( + "MATCH (t:Task {_scoped_key: $task}) RETURN t", task=str(task_sk) + ).records[0]["t"] + assert node["status"] == TaskStatusEnum.waiting.value + assert node.get("reason") is None + + def test_set_task_error_with_reason_finalizes_provenance( + self, n4js, network_tyk2, transformation, scope_test + ): + # the DAG-creation-failure path: error + reason + provenance error, no PDRR + csid = ComputeServiceID.new_from_name("error.reason.prov") + task_sk, _ = self._claimed_task( + n4js, network_tyk2, transformation, scope_test, csid + ) + n4js.set_task_error([task_sk], reason="traceback text", compute_service_id=csid) + tp = _provenance_nodes(n4js, task_sk)[0] + assert tp["outcome"] == TaskOutcomeEnum.error.value + assert tp.get("datetime_end") is not None + + # --- get_task_history -------------------------------------------------- + + def test_get_task_history(self, n4js, network_tyk2, transformation, scope_test): + csid1 = ComputeServiceID.new_from_name("hist.one") + task_sk, taskhub_sk = self._claimed_task( + n4js, network_tyk2, transformation, scope_test, csid1, hostname="h1" + ) + n4js.deregister_computeservice(csid1) # attempt 1 -> expired (no result) + + csid2 = ComputeServiceID.new_from_name("hist.two") + _register(n4js, csid2, hostname="h2") + assert n4js.claim_taskhub_tasks(taskhub_sk, csid2)[0] == task_sk + pdrr = ProtocolDAGResultRef( + scope=task_sk.scope, obj_key=task_sk.gufe_key, ok=True + ) + pdrr_sk = n4js.set_task_result(task_sk, pdrr, compute_service_id=csid2) + + history = n4js.get_task_history(task_sk) + assert len(history) == 2 + assert all(isinstance(a, TaskAttempt) for a in history) + + # most recent first: the completing attempt + assert history[0].compute_service_id == str(csid2) + assert history[0].hostname == "h2" + assert history[0].outcome == TaskOutcomeEnum.complete + assert history[0].protocoldagresultref == pdrr_sk + + # the earlier expired attempt has no result ref + assert history[1].compute_service_id == str(csid1) + assert history[1].outcome == TaskOutcomeEnum.expired + assert history[1].protocoldagresultref is None + + # limit + assert len(n4js.get_task_history(task_sk, limit=1)) == 1 + + # --- get_tasks_details ------------------------------------------------- + + def test_get_tasks_details(self, n4js, network_tyk2, transformation, scope_test): + csid = ComputeServiceID.new_from_name("details.svc") + task_sk, _ = self._claimed_task( + n4js, network_tyk2, transformation, scope_test, csid, hostname="dhost" + ) + + # a second, unclaimed task, plus a nonexistent one + transformation_sk = n4js.get_scoped_key(transformation, scope_test) + task_sk2 = n4js.create_task(transformation_sk) + missing = ScopedKey(gufe_key="Task-doesnotexist", **scope_test.to_dict()) + + details = n4js.get_tasks_details([task_sk, task_sk2, missing]) + assert len(details) == 3 + + d0 = details[0] + assert isinstance(d0, TaskDetails) + assert d0.task == task_sk + assert d0.status == TaskStatusEnum.running + assert d0.datetime_status_changed is not None + assert d0.num_claims == 1 + assert d0.current_claim is not None + assert d0.current_claim.compute_service_id == str(csid) + assert d0.current_claim.hostname == "dhost" + assert d0.most_recent_attempt is not None + assert d0.most_recent_attempt.compute_service_id == str(csid) + + # unclaimed waiting task: no claim, no attempts + d1 = details[1] + assert d1.status == TaskStatusEnum.waiting + assert d1.num_claims == 0 + assert d1.current_claim is None + assert d1.most_recent_attempt is None + + # missing task -> None in place, order preserved + assert details[2] is None + + # --- get_task_tracebacks ---------------------------------------------- + + def test_get_task_tracebacks(self, n4js, network_tyk2, transformation, scope_test): + csid = ComputeServiceID.new_from_name("tb.svc") + task_sk, _ = self._claimed_task( + n4js, network_tyk2, transformation, scope_test, csid + ) + + pdrr = ProtocolDAGResultRef( + scope=task_sk.scope, obj_key=task_sk.gufe_key, ok=False + ) + pdrr_sk = n4js.set_task_result(task_sk, pdrr, compute_service_id=csid) + + pufs = [ + ProtocolUnitFailure( + source_key=f"FakeProtocolUnitKey-{i}", + inputs={}, + outputs={}, + exception=("RuntimeError", ("boom",)), + traceback=f"traceback number {i}", + ) + for i in range(2) + ] + n4js.add_protocol_dag_result_ref_tracebacks(pufs, pdrr_sk) + + tbs = n4js.get_task_tracebacks(task_sk) + assert len(tbs) == 1 + assert isinstance(tbs[0], TaskTracebacks) + assert tbs[0].protocoldagresultref == pdrr_sk + returned = {t.traceback for t in tbs[0].tracebacks} + assert returned == {"traceback number 0", "traceback number 1"} + # no unit refs stored for this synthetic result, so no unit ref link + assert all(t.protocolunitresultref is None for t in tbs[0].tracebacks) + + # --- per-unit result refs (real ProtocolDAGResult) -------------------- + + def test_add_and_get_unit_result_refs( + self, n4js, network_tyk2, transformation, scope_test, protocoldagresults + ): + csid = ComputeServiceID.new_from_name("units.svc") + task_sk, _ = self._claimed_task( + n4js, network_tyk2, transformation, scope_test, csid + ) + + pdr = protocoldagresults[0] + location = ( + f"protocoldagresult/{'/'.join(task_sk.scope.to_tuple())}/" + f"{transformation.key}/results/{pdr.key}/obj.json.zst" + ) + pdrr = ProtocolDAGResultRef( + scope=task_sk.scope, obj_key=pdr.key, ok=True, location=location + ) + pdrr_sk = n4js.set_task_result(task_sk, pdrr, compute_service_id=csid) + + refs_map = n4js.add_protocol_unit_result_refs(pdrr, pdrr_sk, pdr) + assert len(refs_map) == len(pdr.protocol_unit_results) + + # CONTAINS edges from the pdrr to each unit ref + contains = n4js.execute_query( + """ + MATCH (pdrr:ProtocolDAGResultRef {_scoped_key: $pdrr})-[:CONTAINS]->(purr:ProtocolUnitResultRef) + RETURN count(purr) AS n + """, + pdrr=str(pdrr_sk), + ).records[0]["n"] + assert contains == len(pdr.protocol_unit_results) + + # UNIT_DEPENDS_ON edges reproduce the execution topology (a DummyProtocol + # DAG has at least one dependency edge), and are NOT DEPENDS_ON + unit_depends = n4js.execute_query( + """ + MATCH (:ProtocolUnitResultRef)-[r:UNIT_DEPENDS_ON]->(:ProtocolUnitResultRef) + RETURN count(r) AS n + """, + ).records[0]["n"] + assert unit_depends >= 1 + + # records come back one per unit result, in dependency order + recs = n4js.get_result_unit_recs(pdrr_sk) + assert len(recs) == len(pdr.protocol_unit_results) + # location recorded under the unit prefix + assert all( + "units/" in r_location + for r_location in [n4js.get_gufe(rec.scoped_key).location for rec in recs] + ) + + def test_add_protocol_unit_result_refs_idempotent( + self, n4js, network_tyk2, transformation, scope_test, protocoldagresults + ): + csid = ComputeServiceID.new_from_name("units.idem") + task_sk, _ = self._claimed_task( + n4js, network_tyk2, transformation, scope_test, csid + ) + pdr = protocoldagresults[0] + pdrr = ProtocolDAGResultRef(scope=task_sk.scope, obj_key=pdr.key, ok=True) + pdrr_sk = n4js.set_task_result(task_sk, pdrr, compute_service_id=csid) + + refs_map = n4js.add_protocol_unit_result_refs(pdrr, pdrr_sk, pdr) + # a later, separate request flips has_logs on one unit ref + some_purr = next(iter(refs_map.values())) + n4js.set_protocol_unit_result_ref_artifacts(some_purr, has_logs=True) + + # a duplicate result push (same deterministic pdrr) must NOT wipe flags + refs_map2 = n4js.add_protocol_unit_result_refs(pdrr, pdrr_sk, pdr) + assert set(map(str, refs_map2.values())) == set(map(str, refs_map.values())) + + purr = n4js.get_gufe(some_purr) + assert purr.has_logs is True + + def test_get_task_result_recs_ok_filter( + self, n4js, network_tyk2, transformation, scope_test + ): + csid = ComputeServiceID.new_from_name("recs.svc") + task_sk, taskhub_sk = self._claimed_task( + n4js, network_tyk2, transformation, scope_test, csid + ) + ok_ref = ProtocolDAGResultRef( + scope=task_sk.scope, obj_key="ProtocolDAGResult-okresult", ok=True + ) + fail_ref = ProtocolDAGResultRef( + scope=task_sk.scope, obj_key="ProtocolDAGResult-failresult", ok=False + ) + n4js.set_task_result(task_sk, ok_ref, compute_service_id=csid) + n4js.set_task_result(task_sk, fail_ref, compute_service_id=csid) + + all_recs = n4js.get_task_result_recs(task_sk) + assert len(all_recs) == 2 + assert {r.ok for r in all_recs} == {True, False} + + assert all(r.ok for r in n4js.get_task_result_recs(task_sk, ok=True)) + assert all(not r.ok for r in n4js.get_task_result_recs(task_sk, ok=False)) + + # --- live progress ----------------------------------------------------- + + def test_update_and_get_tasks_progress( + self, n4js, network_tyk2, transformation, scope_test + ): + csid = ComputeServiceID.new_from_name("progress.svc") + task_sk, _ = self._claimed_task( + n4js, network_tyk2, transformation, scope_test, csid + ) + + n4js.update_task_progress(csid, {str(task_sk): (3, 10)}) + progress = n4js.get_tasks_progress([task_sk]) + assert progress == [(3, 10)] + + # progress also lands on the open provenance record + tp = _provenance_nodes(n4js, task_sk)[0] + assert tp["units_completed"] == 3 + assert tp["units_total"] == 10 + + def test_progress_dropped_when_claim_absent( + self, n4js, network_tyk2, transformation, scope_test + ): + # a task with no claim by this service: the update is dropped + _, taskhub_sk, _ = n4js.assemble_network(network_tyk2, scope_test) + transformation_sk = n4js.get_scoped_key(transformation, scope_test) + task_sk = n4js.create_task(transformation_sk) + + csid = ComputeServiceID.new_from_name("progress.noclaim") + _register(n4js, csid) + n4js.update_task_progress(csid, {str(task_sk): (1, 5)}) + + # waiting (unclaimed) task reports no progress + assert n4js.get_tasks_progress([task_sk]) == [None] + + # --- compute share ----------------------------------------------------- + + def test_get_scope_compute_share(self, n4js, network_tyk2, transformation): + # set up two orgs with running tasks in a shared campaign/project space + scope_a = Scope("orgA", "camp", "proj") + scope_b = Scope("orgB", "camp", "proj") + + def running_tasks(scope, count, name): + _, taskhub_sk, _ = n4js.assemble_network( + network_tyk2.copy_with_replacements(name=network_tyk2.name + name), + scope, + ) + tf_sk = n4js.get_scoped_key(transformation, scope) + task_sks = n4js.create_tasks([tf_sk] * count) + n4js.action_tasks(task_sks, taskhub_sk) + csid = ComputeServiceID.new_from_name(f"share{name}") + _register(n4js, csid) + claimed = n4js.claim_taskhub_tasks(taskhub_sk, csid, count=count) + assert all(c is not None for c in claimed) + + running_tasks(scope_a, 3, "a") # 3 running in orgA + running_tasks(scope_b, 1, "b") # 1 running in orgB + + # orgA's share of all running tasks across orgs: 3 / (3 + 1) + share = n4js.get_scope_compute_share(Scope(org="orgA")) + assert share == pytest.approx(0.75) + + # empty population -> 0.0 + assert n4js.get_scope_compute_share(Scope(org="orgC")) == 0.0 diff --git a/alchemiscale/tests/unit/compute/__init__.py b/alchemiscale/tests/unit/compute/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/alchemiscale/tests/unit/compute/test_capture.py b/alchemiscale/tests/unit/compute/test_capture.py new file mode 100644 index 00000000..575b2c65 --- /dev/null +++ b/alchemiscale/tests/unit/compute/test_capture.py @@ -0,0 +1,331 @@ +"""Unit tests for :mod:`alchemiscale.compute.capture`. + +These exercise the log-capture handler and the ``SynchronousExecutionHooks`` +seam in isolation --- no Neo4j, no network, no real ``ProtocolUnit`` execution. +We drive the hooks by hand (calling ``on_unit_attempt_start`` / +``on_unit_attempt_end`` / ``on_progress`` directly) and inspect the resulting +``unit_logs`` and the ``gufekey`` logger's handler set. +""" + +import logging + +import pytest + +from gufe.tokenization import GufeKey + +from alchemiscale.compute.capture import ( + GUFEKEY_LOGGER_NAME, + GufeKeyLogHandler, + SynchronousExecutionHooks, +) +from alchemiscale.models import ScopedKey +from alchemiscale.storage.models import ComputeServiceID + +# --------------------------------------------------------------------------- +# fixtures / helpers +# --------------------------------------------------------------------------- + + +class FakeProgressCallback: + """Records every ``(task, units_completed, units_total)`` call.""" + + def __init__(self): + self.calls = [] + + def __call__(self, task, units_completed, units_total): + self.calls.append((task, units_completed, units_total)) + + +class DummyResult: + """Minimal stand-in for a ``ProtocolUnitResult`` --- only exposes ``.key``.""" + + def __init__(self, key): + self.key = key + + +@pytest.fixture +def task(): + return ScopedKey.from_str("Task-x-o-c-p") + + +@pytest.fixture +def compute_service_id(): + return ComputeServiceID("svc-" + "0" * 32) + + +@pytest.fixture +def progress_callback(): + return FakeProgressCallback() + + +@pytest.fixture +def gufekey_logger(): + return logging.getLogger(GUFEKEY_LOGGER_NAME) + + +@pytest.fixture(autouse=True) +def _restore_gufekey_logger(gufekey_logger): + """Snapshot/restore the process-global ``gufekey`` logger around each test. + + The hooks mutate this logger's level and handler set; keep tests hermetic. + """ + saved_level = gufekey_logger.level + saved_handlers = list(gufekey_logger.handlers) + try: + yield + finally: + gufekey_logger.setLevel(saved_level) + gufekey_logger.handlers[:] = saved_handlers + + +def _make_hooks(task, compute_service_id, progress_callback, **kwargs): + return SynchronousExecutionHooks( + task=task, + compute_service_id=compute_service_id, + progress_callback=progress_callback, + **kwargs, + ) + + +def _emit(logger_name, msg, *, gufekey=None): + """Emit one record on ``logger_name``. + + When ``gufekey`` is given, stamp it onto the record via ``extra`` exactly as + ``gufe``'s ``_GufeLoggerAdapter`` would; otherwise emit a bare record with no + ``gufekey`` attribute. + """ + logger = logging.getLogger(logger_name) + extra = {"gufekey": gufekey} if gufekey is not None else None + logger.info(msg, extra=extra) + + +# a syntactically valid gufe key for labeling captured records +_GUFE_KEY = GufeKey("FakeUnit-" + "a" * 32) + + +# --------------------------------------------------------------------------- +# GufeKeyLogHandler +# --------------------------------------------------------------------------- + + +def test_handler_accumulates_lines(): + handler = GufeKeyLogHandler() + rec = logging.LogRecord( + name="gufekey.x", + level=logging.INFO, + pathname=__file__, + lineno=1, + msg="hello", + args=(), + exc_info=None, + ) + rec.gufekey = str(_GUFE_KEY) + handler.emit(rec) + handler.emit(rec) + + assert len(handler.lines) == 2 + # each formatted line carries the gufe key and message + for line in handler.lines: + assert str(_GUFE_KEY) in line + assert "hello" in line + + +def test_handler_missing_gufekey_renders_dash_no_raise(): + """A record with no ``gufekey`` attribute must not raise and renders ``[-]``.""" + handler = GufeKeyLogHandler() + # a bare record, as if ``logging.getLogger("gufekey.x").info(...)`` were + # called without the GufeTokenizable.logger adapter's stamp + rec = logging.LogRecord( + name="gufekey.x", + level=logging.INFO, + pathname=__file__, + lineno=1, + msg="unstamped", + args=(), + exc_info=None, + ) + assert not hasattr(rec, "gufekey") + + handler.emit(rec) # must not raise + + assert len(handler.lines) == 1 + assert "[-]" in handler.lines[0] + assert "unstamped" in handler.lines[0] + + +def test_handler_text_joins_lines(): + handler = GufeKeyLogHandler() + handler.lines = ["a", "b", "c"] + assert handler.text() == "a\nb\nc" + + +def test_handler_text_empty(): + handler = GufeKeyLogHandler() + assert handler.text() == "" + + +# --------------------------------------------------------------------------- +# SynchronousExecutionHooks --- handler attach/detach + capture +# --------------------------------------------------------------------------- + + +def test_start_attaches_end_removes_handler( + task, compute_service_id, progress_callback, gufekey_logger +): + baseline = len(gufekey_logger.handlers) + hooks = _make_hooks(task, compute_service_id, progress_callback) + + hooks.on_unit_attempt_start(unit=None, attempt=0) + assert len(gufekey_logger.handlers) == baseline + 1 + assert isinstance(hooks._handler, GufeKeyLogHandler) + + hooks.on_unit_attempt_end(unit=None, attempt=0, result=DummyResult(_GUFE_KEY)) + # handler count returns to baseline + assert len(gufekey_logger.handlers) == baseline + assert hooks._handler is None + + +def test_record_between_start_and_end_is_captured( + task, compute_service_id, progress_callback +): + hooks = _make_hooks(task, compute_service_id, progress_callback) + result = DummyResult(_GUFE_KEY) + + hooks.on_unit_attempt_start(unit=None, attempt=0) + # a record on a gufekey descendant logger, stamped with a real gufe key + _emit("gufekey.some.Unit", "captured message", gufekey=str(_GUFE_KEY)) + hooks.on_unit_attempt_end(unit=None, attempt=0, result=result) + + stored = hooks.unit_logs[str(result.key)] + assert "captured message" in stored + assert str(_GUFE_KEY) in stored + + +def test_records_before_start_and_after_end_not_captured( + task, compute_service_id, progress_callback +): + hooks = _make_hooks(task, compute_service_id, progress_callback) + result = DummyResult(_GUFE_KEY) + + # before the handler is attached + _emit("gufekey.some.Unit", "before start", gufekey=str(_GUFE_KEY)) + + hooks.on_unit_attempt_start(unit=None, attempt=0) + _emit("gufekey.some.Unit", "during", gufekey=str(_GUFE_KEY)) + hooks.on_unit_attempt_end(unit=None, attempt=0, result=result) + + # after the handler is removed + _emit("gufekey.some.Unit", "after end", gufekey=str(_GUFE_KEY)) + + stored = hooks.unit_logs[str(result.key)] + assert "during" in stored + assert "before start" not in stored + assert "after end" not in stored + + +def test_truncation_keeps_tail(task, compute_service_id, progress_callback): + cap = 200 + hooks = _make_hooks(task, compute_service_id, progress_callback, log_cap_bytes=cap) + result = DummyResult(_GUFE_KEY) + + hooks.on_unit_attempt_start(unit=None, attempt=0) + # emit far more than `cap` bytes, with a distinctive marker at the very end + for i in range(200): + _emit("gufekey.some.Unit", f"line-{i:05d}-padding", gufekey=str(_GUFE_KEY)) + _emit("gufekey.some.Unit", "THE_VERY_LAST_LINE", gufekey=str(_GUFE_KEY)) + hooks.on_unit_attempt_end(unit=None, attempt=0, result=result) + + stored = hooks.unit_logs[str(result.key)] + # kept text is the tail and within the cap + assert len(stored.encode("utf-8")) <= cap + assert "THE_VERY_LAST_LINE" in stored + # an early line should have been dropped from the head + assert "line-00000-padding" not in stored + + +def test_end_with_none_result_closes_handler_stores_nothing( + task, compute_service_id, progress_callback, gufekey_logger +): + baseline = len(gufekey_logger.handlers) + hooks = _make_hooks(task, compute_service_id, progress_callback) + + hooks.on_unit_attempt_start(unit=None, attempt=0) + _emit("gufekey.some.Unit", "should be dropped", gufekey=str(_GUFE_KEY)) + # interrupted attempt: result is None + hooks.on_unit_attempt_end(unit=None, attempt=0, result=None) + + # handler closed (count back to baseline), nothing stored + assert len(gufekey_logger.handlers) == baseline + assert hooks._handler is None + assert hooks.unit_logs == {} + + +def test_empty_capture_not_stored(task, compute_service_id, progress_callback): + """A result with no captured text stores no entry (guarded by ``if text``).""" + hooks = _make_hooks(task, compute_service_id, progress_callback) + result = DummyResult(_GUFE_KEY) + + hooks.on_unit_attempt_start(unit=None, attempt=0) + # emit nothing + hooks.on_unit_attempt_end(unit=None, attempt=0, result=result) + + assert str(result.key) not in hooks.unit_logs + + +# --------------------------------------------------------------------------- +# SynchronousExecutionHooks --- capture_logs=False disables capture entirely +# --------------------------------------------------------------------------- + + +def test_capture_disabled_no_handler_no_logs_no_level_mutation( + task, compute_service_id, progress_callback, gufekey_logger +): + baseline_handlers = len(gufekey_logger.handlers) + baseline_level = gufekey_logger.level + + hooks = _make_hooks(task, compute_service_id, progress_callback, capture_logs=False) + + # constructing the hooks must NOT mutate the (process-global) logger level + assert gufekey_logger.level == baseline_level + + result = DummyResult(_GUFE_KEY) + hooks.on_unit_attempt_start(unit=None, attempt=0) + # no handler attached + assert len(gufekey_logger.handlers) == baseline_handlers + assert hooks._handler is None + + _emit("gufekey.some.Unit", "ignored", gufekey=str(_GUFE_KEY)) + hooks.on_unit_attempt_end(unit=None, attempt=0, result=result) + + # nothing captured, level still untouched + assert hooks.unit_logs == {} + assert gufekey_logger.level == baseline_level + + +def test_capture_enabled_sets_logger_level( + task, compute_service_id, progress_callback, gufekey_logger +): + """With capture on, the ``gufekey`` logger level is set at construction.""" + hooks = _make_hooks( + task, compute_service_id, progress_callback, gufekey_loglevel=logging.INFO + ) + assert gufekey_logger.level == logging.INFO + + +# --------------------------------------------------------------------------- +# SynchronousExecutionHooks --- progress forwarding +# --------------------------------------------------------------------------- + + +def test_on_progress_forwards_to_callback_with_task( + task, compute_service_id, progress_callback +): + hooks = _make_hooks(task, compute_service_id, progress_callback) + + hooks.on_progress(0, 5) + hooks.on_progress(3, 5) + + assert progress_callback.calls == [ + (task, 0, 5), + (task, 3, 5), + ] diff --git a/alchemiscale/tests/unit/compute/test_execute_equivalence.py b/alchemiscale/tests/unit/compute/test_execute_equivalence.py new file mode 100644 index 00000000..ec01b9b9 --- /dev/null +++ b/alchemiscale/tests/unit/compute/test_execute_equivalence.py @@ -0,0 +1,759 @@ +"""Behavioral-equivalence tests for the alchemiscale-owned DAG executor. + +This module is the guard run on every ``gufe`` upgrade. It runs *identical* +:class:`~gufe.protocols.ProtocolDAG`\\ s through both + +* :func:`gufe.protocols.protocoldag.execute_DAG` (the upstream reference), and +* :func:`alchemiscale.compute.execute.execute_DAG` (the alchemiscale fork), + +and asserts the resulting :class:`~gufe.protocols.ProtocolDAGResult`\\ s are +*behaviorally equivalent*. + +Equivalence is compared **structurally**, never by result ``.key``: gufe units +tokenize with a ``uuid4`` per attempt, so two independent runs of the same DAG +produce results with different keys. We instead compare, over the two runs: + +* overall ``.ok()``, +* the number of ``protocol_unit_results``, +* the multiset of result ``source_key``\\ s (which *is* stable --- it is the + key of the originating :class:`ProtocolUnit`, shared by both runs of the + same DAG object), +* the per-result ``.ok()`` sequence (grouped by ``source_key``, since raw + ordering of same-level units is not guaranteed to match), and +* terminal / success / failure counts. + +A separate, non-equivalence test exercises the alchemiscale-only +:class:`~alchemiscale.compute.execute.ExecutionHooks` seam directly. +""" + +from collections import Counter + +import pytest +from rdkit import Chem +from rdkit.Chem import AllChem + +from gufe import ChemicalSystem, SmallMoleculeComponent +from gufe.protocols import Protocol, ProtocolUnit +from gufe.protocols import protocoldag as gufe_protocoldag +from gufe.protocols.errors import ExecutionInterrupt +from gufe.tests.test_protocol import BrokenProtocol, DummyProtocol + +from alchemiscale.compute import execute as alchemiscale_execute +from alchemiscale.compute.execute import ExecutionHooks + +# --------------------------------------------------------------------------- +# chemical-system + DAG fixtures +# +# The unit-test conftest only provides a heavyweight tyk2 network (module +# scope). For executor equivalence we want the smallest possible systems, so +# build minimal ones here (verified working per the task brief). +# --------------------------------------------------------------------------- + + +def _mol(smiles: str, name: str) -> SmallMoleculeComponent: + m = Chem.AddHs(Chem.MolFromSmiles(smiles)) + AllChem.Compute2DCoords(m) + return SmallMoleculeComponent.from_rdkit(m, name=name) + + +@pytest.fixture(scope="module") +def stateA() -> ChemicalSystem: + return ChemicalSystem({"ligand": _mol("CCO", "ethanol")}) + + +@pytest.fixture(scope="module") +def stateB() -> ChemicalSystem: + return ChemicalSystem({"ligand": _mol("CCC", "propane")}) + + +@pytest.fixture +def success_dag(stateA, stateB): + """A DummyProtocol DAG: every unit succeeds (1 init + 21 sims + 1 finish).""" + proto = DummyProtocol(settings=DummyProtocol.default_settings()) + return proto.create(stateA=stateA, stateB=stateB, name="success") + + +@pytest.fixture +def failure_dag(stateA, stateB): + """A BrokenProtocol DAG: exactly one unit always fails, halting the DAG.""" + proto = BrokenProtocol(settings=BrokenProtocol.default_settings()) + return proto.create(stateA=stateA, stateB=stateB, name="failure") + + +# --------------------------------------------------------------------------- +# equivalence helpers +# --------------------------------------------------------------------------- + + +def _make_dirs(base, names): + out = [] + for name in names: + d = base / name + d.mkdir(parents=True, exist_ok=True) + out.append(d) + return out + + +def _ok_by_source(pdr): + """Multiset of ``(source_key, ok)`` pairs across all unit results. + + Keyed on ``source_key`` (the stable originating-``ProtocolUnit`` key) rather + than the per-attempt result ``.key`` (a fresh uuid4 each run). This captures + "which units produced results, and with what ok-status" without depending on + same-level ordering. + """ + return Counter((str(r.source_key), r.ok()) for r in pdr.protocol_unit_results) + + +def assert_equivalent(gufe_pdr, alch_pdr): + """Assert two ``ProtocolDAGResult``\\ s are behaviorally equivalent.""" + # overall success/failure + assert gufe_pdr.ok() == alch_pdr.ok(), "overall .ok() differs" + + # same number of results + assert len(gufe_pdr.protocol_unit_results) == len( + alch_pdr.protocol_unit_results + ), "number of protocol_unit_results differs" + + # same multiset of source_keys (which units produced results) + gufe_sources = Counter(str(r.source_key) for r in gufe_pdr.protocol_unit_results) + alch_sources = Counter(str(r.source_key) for r in alch_pdr.protocol_unit_results) + assert gufe_sources == alch_sources, "multiset of source_keys differs" + + # same per-(source_key) .ok() multiset + assert _ok_by_source(gufe_pdr) == _ok_by_source( + alch_pdr + ), "per-source .ok() sequence differs" + + # same success / failure / terminal counts + assert len(gufe_pdr.protocol_unit_successes) == len( + alch_pdr.protocol_unit_successes + ), "success count differs" + assert len(gufe_pdr.protocol_unit_failures) == len( + alch_pdr.protocol_unit_failures + ), "failure count differs" + assert len(gufe_pdr.terminal_protocol_unit_results) == len( + alch_pdr.terminal_protocol_unit_results + ), "terminal result count differs" + + +def _run_both(dag, tmp_path, *, cache_basedir=None, **kwargs): + """Run ``dag`` through both executors in separate temp dirs. + + Each executor gets its own ``shared``/``scratch`` (and optional ``cache``) + trees so filesystem effects never cross-contaminate. Returns + ``(gufe_pdr, alch_pdr)``. + """ + gufe_dir = tmp_path / "gufe" + alch_dir = tmp_path / "alch" + g_shared, g_scratch = _make_dirs(gufe_dir, ["shared", "scratch"]) + a_shared, a_scratch = _make_dirs(alch_dir, ["shared", "scratch"]) + + g_cache = a_cache = None + if cache_basedir is not None: + (g_cache,) = _make_dirs(gufe_dir, ["cache"]) + (a_cache,) = _make_dirs(alch_dir, ["cache"]) + + gufe_pdr = gufe_protocoldag.execute_DAG( + dag, + shared_basedir=g_shared, + scratch_basedir=g_scratch, + cache_basedir=g_cache, + **kwargs, + ) + alch_pdr = alchemiscale_execute.execute_DAG( + dag, + shared_basedir=a_shared, + scratch_basedir=a_scratch, + cache_basedir=a_cache, + **kwargs, + ) + return gufe_pdr, alch_pdr + + +# --------------------------------------------------------------------------- +# 1. success DAG, n_retries=0 +# --------------------------------------------------------------------------- + + +def test_equivalence_success(success_dag, tmp_path): + gufe_pdr, alch_pdr = _run_both(success_dag, tmp_path, n_retries=0) + + assert gufe_pdr.ok() is True + assert alch_pdr.ok() is True + # all 23 units succeeded once + assert len(alch_pdr.protocol_unit_results) == len(success_dag.protocol_units) + assert_equivalent(gufe_pdr, alch_pdr) + + +# --------------------------------------------------------------------------- +# 2. failure DAG, n_retries=0 (one persistent failure halts the DAG) +# --------------------------------------------------------------------------- + + +def test_equivalence_failure_no_retry(failure_dag, tmp_path): + gufe_pdr, alch_pdr = _run_both( + failure_dag, tmp_path, n_retries=0, raise_error=False + ) + + assert gufe_pdr.ok() is False + assert alch_pdr.ok() is False + # both halt at the same point => same number of results + assert_equivalent(gufe_pdr, alch_pdr) + + +# --------------------------------------------------------------------------- +# 3. failure DAG, n_retries=2 (retries then halts) +# --------------------------------------------------------------------------- + + +def test_equivalence_failure_with_retries(failure_dag, tmp_path): + gufe_pdr, alch_pdr = _run_both( + failure_dag, tmp_path, n_retries=2, raise_error=False + ) + + assert gufe_pdr.ok() is False + assert alch_pdr.ok() is False + + # the single broken unit is attempted n_retries+1 == 3 times in both + assert len(gufe_pdr.protocol_unit_failures) == len(alch_pdr.protocol_unit_failures) + assert len(alch_pdr.protocol_unit_failures) == 3 + assert_equivalent(gufe_pdr, alch_pdr) + + +# --------------------------------------------------------------------------- +# 4. raise_error=True on the failure DAG (both raise the same type) +# --------------------------------------------------------------------------- + + +def test_equivalence_raise_error(failure_dag, tmp_path): + gufe_dir = tmp_path / "gufe" + alch_dir = tmp_path / "alch" + g_shared, g_scratch = _make_dirs(gufe_dir, ["shared", "scratch"]) + a_shared, a_scratch = _make_dirs(alch_dir, ["shared", "scratch"]) + + with pytest.raises(Exception) as gufe_exc: + gufe_protocoldag.execute_DAG( + failure_dag, + shared_basedir=g_shared, + scratch_basedir=g_scratch, + raise_error=True, + n_retries=0, + ) + + with pytest.raises(Exception) as alch_exc: + alchemiscale_execute.execute_DAG( + failure_dag, + shared_basedir=a_shared, + scratch_basedir=a_scratch, + raise_error=True, + n_retries=0, + ) + + # both implementations raise the *same* exception type + assert type(gufe_exc.value) is type(alch_exc.value) + # and it originates from the broken unit's ValueError + assert isinstance(alch_exc.value, ValueError) + assert "I have failed my mission" in str(alch_exc.value) + + +# --------------------------------------------------------------------------- +# 5. cache-resume equivalence +# +# Fair comparison design: each implementation gets its OWN cache dir, but both +# are seeded IDENTICALLY by first running that same implementation once with +# keep_cache=True. We then run each implementation a SECOND time against its +# now-populated cache and assert the two second-runs are equivalent to each +# other (and that both skipped re-execution). +# +# "Skipped execution" is detected structurally: on a fully-cached resume, no +# shared_* directories are created (the per-unit execute() branch is never +# entered), so shared_basedir stays empty. We assert that for both. +# --------------------------------------------------------------------------- + + +def test_equivalence_cache_resume(success_dag, tmp_path): + gufe_dir = tmp_path / "gufe" + alch_dir = tmp_path / "alch" + + def _run(module, base, cache): + shared, scratch = _make_dirs(base, ["shared", "scratch"]) + pdr = module.execute_DAG( + success_dag, + shared_basedir=shared, + scratch_basedir=scratch, + cache_basedir=cache, + keep_cache=True, + n_retries=0, + ) + return pdr, shared + + (g_cache,) = _make_dirs(gufe_dir, ["cache"]) + (a_cache,) = _make_dirs(alch_dir, ["cache"]) + + # first run: populate each cache + gufe_first, _ = _run(gufe_protocoldag, gufe_dir / "run1", g_cache) + alch_first, _ = _run(alchemiscale_execute, alch_dir / "run1", a_cache) + assert gufe_first.ok() and alch_first.ok() + + # second run against the populated cache: should skip execution entirely + gufe_second, gufe_shared2 = _run(gufe_protocoldag, gufe_dir / "run2", g_cache) + alch_second, alch_shared2 = _run(alchemiscale_execute, alch_dir / "run2", a_cache) + + # both second runs succeeded and are equivalent to each other + assert gufe_second.ok() and alch_second.ok() + assert_equivalent(gufe_second, alch_second) + + # both second runs are also equivalent to their respective first runs + assert_equivalent(gufe_first, gufe_second) + assert_equivalent(alch_first, alch_second) + + # cache hit => no unit executed => no shared_* dirs created, in BOTH + assert list(gufe_shared2.iterdir()) == [], "gufe re-executed despite cache" + assert list(alch_shared2.iterdir()) == [], "alchemiscale re-executed despite cache" + + +# --------------------------------------------------------------------------- +# 6. keep_shared / keep_scratch parity +# --------------------------------------------------------------------------- + + +@pytest.mark.parametrize( + "keep_shared,keep_scratch", + [(False, False), (True, False), (False, True), (True, True)], +) +def test_equivalence_keep_dirs(success_dag, tmp_path, keep_shared, keep_scratch): + gufe_dir = tmp_path / "gufe" + alch_dir = tmp_path / "alch" + g_shared, g_scratch = _make_dirs(gufe_dir, ["shared", "scratch"]) + a_shared, a_scratch = _make_dirs(alch_dir, ["shared", "scratch"]) + + gufe_pdr = gufe_protocoldag.execute_DAG( + success_dag, + shared_basedir=g_shared, + scratch_basedir=g_scratch, + keep_shared=keep_shared, + keep_scratch=keep_scratch, + n_retries=0, + ) + alch_pdr = alchemiscale_execute.execute_DAG( + success_dag, + shared_basedir=a_shared, + scratch_basedir=a_scratch, + keep_shared=keep_shared, + keep_scratch=keep_scratch, + n_retries=0, + ) + + assert_equivalent(gufe_pdr, alch_pdr) + + # directory-retention parity: presence/absence of per-unit subdirs must match + def _has_children(d): + return any(d.iterdir()) + + assert _has_children(g_shared) == _has_children(a_shared) == keep_shared + assert _has_children(g_scratch) == _has_children(a_scratch) == keep_scratch + + +# --------------------------------------------------------------------------- +# focused test of the ExecutionHooks seam (alchemiscale-only, non-equivalence) +# --------------------------------------------------------------------------- + + +class RecordingHooks(ExecutionHooks): + """Records every hook invocation for later assertion.""" + + def __init__(self): + self.dag_starts = [] # (units_total,) + self.starts = [] # (source_key, attempt) + self.ends = [] # (source_key, attempt, result_or_None) + self.progress = [] # (units_completed, units_total) + + def on_dag_start(self, protocoldag, units_total): + self.dag_starts.append(units_total) + + def on_unit_attempt_start(self, unit, attempt): + self.starts.append((str(unit.key), attempt)) + + def on_unit_attempt_end(self, unit, attempt, result): + self.ends.append((str(unit.key), attempt, result)) + + def on_progress(self, units_completed, units_total): + self.progress.append((units_completed, units_total)) + + +def test_hooks_success(success_dag, tmp_path): + shared, scratch = _make_dirs(tmp_path, ["shared", "scratch"]) + hooks = RecordingHooks() + + pdr = alchemiscale_execute.execute_DAG( + success_dag, + shared_basedir=shared, + scratch_basedir=scratch, + n_retries=0, + hooks=hooks, + ) + assert pdr.ok() + + units_total = len(success_dag.protocol_units) + + # on_dag_start fired once with the correct total + assert hooks.dag_starts == [units_total] + + # on_progress: first call is (0, N), last is (N, N) + assert hooks.progress[0] == (0, units_total) + assert hooks.progress[-1] == (units_total, units_total) + # progress is monotonically non-decreasing and ends at completion + completed = [c for c, _ in hooks.progress] + assert completed == sorted(completed) + assert completed[-1] == units_total + + # every start has exactly one matching end (balanced) + assert len(hooks.starts) == len(hooks.ends) + start_keys = Counter((k, a) for k, a in hooks.starts) + end_keys = Counter((k, a) for k, a, _ in hooks.ends) + assert start_keys == end_keys + + # every end on a success DAG carries an ok result + assert all(r is not None and r.ok() for _, _, r in hooks.ends) + + +def test_hooks_failure(failure_dag, tmp_path): + shared, scratch = _make_dirs(tmp_path, ["shared", "scratch"]) + hooks = RecordingHooks() + + pdr = alchemiscale_execute.execute_DAG( + failure_dag, + shared_basedir=shared, + scratch_basedir=scratch, + n_retries=0, + raise_error=False, + hooks=hooks, + ) + assert not pdr.ok() + + units_total = len(failure_dag.protocol_units) + assert hooks.dag_starts == [units_total] + + # starts and ends stay balanced even through the failure + assert len(hooks.starts) == len(hooks.ends) + start_keys = Counter((k, a) for k, a in hooks.starts) + end_keys = Counter((k, a) for k, a, _ in hooks.ends) + assert start_keys == end_keys + + # exactly one end carries a non-ok result (the broken unit) + non_ok_ends = [(k, a, r) for k, a, r in hooks.ends if r is not None and not r.ok()] + assert len(non_ok_ends) == 1 + + # progress starts at (0, N) and freezes at the pre-failure completed count: + # the DAG halts before all units finish, so the final completed count is + # strictly less than the total, and equals the number of ok ends. + assert hooks.progress[0] == (0, units_total) + n_ok_ends = sum(1 for _, _, r in hooks.ends if r is not None and r.ok()) + assert hooks.progress[-1] == (n_ok_ends, units_total) + assert hooks.progress[-1][0] < units_total + + +# --------------------------------------------------------------------------- +# custom protocols for the introspection-executor behaviors +# +# Each is a minimal 2-unit DAG (an upstream unit feeding a downstream unit) so +# that "downstream executed" is observable. They subclass DummyProtocol purely +# to inherit its settings/gather machinery and override `_create`. +# --------------------------------------------------------------------------- + + +class _PassUnit(ProtocolUnit): + """A trivially-succeeding unit; used as the downstream in these DAGs.""" + + @staticmethod + def _execute(ctx, **inputs): + return {"ok": True} + + +class _InterruptUnit(ProtocolUnit): + """A unit whose ``_execute`` raises a chosen ``BaseException`` subclass. + + ``gufe``'s ``ProtocolUnit.execute`` lets ``KeyboardInterrupt`` and + ``ExecutionInterrupt`` propagate rather than converting them into a + ``ProtocolUnitFailure``; this unit lets us exercise that path. + """ + + @staticmethod + def _execute(ctx, *, exc_name, **inputs): + if exc_name == "ExecutionInterrupt": + raise ExecutionInterrupt("unrecoverable") + elif exc_name == "KeyboardInterrupt": + raise KeyboardInterrupt("ctrl-c") + raise AssertionError(f"unknown exc_name {exc_name!r}") # pragma: no cover + + +class InterruptProtocol(DummyProtocol): + """A DAG whose first unit raises an interrupt; a downstream unit follows.""" + + exc_name = "ExecutionInterrupt" + + def _create(self, stateA, stateB, mapping=None, extends=None): + head = _InterruptUnit( + settings=self.settings, name="interrupter", exc_name=self.exc_name + ) + tail = _PassUnit(settings=self.settings, name="downstream", upstream=head) + return [head, tail] + + +class KeyboardInterruptProtocol(InterruptProtocol): + exc_name = "KeyboardInterrupt" + + +# module-level attempt counter for the retry-then-success unit, keyed by the +# source unit's gufe key. Units are stateless and re-executed per attempt, so +# state must live outside the unit; the test RESETS this before each executor +# invocation so both runs see identical behavior. +_RETRY_ATTEMPTS: dict[str, int] = {} +# number of leading failures before success +_RETRY_FAIL_COUNT = 2 + + +class _FlakyUnit(ProtocolUnit): + """Fails its first ``_RETRY_FAIL_COUNT`` attempts, then succeeds. + + Attempt bookkeeping is external (``_RETRY_ATTEMPTS``, keyed by the unit's + ``counter_key`` input) since the unit is re-instantiated/re-executed per + attempt. The two executor runs share the same DAG object --- hence the same + ``counter_key`` --- so, after the test resets the counter before each run, + both observe identical flaky behavior. + """ + + @staticmethod + def _execute(ctx, *, counter_key, **inputs): + seen = _RETRY_ATTEMPTS.get(counter_key, 0) + _RETRY_ATTEMPTS[counter_key] = seen + 1 + if seen < _RETRY_FAIL_COUNT: + raise ValueError(f"flaky failure #{seen}") + return {"ok": True, "attempts": seen + 1} + + +class RetryThenSucceedProtocol(DummyProtocol): + """A DAG whose head unit is flaky (N failures then success); tail follows.""" + + def _create(self, stateA, stateB, mapping=None, extends=None): + head = _FlakyUnit(settings=self.settings, name="flaky", counter_key="flaky") + tail = _PassUnit(settings=self.settings, name="downstream", upstream=head) + return [head, tail] + + +class _StreamUnit(ProtocolUnit): + """Writes one file into ``ctx.stdout`` and one into ``ctx.stderr``.""" + + @staticmethod + def _execute(ctx, **inputs): + (ctx.stdout / "out.txt").write_text("hello stdout\n") + (ctx.stderr / "err.txt").write_text("hello stderr\n") + return {"ok": True} + + +class StreamProtocol(DummyProtocol): + """A DAG whose head unit writes stream files; tail follows.""" + + def _create(self, stateA, stateB, mapping=None, extends=None): + head = _StreamUnit(settings=self.settings, name="streamer") + tail = _PassUnit(settings=self.settings, name="downstream", upstream=head) + return [head, tail] + + +class _BalanceHooks(ExecutionHooks): + """Records start/end pairs (and the end ``result``) for balance checks.""" + + def __init__(self): + self.starts = [] # (source_key, attempt) + self.ends = [] # (source_key, attempt, result_or_None) + + def on_unit_attempt_start(self, unit, attempt): + self.starts.append((str(unit.key), attempt)) + + def on_unit_attempt_end(self, unit, attempt, result): + self.ends.append((str(unit.key), attempt, result)) + + +# --------------------------------------------------------------------------- +# TASK A.1 --- interrupt propagation (ExecutionInterrupt / KeyboardInterrupt) +# --------------------------------------------------------------------------- + + +@pytest.mark.parametrize( + "proto_cls,exc_type", + [ + (InterruptProtocol, ExecutionInterrupt), + (KeyboardInterruptProtocol, KeyboardInterrupt), + ], +) +def test_equivalence_interrupt_propagates( + proto_cls, exc_type, stateA, stateB, tmp_path +): + proto = proto_cls(settings=proto_cls.default_settings()) + dag = proto.create(stateA=stateA, stateB=stateB, name="interrupt") + + gufe_dir = tmp_path / "gufe" + alch_dir = tmp_path / "alch" + g_shared, g_scratch = _make_dirs(gufe_dir, ["shared", "scratch"]) + a_shared, a_scratch = _make_dirs(alch_dir, ["shared", "scratch"]) + + # BOTH executors let the interrupt propagate (NOT converted to a failure), + # even with raise_error=False --- interrupts derive from BaseException. + with pytest.raises(exc_type): + gufe_protocoldag.execute_DAG( + dag, + shared_basedir=g_shared, + scratch_basedir=g_scratch, + raise_error=False, + n_retries=0, + ) + + hooks = _BalanceHooks() + with pytest.raises(exc_type): + alchemiscale_execute.execute_DAG( + dag, + shared_basedir=a_shared, + scratch_basedir=a_scratch, + raise_error=False, + n_retries=0, + hooks=hooks, + ) + + # the end hook still fires for the interrupted attempt (so log capture is + # always closed): starts and ends stay balanced... + assert len(hooks.starts) == len(hooks.ends) == 1 + assert Counter(hooks.starts) == Counter((k, a) for k, a, _ in hooks.ends) + # ...and the end hook received result=None for the interrupted attempt + assert hooks.ends[0][2] is None + + +# --------------------------------------------------------------------------- +# TASK A.2 --- retry-then-success (the retry loop's success arm) +# --------------------------------------------------------------------------- + + +def _flaky_source_key(dag): + """The stable source key of the flaky head unit in a RetryThenSucceed DAG.""" + for unit in dag.protocol_units: + if unit.name == "flaky": + return str(unit.key) + raise AssertionError("no flaky unit found") # pragma: no cover + + +def test_equivalence_retry_then_success(stateA, stateB, tmp_path): + proto = RetryThenSucceedProtocol( + settings=RetryThenSucceedProtocol.default_settings() + ) + dag = proto.create(stateA=stateA, stateB=stateB, name="retry") + + flaky_key = _flaky_source_key(dag) + n_retries = _RETRY_FAIL_COUNT + 1 # comfortably >= N + + gufe_dir = tmp_path / "gufe" + alch_dir = tmp_path / "alch" + g_shared, g_scratch = _make_dirs(gufe_dir, ["shared", "scratch"]) + a_shared, a_scratch = _make_dirs(alch_dir, ["shared", "scratch"]) + + # RESET the counter immediately before EACH executor invocation so both + # runs of the same DAG observe identical flaky behavior. + _RETRY_ATTEMPTS.clear() + gufe_pdr = gufe_protocoldag.execute_DAG( + dag, + shared_basedir=g_shared, + scratch_basedir=g_scratch, + raise_error=False, + n_retries=n_retries, + ) + + _RETRY_ATTEMPTS.clear() + alch_pdr = alchemiscale_execute.execute_DAG( + dag, + shared_basedir=a_shared, + scratch_basedir=a_scratch, + raise_error=False, + n_retries=n_retries, + ) + + # both succeed overall (the flaky unit eventually passes; downstream runs) + assert gufe_pdr.ok() is True + assert alch_pdr.ok() is True + + def _for_source(pdr, source_key): + return [r for r in pdr.protocol_unit_results if str(r.source_key) == source_key] + + for pdr in (gufe_pdr, alch_pdr): + flaky_results = _for_source(pdr, flaky_key) + failures = [r for r in flaky_results if not r.ok()] + successes = [r for r in flaky_results if r.ok()] + # exactly N failures + 1 success for the flaky source unit + assert len(failures) == _RETRY_FAIL_COUNT + assert len(successes) == 1 + + # downstream unit executed and produced a (successful) result + downstream = [ + r for r in pdr.protocol_unit_results if str(r.source_key) != flaky_key + ] + assert len(downstream) == 1 + assert downstream[0].ok() + + assert_equivalent(gufe_pdr, alch_pdr) + + +# --------------------------------------------------------------------------- +# TASK A.3 --- stream-dir parity (embedded stdout/stderr; per-attempt cleanup) +# --------------------------------------------------------------------------- + + +def test_equivalence_stream_dirs(stateA, stateB, tmp_path): + proto = StreamProtocol(settings=StreamProtocol.default_settings()) + dag = proto.create(stateA=stateA, stateB=stateB, name="stream") + + gufe_dir = tmp_path / "gufe" + alch_dir = tmp_path / "alch" + g_shared, g_scratch, g_stdout, g_stderr = _make_dirs( + gufe_dir, ["shared", "scratch", "stdout", "stderr"] + ) + a_shared, a_scratch, a_stdout, a_stderr = _make_dirs( + alch_dir, ["shared", "scratch", "stdout", "stderr"] + ) + + gufe_pdr = gufe_protocoldag.execute_DAG( + dag, + shared_basedir=g_shared, + scratch_basedir=g_scratch, + stdout_basedir=g_stdout, + stderr_basedir=g_stderr, + n_retries=0, + ) + alch_pdr = alchemiscale_execute.execute_DAG( + dag, + shared_basedir=a_shared, + scratch_basedir=a_scratch, + stdout_basedir=a_stdout, + stderr_basedir=a_stderr, + n_retries=0, + ) + + assert gufe_pdr.ok() and alch_pdr.ok() + assert_equivalent(gufe_pdr, alch_pdr) + + def _streamer_result(pdr): + (r,) = [r for r in pdr.protocol_unit_results if r.name == "streamer"] + return r + + g_res = _streamer_result(gufe_pdr) + a_res = _streamer_result(alch_pdr) + + # the streamer unit's embedded stdout/stderr (filename -> bytes) match + assert g_res.stdout == a_res.stdout + assert g_res.stderr == a_res.stderr + # and carry the expected content + assert a_res.stdout == {"out.txt": b"hello stdout\n"} + assert a_res.stderr == {"err.txt": b"hello stderr\n"} + + # per-attempt stream subdirs are rmtree'd (gufe/alchemiscale clean them up) + # while the base dirs remain, in BOTH executors. + for base in (g_stdout, g_stderr, a_stdout, a_stderr): + assert base.exists() + assert list(base.iterdir()) == [] diff --git a/docs/compute.rst b/docs/compute.rst index 7a08c2b6..525519c9 100644 --- a/docs/compute.rst +++ b/docs/compute.rst @@ -191,6 +191,58 @@ To scale up the number of compute services on the cluster, increase ``replicas`` A more complete example of this type of deployment can be found in `alchemiscale-k8s`_. +********************************* +Task introspection and capture +********************************* + +The :py:class:`~alchemiscale.compute.settings.ComputeServiceSettings` include several fields that control what execution metadata and artifacts a compute service captures and reports. +These feed the client-side introspection methods described in :ref:`introspection` and the failure-triage tools in :ref:`handling-errors`. +All of them have sensible defaults, so you only need to set them to change the default behavior. + +``hostname`` + Hostname to record on this compute service's registration and copy onto every :py:class:`~alchemiscale.storage.models.TaskProvenance` record it creates. + 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_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. + Captured streams are what :py:meth:`~alchemiscale.interface.client.AlchemiscaleClient.get_result_unit_stdout`, :py:meth:`~alchemiscale.interface.client.AlchemiscaleClient.get_result_unit_stderr`, :py:meth:`~alchemiscale.interface.client.AlchemiscaleClient.get_task_stdout`, and :py:meth:`~alchemiscale.interface.client.AlchemiscaleClient.get_task_stderr` return. + +``capture_logs`` + If ``true`` (the default), log records emitted through ``gufe``'s ``gufekey`` logger namespace (that is, protocol logs written via :external+gufe:py:attr:`~gufe.protocols.protocolunit.ProtocolUnit.logger`) are captured per unit result and uploaded alongside results. + Capture is scoped to the ``gufekey`` namespace only; records from third-party library loggers are **not** captured. + Captured logs are what :py:meth:`~alchemiscale.interface.client.AlchemiscaleClient.get_result_unit_logs` and :py:meth:`~alchemiscale.interface.client.AlchemiscaleClient.get_result_logs` return. + +``gufekey_loglevel`` + The level to set the ``gufekey`` logger to for per-unit log capture (default ``"INFO"``). + The ``gufekey`` logger otherwise inherits the root logger's level (typically ``WARNING``), which would drop protocol ``INFO`` logs before they could be captured. + +``log_cap_bytes`` + Per-unit-result cap, in bytes, on captured log text (default ``1048576``, i.e. 1 MiB). + When a unit's logs exceed the cap, the *tail* is kept — where errors typically live. + +``progress_push_timeout`` + Timeout, in seconds, for a single fire-and-forget progress push (default ``5.0``). + Progress (surfaced through :py:meth:`~alchemiscale.interface.client.AlchemiscaleClient.get_tasks_progress`) is best-effort telemetry: pushes are not retried, and failures are logged and swallowed so that a flaky API never stalls a :external+gufe:py:class:`~gufe.protocols.protocoldag.ProtocolDAG` between units. + +A minimal capture-related config snippet looks like:: + + --- + # override the auto-detected hostname recorded on TaskProvenance + hostname: gpu-node-04.cluster.example.org + + # capture per-unit stdout/stderr (protocol opt-in) and gufekey logs + capture_streams: true + capture_logs: true + + # capture protocol INFO logs, keeping at most 1 MiB of tail per unit + gufekey_loglevel: INFO + log_cap_bytes: 1048576 + + # best-effort progress telemetry + progress_push_timeout: 5.0 + **************** Compute managers **************** diff --git a/docs/user_guide/handling_errors.rst b/docs/user_guide/handling_errors.rst index 15d6c502..3c205948 100644 --- a/docs/user_guide/handling_errors.rst +++ b/docs/user_guide/handling_errors.rst @@ -43,6 +43,112 @@ Note that for some :external+gufe:py:class:`~gufe.protocols.protocol.Protocol`\s * :py:class:`openfe.protocols.openmm_rfe.RelativeHybridTopologyProtocol`: NVIDIA GPU if ``settings.platform == 'CUDA'`` * :py:class:`~feflow.protocols.nonequilibrium_cycling.NonEquilibriumCyclingProtocol`: OpenEye Toolkit license, NVIDIA GPU if ``settings.platform == 'CUDA'`` + +**************************** +Getting tracebacks in bulk +**************************** + +Pulling down full :external+gufe:py:class:`~gufe.protocols.protocoldag.ProtocolDAGResult`\s with :py:meth:`~alchemiscale.interface.client.AlchemiscaleClient.get_transformation_failures` (as above) transfers every failed result object, which can be slow when you only want to read the exceptions. +When your goal is fast failure triage for a single :py:class:`~alchemiscale.storage.models.Task`, use :py:meth:`~alchemiscale.interface.client.AlchemiscaleClient.get_task_tracebacks` instead. +It returns only the traceback text, one :py:class:`~alchemiscale.storage.models.TaskTracebacks` record per failed :external+gufe:py:class:`~gufe.protocols.protocoldag.ProtocolDAGResult` of that :py:class:`~alchemiscale.storage.models.Task`, most recent first:: + + >>> task: ScopedKey + >>> for attempt in asc.get_task_tracebacks(task): + >>> print(attempt.protocoldagresultref, attempt.datetime_created) + >>> for unit in attempt.tracebacks: + >>> print(unit.source_key) + >>> print(unit.traceback) + +Each :py:class:`~alchemiscale.storage.models.TaskTracebacks` carries the :py:class:`~alchemiscale.models.ScopedKey` of the failed result (``protocoldagresultref``) and a list of per-:external+gufe:py:class:`~gufe.protocols.protocolunit.ProtocolUnitFailure` tracebacks; each of those carries the ``source_key`` of the failing :external+gufe:py:class:`~gufe.protocols.protocolunit.ProtocolUnit` and its ``traceback`` string. + +If a :py:class:`~alchemiscale.storage.models.Task` has been attempted many times, you can limit the number of failed results returned to just the most recent ones with the ``limit`` keyword argument:: + + >>> # only the tracebacks from the most recent failed result + >>> asc.get_task_tracebacks(task, limit=1) + + +******************************* +Drilling into per-unit logs +******************************* + +Tracebacks tell you *where* a :external+gufe:py:class:`~gufe.protocols.protocolunit.ProtocolUnit` failed, but often the surrounding logs and captured stdout/stderr are what tell you *why*. +These artifacts are captured per :external+gufe:py:class:`~gufe.protocols.protocolunit.ProtocolUnit` at execution time (see :ref:`compute` for the compute-side settings that control capture), and you can drill into them without transferring the full result objects. + +Start from a :py:class:`~alchemiscale.storage.models.Task` and list records describing its :external+gufe:py:class:`~gufe.protocols.protocoldag.ProtocolDAGResult`\s with :py:meth:`~alchemiscale.interface.client.AlchemiscaleClient.get_task_result_recs`. +Pass ``ok=False`` to look only at failures, ``ok=True`` for successes, or leave it unset for all:: + + >>> task: ScopedKey + >>> pdrrs = asc.get_task_result_recs(task, ok=False) + >>> pdrrs + [, ...] + +Each :py:class:`~alchemiscale.storage.models.ProtocolDAGResultRec` carries the :py:class:`~alchemiscale.models.ScopedKey` of the underlying result as its ``scoped_key`` attribute. +For a given record, list the per-unit records with :py:meth:`~alchemiscale.interface.client.AlchemiscaleClient.get_result_unit_recs`, then pull the captured artifacts for any unit of interest with :py:meth:`~alchemiscale.interface.client.AlchemiscaleClient.get_result_unit_logs`, :py:meth:`~alchemiscale.interface.client.AlchemiscaleClient.get_result_unit_stdout`, and :py:meth:`~alchemiscale.interface.client.AlchemiscaleClient.get_result_unit_stderr`:: + + >>> for pdrr in pdrrs: + >>> for purr in asc.get_result_unit_recs(pdrr): + >>> # skip units with nothing captured + >>> if purr.has_logs: + >>> print(asc.get_result_unit_logs(purr)) + >>> if purr.has_stdout: + >>> print(asc.get_result_unit_stdout(purr)) + >>> if purr.has_stderr: + >>> print(asc.get_result_unit_stderr(purr)) + +Each :py:class:`~alchemiscale.storage.models.ProtocolUnitResultRec` exposes ``has_logs``, ``has_stdout``, and ``has_stderr`` flags so you can skip units that captured nothing. +:py:meth:`~alchemiscale.interface.client.AlchemiscaleClient.get_result_unit_logs` returns the captured log text as a single string (or ``None``), while :py:meth:`~alchemiscale.interface.client.AlchemiscaleClient.get_result_unit_stdout` and :py:meth:`~alchemiscale.interface.client.AlchemiscaleClient.get_result_unit_stderr` return a mapping of filename to captured text (or ``None``). + +.. note:: + The drill-down methods accept either a :py:class:`~alchemiscale.models.ScopedKey` or the corresponding record object (a :py:class:`~alchemiscale.storage.models.ProtocolDAGResultRec` for ``pdrr`` arguments, a :py:class:`~alchemiscale.storage.models.ProtocolUnitResultRec` for ``purr`` arguments), so the chain above composes naturally. + The same :py:class:`~alchemiscale.models.ScopedKey`\s can be copied between Python sessions. + +When you don't need per-unit granularity, three convenience methods render everything for you. +:py:meth:`~alchemiscale.interface.client.AlchemiscaleClient.get_result_logs` returns a single human-readable rendering of all unit logs for one :external+gufe:py:class:`~gufe.protocols.protocoldag.ProtocolDAGResult`; the default ``order='unit'`` groups each unit's logs under a header, while ``order='time'`` interleaves all units' log lines by timestamp:: + + >>> pdrr = asc.get_task_result_recs(task, ok=False)[0] + >>> print(asc.get_result_logs(pdrr)) + >>> # or interleave across units by timestamp + >>> print(asc.get_result_logs(pdrr, order='time')) + +:py:meth:`~alchemiscale.interface.client.AlchemiscaleClient.get_task_stdout` and :py:meth:`~alchemiscale.interface.client.AlchemiscaleClient.get_task_stderr` go one level higher, concatenating the captured stdout/stderr across *all* of a :py:class:`~alchemiscale.storage.models.Task`\'s :external+gufe:py:class:`~gufe.protocols.protocoldag.ProtocolDAGResult`\s (most recent first), with section headers identifying each result, unit, and filename:: + + >>> print(asc.get_task_stdout(task)) + >>> print(asc.get_task_stderr(task)) + +Each returns ``""`` when nothing was captured. + +.. note:: + Logs and stream capture are opt-in on the compute side and depend on what each :external+gufe:py:class:`~gufe.protocols.protocol.Protocol` chooses to emit and archive. + If a :external+gufe:py:class:`~gufe.protocols.protocol.Protocol` writes nothing to its per-unit stdout/stderr, or logs nothing through :external+gufe:py:attr:`~gufe.protocols.protocolunit.ProtocolUnit.logger`, these methods will return empty results even for a :py:class:`~alchemiscale.storage.models.Task` that failed. + See :ref:`compute` for details on the capture mechanism and the settings that govern it. + + +******************************************** +The reason field and DAG-creation failures +******************************************** + +Not every failure produces a :external+gufe:py:class:`~gufe.protocols.protocoldag.ProtocolDAGResult` with tracebacks to inspect. +Before any :external+gufe:py:class:`~gufe.protocols.protocolunit.ProtocolUnit` runs, the compute service must first build the :external+gufe:py:class:`~gufe.protocols.protocoldag.ProtocolDAG` for a :py:class:`~alchemiscale.storage.models.Task` from its :external+gufe:py:class:`~gufe.transformations.transformation.Transformation`. +If that construction itself raises, there is no result to store; instead, the :py:class:`~alchemiscale.storage.models.Task` is set directly to ``error``, and the traceback is recorded on the ``reason`` field of the :py:class:`~alchemiscale.storage.models.Task`. + +You can read this ``reason`` back through :py:meth:`~alchemiscale.interface.client.AlchemiscaleClient.get_tasks_details` (bulk) or :py:meth:`~alchemiscale.interface.client.AlchemiscaleClient.get_task_history` (per-attempt); both are covered in :ref:`introspection`:: + + >>> (detail,) = asc.get_tasks_details([task]) + >>> print(detail.status) # 'error' + >>> print(detail.reason) # the DAG-creation traceback + +.. warning:: + A :external+gufe:py:class:`~gufe.protocols.protocoldag.ProtocolDAG`-creation failure is treated as a *systematic* problem with the :external+gufe:py:class:`~gufe.transformations.transformation.Transformation`, not a transient one. + Such :py:class:`~alchemiscale.storage.models.Task`\s are **not** eligible for automatic retry via :py:class:`~alchemiscale.storage.models.Task` restart patterns (see below), since re-running them would fail again in exactly the same way. + Resolve the underlying problem with the :external+gufe:py:class:`~gufe.transformations.transformation.Transformation` before setting these :py:class:`~alchemiscale.storage.models.Task`\s back to ``waiting``. + +Finally, when you mark :py:class:`~alchemiscale.storage.models.Task`\s ``invalid`` or ``deleted`` (see below), you can attach your own ``reason`` for the record:: + + >>> asc.set_tasks_status(tasks, 'invalid', reason='superseded by re-parameterized transformation') + +The ``reason`` is recorded only for ``invalid`` and ``deleted`` transitions, and surfaces through :py:meth:`~alchemiscale.interface.client.AlchemiscaleClient.get_tasks_details`. + + ************************ Re-running errored Tasks ************************ diff --git a/docs/user_guide/index.rst b/docs/user_guide/index.rst index 2125c0ad..9be7e92f 100644 --- a/docs/user_guide/index.rst +++ b/docs/user_guide/index.rst @@ -11,5 +11,6 @@ If you are making use of an **alchemiscale** instance, this document will famili :maxdepth: 1 ./getting_started + ./introspection ./handling_errors ./strategy_automation diff --git a/docs/user_guide/introspection.rst b/docs/user_guide/introspection.rst new file mode 100644 index 00000000..68723530 --- /dev/null +++ b/docs/user_guide/introspection.rst @@ -0,0 +1,109 @@ +.. _introspection: + +############################ +Understanding Task execution +############################ + +The status counts described in :ref:`getting-started` tell you *how many* of your :py:class:`~alchemiscale.storage.models.Task`\s are ``waiting``, ``running``, ``error``, or ``complete``, but they don't tell you *what happened* to any one :py:class:`~alchemiscale.storage.models.Task`, *where* it ran, or *how far* a currently-running one has gotten. +This document covers the introspection methods on the :py:class:`~alchemiscale.interface.client.AlchemiscaleClient` that answer those questions: per-:py:class:`~alchemiscale.storage.models.Task` execution history, bulk indicators across many :py:class:`~alchemiscale.storage.models.Task`\s, live progress for running :py:class:`~alchemiscale.storage.models.Task`\s, and your compute share within a :py:class:`~alchemiscale.models.Scope`. + +For failure triage specifically — tracebacks, per-unit logs, and captured stdout/stderr — see :ref:`handling-errors`. + + +*********************************** +Task history and execution attempts +*********************************** + +A single :py:class:`~alchemiscale.storage.models.Task` may be executed several times over its lifetime: it can land on a flaky host, be released when you change its status mid-run, or be restarted by a :py:class:`~alchemiscale.storage.models.Task` restart pattern. +Each of these is a distinct *attempt*. +To retrieve the full attempt history of a :py:class:`~alchemiscale.storage.models.Task`, use :py:meth:`~alchemiscale.interface.client.AlchemiscaleClient.get_task_history`, which returns a list of :py:class:`~alchemiscale.storage.models.TaskAttempt` records, most recent first:: + + >>> task: ScopedKey + >>> for attempt in asc.get_task_history(task): + >>> print(attempt.compute_service_id, + >>> attempt.hostname, + >>> attempt.datetime_claimed, + >>> attempt.datetime_end, + >>> attempt.outcome) + +Each :py:class:`~alchemiscale.storage.models.TaskAttempt` records: + +* ``compute_service_id`` and ``hostname`` — which compute service claimed the attempt, and the host it ran on +* ``manager_name`` — the compute manager responsible for the service, if any +* ``datetime_claimed`` and ``datetime_end`` — when the attempt was claimed and when it ended (``datetime_end`` is ``None`` while the attempt is still in flight) +* ``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) + +You can limit the history to the most recent attempts with the ``limit`` keyword argument:: + + >>> # just the most recent attempt + >>> asc.get_task_history(task, limit=1) + +Recall from :ref:`handling-errors` that a :external+gufe:py:class:`~gufe.protocols.protocoldag.ProtocolDAG`-creation failure sets the :py:class:`~alchemiscale.storage.models.Task` to ``error`` and records the traceback on the :py:class:`~alchemiscale.storage.models.Task`\'s ``reason`` (surfaced via :py:meth:`~alchemiscale.interface.client.AlchemiscaleClient.get_tasks_details`, below) rather than producing a result to inspect. + + +**************************** +Bulk indicators for Tasks +**************************** + +When you want a compact status summary across many :py:class:`~alchemiscale.storage.models.Task`\s at once — for a dashboard, a triage sweep, or a quick sanity check — use :py:meth:`~alchemiscale.interface.client.AlchemiscaleClient.get_tasks_details`. +It returns a list of :py:class:`~alchemiscale.storage.models.TaskDetails`, one per input :py:class:`~alchemiscale.storage.models.Task` and in the same order (with ``None`` in place of any :py:class:`~alchemiscale.storage.models.Task` that doesn't exist):: + + >>> tasks = asc.get_network_tasks(an_sk) + >>> for detail in asc.get_tasks_details(tasks): + >>> if detail is None: + >>> continue + >>> print(detail.task, + >>> detail.status, + >>> detail.datetime_status_changed, + >>> detail.num_claims) + +Each :py:class:`~alchemiscale.storage.models.TaskDetails` bundles: + +* ``task`` — the :py:class:`~alchemiscale.models.ScopedKey` of the :py:class:`~alchemiscale.storage.models.Task` +* ``status`` and ``datetime_status_changed`` — the current status and when it last changed +* ``reason`` — the human-readable reason for the current status, where one was recorded; this is where a :external+gufe:py:class:`~gufe.protocols.protocoldag.ProtocolDAG`-creation traceback appears for an errored :py:class:`~alchemiscale.storage.models.Task`, and where a ``reason`` you passed to :py:meth:`~alchemiscale.interface.client.AlchemiscaleClient.set_tasks_status` appears for ``invalid``/``deleted`` :py:class:`~alchemiscale.storage.models.Task`\s +* ``num_claims`` — how many times the :py:class:`~alchemiscale.storage.models.Task` has been claimed for execution +* ``current_claim`` — the live claim on a ``running`` :py:class:`~alchemiscale.storage.models.Task` (compute service, host, claim time, and progress), or ``None`` +* ``most_recent_attempt`` — the most recent :py:class:`~alchemiscale.storage.models.TaskAttempt`, matching the first element of :py:meth:`~alchemiscale.interface.client.AlchemiscaleClient.get_task_history` + +For example, to find all errored :py:class:`~alchemiscale.storage.models.Task`\s on a network and print why each stopped:: + + >>> tasks = asc.get_network_tasks(an_sk, status='error') + >>> for detail in asc.get_tasks_details(tasks): + >>> print(detail.task, detail.reason) + + +****************************** +Live progress of running Tasks +****************************** + +For :py:class:`~alchemiscale.storage.models.Task`\s that are currently ``running``, you can watch how far each has progressed through its :external+gufe:py:class:`~gufe.protocols.protocolunit.ProtocolUnit`\s with :py:meth:`~alchemiscale.interface.client.AlchemiscaleClient.get_tasks_progress`. +It returns a list in the same order as the input, where each element is a ``(units_completed, units_total)`` tuple for a ``running`` :py:class:`~alchemiscale.storage.models.Task` that is reporting progress, or ``None`` otherwise (for example, a :py:class:`~alchemiscale.storage.models.Task` that isn't running, or a running one that hasn't yet reported):: + + >>> tasks = asc.get_network_tasks(an_sk, status='running') + >>> for task, progress in zip(tasks, asc.get_tasks_progress(tasks)): + >>> if progress is None: + >>> print(task, 'no progress reported') + >>> else: + >>> completed, total = progress + >>> print(task, f'{completed}/{total} units') + +Progress is best-effort telemetry pushed by the compute service between :external+gufe:py:class:`~gufe.protocols.protocolunit.ProtocolUnit`\s, so a ``None`` result does not imply anything is wrong — only that no progress datapoint is currently available. + + +**************************** +Compute share within a Scope +**************************** + +When compute is contended, it's useful to know what fraction of it your work is currently receiving. +:py:meth:`~alchemiscale.interface.client.AlchemiscaleClient.get_scope_compute_share` returns, as a ``float``, the fraction of currently-running :py:class:`~alchemiscale.storage.models.Task`\s in a given :py:class:`~alchemiscale.models.Scope` relative to its sibling :py:class:`~alchemiscale.models.Scope`\s at the same level:: + + >>> asc.get_scope_compute_share(Scope('my_org', 'my_campaign', 'my_project')) + 0.42 + +The share is computed server-side as this :py:class:`~alchemiscale.models.Scope`\'s aggregate fraction of running :py:class:`~alchemiscale.storage.models.Task`\s relative to its siblings; only the aggregate fraction is returned. +Your identity must hold the given :py:class:`~alchemiscale.models.Scope`. +A value near ``1.0`` means nearly all currently-running :py:class:`~alchemiscale.storage.models.Task`\s among the sibling :py:class:`~alchemiscale.models.Scope`\s belong to this one; a value near ``0.0`` means the :py:class:`~alchemiscale.models.Scope` is getting little of the available compute right now. +Because it reflects only the *instantaneous* running population, the value fluctuates as :py:class:`~alchemiscale.storage.models.Task`\s are claimed and completed. diff --git a/news/issue-106.rst b/news/issue-106.rst new file mode 100644 index 00000000..22f560ff --- /dev/null +++ b/news/issue-106.rst @@ -0,0 +1,9 @@ +**Added:** + +* 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. +* ``AlchemiscaleClient.get_task_history`` returns the full per-attempt history of a ``Task``, and ``AlchemiscaleClient.get_tasks_details`` returns detailed per-``Task`` information. + +**Changed:** + +* ``Task`` records now expose additional status indicators: ``datetime_status_changed`` and ``reason``. diff --git a/news/issue-195.rst b/news/issue-195.rst new file mode 100644 index 00000000..b72dc699 --- /dev/null +++ b/news/issue-195.rst @@ -0,0 +1,3 @@ +**Fixed:** + +* A ``ProtocolDAG`` creation failure no longer kills the compute service. The affected ``Task`` is now set to ``error`` with an explanatory ``reason``, and the service continues running. diff --git a/news/issue-211.rst b/news/issue-211.rst new file mode 100644 index 00000000..b47baff1 --- /dev/null +++ b/news/issue-211.rst @@ -0,0 +1,3 @@ +**Added:** + +* Durable execution provenance for ``Task``\s via per-attempt ``TaskProvenance`` records, surfaced through ``AlchemiscaleClient.get_task_history`` and ``AlchemiscaleClient.get_tasks_details``. diff --git a/news/issue-295.rst b/news/issue-295.rst new file mode 100644 index 00000000..cdcca654 --- /dev/null +++ b/news/issue-295.rst @@ -0,0 +1,4 @@ +**Added:** + +* Per-unit log capture, including ``stdout`` and ``stderr``, for ``Task`` executions. +* New client methods for retrieving execution logs and captured output: ``AlchemiscaleClient.get_task_result_recs``, ``AlchemiscaleClient.get_result_unit_recs``, ``AlchemiscaleClient.get_result_unit_logs``, ``AlchemiscaleClient.get_result_unit_stdout``, ``AlchemiscaleClient.get_result_unit_stderr``, ``AlchemiscaleClient.get_result_logs``, ``AlchemiscaleClient.get_task_stdout``, and ``AlchemiscaleClient.get_task_stderr``. diff --git a/news/issue-347.rst b/news/issue-347.rst new file mode 100644 index 00000000..ca337df9 --- /dev/null +++ b/news/issue-347.rst @@ -0,0 +1,3 @@ +**Added:** + +* ``AlchemiscaleClient.get_task_tracebacks`` for fast retrieval of tracebacks from failed ``Task`` executions. diff --git a/news/issue-349.rst b/news/issue-349.rst new file mode 100644 index 00000000..ffb6a400 --- /dev/null +++ b/news/issue-349.rst @@ -0,0 +1,3 @@ +**Added:** + +* Retrieval of captured ``stdout``/``stderr`` and per-unit logs from ``Task`` executions via ``AlchemiscaleClient.get_result_unit_stdout``, ``AlchemiscaleClient.get_result_unit_stderr``, ``AlchemiscaleClient.get_task_stdout``, and ``AlchemiscaleClient.get_task_stderr``. diff --git a/news/issue-389.rst b/news/issue-389.rst new file mode 100644 index 00000000..f4e6658f --- /dev/null +++ b/news/issue-389.rst @@ -0,0 +1,3 @@ +**Added:** + +* ``AlchemiscaleClient.get_scope_compute_share`` for reporting the compute share of a given ``Scope``. diff --git a/news/issue-415.rst b/news/issue-415.rst new file mode 100644 index 00000000..ee990fd7 --- /dev/null +++ b/news/issue-415.rst @@ -0,0 +1,7 @@ +**Added:** + +* Live progress reporting for executing ``Task``\s via ``AlchemiscaleClient.get_tasks_progress``. + +**Changed:** + +* ``Task`` history and details now include live progress information. From 682b584fd08b56c180a897b6b4dff52e543ade7f Mon Sep 17 00:00:00 2001 From: David Dotson Date: Fri, 10 Jul 2026 09:13:55 -0600 Subject: [PATCH 02/18] Fix CI failures on the introspection branch - compute API set_task_result: tolerate a missing/`"None"` compute_service_id (the client serializes an absent id as the string "None"); skip provenance finalization instead of raising ValueError -> 500. Fixes the existing test_set_task_result / test_set_task_result_failure integration tests. - equivalence suite: pass `mapping=None` to every `Protocol.create()` call; the pinned gufe requires `mapping` as a keyword-only argument with no default (a newer local gufe had masked this). - docs: replace the unresolvable `gufe...ProtocolUnit.logger` intersphinx attr references with plain literals, so the docs build (fail_on_warning) passes. Co-Authored-By: Claude Opus 4.8 (1M context) --- alchemiscale/compute/api.py | 11 ++++++++++- .../tests/unit/compute/test_execute_equivalence.py | 10 +++++----- docs/compute.rst | 2 +- docs/user_guide/handling_errors.rst | 2 +- 4 files changed, 17 insertions(+), 8 deletions(-) diff --git a/alchemiscale/compute/api.py b/alchemiscale/compute/api.py index 51247346..71ed9c0b 100644 --- a/alchemiscale/compute/api.py +++ b/alchemiscale/compute/api.py @@ -384,6 +384,15 @@ async def set_task_result( protocoldagresult_ = body_["protocoldagresult"] compute_service_id = body_["compute_service_id"] + # the compute client serializes a missing id as the string "None"; treat + # that (and a genuine null) as "no compute service", so provenance + # finalization is simply skipped rather than erroring on an invalid id + compute_service_id_ = ( + ComputeServiceID(compute_service_id) + if compute_service_id and compute_service_id != "None" + else None + ) + task_sk = ScopedKey.from_str(task_scoped_key) validate_scopes(task_sk.scope, token) @@ -409,7 +418,7 @@ async def set_task_result( result_sk: ScopedKey = n4js.set_task_result( task=task_sk, protocoldagresultref=protocoldagresultref, - compute_service_id=ComputeServiceID(compute_service_id), + compute_service_id=compute_service_id_, ) # derive one ProtocolUnitResultRef per unit result, and extract any embedded diff --git a/alchemiscale/tests/unit/compute/test_execute_equivalence.py b/alchemiscale/tests/unit/compute/test_execute_equivalence.py index ec01b9b9..e0f3e7b1 100644 --- a/alchemiscale/tests/unit/compute/test_execute_equivalence.py +++ b/alchemiscale/tests/unit/compute/test_execute_equivalence.py @@ -70,14 +70,14 @@ def stateB() -> ChemicalSystem: def success_dag(stateA, stateB): """A DummyProtocol DAG: every unit succeeds (1 init + 21 sims + 1 finish).""" proto = DummyProtocol(settings=DummyProtocol.default_settings()) - return proto.create(stateA=stateA, stateB=stateB, name="success") + return proto.create(stateA=stateA, stateB=stateB, mapping=None, name="success") @pytest.fixture def failure_dag(stateA, stateB): """A BrokenProtocol DAG: exactly one unit always fails, halting the DAG.""" proto = BrokenProtocol(settings=BrokenProtocol.default_settings()) - return proto.create(stateA=stateA, stateB=stateB, name="failure") + return proto.create(stateA=stateA, stateB=stateB, mapping=None, name="failure") # --------------------------------------------------------------------------- @@ -591,7 +591,7 @@ def test_equivalence_interrupt_propagates( proto_cls, exc_type, stateA, stateB, tmp_path ): proto = proto_cls(settings=proto_cls.default_settings()) - dag = proto.create(stateA=stateA, stateB=stateB, name="interrupt") + dag = proto.create(stateA=stateA, stateB=stateB, mapping=None, name="interrupt") gufe_dir = tmp_path / "gufe" alch_dir = tmp_path / "alch" @@ -645,7 +645,7 @@ def test_equivalence_retry_then_success(stateA, stateB, tmp_path): proto = RetryThenSucceedProtocol( settings=RetryThenSucceedProtocol.default_settings() ) - dag = proto.create(stateA=stateA, stateB=stateB, name="retry") + dag = proto.create(stateA=stateA, stateB=stateB, mapping=None, name="retry") flaky_key = _flaky_source_key(dag) n_retries = _RETRY_FAIL_COUNT + 1 # comfortably >= N @@ -707,7 +707,7 @@ def _for_source(pdr, source_key): def test_equivalence_stream_dirs(stateA, stateB, tmp_path): proto = StreamProtocol(settings=StreamProtocol.default_settings()) - dag = proto.create(stateA=stateA, stateB=stateB, name="stream") + dag = proto.create(stateA=stateA, stateB=stateB, mapping=None, name="stream") gufe_dir = tmp_path / "gufe" alch_dir = tmp_path / "alch" diff --git a/docs/compute.rst b/docs/compute.rst index 525519c9..b24fab18 100644 --- a/docs/compute.rst +++ b/docs/compute.rst @@ -210,7 +210,7 @@ All of them have sensible defaults, so you only need to set them to change the d Captured streams are what :py:meth:`~alchemiscale.interface.client.AlchemiscaleClient.get_result_unit_stdout`, :py:meth:`~alchemiscale.interface.client.AlchemiscaleClient.get_result_unit_stderr`, :py:meth:`~alchemiscale.interface.client.AlchemiscaleClient.get_task_stdout`, and :py:meth:`~alchemiscale.interface.client.AlchemiscaleClient.get_task_stderr` return. ``capture_logs`` - If ``true`` (the default), log records emitted through ``gufe``'s ``gufekey`` logger namespace (that is, protocol logs written via :external+gufe:py:attr:`~gufe.protocols.protocolunit.ProtocolUnit.logger`) are captured per unit result and uploaded alongside results. + If ``true`` (the default), log records emitted through ``gufe``'s ``gufekey`` logger namespace (that is, protocol logs written via ``ProtocolUnit.logger``) are captured per unit result and uploaded alongside results. Capture is scoped to the ``gufekey`` namespace only; records from third-party library loggers are **not** captured. Captured logs are what :py:meth:`~alchemiscale.interface.client.AlchemiscaleClient.get_result_unit_logs` and :py:meth:`~alchemiscale.interface.client.AlchemiscaleClient.get_result_logs` return. diff --git a/docs/user_guide/handling_errors.rst b/docs/user_guide/handling_errors.rst index 3c205948..e957bc51 100644 --- a/docs/user_guide/handling_errors.rst +++ b/docs/user_guide/handling_errors.rst @@ -119,7 +119,7 @@ Each returns ``""`` when nothing was captured. .. note:: Logs and stream capture are opt-in on the compute side and depend on what each :external+gufe:py:class:`~gufe.protocols.protocol.Protocol` chooses to emit and archive. - If a :external+gufe:py:class:`~gufe.protocols.protocol.Protocol` writes nothing to its per-unit stdout/stderr, or logs nothing through :external+gufe:py:attr:`~gufe.protocols.protocolunit.ProtocolUnit.logger`, these methods will return empty results even for a :py:class:`~alchemiscale.storage.models.Task` that failed. + If a :external+gufe:py:class:`~gufe.protocols.protocol.Protocol` writes nothing to its per-unit stdout/stderr, or logs nothing through ``ProtocolUnit.logger``, these methods will return empty results even for a :py:class:`~alchemiscale.storage.models.Task` that failed. See :ref:`compute` for details on the capture mechanism and the settings that govern it. From a8f4b09e45fb3706813b25758787e9d7489ed55b Mon Sep 17 00:00:00 2001 From: David Dotson Date: Fri, 10 Jul 2026 10:07:51 -0600 Subject: [PATCH 03/18] Fix docs build: reword docstrings that broke rST inline markup The Sphinx build (fail_on_warning) flagged autodoc'd docstrings with "inline interpreted text ... start-string without end-string": napoleon's field-list rendering surfaces reStructuredText's markup rules, which reject a closing backtick immediately followed by a letter/apostrophe (`Task`s, `Task`'s) and slash-joined inline markup (`_scoped_key`/`_gufe_key`). Reworded the affected docstrings (get_tasks_details, get_tasks_progress, get_task_tracebacks, get_task_stdout/stderr, get_scope_compute_share, ProtocolUnitResultRef) to keep the code markup while satisfying the markup rules. Verified zero inline-markup warnings across all changed docstrings by running each through napoleon + docutils. Co-Authored-By: Claude Opus 4.8 (1M context) --- alchemiscale/interface/client.py | 18 +++++++++--------- alchemiscale/storage/models.py | 10 ++++++---- alchemiscale/storage/statestore.py | 4 ++-- 3 files changed, 17 insertions(+), 15 deletions(-) diff --git a/alchemiscale/interface/client.py b/alchemiscale/interface/client.py index dbcae1ee..0fee8093 100644 --- a/alchemiscale/interface/client.py +++ b/alchemiscale/interface/client.py @@ -2187,12 +2187,12 @@ def get_task_history( return [TaskAttempt.from_dict(attempt) for attempt in attempts] def get_tasks_details(self, tasks: list[ScopedKey]) -> list[TaskDetails | None]: - """Get summary details for multiple `Task`s. + """Get summary details for multiple Tasks. Parameters ---------- tasks - The `ScopedKey`s of the `Task`s to retrieve details for. + The `ScopedKey` of each `Task` to retrieve details for. Returns ------- @@ -2219,7 +2219,7 @@ def get_task_tracebacks( The `ScopedKey` of the `Task` to retrieve tracebacks for. limit If given, return tracebacks for at most this many of the most - recent failed `ProtocolDAGResult`s. + recent failed `ProtocolDAGResult` objects. Returns ------- @@ -2387,8 +2387,8 @@ def get_result_logs( def get_task_stdout(self, task: ScopedKey) -> str: """Get a human-readable rendering of all captured stdout for a `Task`. - Concatenates stdout across all of the `Task`'s `ProtocolDAGResult`s - (most recent first), with section headers identifying each result, + Concatenates stdout across all `ProtocolDAGResult` objects of the + `Task` (most recent first), with section headers identifying each result, unit, and filename. Parameters @@ -2406,8 +2406,8 @@ def get_task_stdout(self, task: ScopedKey) -> str: def get_task_stderr(self, task: ScopedKey) -> str: """Get a human-readable rendering of all captured stderr for a `Task`. - Concatenates stderr across all of the `Task`'s `ProtocolDAGResult`s - (most recent first), with section headers identifying each result, + Concatenates stderr across all `ProtocolDAGResult` objects of the + `Task` (most recent first), with section headers identifying each result, unit, and filename. Parameters @@ -2425,12 +2425,12 @@ def get_task_stderr(self, task: ScopedKey) -> str: def get_tasks_progress( self, tasks: list[ScopedKey] ) -> list[tuple[int, int] | None]: - """Get execution progress for multiple `Task`s. + """Get execution progress for multiple Tasks. Parameters ---------- tasks - The `ScopedKey`s of the `Task`s to retrieve progress for. + The `ScopedKey` of each `Task` to retrieve progress for. Returns ------- diff --git a/alchemiscale/storage/models.py b/alchemiscale/storage/models.py index 807629d0..b69e7a85 100644 --- a/alchemiscale/storage/models.py +++ b/alchemiscale/storage/models.py @@ -659,12 +659,14 @@ class ProtocolUnitResultRef(ObjectStoreRef): Note ---- - The `has_*` flags and (nothing else) are *mutated in place* via Cypher after - the node is created, as artifacts arrive. The node's `_scoped_key`/`_gufe_key` + The ``has_logs``, ``has_stdout``, and ``has_stderr`` flags (and nothing + else) are mutated in place via Cypher after + the node is created, as artifacts arrive. The node's `_scoped_key` and + `_gufe_key` are computed once at creation and never recomputed, so lookups stay stable even though these tokenizable-contributing fields change. This is safe only - because `ProtocolUnitResultRef`s are an internal state-store detail, never - re-tokenized after creation; keep it that way. + because `ProtocolUnitResultRef` nodes are an internal state-store detail, + never re-tokenized after creation; keep it that way. """ ok: bool diff --git a/alchemiscale/storage/statestore.py b/alchemiscale/storage/statestore.py index b8a89af8..35d24e0f 100644 --- a/alchemiscale/storage/statestore.py +++ b/alchemiscale/storage/statestore.py @@ -3884,8 +3884,8 @@ def get_task_tracebacks( return records def get_scope_compute_share(self, scope: Scope) -> float: - """Return the fraction of currently-`running` `Task`s in `scope` - relative to all `Scope`s at the same level. + """Return the fraction of currently-`running` Tasks in `scope` + relative to all Scopes at the same level. - `Scope('org')` -> the org's running Tasks / all running Tasks; - `Scope('org', 'campaign')` -> the campaign's / all campaigns in that org; From 791f54487e224d05264e07cbd5cb8ef43a903868 Mon Sep 17 00:00:00 2001 From: David Dotson Date: Fri, 10 Jul 2026 10:57:01 -0600 Subject: [PATCH 04/18] Fix remaining docs warning: TaskProvenance Attributes backtick-plural The RTD Sphinx build (older/stricter than local) still flagged TaskProvenance:49 --- backtick-plural (`ProtocolUnit`s) inside a napoleon Attributes section, a context my local docutils reproduction couldn't see (the `.. attribute::` directive masked the inline content). Reworded that and, defensively, the other backtick-plural/apostrophe occurrences in the new docstrings so the render is robust to the stricter build. Pre-existing methods (get_task_results etc.) are untouched; they use the same style in plain summary prose and build cleanly on main. Co-Authored-By: Claude Opus 4.8 (1M context) --- alchemiscale/compute/capture.py | 2 +- alchemiscale/interface/client.py | 16 ++++++++-------- alchemiscale/storage/models.py | 6 +++--- alchemiscale/storage/statestore.py | 6 +++--- 4 files changed, 15 insertions(+), 15 deletions(-) diff --git a/alchemiscale/compute/capture.py b/alchemiscale/compute/capture.py index 0917c897..444fe843 100644 --- a/alchemiscale/compute/capture.py +++ b/alchemiscale/compute/capture.py @@ -65,7 +65,7 @@ def text(self) -> str: class SynchronousExecutionHooks(ExecutionHooks): - """Execution hooks binding one `Task`'s DAG execution to log capture and + """Execution hooks binding the DAG execution of one `Task` to log capture and progress reporting. Parameters diff --git a/alchemiscale/interface/client.py b/alchemiscale/interface/client.py index 0fee8093..392cbbec 100644 --- a/alchemiscale/interface/client.py +++ b/alchemiscale/interface/client.py @@ -2179,7 +2179,7 @@ def get_task_history( Returns ------- list[TaskAttempt] - A list of `TaskAttempt`s, one per execution attempt of the `Task`, + A list of `TaskAttempt` records, one per execution attempt of the `Task`, most recent first. """ params = dict(limit=limit) @@ -2211,7 +2211,7 @@ def get_tasks_details(self, tasks: list[ScopedKey]) -> list[TaskDetails | None]: def get_task_tracebacks( self, task: ScopedKey, limit: int | None = None ) -> list[TaskTracebacks]: - """Get the tracebacks from failed `ProtocolDAGResult`s of a `Task`. + """Get the tracebacks from failed `ProtocolDAGResult` objects of a `Task`. Parameters ---------- @@ -2234,8 +2234,8 @@ def get_task_tracebacks( def get_scope_compute_share(self, scope: Scope) -> float: """Get this identity's fractional compute share within the given `Scope`. - The share is computed server-side as this `Scope`'s aggregate fraction - relative to its sibling `Scope`s; only the aggregate fraction is + The share is computed server-side as the aggregate fraction for this `Scope` + relative to its sibling Scopes; only the aggregate fraction is returned. The identity must hold the given `Scope`. Parameters @@ -2258,7 +2258,7 @@ def _as_scoped_key(obj: ScopedKey | Any) -> ScopedKey: def get_task_result_recs( self, task: ScopedKey, ok: bool | None = None ) -> list[ProtocolDAGResultRec]: - """Get records describing the `ProtocolDAGResult`s of a `Task`. + """Get records describing the `ProtocolDAGResult` objects of a `Task`. Parameters ---------- @@ -2271,7 +2271,7 @@ def get_task_result_recs( Returns ------- list[ProtocolDAGResultRec] - A list of `ProtocolDAGResultRec`s, one per `ProtocolDAGResult` of + A list of `ProtocolDAGResultRec` records, one per `ProtocolDAGResult` of the `Task`, most recent first. """ params = dict(ok=ok) @@ -2281,7 +2281,7 @@ def get_task_result_recs( def get_result_unit_recs( self, pdrr: ScopedKey | ProtocolDAGResultRec ) -> list[ProtocolUnitResultRec]: - """Get records describing the `ProtocolUnitResult`s of a `ProtocolDAGResult`. + """Get records describing the `ProtocolUnitResult` objects of a `ProtocolDAGResult`. Parameters ---------- @@ -2292,7 +2292,7 @@ def get_result_unit_recs( Returns ------- list[ProtocolUnitResultRec] - A list of `ProtocolUnitResultRec`s, one per `ProtocolUnitResult`, + A list of `ProtocolUnitResultRec` records, one per `ProtocolUnitResult`, in dependency order. """ pdrr_sk = self._as_scoped_key(pdrr) diff --git a/alchemiscale/storage/models.py b/alchemiscale/storage/models.py index b69e7a85..94b8e38b 100644 --- a/alchemiscale/storage/models.py +++ b/alchemiscale/storage/models.py @@ -201,10 +201,10 @@ class TaskProvenance(BaseModel): outcome The terminal outcome of the attempt; `None` while the attempt is open. units_completed - The number of distinct `ProtocolUnit`s successfully completed in this + The number of distinct `ProtocolUnit` objects successfully completed in this attempt, as of the last progress update. units_total - The total number of `ProtocolUnit`s in the attempt's `ProtocolDAG`. + The total number of `ProtocolUnit` objects in the attempt's `ProtocolDAG`. """ compute_service_id: ComputeServiceID @@ -967,7 +967,7 @@ def from_dict(cls, d): class TaskUnitTraceback(BaseModel): - """A single `ProtocolUnitFailure`'s traceback within a `TaskTracebacks`.""" + """A single `ProtocolUnitFailure` traceback within a `TaskTracebacks`.""" failure_key: GufeKey source_key: GufeKey diff --git a/alchemiscale/storage/statestore.py b/alchemiscale/storage/statestore.py index 35d24e0f..11d6a5c6 100644 --- a/alchemiscale/storage/statestore.py +++ b/alchemiscale/storage/statestore.py @@ -3747,7 +3747,7 @@ def get_tasks_details(self, tasks: list[ScopedKey]) -> list[TaskDetails | None]: """Return `TaskDetails` for each given `Task`, in input order. `None` is returned in place of any `Task` that does not exist. The - `current_claim`'s live progress fields stay `None` until a compute + The `current_claim` live progress fields stay `None` until a compute service reports progress (section 2 of the design). """ q = """ @@ -3818,11 +3818,11 @@ def get_tasks_details(self, tasks: list[ScopedKey]) -> list[TaskDetails | None]: def get_task_tracebacks( self, task: ScopedKey, limit: int | None = None ) -> list[TaskTracebacks]: - """Return tracebacks for the failed `ProtocolDAGResult`s of a `Task`. + """Return tracebacks for the failed `ProtocolDAGResult` objects of a `Task`. One `TaskTracebacks` record per failed `ProtocolDAGResultRef`, most recent first (by `datetime_created`). Where per-unit - `ProtocolUnitResultRef`s exist (section 3.4), each failure entry carries + `ProtocolUnitResultRef` nodes exist (section 3.4), each failure entry carries the `ScopedKey` of the corresponding unit ref; otherwise it is `None`. """ q = """ From 3f7047e49ac7ef910b931b1348481a7d00d19303 Mon Sep 17 00:00:00 2001 From: David Dotson Date: Fri, 10 Jul 2026 11:29:21 -0600 Subject: [PATCH 05/18] Add tests for introspection record models and object-store artifacts - tests/unit/test_introspection_records.py: to_dict/from_dict round-trips (including None branches) for TaskProvenance and the client-facing record models (TaskAttempt, TaskClaim, TaskDetails, TaskTracebacks, TaskUnitTraceback, ProtocolDAGResultRec, ProtocolUnitResultRec), plus the datetime coercion helpers and the ProtocolUnitResultRef gufe round-trip. - tests/integration/storage/test_objectstore_introspection.py: push/pull round-trips for the per-unit log and stdout/stderr artifact methods (moto-mocked), including UTF-8 errors="replace" decoding, empty-prefix pulls, and invalid-stream-name validation. Co-Authored-By: Claude Opus 4.8 (1M context) --- .../storage/test_objectstore_introspection.py | 98 ++++++ .../tests/unit/test_introspection_records.py | 308 ++++++++++++++++++ 2 files changed, 406 insertions(+) create mode 100644 alchemiscale/tests/integration/storage/test_objectstore_introspection.py create mode 100644 alchemiscale/tests/unit/test_introspection_records.py diff --git a/alchemiscale/tests/integration/storage/test_objectstore_introspection.py b/alchemiscale/tests/integration/storage/test_objectstore_introspection.py new file mode 100644 index 00000000..bb6b4cc9 --- /dev/null +++ b/alchemiscale/tests/integration/storage/test_objectstore_introspection.py @@ -0,0 +1,98 @@ +"""Integration tests for the v0.8.0 per-unit-result artifact methods on +``S3ObjectStore`` (log/stdout/stderr push and pull), exercised against the +moto-mocked ``s3os`` fixture. +""" + +import os + +import pytest + +from alchemiscale.storage.objectstore import ( + S3ObjectStore, + LOGS_FILENAME, + STDOUT_DIRNAME, + STDERR_DIRNAME, +) + + +class TestS3ObjectStoreArtifacts: + + UNIT_LOCATION = "protocoldagresult/o/c/p/T/results/PDR/units/PUR" + + def test_push_pull_logs_roundtrip(self, s3os: S3ObjectStore): + logtext = "[2026-07-10 00:00:00] [gufekey-x] [INFO] line one\nline two\n" + + location = s3os.push_protocol_unit_result_logs(self.UNIT_LOCATION, logtext) + assert location == os.path.join(self.UNIT_LOCATION, LOGS_FILENAME) + + # the object exists under the store prefix + objs = list(s3os.resource.Bucket(s3os.bucket).objects.all()) + assert any(o.key == os.path.join(s3os.prefix, location) for o in objs) + + # round-trip decompresses to the original text + assert s3os.pull_protocol_unit_result_logs(self.UNIT_LOCATION) == logtext + + def test_push_pull_stdout_roundtrip(self, s3os: S3ObjectStore): + files = { + "out.txt": b"hello stdout\n", + "sub.log": b"more stdout bytes", + } + locations = s3os.push_protocol_unit_result_streams( + self.UNIT_LOCATION, STDOUT_DIRNAME, files + ) + assert len(locations) == 2 + for name in files: + assert ( + os.path.join(self.UNIT_LOCATION, STDOUT_DIRNAME, f"{name}.zst") + in locations + ) + + pulled = s3os.pull_protocol_unit_result_streams( + self.UNIT_LOCATION, STDOUT_DIRNAME + ) + assert pulled == { + "out.txt": "hello stdout\n", + "sub.log": "more stdout bytes", + } + + def test_push_pull_stderr_roundtrip(self, s3os: S3ObjectStore): + files = {"err.txt": b"a traceback here"} + s3os.push_protocol_unit_result_streams( + self.UNIT_LOCATION, STDERR_DIRNAME, files + ) + pulled = s3os.pull_protocol_unit_result_streams( + self.UNIT_LOCATION, STDERR_DIRNAME + ) + assert pulled == {"err.txt": "a traceback here"} + + def test_streams_decode_errors_replace(self, s3os: S3ObjectStore): + # invalid UTF-8 bytes are decoded with errors="replace", not raised + files = {"binary.dat": b"\xff\xfe valid tail"} + s3os.push_protocol_unit_result_streams( + self.UNIT_LOCATION, STDOUT_DIRNAME, files + ) + pulled = s3os.pull_protocol_unit_result_streams( + self.UNIT_LOCATION, STDOUT_DIRNAME + ) + assert "valid tail" in pulled["binary.dat"] + # replacement character present for the invalid bytes + assert "�" in pulled["binary.dat"] + + def test_pull_streams_empty_when_absent(self, s3os: S3ObjectStore): + # a location with no stream artifacts yields an empty mapping + assert ( + s3os.pull_protocol_unit_result_streams( + "protocoldagresult/o/c/p/T/results/PDR/units/NOPE", STDOUT_DIRNAME + ) + == {} + ) + + @pytest.mark.parametrize("method", ["push", "pull"]) + def test_invalid_stream_name_raises(self, s3os: S3ObjectStore, method): + with pytest.raises(ValueError, match="stream"): + if method == "push": + s3os.push_protocol_unit_result_streams( + self.UNIT_LOCATION, "notastream", {"x": b"y"} + ) + else: + s3os.pull_protocol_unit_result_streams(self.UNIT_LOCATION, "notastream") diff --git a/alchemiscale/tests/unit/test_introspection_records.py b/alchemiscale/tests/unit/test_introspection_records.py new file mode 100644 index 00000000..db0ccbf1 --- /dev/null +++ b/alchemiscale/tests/unit/test_introspection_records.py @@ -0,0 +1,308 @@ +"""Unit tests for the v0.8.0 introspection data models: the ``TaskProvenance`` +node model and the client-facing ``*Rec``/``Task*`` record models, focusing on +``to_dict``/``from_dict`` round-trips (including the ``None`` branches) that the +API and client rely on for wire transport. +""" + +import datetime + +import pytest +from gufe.tokenization import GufeKey + +from alchemiscale.models import Scope, ScopedKey +from alchemiscale.storage.models import ( + ComputeServiceID, + ProtocolDAGResultRec, + ProtocolUnitResultRec, + ProtocolUnitResultRef, + TaskAttempt, + TaskClaim, + TaskDetails, + TaskOutcomeEnum, + TaskProvenance, + TaskStatusEnum, + TaskTracebacks, + TaskUnitTraceback, + _coerce_datetime, + _iso, +) + +NOW = datetime.datetime(2026, 7, 10, 12, 0, 0, tzinfo=datetime.UTC) +LATER = datetime.datetime(2026, 7, 10, 13, 30, 0, tzinfo=datetime.UTC) +CSID = ComputeServiceID("svc-" + "0" * 32) +PDRR_SK = ScopedKey.from_str("ProtocolDAGResultRef-abc123-org-camp-proj") +PURR_SK = ScopedKey.from_str("ProtocolUnitResultRef-def456-org-camp-proj") +TASK_SK = ScopedKey.from_str("Task-aaa111-org-camp-proj") + + +class TestHelpers: + def test_coerce_datetime_none(self): + assert _coerce_datetime(None) is None + + def test_coerce_datetime_iso_string(self): + assert _coerce_datetime(NOW.isoformat()) == NOW + + def test_coerce_datetime_passthrough(self): + assert _coerce_datetime(NOW) == NOW + + def test_coerce_datetime_neo4j_like(self): + class FakeNeo4jDT: + def to_native(self_inner): + return NOW + + assert _coerce_datetime(FakeNeo4jDT()) == NOW + + def test_iso(self): + assert _iso(None) is None + assert _iso(NOW) == NOW.isoformat() + + +class TestTaskProvenance: + def test_roundtrip_full(self): + tp = TaskProvenance( + compute_service_id=CSID, + hostname="host-a", + manager_name="mgr", + datetime_claimed=NOW, + datetime_end=LATER, + outcome=TaskOutcomeEnum.complete, + units_completed=3, + units_total=5, + ) + d = tp.to_dict() + assert d["compute_service_id"] == str(CSID) + assert d["outcome"] == "complete" + tp2 = TaskProvenance.from_dict(d) + assert isinstance(tp2.compute_service_id, ComputeServiceID) + assert tp2.outcome is TaskOutcomeEnum.complete + assert tp2.units_completed == 3 + assert tp2.datetime_end == LATER + + def test_roundtrip_open(self): + # an open attempt: no end / outcome / progress + tp = TaskProvenance(compute_service_id=CSID, datetime_claimed=NOW) + d = tp.to_dict() + assert d["outcome"] is None + assert d["datetime_end"] is None + tp2 = TaskProvenance.from_dict(d) + assert tp2.outcome is None + assert tp2.datetime_end is None + assert tp2.hostname is None + + +class TestTaskAttempt: + @pytest.mark.parametrize( + "outcome,pdrr", + [ + (TaskOutcomeEnum.complete, PDRR_SK), + (TaskOutcomeEnum.error, PDRR_SK), + (TaskOutcomeEnum.expired, None), + (TaskOutcomeEnum.released, None), + (None, None), + ], + ) + def test_roundtrip(self, outcome, pdrr): + ta = TaskAttempt( + compute_service_id=str(CSID), + hostname="h", + manager_name=None, + datetime_claimed=NOW, + datetime_end=LATER if outcome is not None else None, + outcome=outcome, + units_completed=1 if outcome else None, + units_total=4 if outcome else None, + protocoldagresultref=pdrr, + ) + 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.datetime_claimed == NOW + + +class TestTaskClaim: + def test_roundtrip(self): + tc = TaskClaim( + compute_service_id=str(CSID), + hostname="h", + datetime_claimed=NOW, + units_completed=2, + units_total=6, + ) + tc2 = TaskClaim.from_dict(tc.to_dict()) + assert tc2.hostname == "h" + assert tc2.units_completed == 2 + assert tc2.datetime_claimed == NOW + + def test_roundtrip_minimal(self): + tc = TaskClaim(compute_service_id=str(CSID)) + tc2 = TaskClaim.from_dict(tc.to_dict()) + assert tc2.datetime_claimed is None + assert tc2.units_total is None + + +class TestTaskDetails: + def test_roundtrip_with_claim_and_attempt(self): + claim = TaskClaim( + compute_service_id=str(CSID), + hostname="h", + datetime_claimed=NOW, + units_completed=1, + units_total=3, + ) + attempt = TaskAttempt( + compute_service_id=str(CSID), + datetime_claimed=NOW, + outcome=None, + ) + td = TaskDetails( + task=TASK_SK, + status=TaskStatusEnum.running, + datetime_status_changed=NOW, + reason="because", + num_claims=2, + current_claim=claim, + most_recent_attempt=attempt, + ) + td2 = TaskDetails.from_dict(td.to_dict()) + assert td2.task == TASK_SK + assert td2.status is TaskStatusEnum.running + assert td2.reason == "because" + assert td2.num_claims == 2 + assert td2.current_claim.compute_service_id == str(CSID) + assert td2.most_recent_attempt.compute_service_id == str(CSID) + + def test_roundtrip_minimal(self): + td = TaskDetails(task=TASK_SK, status=TaskStatusEnum.waiting) + td2 = TaskDetails.from_dict(td.to_dict()) + assert td2.status is TaskStatusEnum.waiting + assert td2.num_claims == 0 + assert td2.current_claim is None + assert td2.most_recent_attempt is None + assert td2.reason is None + + +class TestTaskTracebacks: + def test_roundtrip(self): + units = [ + TaskUnitTraceback( + failure_key=GufeKey("ProtocolUnitFailure-f1"), + source_key=GufeKey("ProtocolUnit-u1"), + traceback="boom", + protocolunitresultref=PURR_SK, + ), + TaskUnitTraceback( + failure_key=GufeKey("ProtocolUnitFailure-f2"), + source_key=GufeKey("ProtocolUnit-u2"), + traceback="kaboom", + protocolunitresultref=None, + ), + ] + tt = TaskTracebacks( + protocoldagresultref=PDRR_SK, + datetime_created=NOW, + creator=str(CSID), + tracebacks=units, + ) + tt2 = TaskTracebacks.from_dict(tt.to_dict()) + assert tt2.protocoldagresultref == PDRR_SK + assert tt2.creator == str(CSID) + assert [u.traceback for u in tt2.tracebacks] == ["boom", "kaboom"] + assert tt2.tracebacks[0].protocolunitresultref == PURR_SK + assert tt2.tracebacks[1].protocolunitresultref is None + assert isinstance(tt2.tracebacks[0].failure_key, GufeKey) + + +class TestProtocolDAGResultRec: + @pytest.mark.parametrize("ok", [True, False]) + def test_roundtrip(self, ok): + rec = ProtocolDAGResultRec( + scoped_key=PDRR_SK, + ok=ok, + datetime_created=NOW, + creator=str(CSID), + ) + rec2 = ProtocolDAGResultRec.from_dict(rec.to_dict()) + assert rec2.scoped_key == PDRR_SK + assert rec2.ok is ok + assert rec2.datetime_created == NOW + assert rec2.creator == str(CSID) + + def test_roundtrip_minimal(self): + rec = ProtocolDAGResultRec(scoped_key=PDRR_SK, ok=True) + rec2 = ProtocolDAGResultRec.from_dict(rec.to_dict()) + assert rec2.datetime_created is None + assert rec2.creator is None + + +class TestProtocolUnitResultRec: + def test_roundtrip_full(self): + rec = ProtocolUnitResultRec( + scoped_key=PURR_SK, + obj_key=GufeKey("ProtocolUnitResult-r1"), + source_key=GufeKey("ProtocolUnit-u1"), + name="unit one", + ok=True, + start_time=NOW, + end_time=LATER, + has_logs=True, + has_stdout=True, + has_stderr=False, + ) + rec2 = ProtocolUnitResultRec.from_dict(rec.to_dict()) + assert rec2.scoped_key == PURR_SK + assert isinstance(rec2.obj_key, GufeKey) + assert isinstance(rec2.source_key, GufeKey) + assert rec2.name == "unit one" + assert rec2.ok is True + assert rec2.start_time == NOW + assert rec2.end_time == LATER + assert rec2.has_logs is True + assert rec2.has_stdout is True + assert rec2.has_stderr is False + + def test_roundtrip_minimal(self): + rec = ProtocolUnitResultRec( + scoped_key=PURR_SK, + obj_key=GufeKey("ProtocolUnitFailure-r2"), + source_key=GufeKey("ProtocolUnit-u2"), + ok=False, + ) + rec2 = ProtocolUnitResultRec.from_dict(rec.to_dict()) + assert rec2.ok is False + assert rec2.name is None + assert rec2.start_time is None + assert rec2.end_time is None + assert rec2.has_logs is False + assert rec2.has_stdout is False + assert rec2.has_stderr is False + + +class TestProtocolUnitResultRefNode: + """The state-store node type; verify it tokenizes and its _to_dict/_from_dict + round-trip through gufe's keyed-chain machinery.""" + + def test_gufe_roundtrip(self): + purr = ProtocolUnitResultRef( + location="protocoldagresult/o/c/p/T/results/K/units/R", + obj_key=GufeKey("ProtocolUnitResult-r1"), + source_key=GufeKey("ProtocolUnit-u1"), + scope=Scope("o", "c", "p"), + ok=True, + name="u", + start_time=NOW, + end_time=LATER, + has_logs=True, + ) + # deterministic key computed once at creation + key1 = purr.key + # round-trip through the keyed chain (as the state store does) + from gufe.tokenization import KeyedChain + + purr2 = KeyedChain.from_gufe(purr).to_gufe() + assert purr2.obj_key == purr.obj_key + assert purr2.source_key == purr.source_key + assert purr2.ok is True + assert purr2.has_logs is True + assert purr2.start_time == NOW + assert str(purr2.key) == str(key1) From 07e57238096f4b88e3526b0745262683bd07a333 Mon Sep 17 00:00:00 2001 From: David Dotson Date: Fri, 10 Jul 2026 12:14:59 -0600 Subject: [PATCH 06/18] Add interface-client integration tests for the introspection surface End-to-end (client -> API -> Neo4j/S3) coverage for the new user-facing methods: get_task_history, get_tasks_details, get_task_result_recs, get_result_unit_recs, get_result_unit_logs/stdout/stderr, get_result_logs, get_task_stdout/stderr, get_task_tracebacks, get_tasks_progress, get_scope_compute_share, and set_tasks_status(reason=...). State (claimed provenance, results, per-unit refs, and per-unit artifacts) is set up via the n4js/s3os handles as the existing result tests do, then each method is called through the real client and asserted. Co-Authored-By: Claude Opus 4.8 (1M context) --- .../client/test_client_introspection.py | 296 ++++++++++++++++++ 1 file changed, 296 insertions(+) create mode 100644 alchemiscale/tests/integration/interface/client/test_client_introspection.py diff --git a/alchemiscale/tests/integration/interface/client/test_client_introspection.py b/alchemiscale/tests/integration/interface/client/test_client_introspection.py new file mode 100644 index 00000000..aca12dc1 --- /dev/null +++ b/alchemiscale/tests/integration/interface/client/test_client_introspection.py @@ -0,0 +1,296 @@ +"""Integration tests for the v0.8.0 Task-introspection client surface, exercised +end-to-end through the ``AlchemiscaleClient`` -> API -> Neo4j/S3 round trip. + +State (provenance, results, per-unit refs, artifacts) is set up directly via the +``n4js``/``s3os`` handles (as the existing interface result tests do), then the +new client methods are called and asserted. +""" + +import datetime +from pathlib import Path + +import pytest +from gufe.tokenization import TOKENIZABLE_REGISTRY +from gufe.protocols.protocoldag import execute_DAG + +from alchemiscale.compression import compress_gufe_zstd +from alchemiscale.models import Scope +from alchemiscale.storage.models import ( + ComputeServiceID, + ComputeServiceRegistration, + ProtocolDAGResultRec, + ProtocolUnitResultRec, + TaskAttempt, + TaskDetails, + TaskOutcomeEnum, + TaskStatusEnum, + TaskTracebacks, +) +from alchemiscale.interface import client + + +def _register(n4js, name, hostname="host-a"): + csid = ComputeServiceID.new_from_name(name) + now = datetime.datetime.now(tz=datetime.UTC) + n4js.register_computeservice( + ComputeServiceRegistration( + identifier=csid, + registered=now, + heartbeat=now, + failure_times=[], + hostname=hostname, + ) + ) + return csid + + +class TestClientIntrospection: + + def _claim_and_run( + self, user_client, n4js, s3os, transformation, network_sk, scope_test, csid + ): + """Create+action a task, claim it (creating provenance), execute its DAG, + push the result finalizing provenance, and derive per-unit refs. + + Returns ``(task_sk, pdrr_sk, pdr)``. + """ + transformation_sk = user_client.get_scoped_key(transformation, scope_test) + task_sk = user_client.create_tasks(transformation_sk, count=1)[0] + user_client.action_tasks([task_sk], network_sk) + + taskhub_sk = n4js.get_taskhub(network_sk) + claimed = n4js.claim_taskhub_tasks(taskhub_sk, csid) + assert claimed[0] == task_sk + + # execute the DAG as a compute service would + protocoldag = transformation.create(name=str(task_sk)) + shared = Path("shared").absolute() / str(protocoldag.key) + shared.mkdir(parents=True) + scratch = Path("scratch").absolute() / str(protocoldag.key) + scratch.mkdir(parents=True) + pdr = execute_DAG( + protocoldag, + shared_basedir=shared, + scratch_basedir=scratch, + raise_error=False, + ) + + pdrr = s3os.push_protocoldagresult( + compress_gufe_zstd(pdr), + pdr.ok(), + pdr.key, + transformation=transformation_sk, + ) + pdrr_sk = n4js.set_task_result(task_sk, pdrr, compute_service_id=csid) + n4js.add_protocol_unit_result_refs(pdrr, pdrr_sk, pdr) + return task_sk, pdrr_sk, pdr + + def test_history_details_and_result_recs( + self, + scope_test, + n4js_preloaded, + s3os_server_fresh, + user_client_no_cache: client.AlchemiscaleClient, + network_tyk2, + tmpdir, + ): + user_client = user_client_no_cache + n4js = n4js_preloaded + s3os = s3os_server_fresh + + an = network_tyk2 + transformation = [t for t in an.edges if "_solvent" in t.name][0] + network_sk = user_client.get_scoped_key(an, scope_test) + + csid = _register(n4js, "history.svc", hostname="node-1") + with tmpdir.as_cwd(): + task_sk, pdrr_sk, pdr = self._claim_and_run( + user_client, n4js, s3os, transformation, network_sk, scope_test, csid + ) + n4js.set_task_complete([task_sk]) + + # get_task_history + history = user_client.get_task_history(task_sk) + assert len(history) == 1 + assert isinstance(history[0], TaskAttempt) + assert history[0].compute_service_id == str(csid) + assert history[0].hostname == "node-1" + assert history[0].outcome is TaskOutcomeEnum.complete + assert history[0].protocoldagresultref == pdrr_sk + + # get_tasks_details + details = user_client.get_tasks_details([task_sk]) + assert len(details) == 1 + assert isinstance(details[0], TaskDetails) + assert details[0].task == task_sk + assert details[0].status is TaskStatusEnum.complete + assert details[0].num_claims == 1 + assert details[0].most_recent_attempt is not None + assert details[0].most_recent_attempt.outcome is TaskOutcomeEnum.complete + + # get_task_result_recs + recs = user_client.get_task_result_recs(task_sk) + assert len(recs) == 1 + assert isinstance(recs[0], ProtocolDAGResultRec) + assert recs[0].scoped_key == pdrr_sk + assert recs[0].ok is True + assert user_client.get_task_result_recs(task_sk, ok=True) + assert user_client.get_task_result_recs(task_sk, ok=False) == [] + + # get_result_unit_recs, in dependency order; count matches the PDR + unit_recs = user_client.get_result_unit_recs(recs[0]) + assert len(unit_recs) == len(pdr.protocol_unit_results) + assert all(isinstance(r, ProtocolUnitResultRec) for r in unit_recs) + assert all(r.ok for r in unit_recs) + + def test_unit_artifacts_retrieval( + self, + scope_test, + n4js_preloaded, + s3os_server_fresh, + user_client_no_cache: client.AlchemiscaleClient, + network_tyk2, + tmpdir, + ): + user_client = user_client_no_cache + n4js = n4js_preloaded + s3os = s3os_server_fresh + + an = network_tyk2 + transformation = [t for t in an.edges if "_solvent" in t.name][0] + network_sk = user_client.get_scoped_key(an, scope_test) + + csid = _register(n4js, "artifacts.svc") + with tmpdir.as_cwd(): + task_sk, pdrr_sk, pdr = self._claim_and_run( + user_client, n4js, s3os, transformation, network_sk, scope_test, csid + ) + n4js.set_task_complete([task_sk]) + + # manually attach artifacts to the first unit result and flip its flags + unit_recs = n4js.get_result_unit_recs(pdrr_sk) + purr_sk = unit_recs[0].scoped_key + location = n4js.get_gufe(purr_sk).location + + s3os.push_protocol_unit_result_logs(location, "hello from the unit log\n") + n4js.set_protocol_unit_result_ref_artifacts(purr_sk, has_logs=True) + s3os.push_protocol_unit_result_streams( + location, "stdout", {"out.txt": b"captured stdout\n"} + ) + n4js.set_protocol_unit_result_ref_artifacts(purr_sk, has_stdout=True) + s3os.push_protocol_unit_result_streams( + location, "stderr", {"err.txt": b"captured stderr\n"} + ) + n4js.set_protocol_unit_result_ref_artifacts(purr_sk, has_stderr=True) + + # single-unit retrieval (accepts a ScopedKey or a *Rec) + assert user_client.get_result_unit_logs(purr_sk) == "hello from the unit log\n" + assert user_client.get_result_unit_stdout(purr_sk) == { + "out.txt": "captured stdout\n" + } + assert user_client.get_result_unit_stderr(purr_sk) == { + "err.txt": "captured stderr\n" + } + + # a unit with no logs returns None + no_logs = [r for r in unit_recs if r.scoped_key != purr_sk] + if no_logs: + assert user_client.get_result_unit_logs(no_logs[0].scoped_key) is None + + # rendered aggregations include the captured content + rendered = user_client.get_result_logs(pdrr_sk) + assert "hello from the unit log" in rendered + assert "captured stdout" in user_client.get_task_stdout(task_sk) + assert "captured stderr" in user_client.get_task_stderr(task_sk) + + def test_tracebacks( + self, + scope_test, + n4js_preloaded, + s3os_server_fresh, + user_client_no_cache: client.AlchemiscaleClient, + network_tyk2_failure, + tmpdir, + ): + user_client = user_client_no_cache + n4js = n4js_preloaded + s3os = s3os_server_fresh + + network_sk = user_client.create_network(network_tyk2_failure, scope_test) + transformation = [t for t in network_tyk2_failure.edges if t.name == "broken"][ + 0 + ] + + csid = _register(n4js, "tb.svc") + with tmpdir.as_cwd(): + task_sk, pdrr_sk, pdr = self._claim_and_run( + user_client, n4js, s3os, transformation, network_sk, scope_test, csid + ) + # record the failure's tracebacks (as the compute API does) and error it + n4js.add_protocol_dag_result_ref_tracebacks(pdr.protocol_unit_failures, pdrr_sk) + n4js.set_task_error([task_sk]) + + for pdr_ in [pdr]: + TOKENIZABLE_REGISTRY.pop(pdr_.key, None) + + tracebacks = user_client.get_task_tracebacks(task_sk) + assert len(tracebacks) == 1 + assert isinstance(tracebacks[0], TaskTracebacks) + assert tracebacks[0].protocoldagresultref == pdrr_sk + assert len(tracebacks[0].tracebacks) == len(pdr.protocol_unit_failures) + assert all(tb.traceback for tb in tracebacks[0].tracebacks) + # the unit-ref link is populated (unit refs were derived above) + assert any( + tb.protocolunitresultref is not None for tb in tracebacks[0].tracebacks + ) + + def test_progress_and_compute_share( + self, + scope_test, + n4js_preloaded, + user_client_no_cache: client.AlchemiscaleClient, + network_tyk2, + ): + user_client = user_client_no_cache + n4js = n4js_preloaded + + an = network_tyk2 + transformation = [t for t in an.edges if "_solvent" in t.name][0] + network_sk = user_client.get_scoped_key(an, scope_test) + transformation_sk = user_client.get_scoped_key(transformation, scope_test) + + task_sk = user_client.create_tasks(transformation_sk, count=1)[0] + user_client.action_tasks([task_sk], network_sk) + taskhub_sk = n4js.get_taskhub(network_sk) + csid = _register(n4js, "progress.svc") + assert n4js.claim_taskhub_tasks(taskhub_sk, csid)[0] == task_sk + + # no progress reported yet + assert user_client.get_tasks_progress([task_sk]) == [None] + + n4js.update_task_progress(csid, {str(task_sk): (2, 5)}) + assert user_client.get_tasks_progress([task_sk]) == [(2, 5)] + + # compute share: this scope's running Tasks over same-level siblings. + # only this scope has a running Task, so the org-level share is 1.0 + share = user_client.get_scope_compute_share(Scope(org=scope_test.org)) + assert share == pytest.approx(1.0) + + def test_set_tasks_status_with_reason( + self, + scope_test, + n4js_preloaded, + user_client_no_cache: client.AlchemiscaleClient, + network_tyk2, + ): + user_client = user_client_no_cache + an = network_tyk2 + transformation = [t for t in an.edges if "_solvent" in t.name][0] + transformation_sk = user_client.get_scoped_key(transformation, scope_test) + + task_sk = user_client.create_tasks(transformation_sk, count=1)[0] + user_client.set_tasks_status([task_sk], "invalid", reason="bad inputs") + + details = user_client.get_tasks_details([task_sk]) + assert details[0].status is TaskStatusEnum.invalid + assert details[0].reason == "bad inputs" From f9723ea5d7e8e4491d7a8735165fddf493030f66 Mon Sep 17 00:00:00 2001 From: David Dotson Date: Fri, 10 Jul 2026 12:57:48 -0600 Subject: [PATCH 07/18] Fix compute-share interface test: query at the identity's held scope get_scope_compute_share at the org level requires org-level scope access; the test identity holds only the specific project scope, so the org-level query 401'd (correct authorization behavior). Query at the project scope the identity holds instead; the org-level branch is already covered by the statestore integration test. Co-Authored-By: Claude Opus 4.8 (1M context) --- .../interface/client/test_client_introspection.py | 10 ++++++---- 1 file changed, 6 insertions(+), 4 deletions(-) diff --git a/alchemiscale/tests/integration/interface/client/test_client_introspection.py b/alchemiscale/tests/integration/interface/client/test_client_introspection.py index aca12dc1..307452d0 100644 --- a/alchemiscale/tests/integration/interface/client/test_client_introspection.py +++ b/alchemiscale/tests/integration/interface/client/test_client_introspection.py @@ -14,7 +14,6 @@ from gufe.protocols.protocoldag import execute_DAG from alchemiscale.compression import compress_gufe_zstd -from alchemiscale.models import Scope from alchemiscale.storage.models import ( ComputeServiceID, ComputeServiceRegistration, @@ -271,9 +270,12 @@ def test_progress_and_compute_share( n4js.update_task_progress(csid, {str(task_sk): (2, 5)}) assert user_client.get_tasks_progress([task_sk]) == [(2, 5)] - # compute share: this scope's running Tasks over same-level siblings. - # only this scope has a running Task, so the org-level share is 1.0 - share = user_client.get_scope_compute_share(Scope(org=scope_test.org)) + # compute share: the project's running Tasks over sibling projects in + # the same org-campaign. Query at the fully-specified project scope the + # identity actually holds (an org-level query would require org-level + # access, which this identity lacks); only this project has a running + # Task, so its share is 1.0. + share = user_client.get_scope_compute_share(scope_test) assert share == pytest.approx(1.0) def test_set_tasks_status_with_reason( From eeedfaa7b57fff647226ca0e12ef779de7352629 Mon Sep 17 00:00:00 2001 From: David Dotson Date: Fri, 10 Jul 2026 13:43:06 -0600 Subject: [PATCH 08/18] Cover remaining introspection branches (time-ordered logs, compute-share levels) Nudge patch coverage over the threshold by exercising the last uncovered branches in already-passing tests: - get_result_logs(order='time') interleave path (with a timestamped and an untimestamped log line, hitting both sort keys); - get_scope_compute_share campaign-level and project-level grouping, plus the no-org ValueError. Co-Authored-By: Claude Opus 4.8 (1M context) --- .../client/test_client_introspection.py | 16 +++++++++++----- .../storage/test_statestore_introspection.py | 17 ++++++++++++++--- 2 files changed, 25 insertions(+), 8 deletions(-) diff --git a/alchemiscale/tests/integration/interface/client/test_client_introspection.py b/alchemiscale/tests/integration/interface/client/test_client_introspection.py index 307452d0..dfbcd685 100644 --- a/alchemiscale/tests/integration/interface/client/test_client_introspection.py +++ b/alchemiscale/tests/integration/interface/client/test_client_introspection.py @@ -171,7 +171,11 @@ def test_unit_artifacts_retrieval( purr_sk = unit_recs[0].scoped_key location = n4js.get_gufe(purr_sk).location - s3os.push_protocol_unit_result_logs(location, "hello from the unit log\n") + # a timestamped line plus an untimestamped one, so both the + # timestamp-sorted and unsorted paths of get_result_logs(order='time') + # are exercised + logtext = "[2026-07-10 00:00:00] [gk] [INFO] first line\nsecond line\n" + s3os.push_protocol_unit_result_logs(location, logtext) n4js.set_protocol_unit_result_ref_artifacts(purr_sk, has_logs=True) s3os.push_protocol_unit_result_streams( location, "stdout", {"out.txt": b"captured stdout\n"} @@ -183,7 +187,7 @@ def test_unit_artifacts_retrieval( n4js.set_protocol_unit_result_ref_artifacts(purr_sk, has_stderr=True) # single-unit retrieval (accepts a ScopedKey or a *Rec) - assert user_client.get_result_unit_logs(purr_sk) == "hello from the unit log\n" + assert user_client.get_result_unit_logs(purr_sk) == logtext assert user_client.get_result_unit_stdout(purr_sk) == { "out.txt": "captured stdout\n" } @@ -196,9 +200,11 @@ def test_unit_artifacts_retrieval( if no_logs: assert user_client.get_result_unit_logs(no_logs[0].scoped_key) is None - # rendered aggregations include the captured content - rendered = user_client.get_result_logs(pdrr_sk) - assert "hello from the unit log" in rendered + # rendered aggregations include the captured content, in both orderings + rendered_unit = user_client.get_result_logs(pdrr_sk) + assert "first line" in rendered_unit + rendered_time = user_client.get_result_logs(pdrr_sk, order="time") + assert "first line" in rendered_time assert "captured stdout" in user_client.get_task_stdout(task_sk) assert "captured stderr" in user_client.get_task_stderr(task_sk) diff --git a/alchemiscale/tests/integration/storage/test_statestore_introspection.py b/alchemiscale/tests/integration/storage/test_statestore_introspection.py index 4467a0ac..b10f5e6a 100644 --- a/alchemiscale/tests/integration/storage/test_statestore_introspection.py +++ b/alchemiscale/tests/integration/storage/test_statestore_introspection.py @@ -654,9 +654,20 @@ def running_tasks(scope, count, name): running_tasks(scope_a, 3, "a") # 3 running in orgA running_tasks(scope_b, 1, "b") # 1 running in orgB - # orgA's share of all running tasks across orgs: 3 / (3 + 1) - share = n4js.get_scope_compute_share(Scope(org="orgA")) - assert share == pytest.approx(0.75) + # org-level: orgA's share of all running tasks across orgs: 3 / (3 + 1) + assert n4js.get_scope_compute_share(Scope(org="orgA")) == pytest.approx(0.75) + + # campaign-level: orgA/camp's share of all campaigns in orgA (only one) + assert n4js.get_scope_compute_share(Scope("orgA", "camp")) == pytest.approx(1.0) + + # project-level: orgA/camp/proj's share of all projects in orgA-camp + assert n4js.get_scope_compute_share( + Scope("orgA", "camp", "proj") + ) == pytest.approx(1.0) # empty population -> 0.0 assert n4js.get_scope_compute_share(Scope(org="orgC")) == 0.0 + + # a scope with no org cannot be leveled + with pytest.raises(ValueError): + n4js.get_scope_compute_share(Scope()) From 51a29dd0468be32196ba2bef93a2ea5b85a7d29c Mon Sep 17 00:00:00 2001 From: David Dotson Date: Sat, 11 Jul 2026 14:18:19 -0600 Subject: [PATCH 09/18] Address Fable review: fix dead progress push, migration naming, add wire tests - H1 (HIGH): the /computeservice/{id}/progress route declared Body(embed=True) while the compute client sends the bare map, so every live-progress push 422'd (silently, swallowed at DEBUG) and the whole feature returned None in production. Drop embed so the route accepts the bare map per the design's transport spec; also skip malformed entries instead of 500'ing. - M2: add end-to-end compute-client tests through the API for update_task_progress (would have caught H1), set_task_error (the /error DAG-creation-failure path), and set_task_result_unit_logs (the artifact upload route). - M1: rename the migration v04_to_v05 -> v07_to_v08. Migration names track release versions (v03_to_v04 shipped in v0.4.0); this is the v0.7 -> v0.8 upgrade, so as named an operator upgrading 0.7->0.8 would never run it and the TaskProvenance indexes would never be created. Update the CLI command and document the migration in docs/operations.rst. - N4: constrain the `limit` query param on /history and /tracebacks to ge=1 (422 instead of a 500 from a negative Cypher LIMIT). Co-Authored-By: Claude Opus 4.8 (1M context) --- alchemiscale/cli.py | 8 +- alchemiscale/compute/api.py | 16 ++- alchemiscale/interface/api.py | 6 +- .../{v04_to_v05.py => v07_to_v08.py} | 4 +- .../compute/client/test_compute_client.py | 114 ++++++++++++++++++ docs/operations.rst | 15 +++ 6 files changed, 148 insertions(+), 15 deletions(-) rename alchemiscale/migrations/{v04_to_v05.py => v07_to_v08.py} (93%) diff --git a/alchemiscale/cli.py b/alchemiscale/cli.py index d49fdd7d..db1129b2 100644 --- a/alchemiscale/cli.py +++ b/alchemiscale/cli.py @@ -491,16 +491,16 @@ def v03_to_v04(url, user, password, dbname): @migrate.command() @db_params -def v04_to_v05(url, user, password, dbname): - """Perform migration appropriate for transitioning from alchemiscale v0.4 - to v0.5. +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.v04_to_v05 import migrate + from .migrations.v07_to_v08 import migrate cli_values = url | user | password | dbname settings = get_settings_from_options(cli_values, Neo4jStoreSettings) diff --git a/alchemiscale/compute/api.py b/alchemiscale/compute/api.py index 71ed9c0b..ea9b1e4d 100644 --- a/alchemiscale/compute/api.py +++ b/alchemiscale/compute/api.py @@ -506,20 +506,24 @@ async def set_task_error( def update_task_progress( compute_service_id, *, - progress: dict[str, dict[str, int]] = Body(..., embed=True), + progress: dict[str, dict[str, int]] = Body(...), n4js: Neo4jStore = Depends(get_n4js_depends), ): """Record live progress counts for a service's claimed Tasks. - Body maps `Task` ScopedKey strings to - ``{"units_completed": int, "units_total": int}`` --- one batched request per - push event, regardless of claim count. Like a heartbeat, this route never - rejects: updates for Tasks the service no longer claims are silently dropped - server-side (the claim expired mid-flight). + The request body is the bare map from `Task` ScopedKey string to + ``{"units_completed": int, "units_total": int}`` (per the design's transport + spec) --- one batched request per push event, regardless of claim count. + NOTE: this is deliberately NOT ``embed``ed; the compute client sends the map + as the whole body. Like a heartbeat, this route never rejects: updates for + Tasks the service no longer claims are silently dropped server-side (the + claim expired mid-flight), and malformed entries are skipped rather than + 500'd. """ progress_ = { task_sk: (counts["units_completed"], counts["units_total"]) for task_sk, counts in progress.items() + if "units_completed" in counts and "units_total" in counts } n4js.update_task_progress(ComputeServiceID(compute_service_id), progress_) return None diff --git a/alchemiscale/interface/api.py b/alchemiscale/interface/api.py index 07963bfa..8c557112 100644 --- a/alchemiscale/interface/api.py +++ b/alchemiscale/interface/api.py @@ -7,7 +7,7 @@ import re from collections import Counter -from fastapi import FastAPI, APIRouter, Body, Depends, HTTPException, Request +from fastapi import FastAPI, APIRouter, Body, Depends, HTTPException, Query, Request from fastapi import status as http_status from fastapi.middleware.gzip import GZipMiddleware @@ -1199,7 +1199,7 @@ def get_task_failures( def get_task_history( task_scoped_key, *, - limit: int | None = None, + limit: int | None = Query(None, ge=1), n4js: Neo4jStore = Depends(get_n4js_depends), token: TokenData = Depends(get_token_data_depends), ): @@ -1230,7 +1230,7 @@ def get_tasks_details( def get_task_tracebacks( task_scoped_key, *, - limit: int | None = None, + limit: int | None = Query(None, ge=1), n4js: Neo4jStore = Depends(get_n4js_depends), token: TokenData = Depends(get_token_data_depends), ): diff --git a/alchemiscale/migrations/v04_to_v05.py b/alchemiscale/migrations/v07_to_v08.py similarity index 93% rename from alchemiscale/migrations/v04_to_v05.py rename to alchemiscale/migrations/v07_to_v08.py index b1ba9fb1..929817cc 100644 --- a/alchemiscale/migrations/v04_to_v05.py +++ b/alchemiscale/migrations/v07_to_v08.py @@ -1,5 +1,5 @@ """ -:mod:`alchemiscale.migrations.v04_to_v05` --- migration for v0.4 to v0.5 +:mod:`alchemiscale.migrations.v07_to_v08` --- migration for v0.7 to v0.8 ======================================================================== """ @@ -8,7 +8,7 @@ def migrate(n4js: Neo4jStore): - """Migrate state store from alchemiscale v0.4 to v0.5. + """Migrate state store from alchemiscale v0.7 to v0.8. Changes: - adds indexes on the new ``TaskProvenance`` node label to support the diff --git a/alchemiscale/tests/integration/compute/client/test_compute_client.py b/alchemiscale/tests/integration/compute/client/test_compute_client.py index 21bd591c..7ff60811 100644 --- a/alchemiscale/tests/integration/compute/client/test_compute_client.py +++ b/alchemiscale/tests/integration/compute/client/test_compute_client.py @@ -6,10 +6,13 @@ from gufe.tokenization import JSON_HANDLER +from gufe.tokenization import GufeKey + from alchemiscale.compute import client from alchemiscale.models import ScopedKey from alchemiscale.storage.models import ( TaskStatusEnum, + TaskOutcomeEnum, ProtocolDAGResultRef, ) from alchemiscale.tests.integration.compute.utils import get_compute_settings_override @@ -503,3 +506,114 @@ def test_set_task_result_failure( _ = compute_client.set_task_result(task_sks[0], protocoldagresults_failure[0]) assert n4js_preloaded.get_task_status(task_sks)[0] == TaskStatusEnum.error + + def test_update_task_progress( + self, + scope_test, + n4js_preloaded, + compute_client: client.AlchemiscaleComputeClient, + compute_service_id, + network_tyk2, + uvicorn_server, + ): + # exercises the live-progress push over the wire (client -> /progress + # route -> state store); a body-shape mismatch here would silently 422 + compute_client.register(compute_service_id) + an_sk = ScopedKey(gufe_key=network_tyk2.key, **scope_test.to_dict()) + taskhub_sk = n4js_preloaded.get_taskhub(an_sk) + + task_sk = compute_client.claim_taskhub_tasks( + taskhub_sk, compute_service_id=compute_service_id + )[0] + assert task_sk is not None + + # no progress reported yet for the running Task + assert n4js_preloaded.get_tasks_progress([task_sk]) == [None] + + compute_client.update_task_progress( + compute_service_id, + {str(task_sk): {"units_completed": 2, "units_total": 5}}, + ) + + assert n4js_preloaded.get_tasks_progress([task_sk]) == [(2, 5)] + + def test_set_task_error( + self, + scope_test, + n4js_preloaded, + compute_client: client.AlchemiscaleComputeClient, + compute_service_id, + network_tyk2, + uvicorn_server, + ): + # the ProtocolDAG creation-failure path (client -> /error route) + compute_client.register(compute_service_id) + an_sk = ScopedKey(gufe_key=network_tyk2.key, **scope_test.to_dict()) + taskhub_sk = n4js_preloaded.get_taskhub(an_sk) + + task_sk = compute_client.claim_taskhub_tasks( + taskhub_sk, compute_service_id=compute_service_id + )[0] + assert task_sk is not None + + returned = compute_client.set_task_error( + task_sk, reason="boom during create", compute_service_id=compute_service_id + ) + assert returned == task_sk + + assert n4js_preloaded.get_task_status([task_sk])[0] == TaskStatusEnum.error + + # reason recorded on the Task; open provenance finalized as error + details = n4js_preloaded.get_tasks_details([task_sk])[0] + assert details.reason == "boom during create" + + tp = n4js_preloaded.execute_query( + """ + MATCH (tp:TaskProvenance {compute_service_id: $csid})-[:PROVENANCE_OF]->(t:Task {_scoped_key: $task}) + RETURN tp + """, + csid=str(compute_service_id), + task=str(task_sk), + ).records[0]["tp"] + assert tp["outcome"] == TaskOutcomeEnum.error.value + + def test_set_task_result_unit_logs( + self, + scope_test, + n4js_preloaded, + s3os_server_fresh, + compute_client: client.AlchemiscaleComputeClient, + compute_service_id, + network_tyk2, + protocoldagresults, + uvicorn_server, + ): + # push a result (server derives unit refs), then upload logs for one + # unit result over the wire (client -> /artifacts/logs route) + compute_client.register(compute_service_id) + an_sk = ScopedKey(gufe_key=network_tyk2.key, **scope_test.to_dict()) + taskhub_sk = n4js_preloaded.get_taskhub(an_sk) + + task_sk = compute_client.claim_taskhub_tasks( + taskhub_sk, compute_service_id=compute_service_id + )[0] + assert task_sk is not None + + pdr = protocoldagresults[0] + pdrr_sk = compute_client.set_task_result(task_sk, pdr, compute_service_id) + + unit_key = str(pdr.protocol_unit_results[0].key) + compute_client.set_task_result_unit_logs( + task_sk, pdrr_sk, unit_key, "captured log line\n" + ) + + purr_sk = n4js_preloaded.get_protocol_unit_result_ref_scoped_key( + pdrr_sk, GufeKey(unit_key) + ) + assert purr_sk is not None + purr = n4js_preloaded.get_gufe(purr_sk) + assert purr.has_logs is True + assert ( + s3os_server_fresh.pull_protocol_unit_result_logs(purr.location) + == "captured log line\n" + ) diff --git a/docs/operations.rst b/docs/operations.rst index 7c18b2b6..a4fe7f30 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 and related +introspection features. +This requires a lightweight schema migration that adds ``neo4j`` indexes for the +new ``TaskProvenance`` node label; it is idempotent and requires no data +migration (pre-existing ``Task``\ s simply have empty execution history). + +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 From 4a8195c38279f96c81769765411c6fc5fde4c7c4 Mon Sep 17 00:00:00 2001 From: David Dotson Date: Thu, 16 Jul 2026 21:59:54 -0600 Subject: [PATCH 10/18] Remove the v07_to_v08 index migration Every TaskProvenance access in the introspection queries is Task-anchored -- `(t)<-[:PROVENANCE_OF]-(tp:TaskProvenance ...)` -- starting from a Task pinned by its `_scoped_key` (already covered by the GufeTokenizable uniqueness constraint) and expanding the relationship to that Task's handful of attempt records. The `compute_service_id`/`datetime_claimed` label indexes are never consulted by these anchored traversals; they'd only help a query that scans the TaskProvenance label globally (e.g. future pruning or #180), of which there are none yet. Drop the migration entirely (module, CLI command, and operations docs); we'll add indexes if a global provenance query ever proves them critical. No data migration is required for this release regardless. Co-Authored-By: Claude Opus 4.8 (1M context) --- alchemiscale/cli.py | 23 ------------ alchemiscale/migrations/v07_to_v08.py | 52 --------------------------- docs/operations.rst | 15 -------- 3 files changed, 90 deletions(-) delete mode 100644 alchemiscale/migrations/v07_to_v08.py diff --git a/alchemiscale/cli.py b/alchemiscale/cli.py index db1129b2..c87caa5e 100644 --- a/alchemiscale/cli.py +++ b/alchemiscale/cli.py @@ -489,29 +489,6 @@ 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/migrations/v07_to_v08.py b/alchemiscale/migrations/v07_to_v08.py deleted file mode 100644 index 929817cc..00000000 --- a/alchemiscale/migrations/v07_to_v08.py +++ /dev/null @@ -1,52 +0,0 @@ -""" -: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. - - Changes: - - adds indexes on the new ``TaskProvenance`` node label to support the - introspection queries introduced in v0.5. ``TaskProvenance`` is a plain - labeled node (not a ``GufeTokenizable``), identified by its properties and - reached from a ``Task`` via the ``PROVENANCE_OF`` relationship: - - ``TaskProvenance.compute_service_id``: provenance records are matched - by the id of the compute service that produced them, both when - finalizing an attempt (``set_task_result``, expiry/deregistration) and - when reading live progress for the current claimant. - - ``TaskProvenance.datetime_claimed``: attempt histories and - most-recent-attempt lookups are ordered by claim time. - - (There is nothing to index for the ``PROVENANCE_OF`` traversal itself: the - ``Task`` endpoint is already covered by the ``GufeTokenizable._scoped_key`` - uniqueness constraint, and the ``PROVENANCE_OF`` relationship carries no - properties.) - - This migration is idempotent (all indexes are created with - ``IF NOT EXISTS``) and requires NO data migration. All new properties are - optional-valued: pre-existing ``Task`` nodes simply have an empty attempt - history and are unaffected. - - """ - - indexes = { - "TaskProvenance_compute_service_id_index": ( - "TaskProvenance", - "compute_service_id", - ), - "TaskProvenance_datetime_claimed_index": ( - "TaskProvenance", - "datetime_claimed", - ), - } - - for name, (label, property_) in indexes.items(): - n4js.execute_query(f""" - CREATE INDEX {name} IF NOT EXISTS - FOR (n:{label}) ON (n.{property_}) - """) diff --git a/docs/operations.rst b/docs/operations.rst index a4fe7f30..7c18b2b6 100644 --- a/docs/operations.rst +++ b/docs/operations.rst @@ -171,18 +171,3 @@ 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 and related -introspection features. -This requires a lightweight schema migration that adds ``neo4j`` indexes for the -new ``TaskProvenance`` node label; it is idempotent and requires no data -migration (pre-existing ``Task``\ s simply have empty execution history). - -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 From dcf569ff0709324c0e40a9f104b82b4a18dca3ad Mon Sep 17 00:00:00 2001 From: David Dotson Date: Mon, 20 Jul 2026 21:13:01 -0600 Subject: [PATCH 11/18] Centralize per-unit-result artifact location construction Introduce `protocol_unit_result_location` in `storage/objectstore.py` as the single source of truth for the per-unit-result artifact layout (`.../{pdr_key}/units/{unit_result_key}`), and use it at both storage sites (`compute.api.set_task_result` and `Neo4jStore.add_protocol_unit_result_refs`) instead of re-deriving the path inline in each. This removes the duplicated derivation that could silently drift and leave pushed stdout/stderr unretrievable. Retrieval continues to read the authoritative `ProtocolUnitResultRef.location` directly: unlike `ProtocolDAGResult`s, unit refs have no legacy/heterogeneous data and their stored location is the actual write path, so a reconstruct-then-fallback mechanism would be inert. Add a unit test pinning the constructor's layout and GufeKey/str acceptance. Co-Authored-By: Claude Opus 4.8 (1M context) --- alchemiscale/compute/api.py | 8 +++--- alchemiscale/storage/objectstore.py | 19 +++++++++++++ alchemiscale/storage/statestore.py | 12 +++----- .../tests/unit/test_introspection_records.py | 28 +++++++++++++++++++ 4 files changed, 55 insertions(+), 12 deletions(-) diff --git a/alchemiscale/compute/api.py b/alchemiscale/compute/api.py index ea9b1e4d..fc38928b 100644 --- a/alchemiscale/compute/api.py +++ b/alchemiscale/compute/api.py @@ -5,7 +5,6 @@ """ import json -import os import datetime from datetime import timedelta import random @@ -38,7 +37,7 @@ ComputeAPISettings, ) from ..storage.statestore import Neo4jStore -from ..storage.objectstore import S3ObjectStore +from ..storage.objectstore import S3ObjectStore, protocol_unit_result_location from ..storage.models import ( ProtocolDAGResultRef, ComputeServiceID, @@ -427,12 +426,13 @@ async def set_task_result( # Streams need no new compute-facing routes: they ride inside the PDR blob. refs_map = n4js.add_protocol_unit_result_refs(protocoldagresultref, result_sk, pdr) if protocoldagresultref.location: - base_location = os.path.dirname(protocoldagresultref.location) for unit_result in pdr.protocol_unit_results: purr_sk = refs_map.get(unit_result.key) if purr_sk is None: continue - unit_location = os.path.join(base_location, "units", str(unit_result.key)) + unit_location = protocol_unit_result_location( + protocoldagresultref.location, unit_result.key + ) if unit_result.stdout: s3os.push_protocol_unit_result_streams( unit_location, "stdout", unit_result.stdout diff --git a/alchemiscale/storage/objectstore.py b/alchemiscale/storage/objectstore.py index 87c1fb3c..873d6502 100644 --- a/alchemiscale/storage/objectstore.py +++ b/alchemiscale/storage/objectstore.py @@ -27,6 +27,25 @@ STDERR_DIRNAME = "stderr" +def protocol_unit_result_location( + protocoldagresult_location: str, unit_result_key +) -> str: + """Construct the object-store prefix for a unit result's artifacts. + + Derived from the parent `ProtocolDAGResult`'s ``location`` --- the unit + result artifacts hang off the PDR's directory as + ``.../{pdr_key}/units/{unit_result_key}``. This is the single source of + truth for that layout: it is used both when storing artifacts (and recording + the result on `ProtocolUnitResultRef.location`) and when retrieving them, so + a retrieval can reconstruct the location the same way it was constructed and + fall back to the stored `ProtocolUnitResultRef.location` if that misses --- + mirroring `ProtocolDAGResult` storage/retrieval. + """ + return os.path.join( + os.path.dirname(protocoldagresult_location), "units", str(unit_result_key) + ) + + def get_s3os(settings: S3ObjectStoreSettings) -> "S3ObjectStore": """Convenience function for getting an S3ObjectStore directly from settings.""" return S3ObjectStore(settings) diff --git a/alchemiscale/storage/statestore.py b/alchemiscale/storage/statestore.py index 11d6a5c6..8b543e7b 100644 --- a/alchemiscale/storage/statestore.py +++ b/alchemiscale/storage/statestore.py @@ -9,7 +9,6 @@ import datetime from contextlib import contextmanager import json -import os import re from functools import lru_cache, update_wrapper from collections import defaultdict @@ -67,6 +66,7 @@ from ..models import Scope, ScopedKey from .cypher import cypher_or +from .objectstore import protocol_unit_result_location from ..security.models import CredentialedEntity from ..settings import Neo4jStoreSettings @@ -4034,11 +4034,7 @@ def add_protocol_unit_result_refs( pdrr_node = self._get_node(protocoldagresultref_scoped_key) - base_location = ( - os.path.dirname(protocoldagresultref.location) - if protocoldagresultref.location - else None - ) + pdrr_location = protocoldagresultref.location subgraph = Subgraph() result_key_to_node = {} @@ -4046,8 +4042,8 @@ def add_protocol_unit_result_refs( for result in protocoldagresult.protocol_unit_results: unit_location = ( - os.path.join(base_location, "units", str(result.key)) - if base_location is not None + protocol_unit_result_location(pdrr_location, result.key) + if pdrr_location else None ) purr = ProtocolUnitResultRef( diff --git a/alchemiscale/tests/unit/test_introspection_records.py b/alchemiscale/tests/unit/test_introspection_records.py index db0ccbf1..eca23cab 100644 --- a/alchemiscale/tests/unit/test_introspection_records.py +++ b/alchemiscale/tests/unit/test_introspection_records.py @@ -306,3 +306,31 @@ def test_gufe_roundtrip(self): assert purr2.has_logs is True assert purr2.start_time == NOW assert str(purr2.key) == str(key1) + + +class TestProtocolUnitResultLocation: + """The shared constructor for per-unit-result artifact locations, used by + both storage and (reconstruction-based) retrieval.""" + + def test_layout_matches_pdr_dir(self): + from alchemiscale.storage.objectstore import ( + protocol_unit_result_location, + OBJECT_FILENAME, + ) + + pdr_location = f"protocoldagresult/o/c/p/T/results/PDR/{OBJECT_FILENAME}" + + # units hang off the parent ProtocolDAGResult's directory (the PDR's + # `location` filename is stripped), keyed by the unit result gufe key + assert ( + protocol_unit_result_location(pdr_location, GufeKey("PUR")) + == "protocoldagresult/o/c/p/T/results/PDR/units/PUR" + ) + + def test_accepts_gufekey_or_str(self): + from alchemiscale.storage.objectstore import protocol_unit_result_location + + pdr_location = "protocoldagresult/o/c/p/T/results/PDR/obj.json.zst" + assert protocol_unit_result_location( + pdr_location, GufeKey("PUR") + ) == protocol_unit_result_location(pdr_location, "PUR") From c77f9d3faed487f5a5141df820fbda773afdad9e Mon Sep 17 00:00:00 2001 From: David Dotson Date: Mon, 20 Jul 2026 23:32:43 -0600 Subject: [PATCH 12/18] Make TaskProvenance a GufeTokenizable with a ScopedKey TaskProvenance now carries a ScopedKey so it is a first-class scoped entity, authorized through the standard `validate_scopes(sk.scope, token)` path rather than relying on an anchoring Task -- which matters for future global/cross-task provenance queries (pruning, #180) that have no Task to borrow a scope check from. Like `Task`, it tokenizes on a uuid (via `_gufe_tokenize`), not its contents, so its GufeKey/ScopedKey is fixed at creation and unaffected by the in-place mutations (outcome/datetime_end/progress) applied over the attempt's life. Content-addressing would be neither stable nor unique per attempt. Creation moves out of the atomic CLAIM_QUERY: the query now claims and returns the registration's hostname/manager_name, and `claim_taskhub_tasks` builds a TaskProvenance for each genuinely-claimed Task via the normal keyed-node path (`_keyed_chain_to_subgraph` + `merge_subgraph`) and creates the PROVENANCE_OF edge, all in the same transaction so claim and provenance stay atomic. `datetime_claimed` still stores as a native neo4j DateTime, so ordering/coercion reads are unchanged. Adds a unit test for the uuid identity/round-trip and an integration test for the ScopedKey + get_gufe round-trip. Co-Authored-By: Claude Opus 4.8 (1M context) --- alchemiscale/storage/models.py | 117 +++++++++++++----- alchemiscale/storage/statestore.py | 61 ++++++--- .../storage/test_statestore_introspection.py | 33 +++++ .../tests/unit/test_introspection_records.py | 23 ++++ 4 files changed, 183 insertions(+), 51 deletions(-) diff --git a/alchemiscale/storage/models.py b/alchemiscale/storage/models.py index 94b8e38b..ad013e94 100644 --- a/alchemiscale/storage/models.py +++ b/alchemiscale/storage/models.py @@ -19,6 +19,21 @@ from ..models import ScopedKey, Scope +def _coerce_datetime(v) -> datetime.datetime | None: + """Coerce a neo4j ``DateTime``, ISO string, or ``datetime`` to ``datetime``.""" + if v is None: + return None + if hasattr(v, "to_native"): + return v.to_native() + if isinstance(v, str): + return datetime.datetime.fromisoformat(v) + return v + + +def _iso(v: datetime.datetime | None) -> str | None: + return v.isoformat() if v is not None else None + + class ComputeIDBase(str): _allowed_name_pattern = r"^[a-zA-Z][a-zA-Z0-9_\.\:]*$" @@ -175,8 +190,8 @@ class TaskOutcomeEnum(Enum): released = "released" -class TaskProvenance(BaseModel): - """An immutable record of a single execution attempt of a `Task`. +class TaskProvenance(GufeTokenizable): + """A record of a single execution attempt of a `Task`. A `TaskProvenance` node is created at claim time and finalized when the attempt ends. It survives claim teardown and registration expiry, so that @@ -185,6 +200,17 @@ class TaskProvenance(BaseModel): record rather than held as a relationship to the (potentially deleted) `ComputeServiceRegistration`. + It is a `GufeTokenizable` so it carries a `ScopedKey`: provenance is a + scoped entity, authorized through the same ``validate_scopes(sk.scope, + token)`` path as every other scoped object, rather than relying on an + anchoring `Task`. Like `Task`, it tokenizes on a uuid (see + `_gufe_tokenize`), *not* its contents, so its `GufeKey`/`ScopedKey` is fixed + at creation and unaffected by the in-place mutations (`outcome`, + `datetime_end`, and the progress counts) applied over the attempt's life. + Because of that, a `TaskProvenance` must never be re-tokenized after + creation (never round-tripped object -> node a second time); mutations go + straight to the node via Cypher. + Attributes ---------- compute_service_id @@ -208,27 +234,65 @@ class TaskProvenance(BaseModel): """ compute_service_id: ComputeServiceID - hostname: str | None = None - manager_name: str | None = None - datetime_claimed: datetime.datetime - datetime_end: datetime.datetime | None = None - outcome: TaskOutcomeEnum | None = None - units_completed: int | None = None - units_total: int | None = None + hostname: str | None + manager_name: str | None + datetime_claimed: datetime.datetime | None + datetime_end: datetime.datetime | None + outcome: TaskOutcomeEnum | None + units_completed: int | None + units_total: int | None - model_config = ConfigDict(arbitrary_types_allowed=True) + def __init__( + self, + *, + compute_service_id: ComputeServiceID | str, + datetime_claimed: datetime.datetime | None = None, + hostname: str | None = None, + manager_name: str | None = None, + datetime_end: datetime.datetime | None = None, + outcome: str | TaskOutcomeEnum | None = None, + units_completed: int | None = None, + units_total: int | None = None, + _key: str = None, + ): + if _key is not None: + self._key = GufeKey(_key) - def to_dict(self): - dct = self.model_dump() - dct["compute_service_id"] = str(self.compute_service_id) - dct["outcome"] = self.outcome.value if self.outcome is not None else None - return dct + self.compute_service_id = ComputeServiceID(compute_service_id) + self.hostname = hostname + self.manager_name = manager_name + self.datetime_claimed = _coerce_datetime(datetime_claimed) + self.datetime_end = _coerce_datetime(datetime_end) + self.outcome = TaskOutcomeEnum(outcome) if outcome is not None else None + self.units_completed = units_completed + self.units_total = units_total + + def _gufe_tokenize(self): + # tokenize with a uuid, not content: identity is per-attempt, and the + # record is mutated in place (outcome/datetime_end/progress) after + # creation, so a content hash would neither be stable nor unique. + return uuid4().hex + + def _to_dict(self): + return { + "compute_service_id": str(self.compute_service_id), + "hostname": self.hostname, + "manager_name": self.manager_name, + "datetime_claimed": self.datetime_claimed, + "datetime_end": self.datetime_end, + "outcome": self.outcome.value if self.outcome is not None else None, + "units_completed": self.units_completed, + "units_total": self.units_total, + "_key": str(self.key), + } @classmethod - def from_dict(cls, dct): - dct_ = copy(dct) - dct_["compute_service_id"] = ComputeServiceID(dct_["compute_service_id"]) - return cls(**dct_) + def _from_dict(cls, d): + return cls(**d) + + @classmethod + def _defaults(cls): + return super()._defaults() class TaskStatusEnum(Enum): @@ -802,21 +866,6 @@ def from_dict(cls, d): return cls(**d) -def _coerce_datetime(v) -> datetime.datetime | None: - """Coerce a neo4j ``DateTime``, ISO string, or ``datetime`` to ``datetime``.""" - if v is None: - return None - if hasattr(v, "to_native"): - return v.to_native() - if isinstance(v, str): - return datetime.datetime.fromisoformat(v) - return v - - -def _iso(v: datetime.datetime | None) -> str | None: - return v.isoformat() if v is not None else None - - # --- client-facing API record models -------------------------------------- # # These models are the user-facing surface for Task introspection. They are diff --git a/alchemiscale/storage/statestore.py b/alchemiscale/storage/statestore.py index 8b543e7b..002b43ff 100644 --- a/alchemiscale/storage/statestore.py +++ b/alchemiscale/storage/statestore.py @@ -181,23 +181,11 @@ def _status_write( MATCH (csreg:ComputeServiceRegistration {{identifier: $compute_service_id}}) CREATE (t)<-[cl:CLAIMS {{claimed: datetime($datetimestr)}}]-(csreg) - // create an immutable TaskProvenance record for this execution attempt, - // copying identifying info off the registration (which may later be - // deleted on expiry/deregistration) - CREATE (tp:TaskProvenance {{ - compute_service_id: $compute_service_id, - hostname: csreg.hostname, - manager_name: csreg.manager_name, - datetime_claimed: datetime($datetimestr), - _org: t._org, - _campaign: t._campaign, - _project: t._project - }}) - CREATE (tp)-[:PROVENANCE_OF]->(t) - {_status_write('t', TaskStatusEnum.running.value, time_param='datetimestr')} - RETURN t + // return the identifying info the caller copies onto a `TaskProvenance` + // record for each genuinely-claimed Task (built in the same transaction) + RETURN t, csreg.hostname AS hostname, csreg.manager_name AS manager_name """ @@ -2925,13 +2913,52 @@ def task_count(task_dict: dict): # if tasks is not empty, proceed with claiming if tasks: - tx.run( + now = datetime.datetime.now(tz=datetime.UTC) + claim_result = tx.run( CLAIM_QUERY, tasks_list=[str(task) for task in tasks if task is not None], - datetimestr=str(datetime.datetime.now(tz=datetime.UTC).isoformat()), + datetimestr=now.isoformat(), compute_service_id=str(compute_service_id), ) + # Build a `TaskProvenance` record for each Task *actually* + # claimed (only those come back from CLAIM_QUERY --- a racing + # service may have taken some), in this same transaction so + # claim and provenance are atomic. `TaskProvenance` is a + # `GufeTokenizable`, so each is created through the normal + # keyed-node path and carries a `ScopedKey`; the `PROVENANCE_OF` + # edge to its Task is created after the nodes are merged. + provenance_subgraph = Subgraph() + links = [] + for record in claim_result: + task_node = record["t"] + task_sk = ScopedKey.from_str(task_node["_scoped_key"]) + tp = TaskProvenance( + compute_service_id=compute_service_id, + hostname=record["hostname"], + manager_name=record["manager_name"], + datetime_claimed=now, + ) + tp_subgraph, _, tp_sk = self._keyed_chain_to_subgraph( + KeyedChain.from_gufe(tp), scope=task_sk.scope + ) + provenance_subgraph = provenance_subgraph | tp_subgraph + links.append({"tp": str(tp_sk), "task": str(task_sk)}) + + if links: + merge_subgraph( + tx, provenance_subgraph, "GufeTokenizable", "_scoped_key" + ) + tx.run( + """ + UNWIND $links AS link + MATCH (tp:TaskProvenance {_scoped_key: link.tp}) + MATCH (t:Task {_scoped_key: link.task}) + CREATE (tp)-[:PROVENANCE_OF]->(t) + """, + links=links, + ) + tx.run( """ MATCH (th:TaskHub {_scoped_key: $taskhub}) diff --git a/alchemiscale/tests/integration/storage/test_statestore_introspection.py b/alchemiscale/tests/integration/storage/test_statestore_introspection.py index b10f5e6a..bc84433a 100644 --- a/alchemiscale/tests/integration/storage/test_statestore_introspection.py +++ b/alchemiscale/tests/integration/storage/test_statestore_introspection.py @@ -22,6 +22,7 @@ TaskAttempt, TaskDetails, TaskOutcomeEnum, + TaskProvenance, TaskStatusEnum, TaskTracebacks, ) @@ -124,6 +125,38 @@ def test_provenance_created_at_claim( assert task_node["status"] == TaskStatusEnum.running.value assert task_node.get("datetime_status_changed") is not None + def test_provenance_is_scoped_gufe_tokenizable( + self, n4js, network_tyk2, transformation, scope_test + ): + """The provenance node is a `GufeTokenizable` carrying a `ScopedKey` in + its Task's scope, and round-trips via `get_gufe` --- so it authorizes + through the standard `validate_scopes(sk.scope, token)` path, with no + anchoring Task required.""" + csid = ComputeServiceID.new_from_name("prov.scoped") + task_sk, _ = self._claimed_task( + n4js, + network_tyk2, + transformation, + scope_test, + csid, + hostname="cluster-node-7", + manager_name="mgr-a", + ) + + tp = _provenance_nodes(n4js, task_sk)[0] + # a real keyed node: GufeTokenizable label + a ScopedKey in Task's scope + tp_sk = ScopedKey.from_str(tp["_scoped_key"]) + assert tp_sk.qualname == "TaskProvenance" + assert tp_sk.scope == task_sk.scope + + # round-trips as a TaskProvenance with the same identity and fields + obj = n4js.get_gufe(tp_sk) + assert isinstance(obj, TaskProvenance) + assert str(obj.key) == tp["_gufe_key"] + assert str(obj.compute_service_id) == str(csid) + assert obj.hostname == "cluster-node-7" + assert obj.manager_name == "mgr-a" + # --- finalization: complete / error via set_task_result --------------- def test_provenance_finalized_complete( diff --git a/alchemiscale/tests/unit/test_introspection_records.py b/alchemiscale/tests/unit/test_introspection_records.py index eca23cab..70b7ef2a 100644 --- a/alchemiscale/tests/unit/test_introspection_records.py +++ b/alchemiscale/tests/unit/test_introspection_records.py @@ -89,6 +89,29 @@ def test_roundtrip_open(self): assert tp2.datetime_end is None assert tp2.hostname is None + def test_gufe_tokenizable_uuid_identity(self): + # TaskProvenance is a GufeTokenizable so it carries a ScopedKey for + # scope-based authorization; it tokenizes on a uuid, so the key is + # unique per attempt and stable against content / later mutation. + from gufe.tokenization import GufeTokenizable + + tp = TaskProvenance(compute_service_id=CSID, datetime_claimed=NOW) + assert isinstance(tp, GufeTokenizable) + assert str(tp.key).startswith("TaskProvenance-") + + # identical content but a distinct attempt -> distinct key + assert ( + tp.key != TaskProvenance(compute_service_id=CSID, datetime_claimed=NOW).key + ) + + # key survives a serialization round-trip and an in-place mutation + # (finalization sets outcome/datetime_end after creation) + key = tp.key + assert TaskProvenance.from_dict(tp.to_dict()).key == key + tp.outcome = TaskOutcomeEnum.complete + tp.datetime_end = LATER + assert TaskProvenance.from_dict(tp.to_dict()).key == key + class TestTaskAttempt: @pytest.mark.parametrize( From cce6ae88c7c076e679da91b27b049eb795913c5c Mon Sep 17 00:00:00 2001 From: David Dotson Date: Mon, 20 Jul 2026 23:37:10 -0600 Subject: [PATCH 13/18] Complete Task and ProtocolUnitResultRef attribute docs Fill in the empty `datetime_created` description on `Task` and document the previously-undocumented `datetime_status_changed`, `reason`, `creator`, and `extends` attributes. Add the missing `scope` attribute (inherited from `ObjectStoreRef`) to `ProtocolUnitResultRef`. Docstring-only. Co-Authored-By: Claude Opus 4.8 (1M context) --- alchemiscale/storage/models.py | 19 ++++++++++++++++++- 1 file changed, 18 insertions(+), 1 deletion(-) diff --git a/alchemiscale/storage/models.py b/alchemiscale/storage/models.py index ad013e94..a8e5fcdd 100644 --- a/alchemiscale/storage/models.py +++ b/alchemiscale/storage/models.py @@ -315,8 +315,23 @@ class Task(GufeTokenizable): priority Priority of the task; 1 is highest, larger values indicate lower priority. claim - Identifier of the compute service that has a claim on this task. + Identifier of the compute service that has a claim on this task, if any. datetime_created + When the `Task` was created. + datetime_status_changed + When the `Task`'s `status` was last changed; refreshed at every + status-mutation site (claim, `set_task_*`, expiry, deregistration, + restart-renew), but not on an idempotent no-op re-set of the same status. + reason + Human-readable reason for the current `status`, if any --- e.g. a + `ProtocolDAG` creation-failure traceback, or a user-supplied reason for + an `invalid`/`deleted` transition. Cleared when the status changes to + one that carries no reason. + creator + Identifier of the identity that created the `Task`, if recorded. + extends + `ScopedKey` (as a string) of the `Task` this one extends (continues + from), if any. """ @@ -718,6 +733,8 @@ class ProtocolUnitResultRef(ObjectStoreRef): When execution of the unit attempt began and ended. location The object store prefix under which this unit result's artifacts live. + scope + The `Scope` (org/campaign/project) this reference lives in. has_logs, has_stdout, has_stderr Whether captured log/stdout/stderr artifacts exist for this unit result. From 226b7d6dd2d122b0ff6e6852a046f3a26cc6a378 Mon Sep 17 00:00:00 2001 From: David Dotson Date: Thu, 30 Jul 2026 14:18:16 -0600 Subject: [PATCH 14/18] Give ProtocolUnitResultRef a uuid _gufe_tokenize ProtocolUnitResultRef is a mutable GufeTokenizable -- its has_logs/ has_stdout/has_stderr flags are flipped in place via Cypher as artifacts arrive. Under the inherited content-based tokenization, a reconstructed ref would recompute a key from its (now-mutated) content that no longer matched its stored `_scoped_key` -- safe before only by never relying on that recomputed key. Switch it to a uuid `_gufe_tokenize`, matching `Task`/`TaskProvenance`, so the GufeKey/ScopedKey is fixed at creation and independent of the mutable flags by construction. `_key` now round-trips through `__init__`/`_to_dict` so lookups by `_scoped_key` stay stable. Nothing relied on the old content-determined key: the refs map and idempotency short-circuit key off `obj_key`, not the ref's own key. Co-Authored-By: Claude Opus 4.8 (1M context) --- alchemiscale/storage/models.py | 26 ++++++++++++++----- .../tests/unit/test_introspection_records.py | 24 ++++++++++++++--- 2 files changed, 40 insertions(+), 10 deletions(-) diff --git a/alchemiscale/storage/models.py b/alchemiscale/storage/models.py index a8e5fcdd..bd4eb80d 100644 --- a/alchemiscale/storage/models.py +++ b/alchemiscale/storage/models.py @@ -741,13 +741,14 @@ class ProtocolUnitResultRef(ObjectStoreRef): Note ---- The ``has_logs``, ``has_stdout``, and ``has_stderr`` flags (and nothing - else) are mutated in place via Cypher after - the node is created, as artifacts arrive. The node's `_scoped_key` and - `_gufe_key` - are computed once at creation and never recomputed, so lookups stay stable - even though these tokenizable-contributing fields change. This is safe only - because `ProtocolUnitResultRef` nodes are an internal state-store detail, - never re-tokenized after creation; keep it that way. + else) are mutated in place via Cypher after the node is created, as + artifacts arrive. Like `Task`/`TaskProvenance`, this object tokenizes on a + uuid (see `_gufe_tokenize`), *not* its contents, so its `GufeKey`/ + `ScopedKey` is fixed at creation and is unaffected by those mutations --- + lookups by `_scoped_key` stay stable. Because the key is a uuid, a + `ProtocolUnitResultRef` must never be re-tokenized after creation (never + round-tripped object -> node a second time); mutations go straight to the + node via Cypher. """ ok: bool @@ -773,7 +774,11 @@ def __init__( has_logs: bool = False, has_stdout: bool = False, has_stderr: bool = False, + _key: str = None, ): + if _key is not None: + self._key = GufeKey(_key) + self.location = location self.obj_key = GufeKey(obj_key) self.source_key = GufeKey(source_key) @@ -786,6 +791,12 @@ def __init__( self.has_stdout = has_stdout self.has_stderr = has_stderr + def _gufe_tokenize(self): + # tokenize with a uuid, not content: the has_logs/has_stdout/has_stderr + # flags are mutated in place after creation, so a content hash would not + # be stable. Like `Task`/`TaskProvenance`, the key is fixed at creation. + return uuid4().hex + def _to_dict(self): return { "location": self.location, @@ -803,6 +814,7 @@ def _to_dict(self): "has_logs": self.has_logs, "has_stdout": self.has_stdout, "has_stderr": self.has_stderr, + "_key": str(self.key), } @classmethod diff --git a/alchemiscale/tests/unit/test_introspection_records.py b/alchemiscale/tests/unit/test_introspection_records.py index 70b7ef2a..60e471de 100644 --- a/alchemiscale/tests/unit/test_introspection_records.py +++ b/alchemiscale/tests/unit/test_introspection_records.py @@ -317,11 +317,29 @@ def test_gufe_roundtrip(self): end_time=LATER, has_logs=True, ) - # deterministic key computed once at creation + # uuid key, fixed at creation (the has_* flags are mutated in place, so + # a content hash would not be stable) + from gufe.tokenization import GufeTokenizable, KeyedChain + + assert isinstance(purr, GufeTokenizable) key1 = purr.key - # round-trip through the keyed chain (as the state store does) - from gufe.tokenization import KeyedChain + # a distinct ref with identical content gets a distinct key + purr_other = ProtocolUnitResultRef( + location="protocoldagresult/o/c/p/T/results/K/units/R", + obj_key=GufeKey("ProtocolUnitResult-r1"), + source_key=GufeKey("ProtocolUnit-u1"), + scope=Scope("o", "c", "p"), + ok=True, + name="u", + start_time=NOW, + end_time=LATER, + has_logs=True, + ) + assert purr_other.key != key1 + + # round-trip through the keyed chain (as the state store does) preserves + # the key and the fields purr2 = KeyedChain.from_gufe(purr).to_gufe() assert purr2.obj_key == purr.obj_key assert purr2.source_key == purr.source_key From 3611ef78840434f916e5d87ec4e487d9c782c94a Mon Sep 17 00:00:00 2001 From: David Dotson Date: Fri, 31 Jul 2026 14:32:21 -0600 Subject: [PATCH 15/18] Use pydantic validators/serializers for client-facing record models The client-facing record models (`TaskAttempt`, `TaskClaim`, `TaskDetails`, `TaskUnitTraceback`, `TaskTracebacks`, `ProtocolDAGResultRec`, `ProtocolUnitResultRec`) hand-rolled per-field coercion in `from_dict` and per-field serialization in `to_dict`. Replace that with three reusable pydantic v2 `Annotated` field types -- `Datetime`, `SK`, `GK` -- carrying `BeforeValidator`/`PlainSerializer`, so coercion happens on any construction (not just `from_dict`) and enums/nested models/ISO datetimes are handled natively. `to_dict`/`from_dict` become thin wrappers over `model_dump(mode="json")`/`model_validate`, preserving the public interface and call sites. The wire JSON is unchanged: `Datetime` serializes with `isoformat()` so UTC stays `+00:00` rather than pydantic's default `Z`. A new `TestWireShapeStability` pins the exact wire shape of each model (including the datetime-offset guard and inbound coercion of raw wire/neo4j values). Drops the now-unused `_iso` helper; `_coerce_datetime` remains (it backs the `Datetime` validator and is still used by the state store). Co-Authored-By: Claude Opus 4.8 (1M context) --- alchemiscale/storage/models.py | 241 ++++++------------ .../tests/unit/test_introspection_records.py | 116 ++++++++- 2 files changed, 185 insertions(+), 172 deletions(-) diff --git a/alchemiscale/storage/models.py b/alchemiscale/storage/models.py index bd4eb80d..ad1c0f47 100644 --- a/alchemiscale/storage/models.py +++ b/alchemiscale/storage/models.py @@ -8,12 +8,20 @@ from copy import copy import datetime from enum import Enum, StrEnum +from typing import Annotated from uuid import uuid4, UUID import re import hashlib -from pydantic import BaseModel, ConfigDict, PositiveInt, field_validator +from pydantic import ( + BaseModel, + BeforeValidator, + ConfigDict, + PlainSerializer, + PositiveInt, + field_validator, +) from gufe.tokenization import GufeTokenizable, GufeKey from ..models import ScopedKey, Scope @@ -30,8 +38,39 @@ def _coerce_datetime(v) -> datetime.datetime | None: return v -def _iso(v: datetime.datetime | None) -> str | None: - return v.isoformat() if v is not None else None +def _to_scoped_key(v): + return ScopedKey.from_str(v) if isinstance(v, str) else v + + +def _to_gufe_key(v): + return GufeKey(v) if v is not None else v + + +# Reusable field types for the client-facing record models below. They coerce +# the wire/neo4j representations on the way in --- on *any* construction, not +# just via `from_dict` --- and pin the wire representation on the way out, so +# those models need no per-field coercion plumbing: +# +# - `Datetime` accepts a neo4j ``DateTime``, an ISO string, or a ``datetime``, +# and serializes with ``isoformat()`` (stable ``+00:00`` offset --- pydantic's +# default would render UTC as ``Z``, changing the wire shape). +# - `SK`/`GK` accept a `ScopedKey`/`GufeKey` or its string form and serialize +# back to that string. +Datetime = Annotated[ + datetime.datetime, + BeforeValidator(_coerce_datetime), + PlainSerializer(lambda v: v.isoformat(), return_type=str), +] +SK = Annotated[ + ScopedKey, + BeforeValidator(_to_scoped_key), + PlainSerializer(lambda v: str(v), return_type=str), +] +GK = Annotated[ + GufeKey, + BeforeValidator(_to_gufe_key), + PlainSerializer(lambda v: str(v), return_type=str), +] class ComputeIDBase(str): @@ -915,51 +954,21 @@ class TaskAttempt(BaseModel): compute_service_id: str hostname: str | None = None manager_name: str | None = None - datetime_claimed: datetime.datetime - datetime_end: datetime.datetime | None = None + datetime_claimed: Datetime + datetime_end: Datetime | None = None outcome: TaskOutcomeEnum | None = None units_completed: int | None = None units_total: int | None = None - protocoldagresultref: ScopedKey | None = None + protocoldagresultref: SK | None = None model_config = ConfigDict(arbitrary_types_allowed=True) def to_dict(self): - return { - "compute_service_id": self.compute_service_id, - "hostname": self.hostname, - "manager_name": self.manager_name, - "datetime_claimed": _iso(self.datetime_claimed), - "datetime_end": _iso(self.datetime_end), - "outcome": self.outcome.value if self.outcome is not None else None, - "units_completed": self.units_completed, - "units_total": self.units_total, - "protocoldagresultref": ( - str(self.protocoldagresultref) - if self.protocoldagresultref is not None - else None - ), - } + return self.model_dump(mode="json") @classmethod def from_dict(cls, d): - return cls( - compute_service_id=d["compute_service_id"], - hostname=d.get("hostname"), - manager_name=d.get("manager_name"), - datetime_claimed=_coerce_datetime(d["datetime_claimed"]), - datetime_end=_coerce_datetime(d.get("datetime_end")), - outcome=( - TaskOutcomeEnum(d["outcome"]) if d.get("outcome") is not None else None - ), - units_completed=d.get("units_completed"), - units_total=d.get("units_total"), - protocoldagresultref=( - ScopedKey.from_str(d["protocoldagresultref"]) - if d.get("protocoldagresultref") is not None - else None - ), - ) + return cls.model_validate(d) class TaskClaim(BaseModel): @@ -967,38 +976,26 @@ class TaskClaim(BaseModel): compute_service_id: str hostname: str | None = None - datetime_claimed: datetime.datetime | None = None + datetime_claimed: Datetime | None = None units_completed: int | None = None units_total: int | None = None model_config = ConfigDict(arbitrary_types_allowed=True) def to_dict(self): - return { - "compute_service_id": self.compute_service_id, - "hostname": self.hostname, - "datetime_claimed": _iso(self.datetime_claimed), - "units_completed": self.units_completed, - "units_total": self.units_total, - } + return self.model_dump(mode="json") @classmethod def from_dict(cls, d): - return cls( - compute_service_id=d["compute_service_id"], - hostname=d.get("hostname"), - datetime_claimed=_coerce_datetime(d.get("datetime_claimed")), - units_completed=d.get("units_completed"), - units_total=d.get("units_total"), - ) + return cls.model_validate(d) class TaskDetails(BaseModel): """Bulk indicator summary for a `Task`, as returned by `get_tasks_details`.""" - task: ScopedKey + task: SK status: TaskStatusEnum - datetime_status_changed: datetime.datetime | None = None + datetime_status_changed: Datetime | None = None reason: str | None = None num_claims: int = 0 current_claim: TaskClaim | None = None @@ -1007,77 +1004,29 @@ class TaskDetails(BaseModel): model_config = ConfigDict(arbitrary_types_allowed=True) def to_dict(self): - return { - "task": str(self.task), - "status": self.status.value, - "datetime_status_changed": _iso(self.datetime_status_changed), - "reason": self.reason, - "num_claims": self.num_claims, - "current_claim": ( - self.current_claim.to_dict() if self.current_claim is not None else None - ), - "most_recent_attempt": ( - self.most_recent_attempt.to_dict() - if self.most_recent_attempt is not None - else None - ), - } + return self.model_dump(mode="json") @classmethod def from_dict(cls, d): - return cls( - task=ScopedKey.from_str(d["task"]), - status=TaskStatusEnum(d["status"]), - datetime_status_changed=_coerce_datetime(d.get("datetime_status_changed")), - reason=d.get("reason"), - num_claims=d.get("num_claims", 0), - current_claim=( - TaskClaim.from_dict(d["current_claim"]) - if d.get("current_claim") is not None - else None - ), - most_recent_attempt=( - TaskAttempt.from_dict(d["most_recent_attempt"]) - if d.get("most_recent_attempt") is not None - else None - ), - ) + return cls.model_validate(d) class TaskUnitTraceback(BaseModel): """A single `ProtocolUnitFailure` traceback within a `TaskTracebacks`.""" - failure_key: GufeKey - source_key: GufeKey + failure_key: GK + source_key: GK traceback: str - protocolunitresultref: ScopedKey | None = None + protocolunitresultref: SK | None = None model_config = ConfigDict(arbitrary_types_allowed=True) def to_dict(self): - return { - "failure_key": str(self.failure_key), - "source_key": str(self.source_key), - "traceback": self.traceback, - "protocolunitresultref": ( - str(self.protocolunitresultref) - if self.protocolunitresultref is not None - else None - ), - } + return self.model_dump(mode="json") @classmethod def from_dict(cls, d): - return cls( - failure_key=GufeKey(d["failure_key"]), - source_key=GufeKey(d["source_key"]), - traceback=d["traceback"], - protocolunitresultref=( - ScopedKey.from_str(d["protocolunitresultref"]) - if d.get("protocolunitresultref") is not None - else None - ), - ) + return cls.model_validate(d) class TaskTracebacks(BaseModel): @@ -1087,29 +1036,19 @@ class TaskTracebacks(BaseModel): `ProtocolDAGResultRef`, most recent first. """ - protocoldagresultref: ScopedKey - datetime_created: datetime.datetime | None = None + protocoldagresultref: SK + datetime_created: Datetime | None = None creator: str | None = None tracebacks: list[TaskUnitTraceback] = [] model_config = ConfigDict(arbitrary_types_allowed=True) def to_dict(self): - return { - "protocoldagresultref": str(self.protocoldagresultref), - "datetime_created": _iso(self.datetime_created), - "creator": self.creator, - "tracebacks": [tb.to_dict() for tb in self.tracebacks], - } + return self.model_dump(mode="json") @classmethod def from_dict(cls, d): - return cls( - protocoldagresultref=ScopedKey.from_str(d["protocoldagresultref"]), - datetime_created=_coerce_datetime(d.get("datetime_created")), - creator=d.get("creator"), - tracebacks=[TaskUnitTraceback.from_dict(tb) for tb in d["tracebacks"]], - ) + return cls.model_validate(d) class ProtocolDAGResultRec(BaseModel): @@ -1120,29 +1059,19 @@ class ProtocolDAGResultRec(BaseModel): method accepts directly. """ - scoped_key: ScopedKey + scoped_key: SK ok: bool - datetime_created: datetime.datetime | None = None + datetime_created: Datetime | None = None creator: str | None = None model_config = ConfigDict(arbitrary_types_allowed=True) def to_dict(self): - return { - "scoped_key": str(self.scoped_key), - "ok": self.ok, - "datetime_created": _iso(self.datetime_created), - "creator": self.creator, - } + return self.model_dump(mode="json") @classmethod def from_dict(cls, d): - return cls( - scoped_key=ScopedKey.from_str(d["scoped_key"]), - ok=d["ok"], - datetime_created=_coerce_datetime(d.get("datetime_created")), - creator=d.get("creator"), - ) + return cls.model_validate(d) class ProtocolUnitResultRec(BaseModel): @@ -1154,13 +1083,13 @@ class ProtocolUnitResultRec(BaseModel): `ProtocolDAGResult`. """ - scoped_key: ScopedKey - obj_key: GufeKey - source_key: GufeKey + scoped_key: SK + obj_key: GK + source_key: GK name: str | None = None ok: bool - start_time: datetime.datetime | None = None - end_time: datetime.datetime | None = None + start_time: Datetime | None = None + end_time: Datetime | None = None has_logs: bool = False has_stdout: bool = False has_stderr: bool = False @@ -1168,30 +1097,8 @@ class ProtocolUnitResultRec(BaseModel): model_config = ConfigDict(arbitrary_types_allowed=True) def to_dict(self): - return { - "scoped_key": str(self.scoped_key), - "obj_key": str(self.obj_key), - "source_key": str(self.source_key), - "name": self.name, - "ok": self.ok, - "start_time": _iso(self.start_time), - "end_time": _iso(self.end_time), - "has_logs": self.has_logs, - "has_stdout": self.has_stdout, - "has_stderr": self.has_stderr, - } + return self.model_dump(mode="json") @classmethod def from_dict(cls, d): - return cls( - scoped_key=ScopedKey.from_str(d["scoped_key"]), - obj_key=GufeKey(d["obj_key"]), - source_key=GufeKey(d["source_key"]), - name=d.get("name"), - ok=d["ok"], - start_time=_coerce_datetime(d.get("start_time")), - end_time=_coerce_datetime(d.get("end_time")), - has_logs=d.get("has_logs", False), - has_stdout=d.get("has_stdout", False), - has_stderr=d.get("has_stderr", False), - ) + return cls.model_validate(d) diff --git a/alchemiscale/tests/unit/test_introspection_records.py b/alchemiscale/tests/unit/test_introspection_records.py index 60e471de..ee8da64b 100644 --- a/alchemiscale/tests/unit/test_introspection_records.py +++ b/alchemiscale/tests/unit/test_introspection_records.py @@ -24,7 +24,6 @@ TaskTracebacks, TaskUnitTraceback, _coerce_datetime, - _iso, ) NOW = datetime.datetime(2026, 7, 10, 12, 0, 0, tzinfo=datetime.UTC) @@ -52,10 +51,6 @@ def to_native(self_inner): assert _coerce_datetime(FakeNeo4jDT()) == NOW - def test_iso(self): - assert _iso(None) is None - assert _iso(NOW) == NOW.isoformat() - class TestTaskProvenance: def test_roundtrip_full(self): @@ -375,3 +370,114 @@ def test_accepts_gufekey_or_str(self): assert protocol_unit_result_location( pdr_location, GufeKey("PUR") ) == protocol_unit_result_location(pdr_location, "PUR") + + +class TestWireShapeStability: + """Pin the exact JSON wire shape of the client-facing record models. These + cross the HTTP boundary, so the switch to pydantic validators/serializers + must not change what a client sees --- most importantly the datetime format + (isoformat ``+00:00``, not pydantic's default ``Z``).""" + + def test_task_attempt_shape(self): + ta = TaskAttempt( + compute_service_id=str(CSID), + hostname="h", + datetime_claimed=NOW, + datetime_end=LATER, + outcome=TaskOutcomeEnum.complete, + units_completed=3, + units_total=5, + protocoldagresultref=PDRR_SK, + ) + assert ta.to_dict() == { + "compute_service_id": str(CSID), + "hostname": "h", + "manager_name": None, + "datetime_claimed": "2026-07-10T12:00:00+00:00", + "datetime_end": "2026-07-10T13:30:00+00:00", + "outcome": "complete", + "units_completed": 3, + "units_total": 5, + "protocoldagresultref": str(PDRR_SK), + } + + def test_unit_result_rec_shape(self): + pur = ProtocolUnitResultRec( + scoped_key=PURR_SK, + obj_key=GufeKey("ProtocolUnitResult-r1"), + source_key=GufeKey("ProtocolUnit-u1"), + name="u", + ok=True, + start_time=NOW, + end_time=LATER, + has_logs=True, + ) + assert pur.to_dict() == { + "scoped_key": str(PURR_SK), + "obj_key": "ProtocolUnitResult-r1", + "source_key": "ProtocolUnit-u1", + "name": "u", + "ok": True, + "start_time": "2026-07-10T12:00:00+00:00", + "end_time": "2026-07-10T13:30:00+00:00", + "has_logs": True, + "has_stdout": False, + "has_stderr": False, + } + + def test_task_details_nested_shape(self): + tc = TaskClaim(compute_service_id=str(CSID), hostname="h", datetime_claimed=NOW) + ta = TaskAttempt( + compute_service_id=str(CSID), datetime_claimed=NOW, outcome=None + ) + td = TaskDetails( + task=TASK_SK, + status=TaskStatusEnum.running, + datetime_status_changed=NOW, + num_claims=2, + current_claim=tc, + most_recent_attempt=ta, + ) + d = td.to_dict() + assert d["task"] == str(TASK_SK) + assert d["status"] == "running" + assert d["datetime_status_changed"] == "2026-07-10T12:00:00+00:00" + # nested models serialize identically to their own to_dict() + assert d["current_claim"] == tc.to_dict() + assert d["most_recent_attempt"] == ta.to_dict() + + def test_datetime_is_isoformat_offset_not_z(self): + # the one real regression risk of pydantic serialization + d = ProtocolDAGResultRec( + scoped_key=PDRR_SK, ok=True, datetime_created=NOW + ).to_dict() + assert d["datetime_created"] == "2026-07-10T12:00:00+00:00" + assert not d["datetime_created"].endswith("Z") + + def test_validators_coerce_wire_values_on_any_construction(self): + # coercion applies on model_validate (and any construction), not only + # via from_dict: raw wire strings become the proper Python types + pur = ProtocolUnitResultRec.model_validate( + { + "scoped_key": str(PURR_SK), + "obj_key": "ProtocolUnitResult-r1", + "source_key": "ProtocolUnit-u1", + "ok": True, + "start_time": "2026-07-10T12:00:00+00:00", + } + ) + assert isinstance(pur.scoped_key, ScopedKey) and pur.scoped_key == PURR_SK + assert isinstance(pur.obj_key, GufeKey) + assert pur.start_time == NOW + + def test_neo4j_datetime_coerced(self): + # a neo4j DateTime-like object (has .to_native()) is coerced inbound + class FakeNeo4jDT: + def to_native(self): + return NOW + + rec = ProtocolDAGResultRec( + scoped_key=PDRR_SK, ok=True, datetime_created=FakeNeo4jDT() + ) + assert rec.datetime_created == NOW + assert rec.to_dict()["datetime_created"] == "2026-07-10T12:00:00+00:00" From 5706869ed9e97e3709db81b460378b4a3d4b275c Mon Sep 17 00:00:00 2001 From: David Dotson Date: Fri, 31 Jul 2026 14:32:30 -0600 Subject: [PATCH 16/18] Reuse iter_contents for stream iteration; docstring cleanups `pull_protocol_unit_result_streams` used a redundant private `_get_filename_prefix_contents` helper identical to the existing public `S3ObjectStore.iter_contents`; use `iter_contents` and drop the helper. Minor docstring clarifications in the state store and interface client. Co-Authored-By: Claude Opus 4.8 (1M context) --- alchemiscale/interface/client.py | 2 +- alchemiscale/storage/objectstore.py | 10 ++-------- alchemiscale/storage/statestore.py | 5 ++--- 3 files changed, 5 insertions(+), 12 deletions(-) diff --git a/alchemiscale/interface/client.py b/alchemiscale/interface/client.py index 392cbbec..207d43b7 100644 --- a/alchemiscale/interface/client.py +++ b/alchemiscale/interface/client.py @@ -2236,7 +2236,7 @@ def get_scope_compute_share(self, scope: Scope) -> float: The share is computed server-side as the aggregate fraction for this `Scope` relative to its sibling Scopes; only the aggregate fraction is - returned. The identity must hold the given `Scope`. + returned. The identity must be able to access the given `Scope`. Parameters ---------- diff --git a/alchemiscale/storage/objectstore.py b/alchemiscale/storage/objectstore.py index 873d6502..3cb7afbf 100644 --- a/alchemiscale/storage/objectstore.py +++ b/alchemiscale/storage/objectstore.py @@ -379,15 +379,14 @@ def pull_protocol_unit_result_streams( """Return a unit result's captured stream files, filename -> decoded text. Bytes are decoded as UTF-8 with ``errors="replace"``; protocols - overwhelmingly archive text (binary outputs are #180 `ResultFile` - territory). + overwhelmingly archive text. """ if stream not in (STDOUT_DIRNAME, STDERR_DIRNAME): raise ValueError("`stream` must be 'stdout' or 'stderr'") prefix = os.path.join(unit_location, stream) + "/" decompressor = zstd.ZstdDecompressor() out = {} - for obj in self._get_filename_prefix_contents(prefix): + for obj in self.iter_contents(prefix): # key includes self.prefix and the full location; recover the # filename relative to the stream directory, dropping the .zst suffix key = obj.key @@ -399,8 +398,3 @@ def pull_protocol_unit_result_streams( "utf-8", errors="replace" ) return out - - def _get_filename_prefix_contents(self, prefix: str): - """Iterate S3 objects under a location prefix (excluding ``self.prefix``).""" - filter_prefix = os.path.join(self.prefix, prefix) - return self.resource.Bucket(self.bucket).objects.filter(Prefix=filter_prefix) diff --git a/alchemiscale/storage/statestore.py b/alchemiscale/storage/statestore.py index 002b43ff..97de910d 100644 --- a/alchemiscale/storage/statestore.py +++ b/alchemiscale/storage/statestore.py @@ -3775,7 +3775,7 @@ def get_tasks_details(self, tasks: list[ScopedKey]) -> list[TaskDetails | None]: `None` is returned in place of any `Task` that does not exist. The The `current_claim` live progress fields stay `None` until a compute - service reports progress (section 2 of the design). + service reports progress. """ q = """ UNWIND $tasks AS task_sk @@ -4284,8 +4284,7 @@ def set_task_status( raise_error If `True`, raise a `ValueError` if the status of a given Task cannot be changed. reason - Optional human-readable reason for the status change; only recorded - for `invalid`/`deleted` transitions (ignored otherwise). + Optional human-readable reason for the status change. Returns ------- From 00f3c83b52cdac8938f73ffc8bf9bce98eefcd9a Mon Sep 17 00:00:00 2001 From: David Dotson Date: Fri, 31 Jul 2026 14:55:05 -0600 Subject: [PATCH 17/18] Expose stdout/stderr at the result tier, not the task tier Logs were retrievable at the unit and result (ProtocolDAGResult) tiers, but stdout/stderr only at the unit and task tiers -- an accidental asymmetry. The result is the coherent aggregation boundary (one execution attempt's DAG), which is why logs uses it; task-level aggregation instead concatenates raw streams across independent attempts, which is noisier and less useful. Replace `get_task_stdout`/`get_task_stderr` (and their `/tasks/{task}/stdout|stderr` routes) with `get_result_stdout`/ `get_result_stderr` (`/protocoldagresultrefs/{pdrr}/stdout|stderr`), mirroring `get_result_logs`. This yields a consistent unit + result matrix for logs/stdout/stderr; a caller wanting everything for a Task iterates `get_task_result_recs`, the same drill-down logs already requires. Removal is safe: the introspection surface is unreleased (0.8.0). Updates the integration test and the docs accordingly. Co-Authored-By: Claude Opus 4.8 (1M context) --- alchemiscale/interface/api.py | 48 ++++++++++--------- alchemiscale/interface/client.py | 34 ++++++------- .../client/test_client_introspection.py | 4 +- docs/compute.rst | 2 +- docs/user_guide/handling_errors.rst | 6 +-- 5 files changed, 50 insertions(+), 44 deletions(-) diff --git a/alchemiscale/interface/api.py b/alchemiscale/interface/api.py index 8c557112..61a21ebd 100644 --- a/alchemiscale/interface/api.py +++ b/alchemiscale/interface/api.py @@ -1387,48 +1387,52 @@ def get_result_logs( return "\n".join(line for _, _, line in entries) -def _render_task_stream(task_scoped_key, stream, n4js, s3os, token) -> str: +def _render_result_stream( + protocoldagresultref_scoped_key, stream, n4js, s3os, token +) -> str: has_attr = "has_stdout" if stream == "stdout" else "has_stderr" - sk = ScopedKey.from_str(task_scoped_key) - validate_scopes(sk.scope, token) + pdrr_sk = ScopedKey.from_str(protocoldagresultref_scoped_key) + validate_scopes(pdrr_sk.scope, token) sections = [] - for pdrr_rec in n4js.get_task_result_recs(sk): - for unit_rec in n4js.get_result_unit_recs(pdrr_rec.scoped_key): - if not getattr(unit_rec, has_attr): - continue - purr = n4js.get_gufe(unit_rec.scoped_key) - files = s3os.pull_protocol_unit_result_streams(purr.location, stream) - for filename, text in files.items(): - sections.append( - f"=== result {pdrr_rec.scoped_key} :: " - f"unit {_unit_label(unit_rec)} :: {filename} ===\n{text}" - ) + for unit_rec in n4js.get_result_unit_recs(pdrr_sk): + if not getattr(unit_rec, has_attr): + continue + purr = n4js.get_gufe(unit_rec.scoped_key) + files = s3os.pull_protocol_unit_result_streams(purr.location, stream) + for filename, text in files.items(): + sections.append( + f"=== unit {_unit_label(unit_rec)} :: {filename} ===\n{text}" + ) return "\n".join(sections) -@router.get("/tasks/{task_scoped_key}/stdout") -def get_task_stdout( - task_scoped_key, +@router.get("/protocoldagresultrefs/{protocoldagresultref_scoped_key}/stdout") +def get_result_stdout( + protocoldagresultref_scoped_key, *, n4js: Neo4jStore = Depends(get_n4js_depends), s3os: S3ObjectStore = Depends(get_s3os_depends), token: TokenData = Depends(get_token_data_depends), ) -> str: - return _render_task_stream(task_scoped_key, "stdout", n4js, s3os, token) + return _render_result_stream( + protocoldagresultref_scoped_key, "stdout", n4js, s3os, token + ) -@router.get("/tasks/{task_scoped_key}/stderr") -def get_task_stderr( - task_scoped_key, +@router.get("/protocoldagresultrefs/{protocoldagresultref_scoped_key}/stderr") +def get_result_stderr( + protocoldagresultref_scoped_key, *, n4js: Neo4jStore = Depends(get_n4js_depends), s3os: S3ObjectStore = Depends(get_s3os_depends), token: TokenData = Depends(get_token_data_depends), ) -> str: - return _render_task_stream(task_scoped_key, "stderr", n4js, s3os, token) + return _render_result_stream( + protocoldagresultref_scoped_key, "stderr", n4js, s3os, token + ) @router.post("/bulk/tasks/progress") diff --git a/alchemiscale/interface/client.py b/alchemiscale/interface/client.py index 207d43b7..377a9b8f 100644 --- a/alchemiscale/interface/client.py +++ b/alchemiscale/interface/client.py @@ -2384,43 +2384,45 @@ def get_result_logs( f"/protocoldagresultrefs/{pdrr_sk}/logs", params=params ) - def get_task_stdout(self, task: ScopedKey) -> str: - """Get a human-readable rendering of all captured stdout for a `Task`. + def get_result_stdout(self, pdrr: ScopedKey | ProtocolDAGResultRec) -> str: + """Get a human-readable rendering of all captured stdout of a `ProtocolDAGResult`. - Concatenates stdout across all `ProtocolDAGResult` objects of the - `Task` (most recent first), with section headers identifying each result, - unit, and filename. + Concatenates stdout across the unit results of the `ProtocolDAGResult`, + with section headers identifying each unit and filename. Parameters ---------- - task - The `ScopedKey` of the `Task` to retrieve stdout for. + pdrr + The `ScopedKey` of the `ProtocolDAGResultRef` (or the + `ProtocolDAGResultRec` describing it) to retrieve stdout for. Returns ------- str The rendered stdout, or ``""`` if none was captured. """ - return self._get_resource(f"/tasks/{task}/stdout") + pdrr_sk = self._as_scoped_key(pdrr) + return self._get_resource(f"/protocoldagresultrefs/{pdrr_sk}/stdout") - def get_task_stderr(self, task: ScopedKey) -> str: - """Get a human-readable rendering of all captured stderr for a `Task`. + def get_result_stderr(self, pdrr: ScopedKey | ProtocolDAGResultRec) -> str: + """Get a human-readable rendering of all captured stderr of a `ProtocolDAGResult`. - Concatenates stderr across all `ProtocolDAGResult` objects of the - `Task` (most recent first), with section headers identifying each result, - unit, and filename. + Concatenates stderr across the unit results of the `ProtocolDAGResult`, + with section headers identifying each unit and filename. Parameters ---------- - task - The `ScopedKey` of the `Task` to retrieve stderr for. + pdrr + The `ScopedKey` of the `ProtocolDAGResultRef` (or the + `ProtocolDAGResultRec` describing it) to retrieve stderr for. Returns ------- str The rendered stderr, or ``""`` if none was captured. """ - return self._get_resource(f"/tasks/{task}/stderr") + pdrr_sk = self._as_scoped_key(pdrr) + return self._get_resource(f"/protocoldagresultrefs/{pdrr_sk}/stderr") def get_tasks_progress( self, tasks: list[ScopedKey] diff --git a/alchemiscale/tests/integration/interface/client/test_client_introspection.py b/alchemiscale/tests/integration/interface/client/test_client_introspection.py index dfbcd685..a8d71f92 100644 --- a/alchemiscale/tests/integration/interface/client/test_client_introspection.py +++ b/alchemiscale/tests/integration/interface/client/test_client_introspection.py @@ -205,8 +205,8 @@ def test_unit_artifacts_retrieval( assert "first line" in rendered_unit rendered_time = user_client.get_result_logs(pdrr_sk, order="time") assert "first line" in rendered_time - assert "captured stdout" in user_client.get_task_stdout(task_sk) - assert "captured stderr" in user_client.get_task_stderr(task_sk) + assert "captured stdout" in user_client.get_result_stdout(pdrr_sk) + assert "captured stderr" in user_client.get_result_stderr(pdrr_sk) def test_tracebacks( self, diff --git a/docs/compute.rst b/docs/compute.rst index b24fab18..04cde328 100644 --- a/docs/compute.rst +++ b/docs/compute.rst @@ -207,7 +207,7 @@ All of them have sensible defaults, so you only need to set them to change the d ``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. - Captured streams are what :py:meth:`~alchemiscale.interface.client.AlchemiscaleClient.get_result_unit_stdout`, :py:meth:`~alchemiscale.interface.client.AlchemiscaleClient.get_result_unit_stderr`, :py:meth:`~alchemiscale.interface.client.AlchemiscaleClient.get_task_stdout`, and :py:meth:`~alchemiscale.interface.client.AlchemiscaleClient.get_task_stderr` return. + Captured streams are what :py:meth:`~alchemiscale.interface.client.AlchemiscaleClient.get_result_unit_stdout`, :py:meth:`~alchemiscale.interface.client.AlchemiscaleClient.get_result_unit_stderr`, :py:meth:`~alchemiscale.interface.client.AlchemiscaleClient.get_result_stdout`, and :py:meth:`~alchemiscale.interface.client.AlchemiscaleClient.get_result_stderr` return. ``capture_logs`` If ``true`` (the default), log records emitted through ``gufe``'s ``gufekey`` logger namespace (that is, protocol logs written via ``ProtocolUnit.logger``) are captured per unit result and uploaded alongside results. diff --git a/docs/user_guide/handling_errors.rst b/docs/user_guide/handling_errors.rst index e957bc51..5f2e3905 100644 --- a/docs/user_guide/handling_errors.rst +++ b/docs/user_guide/handling_errors.rst @@ -110,10 +110,10 @@ When you don't need per-unit granularity, three convenience methods render every >>> # or interleave across units by timestamp >>> print(asc.get_result_logs(pdrr, order='time')) -:py:meth:`~alchemiscale.interface.client.AlchemiscaleClient.get_task_stdout` and :py:meth:`~alchemiscale.interface.client.AlchemiscaleClient.get_task_stderr` go one level higher, concatenating the captured stdout/stderr across *all* of a :py:class:`~alchemiscale.storage.models.Task`\'s :external+gufe:py:class:`~gufe.protocols.protocoldag.ProtocolDAGResult`\s (most recent first), with section headers identifying each result, unit, and filename:: +:py:meth:`~alchemiscale.interface.client.AlchemiscaleClient.get_result_stdout` and :py:meth:`~alchemiscale.interface.client.AlchemiscaleClient.get_result_stderr` do the same for the captured stdout/stderr of one :external+gufe:py:class:`~gufe.protocols.protocoldag.ProtocolDAGResult`, concatenating across its unit results with section headers identifying each unit and filename:: - >>> print(asc.get_task_stdout(task)) - >>> print(asc.get_task_stderr(task)) + >>> print(asc.get_result_stdout(pdrr)) + >>> print(asc.get_result_stderr(pdrr)) Each returns ``""`` when nothing was captured. From a4311779531700f7ce287a4553e0a4a3f8514de2 Mon Sep 17 00:00:00 2001 From: David Dotson Date: Fri, 31 Jul 2026 15:48:51 -0600 Subject: [PATCH 18/18] Add visualize kwarg to Task introspection client methods Give `get_task_history`, `get_tasks_details`, `get_tasks_progress`, and `get_task_tracebacks` a `visualize: bool = True` kwarg that renders a rich-formatted view (side-effect display; the raw records are still returned), matching the existing `get_*_status` convention. - history: per-attempt table with color-coded outcome, hostname, claimed, a compact duration ("how long did it run"), progress, and a result flag. - details: table across tasks; a not-found (None) entry is shown, not dropped. - progress: real rich progress bars; non-reporting tasks marked. - tracebacks: each unit traceback in a bordered panel. Consistency/cleanup: extract a shared status color palette and an outcome palette, and refactor `_visualize_status` onto it. Drop `expand=True` on the wide tables (it starved flexible columns to zero width) and use minute precision + a duration column so they read cleanly at typical widths. All user-derived text (reasons, tracebacks) is wrapped in `rich.text.Text` so bracketed content (e.g. `[Errno 2]`) renders literally instead of being parsed as markup. Adds tests/unit/test_client_visualize.py smoke-testing every renderer against edge cases (open attempts, not-found tasks, zero-total progress, empty input, markup-hostile tracebacks). Co-Authored-By: Claude Opus 4.8 (1M context) --- alchemiscale/interface/client.py | 245 ++++++++++++++++-- .../tests/unit/test_client_visualize.py | 123 +++++++++ 2 files changed, 350 insertions(+), 18 deletions(-) create mode 100644 alchemiscale/tests/unit/test_client_visualize.py diff --git a/alchemiscale/interface/client.py b/alchemiscale/interface/client.py index 377a9b8f..edee7523 100644 --- a/alchemiscale/interface/client.py +++ b/alchemiscale/interface/client.py @@ -1010,23 +1010,33 @@ def get_task_transformation(self, task: ScopedKey) -> ScopedKey: transformation = self._get_resource(f"/tasks/{task}/transformation") return ScopedKey.from_str(transformation) + # rich styles for Task status / attempt-outcome cells, shared across the + # ``visualize=True`` renderings for a consistent color scheme + _STATUS_STYLES = { + "complete": "green", + "running": "orange3", + "waiting": "#1793d0", + "error": "#ff073a", + "invalid": "magenta1", + "deleted": "purple", + } + _OUTCOME_STYLES = { + "complete": "green", + "error": "#ff073a", + "expired": "orange3", + "released": "grey62", + } + def _visualize_status(self, status_counts, status_object): from rich import print as rprint from rich.table import Table - title = f"{status_object}" - table = Table(title=title, title_justify="left", expand=True) - # table = Table() - + table = Table(title=f"{status_object}", title_justify="left", expand=True) table.add_column("status", justify="left", no_wrap=True) table.add_column("count", justify="right") - table.add_row("complete", f"{status_counts.get('complete', 0)}", style="green") - table.add_row("running", f"{status_counts.get('running', 0)}", style="orange3") - table.add_row("waiting", f"{status_counts.get('waiting', 0)}", style="#1793d0") - table.add_row("error", f"{status_counts.get('error', 0)}", style="#ff073a") - table.add_row("invalid", f"{status_counts.get('invalid', 0)}", style="magenta1") - table.add_row("deleted", f"{status_counts.get('deleted', 0)}", style="purple") + for status, style in self._STATUS_STYLES.items(): + table.add_row(status, f"{status_counts.get(status, 0)}", style=style) rprint(table) @@ -2164,8 +2174,173 @@ def get_task_failures( return pdrs + @staticmethod + def _fmt_dt(dt) -> str: + # minute precision keeps the glance tables narrow; the returned records + # carry full-precision datetimes + return dt.strftime("%Y-%m-%d %H:%M") if dt is not None else "—" + + @staticmethod + def _fmt_duration(start, end) -> str: + """Compact ``start``..``end`` duration (e.g. ``1h30m``), or ``—``.""" + if start is None or end is None: + return "—" + seconds = int((end - start).total_seconds()) + if seconds < 0: + return "—" + hours, rem = divmod(seconds, 3600) + minutes, secs = divmod(rem, 60) + if hours: + return f"{hours}h{minutes}m" + if minutes: + return f"{minutes}m{secs}s" + return f"{secs}s" + + def _visualize_task_history(self, task, attempts): + from rich import print as rprint + from rich.table import Table + from rich.text import Text + + table = Table(title=f"Task history: {task}", title_justify="left") + table.add_column("#", justify="right", no_wrap=True) + table.add_column("outcome", no_wrap=True) + table.add_column("hostname", overflow="fold") + table.add_column("claimed", no_wrap=True) + table.add_column("duration", justify="right", no_wrap=True) + table.add_column("progress", justify="right", no_wrap=True) + table.add_column("result", justify="center", no_wrap=True) + + # `compute_service_id` is intentionally omitted (long, and present on the + # returned records); `hostname` is the human-friendly "where". `duration` + # (rather than an end timestamp) directly answers "how long did it run". + n = len(attempts) + for i, a in enumerate(attempts): + outcome = a.outcome.value if a.outcome is not None else "running" + progress = ( + f"{a.units_completed}/{a.units_total}" + if a.units_total is not None + else "—" + ) + table.add_row( + str(n - i), # 1-based, oldest = 1 (attempts are most-recent-first) + Text(outcome, style=self._OUTCOME_STYLES.get(outcome, "grey62")), + Text(a.hostname or "—"), + self._fmt_dt(a.datetime_claimed), + self._fmt_duration(a.datetime_claimed, a.datetime_end), + progress, + "✓" if a.protocoldagresultref is not None else "—", + ) + + rprint(table) + + def _visualize_tasks_details(self, tasks, details): + from rich import print as rprint + from rich.table import Table + from rich.text import Text + + table = Table(title="Task details", title_justify="left") + table.add_column("task", overflow="fold") + table.add_column("status", no_wrap=True) + table.add_column("changed", no_wrap=True) + table.add_column("reason", overflow="fold") + table.add_column("claims", justify="right", no_wrap=True) + table.add_column("current claim", overflow="fold") + table.add_column("last outcome", no_wrap=True) + + for task, d in zip(tasks, details): + if d is None: + table.add_row( + str(task), Text("(not found)", style="grey62"), *(["—"] * 5) + ) + continue + + status = d.status.value + claim = "—" + if d.current_claim is not None: + claim = d.current_claim.hostname or d.current_claim.compute_service_id + + last_outcome = "—" + attempt = d.most_recent_attempt + if attempt is not None and attempt.outcome is not None: + oc = attempt.outcome.value + last_outcome = Text(oc, style=self._OUTCOME_STYLES.get(oc, "grey62")) + + reason = d.reason or "" + if len(reason) > 30: + reason = reason[:29] + "…" + + table.add_row( + str(d.task), + Text(status, style=self._STATUS_STYLES.get(status, "")), + self._fmt_dt(d.datetime_status_changed), + Text(reason), + str(d.num_claims), + Text(claim), + last_outcome, + ) + + rprint(table) + + def _visualize_tasks_progress(self, tasks, progress): + from rich import print as rprint + from rich.progress_bar import ProgressBar + from rich.table import Table + from rich.text import Text + + table = Table(title="Task progress", title_justify="left") + table.add_column("task", overflow="fold") + table.add_column("progress", ratio=1) + table.add_column("units", justify="right", no_wrap=True) + table.add_column("%", justify="right", no_wrap=True) + + for task, p in zip(tasks, progress): + if p is None: + table.add_row( + str(task), Text("— not reporting —", style="grey62"), "—", "—" + ) + continue + + completed, total = p + pct = (100 * completed / total) if total else 0.0 + table.add_row( + str(task), + ProgressBar(total=total or 1, completed=completed, width=40), + f"{completed}/{total}", + f"{pct:.0f}%", + ) + + rprint(table) + + def _visualize_task_tracebacks(self, task, tracebacks): + from rich import print as rprint + from rich.panel import Panel + from rich.text import Text + + if not tracebacks: + rprint(f"[grey62]No tracebacks for task {task}.[/grey62]") + return + + for tb in tracebacks: + header = f"[bold]{tb.protocoldagresultref}[/bold]" + if tb.datetime_created is not None: + header += f" · {self._fmt_dt(tb.datetime_created)}" + rprint(header) + for ut in tb.tracebacks: + # wrap traceback text in `Text` so its bracketed content (e.g. + # ``[Errno 2]``) is never interpreted as rich markup + rprint( + Panel( + Text(ut.traceback), + title=str(ut.source_key), + subtitle=str(ut.failure_key), + title_align="left", + subtitle_align="right", + border_style="#ff073a", + ) + ) + def get_task_history( - self, task: ScopedKey, limit: int | None = None + self, task: ScopedKey, limit: int | None = None, visualize: bool = True ) -> list[TaskAttempt]: """Get the execution history of a `Task`. @@ -2175,6 +2350,9 @@ def get_task_history( The `ScopedKey` of the `Task` to retrieve the history for. limit If given, return at most this many of the most recent attempts. + visualize + If ``True`` (default), also print a rich-formatted table of the + attempts. Returns ------- @@ -2184,15 +2362,25 @@ def get_task_history( """ params = dict(limit=limit) attempts = self._get_resource(f"/tasks/{task}/history", params=params) - return [TaskAttempt.from_dict(attempt) for attempt in attempts] + attempts = [TaskAttempt.from_dict(attempt) for attempt in attempts] + + if visualize: + self._visualize_task_history(task, attempts) + + return attempts - def get_tasks_details(self, tasks: list[ScopedKey]) -> list[TaskDetails | None]: + def get_tasks_details( + self, tasks: list[ScopedKey], visualize: bool = True + ) -> list[TaskDetails | None]: """Get summary details for multiple Tasks. Parameters ---------- tasks The `ScopedKey` of each `Task` to retrieve details for. + visualize + If ``True`` (default), also print a rich-formatted table of the + details. Returns ------- @@ -2203,13 +2391,18 @@ def get_tasks_details(self, tasks: list[ScopedKey]) -> list[TaskDetails | None]: """ data = dict(tasks=[str(task) for task in tasks]) details = self._post_resource("/bulk/tasks/details", data=data) - return [ + details = [ TaskDetails.from_dict(detail) if detail is not None else None for detail in details ] + if visualize: + self._visualize_tasks_details(tasks, details) + + return details + def get_task_tracebacks( - self, task: ScopedKey, limit: int | None = None + self, task: ScopedKey, limit: int | None = None, visualize: bool = True ) -> list[TaskTracebacks]: """Get the tracebacks from failed `ProtocolDAGResult` objects of a `Task`. @@ -2220,6 +2413,9 @@ def get_task_tracebacks( limit If given, return tracebacks for at most this many of the most recent failed `ProtocolDAGResult` objects. + visualize + If ``True`` (default), also print each traceback in a rich-formatted + panel. Returns ------- @@ -2229,7 +2425,12 @@ def get_task_tracebacks( """ params = dict(limit=limit) tracebacks = self._get_resource(f"/tasks/{task}/tracebacks", params=params) - return [TaskTracebacks.from_dict(tb) for tb in tracebacks] + tracebacks = [TaskTracebacks.from_dict(tb) for tb in tracebacks] + + if visualize: + self._visualize_task_tracebacks(task, tracebacks) + + return tracebacks def get_scope_compute_share(self, scope: Scope) -> float: """Get this identity's fractional compute share within the given `Scope`. @@ -2425,7 +2626,7 @@ def get_result_stderr(self, pdrr: ScopedKey | ProtocolDAGResultRec) -> str: return self._get_resource(f"/protocoldagresultrefs/{pdrr_sk}/stderr") def get_tasks_progress( - self, tasks: list[ScopedKey] + self, tasks: list[ScopedKey], visualize: bool = True ) -> list[tuple[int, int] | None]: """Get execution progress for multiple Tasks. @@ -2433,6 +2634,9 @@ def get_tasks_progress( ---------- tasks The `ScopedKey` of each `Task` to retrieve progress for. + visualize + If ``True`` (default), also print a rich-formatted table of progress + bars. Returns ------- @@ -2443,7 +2647,12 @@ def get_tasks_progress( """ data = dict(tasks=[str(task) for task in tasks]) progress = self._post_resource("/bulk/tasks/progress", data=data) - return [tuple(p) if p is not None else None for p in progress] + progress = [tuple(p) if p is not None else None for p in progress] + + if visualize: + self._visualize_tasks_progress(tasks, progress) + + return progress def add_task_restart_patterns( self, diff --git a/alchemiscale/tests/unit/test_client_visualize.py b/alchemiscale/tests/unit/test_client_visualize.py new file mode 100644 index 00000000..b858777a --- /dev/null +++ b/alchemiscale/tests/unit/test_client_visualize.py @@ -0,0 +1,123 @@ +"""Robustness tests for the `AlchemiscaleClient` introspection ``visualize`` +renderings (``get_task_history``/``get_tasks_details``/``get_tasks_progress``/ +``get_task_tracebacks``). + +These render via ``rich``; the concern is that they never raise on edge cases +(open/running attempts, not-found Tasks, zero-total progress, empty input) and +never mis-interpret bracketed content (e.g. ``[Errno 2]``) as ``rich`` markup. +The visualizers only use class-level state, so an un-``__init__``-ed client +instance suffices --- no server needed. ``rich`` strips ANSI when stdout is not +a TTY (as under ``capsys``), so plain-text assertions hold. +""" + +import datetime + +import pytest +from gufe.tokenization import GufeKey + +from alchemiscale.models import ScopedKey, Scope +from alchemiscale.storage.models import ( + TaskAttempt, + TaskClaim, + TaskDetails, + TaskOutcomeEnum, + TaskStatusEnum, + TaskTracebacks, + TaskUnitTraceback, +) +from alchemiscale.interface.client import AlchemiscaleClient + +NOW = datetime.datetime(2026, 7, 10, 12, 0, 0, tzinfo=datetime.UTC) +LATER = datetime.datetime(2026, 7, 10, 13, 30, 0, tzinfo=datetime.UTC) +CSID = "compute-a.svc-" + "0" * 32 +T = ScopedKey.from_str("Task-aaa111-org-camp-proj") +T2 = ScopedKey.from_str("Task-bbb222-org-camp-proj") +PDRR = ScopedKey.from_str("ProtocolDAGResultRef-abc123-org-camp-proj") + + +@pytest.fixture +def client(): + # bypass __init__ (no server): the visualizers only touch class-level state + return AlchemiscaleClient.__new__(AlchemiscaleClient) + + +class TestVisualizeIntrospection: + def test_task_history(self, client, capsys): + attempts = [ + TaskAttempt( + compute_service_id=CSID, + hostname="node-7", + datetime_claimed=NOW, + datetime_end=LATER, + outcome=TaskOutcomeEnum.complete, + units_completed=5, + units_total=5, + protocoldagresultref=PDRR, + ), + # an open/running attempt: no end, no outcome, no result + TaskAttempt( + compute_service_id=CSID, + hostname="node-9", + datetime_claimed=NOW, + outcome=None, + units_completed=1, + units_total=5, + ), + ] + client._visualize_task_history(T, attempts) + out = capsys.readouterr().out + assert "complete" in out + assert "running" in out # open attempt shown as running + assert "1h30m" in out # duration rendering + + def test_tasks_details_handles_missing(self, client, capsys): + details = TaskDetails( + task=T, + status=TaskStatusEnum.error, + datetime_status_changed=NOW, + reason="something went wrong", + num_claims=1, + current_claim=TaskClaim(compute_service_id=CSID, hostname="node-9"), + ) + client._visualize_tasks_details([T, T2], [details, None]) + out = capsys.readouterr().out + assert "error" in out + assert "not found" in out # the None entry is rendered, not skipped + + def test_tasks_progress_edges(self, client, capsys): + # a reporting Task, a non-reporting (None) Task, and a zero-total Task + client._visualize_tasks_progress([T, T2, T], [(3, 10), None, (0, 0)]) + out = capsys.readouterr().out + assert "3/10" in out + assert "reporting" in out # "— not reporting —" for the None entry + + def test_tracebacks_markup_safe(self, client, capsys): + # bracketed content must render literally, never as rich markup + tut = TaskUnitTraceback( + failure_key=GufeKey("ProtocolUnitFailure-f1"), + source_key=GufeKey("ProtocolUnit-u1"), + traceback='raise ValueError("[boom]") # [Errno 2]', + ) + tt = TaskTracebacks( + protocoldagresultref=PDRR, + datetime_created=NOW, + creator=CSID, + tracebacks=[tut], + ) + client._visualize_task_tracebacks(T, [tt]) + out = capsys.readouterr().out + assert "boom" in out + assert "Errno" in out + + def test_tracebacks_empty(self, client, capsys): + client._visualize_task_tracebacks(T, []) + out = capsys.readouterr().out + assert "No tracebacks" in out + + def test_status_palette_unchanged(self, client, capsys): + # the refactor of _visualize_status onto the shared palette keeps all + # six statuses in order + client._visualize_status({"complete": 2, "error": 1}, Scope("o", "c", "p")) + out = capsys.readouterr().out + for status in ("complete", "running", "waiting", "error", "invalid", "deleted"): + assert status in out