From 6a1fe2422a33939328c8891b15ba8ecf41a6a1d4 Mon Sep 17 00:00:00 2001 From: Claude Date: Fri, 9 Jan 2026 21:27:52 +0000 Subject: [PATCH 1/2] Implement stdout/stderr log capture for ProtocolDAGResults (#295) This commit implements the feature requested in issue #295 to capture and retain stdout/stderr logs from ProtocolDAG executions and expose them via the AlchemiscaleClient. Changes: - Added ProtocolDAGResultLog model to store log references in the state store - Modified compute service to capture stdout/stderr during DAG execution - Updated compute client to send logs along with results - Updated compute API to receive and store logs in object store - Added state store methods to link logs to ProtocolDAGResults - Added object store methods to persist and retrieve logs - Implemented get_task_stdout/get_task_stderr client methods - Added API endpoint to retrieve task logs The implementation captures logs during execution, stores them as artifacts in S3, and provides easy access through new client methods for debugging failed or completed tasks. --- alchemiscale/compute/api.py | 25 +++++++ alchemiscale/compute/client.py | 8 ++ alchemiscale/compute/service.py | 45 ++++++++---- alchemiscale/interface/api.py | 59 +++++++++++++++ alchemiscale/interface/client.py | 36 +++++++++ alchemiscale/storage/models.py | 48 ++++++++++++ alchemiscale/storage/objectstore.py | 110 ++++++++++++++++++++++++++++ alchemiscale/storage/statestore.py | 97 ++++++++++++++++++++++++ 8 files changed, 415 insertions(+), 13 deletions(-) diff --git a/alchemiscale/compute/api.py b/alchemiscale/compute/api.py index ada77df4..876f42d8 100644 --- a/alchemiscale/compute/api.py +++ b/alchemiscale/compute/api.py @@ -367,6 +367,8 @@ async def set_task_result( protocoldagresult_ = body_["protocoldagresult"] compute_service_id = body_["compute_service_id"] + stdout = body_.get("stdout") + stderr = body_.get("stderr") task_sk = ScopedKey.from_str(task_scoped_key) validate_scopes(task_sk.scope, token) @@ -392,6 +394,29 @@ async def set_task_result( task=task_sk, protocoldagresultref=protocoldagresultref ) + # Store stdout and stderr logs if provided + if stdout: + stdout_log_ref = s3os.push_protocoldagresult_log( + log_content=stdout, + stream="stdout", + protocoldagresult_gufekey=pdr.key, + transformation=tf_sk, + protocoldagresult_ok=pdr.ok(), + creator=compute_service_id, + ) + n4js.set_protocoldagresult_log(result_sk, stdout_log_ref) + + if stderr: + stderr_log_ref = s3os.push_protocoldagresult_log( + log_content=stderr, + stream="stderr", + protocoldagresult_gufekey=pdr.key, + transformation=tf_sk, + protocoldagresult_ok=pdr.ok(), + creator=compute_service_id, + ) + n4js.set_protocoldagresult_log(result_sk, stderr_log_ref) + # if success, set task complete, remove from all hubs # otherwise, set as errored, leave in hubs if protocoldagresultref.ok: diff --git a/alchemiscale/compute/client.py b/alchemiscale/compute/client.py index 7a685925..8f047658 100644 --- a/alchemiscale/compute/client.py +++ b/alchemiscale/compute/client.py @@ -148,6 +148,8 @@ def set_task_result( task: ScopedKey, protocoldagresult: ProtocolDAGResult, compute_service_id: ComputeServiceID | None = None, + stdout: str | None = None, + stderr: str | None = None, ) -> ScopedKey: data = dict( @@ -155,6 +157,12 @@ def set_task_result( compute_service_id=str(compute_service_id), ) + # Add logs if provided + if stdout is not None: + data["stdout"] = stdout + if stderr is not None: + data["stderr"] = stderr + pdr_sk = self._post_resource(f"/tasks/{task}/results", data) return ScopedKey.from_dict(pdr_sk) diff --git a/alchemiscale/compute/service.py b/alchemiscale/compute/service.py index dc4ee30c..259dd41e 100644 --- a/alchemiscale/compute/service.py +++ b/alchemiscale/compute/service.py @@ -12,6 +12,9 @@ import threading from pathlib import Path import shutil +import sys +import io +from contextlib import redirect_stdout, redirect_stderr from gufe import Transformation from gufe.protocols.protocoldag import execute_DAG, ProtocolDAG, ProtocolDAGResult @@ -198,16 +201,20 @@ def task_to_protocoldag( return protocoldag, transformation, extends_protocoldagresult def push_result( - self, task: ScopedKey, protocoldagresult: ProtocolDAGResult + self, + task: ScopedKey, + protocoldagresult: ProtocolDAGResult, + stdout: str | None = None, + stderr: str | None = None, ) -> ScopedKey: # TODO: this method should postprocess any paths, # leaf nodes in DAG for blob results that should go to object store # TODO: ship paths to object store - # finally, push ProtocolDAGResult + # finally, push ProtocolDAGResult with logs sk: ScopedKey = self.client.set_task_result( - task, protocoldagresult, self.compute_service_id + task, protocoldagresult, self.compute_service_id, stdout=stdout, stderr=stderr ) return sk @@ -237,15 +244,21 @@ def execute(self, task: ScopedKey) -> ScopedKey: scratch.mkdir() self.logger.info("Executing '%s'...", protocoldag) + + # Capture stdout and stderr during execution + stdout_capture = io.StringIO() + stderr_capture = io.StringIO() + try: - protocoldagresult = execute_DAG( - protocoldag, - shared_basedir=shared, - scratch_basedir=scratch, - keep_scratch=self.keep_scratch, - raise_error=False, - n_retries=self.settings.n_retries, - ) + with redirect_stdout(stdout_capture), redirect_stderr(stderr_capture): + protocoldagresult = execute_DAG( + protocoldag, + shared_basedir=shared, + scratch_basedir=scratch, + keep_scratch=self.keep_scratch, + raise_error=False, + n_retries=self.settings.n_retries, + ) finally: if not self.keep_shared: shutil.rmtree(shared) @@ -253,6 +266,10 @@ def execute(self, task: ScopedKey) -> ScopedKey: if not self.keep_scratch: shutil.rmtree(scratch) + # Get captured output + stdout_content = stdout_capture.getvalue() + stderr_content = stderr_capture.getvalue() + if protocoldagresult.ok(): self.logger.info("'%s' -> '%s' : SUCCESS", protocoldag, protocoldagresult) else: @@ -265,8 +282,10 @@ def execute(self, task: ScopedKey) -> ScopedKey: failure.exception, ) - # push the result (or failure) back to the compute API - result_sk = self.push_result(task, protocoldagresult) + # push the result (or failure) back to the compute API with captured logs + result_sk = self.push_result( + task, protocoldagresult, stdout=stdout_content, stderr=stderr_content + ) self.logger.info("Pushed result `%s'", protocoldagresult) return result_sk diff --git a/alchemiscale/interface/api.py b/alchemiscale/interface/api.py index e41a1c87..e95dd9b3 100644 --- a/alchemiscale/interface/api.py +++ b/alchemiscale/interface/api.py @@ -1177,6 +1177,65 @@ def get_task_failures( return [str(sk) for sk in n4js.get_task_failures(sk)] +@router.get("/tasks/{task_scoped_key}/logs/{stream}") +def get_task_logs( + task_scoped_key, + stream: str, + *, + n4js: Neo4jStore = Depends(get_n4js_depends), + s3os: S3ObjectStore = Depends(get_s3os_depends), + token: TokenData = Depends(get_token_data_depends), +): + """Get log content for a Task. + + Parameters + ---------- + task_scoped_key + The ScopedKey of the Task. + stream + Either "stdout" or "stderr". + + Returns + ------- + list[str] + List of log contents from all ProtocolDAGResults for this Task. + """ + if stream not in ["stdout", "stderr"]: + raise HTTPException( + status_code=400, detail="stream must be 'stdout' or 'stderr'" + ) + + sk = ScopedKey.from_str(task_scoped_key) + validate_scopes(sk.scope, token) + + # Get all log references for this task and stream + log_refs = n4js.get_task_logs(sk, stream=stream) + + # Retrieve the actual log content from the object store + logs = [] + for log_ref_sk in log_refs: + # Get the log reference node to extract location and other metadata + log_ref_data = n4js.get_scoped_key(log_ref_sk, resolve_gufe=False) + from ..storage.models import ProtocolDAGResultLog + + log_ref = ProtocolDAGResultLog._from_dict(log_ref_data) + + # Get transformation for this task + tf_sk, _ = n4js.get_task_transformation(task=task_scoped_key, return_gufe=False) + + # Pull the log content from S3 + log_content = s3os.pull_protocoldagresult_log( + protocoldagresult=ScopedKey(log_ref.scope, log_ref.obj_key), + transformation=tf_sk, + stream=stream, + ok=True, # We'll try both success and failure + location=log_ref.location, + ) + logs.append(log_content) + + return logs + + ### strategies diff --git a/alchemiscale/interface/client.py b/alchemiscale/interface/client.py index 9a321a39..e59eac59 100644 --- a/alchemiscale/interface/client.py +++ b/alchemiscale/interface/client.py @@ -1920,6 +1920,42 @@ def get_task_failures( return pdrs + def get_task_stdout(self, task: ScopedKey) -> list[str]: + """Get stdout logs from all `ProtocolDAGResult`s for the given `Task`. + + Parameters + ---------- + task + The `ScopedKey` of the `Task` to retrieve stdout logs for. + + Returns + ------- + list[str] + List of stdout log contents from all ProtocolDAGResults for this Task. + Each element corresponds to one execution attempt. + + """ + logs = self._get_resource(f"/tasks/{task}/logs/stdout") + return logs + + def get_task_stderr(self, task: ScopedKey) -> list[str]: + """Get stderr logs from all `ProtocolDAGResult`s for the given `Task`. + + Parameters + ---------- + task + The `ScopedKey` of the `Task` to retrieve stderr logs for. + + Returns + ------- + list[str] + List of stderr log contents from all ProtocolDAGResults for this Task. + Each element corresponds to one execution attempt. + + """ + logs = self._get_resource(f"/tasks/{task}/logs/stderr") + return logs + def add_task_restart_patterns( self, network_scoped_key: ScopedKey, diff --git a/alchemiscale/storage/models.py b/alchemiscale/storage/models.py index a9aaf8bd..de7aaed3 100644 --- a/alchemiscale/storage/models.py +++ b/alchemiscale/storage/models.py @@ -550,6 +550,54 @@ def _from_dict(cls, d): return super()._from_dict(d_) +class ProtocolDAGResultLog(ObjectStoreRef): + """Reference to stdout or stderr logs from a ProtocolDAGResult execution.""" + + stream: str # "stdout" or "stderr" + + def __init__( + self, + *, + location: str | None = None, + obj_key: GufeKey, + scope: Scope, + stream: str, + datetime_created: datetime.datetime | None = None, + creator: str | None = None, + ): + self.location = location + self.obj_key = GufeKey(obj_key) + self.scope = scope + self.stream = stream + self.datetime_created = datetime_created + self.creator = creator + + def _to_dict(self): + return { + "location": self.location, + "obj_key": str(self.obj_key), + "scope": str(self.scope), + "stream": self.stream, + "datetime_created": ( + self.datetime_created.isoformat() + if self.datetime_created is not None + else None + ), + "creator": self.creator, + } + + @classmethod + def _from_dict(cls, d): + d_ = copy(d) + d_["datetime_created"] = ( + datetime.datetime.fromisoformat(d["datetime_created"]) + if d.get("datetime_created") is not None + else None + ) + + return super()._from_dict(d_) + + class StrategyModeEnum(StrEnum): full = "full" partial = "partial" diff --git a/alchemiscale/storage/objectstore.py b/alchemiscale/storage/objectstore.py index 352b5bee..238c3c7c 100644 --- a/alchemiscale/storage/objectstore.py +++ b/alchemiscale/storage/objectstore.py @@ -297,3 +297,113 @@ def pull_protocoldagresult( pdr_bytes = self._get_bytes(location) return pdr_bytes + + def push_protocoldagresult_log( + self, + log_content: str, + stream: str, + protocoldagresult_gufekey: GufeKey, + transformation: ScopedKey, + protocoldagresult_ok: bool, + creator: str | None = None, + ) -> ProtocolDAGResultLog: + """Push stdout or stderr log for a ProtocolDAGResult to this ObjectStore. + + Parameters + ---------- + log_content + The log content (stdout or stderr) as a string. + stream + Either "stdout" or "stderr". + protocoldagresult_gufekey + The GufeKey of the ProtocolDAGResult this log belongs to. + transformation + The ScopedKey of the Transformation this log corresponds to. + protocoldagresult_ok + ``True`` if ProtocolDAGResult completed successfully; ``False`` if failed. + creator + Identifier of the entity creating this log (usually compute_service_id). + + Returns + ------- + ProtocolDAGResultLog + Reference to the log in the object store. + + """ + from ..storage.models import ProtocolDAGResultLog + + ok = protocoldagresult_ok + route = "results" if ok else "failures" + + # build `location` based on gufe key + location = os.path.join( + "protocoldagresult", + *transformation.scope.to_tuple(), + transformation.gufe_key, + route, + protocoldagresult_gufekey, + f"{stream}.log", + ) + + # encode log content to bytes (UTF-8) + log_bytes = log_content.encode("utf-8") + self._store_bytes(location, log_bytes) + + return ProtocolDAGResultLog( + location=location, + obj_key=protocoldagresult_gufekey, + scope=transformation.scope, + stream=stream, + datetime_created=datetime.datetime.now(tz=datetime.UTC), + creator=creator, + ) + + def pull_protocoldagresult_log( + self, + protocoldagresult: ScopedKey, + transformation: ScopedKey, + stream: str, + ok: bool = True, + location: str | None = None, + ) -> str: + """Pull the log content for a ProtocolDAGResult. + + Parameters + ---------- + protocoldagresult + ScopedKey for ProtocolDAGResult in the object store. + transformation + The ScopedKey of the Transformation this log corresponds to. + stream + Either "stdout" or "stderr". + ok + ``True`` if ProtocolDAGResult completed successfully; ``False`` if failed. + location + The full path in the object store to the log. If provided, this will be used. + + Returns + ------- + str + The log content as a string. + + """ + route = "results" if ok else "failures" + + # build `location` based on provided ScopedKey if not provided + if location is None: + if transformation.scope != protocoldagresult.scope: + raise ValueError( + f"transformation scope '{transformation.scope}' differs from protocoldagresult scope '{protocoldagresult.scope}'" + ) + + location = os.path.join( + "protocoldagresult", + *protocoldagresult.scope.to_tuple(), + transformation.gufe_key, + route, + protocoldagresult.gufe_key, + f"{stream}.log", + ) + + log_bytes = self._get_bytes(location) + return log_bytes.decode("utf-8") diff --git a/alchemiscale/storage/statestore.py b/alchemiscale/storage/statestore.py index ce6de601..9838fc38 100644 --- a/alchemiscale/storage/statestore.py +++ b/alchemiscale/storage/statestore.py @@ -3551,6 +3551,103 @@ def add_protocol_dag_result_ref_tracebacks( merge_subgraph(tx, subgraph, "GufeTokenizable", "_scoped_key") + def set_protocoldagresult_log( + self, + protocoldagresultref: ScopedKey, + protocoldagresultlog: "ProtocolDAGResultLog", + ) -> ScopedKey: + """Set a `ProtocolDAGResultLog` for the given `ProtocolDAGResultRef`. + + Parameters + ---------- + protocoldagresultref + ScopedKey of the ProtocolDAGResultRef this log belongs to. + protocoldagresultlog + The ProtocolDAGResultLog reference to store. + + Returns + ------- + ScopedKey + The ScopedKey of the stored ProtocolDAGResultLog. + """ + from ..storage.models import ProtocolDAGResultLog + + if protocoldagresultref.qualname != "ProtocolDAGResultRef": + raise ValueError( + "`protocoldagresultref` ScopedKey does not correspond to a `ProtocolDAGResultRef`" + ) + + scope = protocoldagresultref.scope + protocoldagresultref_node = self._get_node(protocoldagresultref) + + subgraph, log_node, scoped_key = self._keyed_chain_to_subgraph( + KeyedChain.from_gufe(protocoldagresultlog), + scope=scope, + ) + + subgraph = subgraph | Relationship.type("HAS_LOG")( + protocoldagresultref_node, + log_node, + _org=scope.org, + _campaign=scope.campaign, + _project=scope.project, + ) + + with self.transaction() as tx: + merge_subgraph(tx, subgraph, "GufeTokenizable", "_scoped_key") + + return scoped_key + + def _get_protocoldagresultlogs(self, q: str, scoped_key: ScopedKey): + """Helper method to retrieve ProtocolDAGResultLog ScopedKeys.""" + sks = [] + with self.transaction() as tx: + res = tx.run(q, scoped_key=str(scoped_key)) + for rec in res: + sks.append(rec["sk"]) + + return [ScopedKey.from_str(sk) for sk in sks] + + def get_task_logs( + self, task: ScopedKey, stream: str | None = None + ) -> list[ScopedKey]: + """Get all log references for a given Task. + + Parameters + ---------- + task + ScopedKey of the Task. + stream + Optional filter for "stdout" or "stderr". If None, returns all logs. + + Returns + ------- + list[ScopedKey] + List of ScopedKeys for ProtocolDAGResultLog objects. + """ + if stream is not None and stream not in ["stdout", "stderr"]: + raise ValueError("`stream` must be 'stdout', 'stderr', or None") + + if stream is None: + q = """ + MATCH (task:Task {_scoped_key: $scoped_key}), + (task)-[:RESULTS_IN]->(res:ProtocolDAGResultRef), + (res)-[:HAS_LOG]->(log:ProtocolDAGResultLog) + WITH log._scoped_key as sk + RETURN DISTINCT sk + """ + else: + q = f""" + MATCH (task:Task {{_scoped_key: $scoped_key}}), + (task)-[:RESULTS_IN]->(res:ProtocolDAGResultRef), + (res)-[:HAS_LOG]->(log:ProtocolDAGResultLog) + WHERE log.stream = '{stream}' + WITH log._scoped_key as sk + RETURN DISTINCT sk + """ + + return self._get_protocoldagresultlogs(q, task) + def set_task_status( self, tasks: list[ScopedKey], status: TaskStatusEnum, raise_error: bool = False ) -> list[ScopedKey | None]: From 34754a65e668e204a705d5fa30e70121de3e3daf Mon Sep 17 00:00:00 2001 From: David Dotson Date: Tue, 17 Feb 2026 19:42:01 -0700 Subject: [PATCH 2/2] Fix multiple bugs in stdout/stderr log capture implementation - interface/api.py: Rewrite get_task_logs endpoint; the original called n4js.get_scoped_key() with wrong argument types and constructed ScopedKey with invalid positional args. Now retrieves S3 locations directly from statestore and passes them to objectstore. Use http_status constant instead of bare 400. - compute/api.py: Use `is not None` checks instead of truthiness for stdout/stderr to preserve empty string logs. - statestore.py: Replace f-string interpolation in Cypher query with parameterized $stream to prevent injection. Return S3 locations directly instead of ScopedKeys to avoid unnecessary round-trips. - objectstore.py: Simplify pull_protocoldagresult_log to accept only location (always available from statestore). Move ProtocolDAGResultLog import to module level. - compute/service.py: Use _TeeStream for stderr capture so that logging output from the StreamHandler is both captured and still emitted to the original stderr, preventing log suppression during DAG execution. Co-Authored-By: Claude Opus 4.6 --- alchemiscale/compute/api.py | 4 +-- alchemiscale/compute/service.py | 23 +++++++++++-- alchemiscale/interface/api.py | 27 ++++----------- alchemiscale/storage/objectstore.py | 38 ++------------------- alchemiscale/storage/statestore.py | 52 +++++++++++------------------ 5 files changed, 51 insertions(+), 93 deletions(-) diff --git a/alchemiscale/compute/api.py b/alchemiscale/compute/api.py index 9d624003..12de2f19 100644 --- a/alchemiscale/compute/api.py +++ b/alchemiscale/compute/api.py @@ -394,7 +394,7 @@ async def set_task_result( ) # Store stdout and stderr logs if provided - if stdout: + if stdout is not None: stdout_log_ref = s3os.push_protocoldagresult_log( log_content=stdout, stream="stdout", @@ -405,7 +405,7 @@ async def set_task_result( ) n4js.set_protocoldagresult_log(result_sk, stdout_log_ref) - if stderr: + if stderr is not None: stderr_log_ref = s3os.push_protocoldagresult_log( log_content=stderr, stream="stderr", diff --git a/alchemiscale/compute/service.py b/alchemiscale/compute/service.py index 259dd41e..1a16826a 100644 --- a/alchemiscale/compute/service.py +++ b/alchemiscale/compute/service.py @@ -16,6 +16,22 @@ import io from contextlib import redirect_stdout, redirect_stderr + +class _TeeStream(io.TextIOBase): + """A stream wrapper that writes to both a capture buffer and the original stream.""" + + def __init__(self, capture: io.StringIO, original: io.TextIOBase): + self._capture = capture + self._original = original + + def write(self, s): + self._capture.write(s) + return self._original.write(s) + + def flush(self): + self._capture.flush() + self._original.flush() + from gufe import Transformation from gufe.protocols.protocoldag import execute_DAG, ProtocolDAG, ProtocolDAGResult @@ -245,12 +261,15 @@ def execute(self, task: ScopedKey) -> ScopedKey: self.logger.info("Executing '%s'...", protocoldag) - # Capture stdout and stderr during execution + # Capture stdout and stderr during execution. + # Use _TeeStream for stderr so that logging output (which goes to + # stderr via the StreamHandler) is both captured and still emitted. stdout_capture = io.StringIO() stderr_capture = io.StringIO() + stderr_tee = _TeeStream(stderr_capture, sys.stderr) try: - with redirect_stdout(stdout_capture), redirect_stderr(stderr_capture): + with redirect_stdout(stdout_capture), redirect_stderr(stderr_tee): protocoldagresult = execute_DAG( protocoldag, shared_basedir=shared, diff --git a/alchemiscale/interface/api.py b/alchemiscale/interface/api.py index ea3a7318..6fd954e4 100644 --- a/alchemiscale/interface/api.py +++ b/alchemiscale/interface/api.py @@ -1201,35 +1201,20 @@ def get_task_logs( """ if stream not in ["stdout", "stderr"]: raise HTTPException( - status_code=400, detail="stream must be 'stdout' or 'stderr'" + status_code=http_status.HTTP_400_BAD_REQUEST, + detail="stream must be 'stdout' or 'stderr'", ) sk = ScopedKey.from_str(task_scoped_key) validate_scopes(sk.scope, token) - # Get all log references for this task and stream - log_refs = n4js.get_task_logs(sk, stream=stream) + # Get all log S3 locations for this task and stream + log_locations = n4js.get_task_log_locations(sk, stream=stream) # Retrieve the actual log content from the object store logs = [] - for log_ref_sk in log_refs: - # Get the log reference node to extract location and other metadata - log_ref_data = n4js.get_scoped_key(log_ref_sk, resolve_gufe=False) - from ..storage.models import ProtocolDAGResultLog - - log_ref = ProtocolDAGResultLog._from_dict(log_ref_data) - - # Get transformation for this task - tf_sk, _ = n4js.get_task_transformation(task=task_scoped_key, return_gufe=False) - - # Pull the log content from S3 - log_content = s3os.pull_protocoldagresult_log( - protocoldagresult=ScopedKey(log_ref.scope, log_ref.obj_key), - transformation=tf_sk, - stream=stream, - ok=True, # We'll try both success and failure - location=log_ref.location, - ) + for location in log_locations: + log_content = s3os.pull_protocoldagresult_log(location=location) logs.append(log_content) return logs diff --git a/alchemiscale/storage/objectstore.py b/alchemiscale/storage/objectstore.py index 238c3c7c..e4eb8ae5 100644 --- a/alchemiscale/storage/objectstore.py +++ b/alchemiscale/storage/objectstore.py @@ -12,7 +12,7 @@ from gufe.tokenization import GufeKey from ..models import ScopedKey -from .models import ProtocolDAGResultRef +from .models import ProtocolDAGResultRef, ProtocolDAGResultLog from ..settings import S3ObjectStoreSettings # default filename for object store files @@ -330,8 +330,6 @@ def push_protocoldagresult_log( Reference to the log in the object store. """ - from ..storage.models import ProtocolDAGResultLog - ok = protocoldagresult_ok route = "results" if ok else "failures" @@ -360,26 +358,14 @@ def push_protocoldagresult_log( def pull_protocoldagresult_log( self, - protocoldagresult: ScopedKey, - transformation: ScopedKey, - stream: str, - ok: bool = True, - location: str | None = None, + location: str, ) -> str: """Pull the log content for a ProtocolDAGResult. Parameters ---------- - protocoldagresult - ScopedKey for ProtocolDAGResult in the object store. - transformation - The ScopedKey of the Transformation this log corresponds to. - stream - Either "stdout" or "stderr". - ok - ``True`` if ProtocolDAGResult completed successfully; ``False`` if failed. location - The full path in the object store to the log. If provided, this will be used. + The full path in the object store to the log file. Returns ------- @@ -387,23 +373,5 @@ def pull_protocoldagresult_log( The log content as a string. """ - route = "results" if ok else "failures" - - # build `location` based on provided ScopedKey if not provided - if location is None: - if transformation.scope != protocoldagresult.scope: - raise ValueError( - f"transformation scope '{transformation.scope}' differs from protocoldagresult scope '{protocoldagresult.scope}'" - ) - - location = os.path.join( - "protocoldagresult", - *protocoldagresult.scope.to_tuple(), - transformation.gufe_key, - route, - protocoldagresult.gufe_key, - f"{stream}.log", - ) - log_bytes = self._get_bytes(location) return log_bytes.decode("utf-8") diff --git a/alchemiscale/storage/statestore.py b/alchemiscale/storage/statestore.py index 5abbdf60..ed7c59c0 100644 --- a/alchemiscale/storage/statestore.py +++ b/alchemiscale/storage/statestore.py @@ -3594,20 +3594,10 @@ def set_protocoldagresult_log( return scoped_key - def _get_protocoldagresultlogs(self, q: str, scoped_key: ScopedKey): - """Helper method to retrieve ProtocolDAGResultLog ScopedKeys.""" - sks = [] - with self.transaction() as tx: - res = tx.run(q, scoped_key=str(scoped_key)) - for rec in res: - sks.append(rec["sk"]) - - return [ScopedKey.from_str(sk) for sk in sks] - - def get_task_logs( + def get_task_log_locations( self, task: ScopedKey, stream: str | None = None - ) -> list[ScopedKey]: - """Get all log references for a given Task. + ) -> list[str]: + """Get all log S3 locations for a given Task. Parameters ---------- @@ -3618,31 +3608,27 @@ def get_task_logs( Returns ------- - list[ScopedKey] - List of ScopedKeys for ProtocolDAGResultLog objects. + list[str] + List of S3 location strings for ProtocolDAGResultLog objects. """ if stream is not None and stream not in ["stdout", "stderr"]: raise ValueError("`stream` must be 'stdout', 'stderr', or None") - if stream is None: - q = """ - MATCH (task:Task {_scoped_key: $scoped_key}), - (task)-[:RESULTS_IN]->(res:ProtocolDAGResultRef), - (res)-[:HAS_LOG]->(log:ProtocolDAGResultLog) - WITH log._scoped_key as sk - RETURN DISTINCT sk - """ - else: - q = f""" - MATCH (task:Task {{_scoped_key: $scoped_key}}), - (task)-[:RESULTS_IN]->(res:ProtocolDAGResultRef), - (res)-[:HAS_LOG]->(log:ProtocolDAGResultLog) - WHERE log.stream = '{stream}' - WITH log._scoped_key as sk - RETURN DISTINCT sk - """ + q = """ + MATCH (task:Task {_scoped_key: $scoped_key}), + (task)-[:RESULTS_IN]->(res:ProtocolDAGResultRef), + (res)-[:HAS_LOG]->(log:ProtocolDAGResultLog) + WHERE ($stream IS NULL OR log.stream = $stream) + RETURN DISTINCT log.location as location + """ + + locations = [] + with self.transaction() as tx: + res = tx.run(q, scoped_key=str(task), stream=stream) + for rec in res: + locations.append(rec["location"]) - return self._get_protocoldagresultlogs(q, task) + return locations def set_task_status( self, tasks: list[ScopedKey], status: TaskStatusEnum, raise_error: bool = False