Add runtime heartbeat and stall detection

Author: LIHUA919

research_agent/run_infer_idea.py

@@ -79,6 +79,12 @@ async def forward(self, instance_path: str, task_level: str, local_root: str, wo
         metadata = self.load_ins({"instance_path": instance_path, "task_level": task_level})
         runtime = MasterRuntime(self.cache_path)
         runtime.sync_stage_state()
+        run_id = metadata.get("instance_id", task_level)
+        runtime.write_runtime_status(
+            run_id=run_id,
+            status="running",
+            metadata={"entrypoint": "run_infer_idea", "task_level": task_level},
+        )
         stage_state = runtime.load_state()
         context_variables = {
             "working_dir": f"/{workplace_name}", # TODO: change to the codebase path
@@ -634,6 +640,11 @@ async def forward(self, instance_path: str, task_level: str, local_root: str, wo
         print(refine_res)
         runtime.sync_stage_state()
         goal_evaluation = runtime.evaluate_goal()
+        runtime.write_runtime_status(
+            run_id=run_id,
+            status="completed" if goal_evaluation.all_criteria_met else "running",
+            metadata={"entrypoint": "run_infer_idea", "task_level": task_level},
+        )
         return {
             "task_id": metadata.get("instance_id", task_level),
             "query": idea_query,
@@ -721,9 +732,18 @@ def main(args, references):
     web_env = BrowserEnv(browsergym_eval_env = None, local_root=env_config.local_root, workplace_name=env_config.workplace_name)
     file_env = RequestsMarkdownBrowser(viewport_size=1024 * 4, local_root=env_config.local_root, workplace_name=env_config.workplace_name, downloads_folder=os.path.join(env_config.local_root, env_config.workplace_name, "downloads"))
     flow = InnoFlow(cache_path=cache_path, log_path="log_" + instance_id, code_env=code_env, web_env=web_env, file_env=file_env, model=args.model)
-    # ml_result = await flow(instance_path=instance_path)
-    result = asyncio.run(flow(instance_path=args.instance_path, task_level=args.task_level, local_root=local_root, workplace_name=args.workplace_name, max_iter_times=args.max_iter_times, category=args.category, references = references))
-    return build_and_save_eval_result(result, cache_path)
+    runtime = MasterRuntime(cache_path)
+    try:
+        result = asyncio.run(flow(instance_path=args.instance_path, task_level=args.task_level, local_root=local_root, workplace_name=args.workplace_name, max_iter_times=args.max_iter_times, category=args.category, references = references))
+        return build_and_save_eval_result(result, cache_path)
+    except Exception as exc:
+        runtime.write_failure_status(
+            run_id=instance_id,
+            error_message=str(exc),
+            stage_name=runtime.next_stage(),
+            metadata={"entrypoint": "run_infer_idea", "task_level": args.task_level},
+        )
+        raise
     # print(judge_result)
 
 

research_agent/run_infer_plan.py

@@ -77,6 +77,12 @@ async def forward(self, instance_path: str, task_level: str, local_root: str, wo
         metadata = self.load_ins({"instance_path": instance_path, "task_level": task_level})
         runtime = MasterRuntime(self.cache_path)
         runtime.sync_stage_state()
+        run_id = metadata.get("instance_id", task_level)
+        runtime.write_runtime_status(
+            run_id=run_id,
+            status="running",
+            metadata={"entrypoint": "run_infer_plan", "task_level": task_level},
+        )
         stage_state = runtime.load_state()
         context_variables = {
             "working_dir": workplace_name,  # Agent instructions already prepend "/"
@@ -603,6 +609,11 @@ async def forward(self, instance_path: str, task_level: str, local_root: str, wo
 
         runtime.sync_stage_state()
         goal_evaluation = runtime.evaluate_goal()
+        runtime.write_runtime_status(
+            run_id=run_id,
+            status="completed" if goal_evaluation.all_criteria_met else "running",
+            metadata={"entrypoint": "run_infer_plan", "task_level": task_level},
+        )
         return {
             "task_id": metadata.get("instance_id", task_level),
             "query": plan_query,
@@ -691,9 +702,18 @@ def main(args, ideas, references):
     web_env = BrowserEnv(browsergym_eval_env = None, local_root=env_config.local_root, workplace_name=env_config.workplace_name)
     file_env = RequestsMarkdownBrowser(viewport_size=1024 * 4, local_root=env_config.local_root, workplace_name=env_config.workplace_name, downloads_folder=os.path.join(env_config.local_root, env_config.workplace_name, "downloads"))
     flow = InnoFlow(cache_path=cache_path, log_path="log_" + instance_id, code_env=code_env, web_env=web_env, file_env=file_env, model=args.model)
-    # ml_result = await flow(instance_path=instance_path)
-    result = asyncio.run(flow(instance_path=args.instance_path, task_level=args.task_level, local_root=local_root, workplace_name=args.workplace_name, max_iter_times=args.max_iter_times, category=args.category, ideas = ideas, references = references))
-    return build_and_save_eval_result(result, cache_path)
+    runtime = MasterRuntime(cache_path)
+    try:
+        result = asyncio.run(flow(instance_path=args.instance_path, task_level=args.task_level, local_root=local_root, workplace_name=args.workplace_name, max_iter_times=args.max_iter_times, category=args.category, ideas = ideas, references = references))
+        return build_and_save_eval_result(result, cache_path)
+    except Exception as exc:
+        runtime.write_failure_status(
+            run_id=instance_id,
+            error_message=str(exc),
+            stage_name=runtime.next_stage(),
+            metadata={"entrypoint": "run_infer_plan", "task_level": args.task_level},
+        )
+        raise
     # print(judge_result)
 
 

research_agent/runtime/__init__.py

@@ -3,6 +3,12 @@
     StageCriteria,
     validate_stage_artifacts,
 )
+from research_agent.runtime.heartbeat import (
+    parse_runtime_timestamp,
+    read_runtime_json,
+    write_heartbeat,
+    write_run_status,
+)
 from research_agent.runtime.master import GoalEvaluation, MasterRuntime
 
 __all__ = [
@@ -11,4 +17,8 @@
     "MasterRuntime",
     "StageCriteria",
     "validate_stage_artifacts",
+    "parse_runtime_timestamp",
+    "read_runtime_json",
+    "write_heartbeat",
+    "write_run_status",
 ]

research_agent/runtime/heartbeat.py

@@ -0,0 +1,87 @@
+from __future__ import annotations
+
+import json
+import os
+from datetime import datetime, timezone
+from typing import Any
+
+
+def utc_now_iso() -> str:
+    return datetime.now(timezone.utc).isoformat()
+
+
+def write_runtime_json(cache_path: str, file_name: str, payload: dict[str, Any]) -> str:
+    os.makedirs(cache_path, exist_ok=True)
+    output_path = os.path.join(cache_path, file_name)
+    with open(output_path, "w", encoding="utf-8") as f:
+        json.dump(payload, f, ensure_ascii=False, indent=4)
+    return output_path
+
+
+def read_runtime_json(cache_path: str, file_name: str) -> dict[str, Any]:
+    output_path = os.path.join(cache_path, file_name)
+    if not os.path.exists(output_path):
+        return {}
+    try:
+        with open(output_path, "r", encoding="utf-8") as f:
+            payload = json.load(f)
+        return payload if isinstance(payload, dict) else {}
+    except (OSError, json.JSONDecodeError):
+        return {}
+
+
+def parse_runtime_timestamp(value: str | None) -> datetime | None:
+    if not value:
+        return None
+    try:
+        return datetime.fromisoformat(value)
+    except ValueError:
+        return None
+
+
+def write_heartbeat(
+    cache_path: str,
+    *,
+    current_stage: str | None,
+    status: str,
+    last_error: str | None = None,
+    metadata: dict[str, Any] | None = None,
+) -> str:
+    return write_runtime_json(
+        cache_path,
+        "heartbeat.json",
+        {
+            "current_stage": current_stage,
+            "status": status,
+            "last_error": last_error,
+            "updated_at": utc_now_iso(),
+            "metadata": metadata or {},
+        },
+    )
+
+
+def write_run_status(
+    cache_path: str,
+    *,
+    run_id: str,
+    status: str,
+    current_stage: str | None,
+    completed_stages: list[str],
+    incomplete_stages: list[str],
+    latest_artifact: str | None = None,
+    last_error: str | None = None,
+) -> str:
+    return write_runtime_json(
+        cache_path,
+        "run_status.json",
+        {
+            "run_id": run_id,
+            "status": status,
+            "current_stage": current_stage,
+            "completed_stages": completed_stages,
+            "incomplete_stages": incomplete_stages,
+            "latest_artifact": latest_artifact,
+            "last_error": last_error,
+            "updated_at": utc_now_iso(),
+        },
+    )

research_agent/runtime/master.py

@@ -1,6 +1,7 @@
 from __future__ import annotations
 
 from dataclasses import dataclass
+from datetime import datetime, timezone
 
 from research_agent.inno_common import load_stage_state, update_stage_state
 from research_agent.runtime.criteria import (
@@ -9,6 +10,12 @@
     StageCriteria,
     validate_stage_artifacts,
 )
+from research_agent.runtime.heartbeat import (
+    parse_runtime_timestamp,
+    read_runtime_json,
+    write_heartbeat,
+    write_run_status,
+)
 
 
 @dataclass(frozen=True)
@@ -19,6 +26,14 @@ class GoalEvaluation:
     all_criteria_met: bool
 
 
+@dataclass(frozen=True)
+class StallEvaluation:
+    stalled: bool
+    current_stage: str | None
+    age_seconds: float | None
+    reason: str
+
+
 class MasterRuntime:
     def __init__(
         self,
@@ -33,6 +48,22 @@ def __init__(
     def load_state(self) -> dict:
         return load_stage_state(self.cache_path)
 
+    def load_heartbeat(self) -> dict:
+        return read_runtime_json(self.cache_path, "heartbeat.json")
+
+    def load_run_status(self) -> dict:
+        return read_runtime_json(self.cache_path, "run_status.json")
+
+    def latest_artifact(self) -> str | None:
+        state = self.load_state()
+        for stage_name in reversed(self.stage_order):
+            stage_state = state.get(stage_name, {})
+            artifacts = stage_state.get("artifacts") or {}
+            for artifact_path in artifacts.values():
+                if artifact_path:
+                    return artifact_path
+        return None
+
     def stage_index(self, stage_name: str) -> int:
         return self.stage_order.index(stage_name)
 
@@ -98,3 +129,126 @@ def record_stage_completion(
             metadata=metadata or {},
         )
         return self.sync_stage_state()
+
+    def record_stage_failure(
+        self,
+        stage_name: str,
+        error_message: str,
+        metadata: dict | None = None,
+    ) -> dict:
+        state = self.load_state()
+        stage_state = state.get(stage_name, {})
+        merged_metadata = dict(stage_state.get("metadata", {}))
+        merged_metadata.update(metadata or {})
+        merged_metadata["last_error"] = error_message
+        update_stage_state(
+            self.cache_path,
+            stage_name,
+            "failed",
+            artifacts=stage_state.get("artifacts", {}),
+            metadata=merged_metadata,
+        )
+        return self.load_state()
+
+    def write_runtime_status(
+        self,
+        *,
+        run_id: str,
+        status: str,
+        last_error: str | None = None,
+        metadata: dict | None = None,
+    ) -> dict[str, str]:
+        goal = self.evaluate_goal()
+        current_stage = goal.current_stage
+        heartbeat_path = write_heartbeat(
+            self.cache_path,
+            current_stage=current_stage,
+            status=status,
+            last_error=last_error,
+            metadata=metadata,
+        )
+        run_status_path = write_run_status(
+            self.cache_path,
+            run_id=run_id,
+            status=status,
+            current_stage=current_stage,
+            completed_stages=goal.completed_stages,
+            incomplete_stages=goal.incomplete_stages,
+            latest_artifact=self.latest_artifact(),
+            last_error=last_error,
+        )
+        return {
+            "heartbeat": heartbeat_path,
+            "run_status": run_status_path,
+        }
+
+    def write_failure_status(
+        self,
+        *,
+        run_id: str,
+        error_message: str,
+        stage_name: str | None = None,
+        metadata: dict | None = None,
+    ) -> dict[str, str]:
+        if stage_name:
+            self.record_stage_failure(stage_name, error_message, metadata=metadata)
+        return self.write_runtime_status(
+            run_id=run_id,
+            status="failed",
+            last_error=error_message,
+            metadata=metadata,
+        )
+
+    def evaluate_stall(
+        self,
+        *,
+        max_stale_seconds: int,
+        now: datetime | None = None,
+    ) -> StallEvaluation:
+        heartbeat = self.load_heartbeat()
+        if not heartbeat:
+            return StallEvaluation(False, None, None, "missing_heartbeat")
+
+        status = heartbeat.get("status")
+        current_stage = heartbeat.get("current_stage")
+        if status != "running":
+            return StallEvaluation(False, current_stage, None, f"status_{status or 'unknown'}")
+
+        updated_at = parse_runtime_timestamp(heartbeat.get("updated_at"))
+        if updated_at is None:
+            return StallEvaluation(False, current_stage, None, "invalid_timestamp")
+
+        effective_now = now or datetime.now(timezone.utc)
+        age_seconds = max((effective_now - updated_at).total_seconds(), 0.0)
+        if age_seconds > max_stale_seconds:
+            return StallEvaluation(True, current_stage, age_seconds, "heartbeat_stale")
+
+        return StallEvaluation(False, current_stage, age_seconds, "healthy")
+
+    def write_stalled_status(
+        self,
+        *,
+        run_id: str,
+        reason: str,
+        stage_name: str | None = None,
+        metadata: dict | None = None,
+    ) -> dict[str, str]:
+        if stage_name:
+            state = self.load_state()
+            stage_state = state.get(stage_name, {})
+            merged_metadata = dict(stage_state.get("metadata", {}))
+            merged_metadata.update(metadata or {})
+            merged_metadata["stalled_reason"] = reason
+            update_stage_state(
+                self.cache_path,
+                stage_name,
+                "stalled",
+                artifacts=stage_state.get("artifacts", {}),
+                metadata=merged_metadata,
+            )
+        return self.write_runtime_status(
+            run_id=run_id,
+            status="stalled",
+            last_error=reason,
+            metadata=metadata,
+        )