diff --git a/alchemiscale/compute/api.py b/alchemiscale/compute/api.py index 56adf285..12de2f19 100644 --- a/alchemiscale/compute/api.py +++ b/alchemiscale/compute/api.py @@ -366,6 +366,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) @@ -391,6 +393,29 @@ async def set_task_result( task=task_sk, protocoldagresultref=protocoldagresultref ) + # Store stdout and stderr logs if provided + if stdout is not None: + 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 is not None: + 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..1a16826a 100644 --- a/alchemiscale/compute/service.py +++ b/alchemiscale/compute/service.py @@ -12,6 +12,25 @@ import threading from pathlib import Path import shutil +import sys +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 @@ -198,16 +217,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 +260,24 @@ def execute(self, task: ScopedKey) -> ScopedKey: scratch.mkdir() self.logger.info("Executing '%s'...", protocoldag) + + # 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: - 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_tee): + 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 +285,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 +301,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 341069e9..6fd954e4 100644 --- a/alchemiscale/interface/api.py +++ b/alchemiscale/interface/api.py @@ -1176,6 +1176,50 @@ 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=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 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 location in log_locations: + log_content = s3os.pull_protocoldagresult_log(location=location) + logs.append(log_content) + + return logs + + ### strategies diff --git a/alchemiscale/interface/client.py b/alchemiscale/interface/client.py index f53859ba..9952dd8e 100644 --- a/alchemiscale/interface/client.py +++ b/alchemiscale/interface/client.py @@ -1973,6 +1973,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..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 @@ -297,3 +297,81 @@ 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. + + """ + 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, + location: str, + ) -> str: + """Pull the log content for a ProtocolDAGResult. + + Parameters + ---------- + location + The full path in the object store to the log file. + + Returns + ------- + str + The log content as a string. + + """ + 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 343bdcde..ed7c59c0 100644 --- a/alchemiscale/storage/statestore.py +++ b/alchemiscale/storage/statestore.py @@ -3547,6 +3547,89 @@ 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_task_log_locations( + self, task: ScopedKey, stream: str | None = None + ) -> list[str]: + """Get all log S3 locations for a given Task. + + Parameters + ---------- + task + ScopedKey of the Task. + stream + Optional filter for "stdout" or "stderr". If None, returns all logs. + + Returns + ------- + 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") + + 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 locations + def set_task_status( self, tasks: list[ScopedKey], status: TaskStatusEnum, raise_error: bool = False ) -> list[ScopedKey | None]: