feat: fallback to acceptable results (#1129)

* refactor: add is_acceptable, fallback logic and generify evolving agent

* refine lint

* small

* lint

* lint

* lint

* feat: add is_acceptable to CoSTEERMultiFeedback

* feat: add in-memory workspace checkpoint and recovery

* feat: preserve symbolic links in workspace checkpoints and recovery

* lint

* lint

* feat: limit workspace checkpoint to files under 100KB

* feat: add workspace checkpoint size limit setting

* prompt

* lint

rdagent/core/conf.py

@@ -55,6 +55,11 @@ class RDAgentSettings(ExtendedBaseSettings):
 
     # workspace conf
     workspace_path: Path = Path.cwd() / "git_ignore_folder" / "RD-Agent_workspace"
+    workspace_ckp_size_limit: int = 0
+    """
+    the checkpoint for the workspace is a zip file.
+    0 (or any value <=0) means *no* size limit for files in workspace checkpoints
+    """
 
     # multi processing conf
     multi_proc_n: int = 1

rdagent/core/evaluation.py

@@ -12,6 +12,13 @@ class Feedback:
         The building process of feedback will should be in evaluator
     """
 
+    def is_acceptable(self) -> bool:
+        """
+        Sometimes, the solution is already acceptable, but we still want to refine it.
+        So we use different logic to determine whether the solution is acceptable or finished.
+        """
+        return self.__bool__()
+
     def finished(self) -> bool:
         """
         In some implementations, tasks may fail multiple times, leading agents to skip the implementation.

rdagent/core/evolving_agent.py

@@ -2,24 +2,21 @@ from __future__ import annotations
 
 from abc import ABC, abstractmethod
 from collections.abc import Generator
-from typing import TYPE_CHECKING, Any, Generic, TypeVar
+from contextlib import nullcontext
+from typing import Any, Generic, TypeVar
 
 from filelock import FileLock
 from tqdm import tqdm
 
-if TYPE_CHECKING:
-    from rdagent.core.evolving_framework import EvolvableSubjects
-
-from contextlib import nullcontext
-
 from rdagent.core.evaluation import EvaluableObj, Evaluator, Feedback
-from rdagent.core.evolving_framework import EvolvingStrategy, EvoStep
+from rdagent.core.evolving_framework import EvolvableSubjects, EvolvingStrategy, EvoStep
 from rdagent.log import rdagent_logger as logger
 
 ASpecificEvaluator = TypeVar("ASpecificEvaluator", bound=Evaluator)
+ASpecificEvolvableSubjects = TypeVar("ASpecificEvolvableSubjects", bound=EvolvableSubjects)
 
 
-class EvoAgent(ABC, Generic[ASpecificEvaluator]):
+class EvoAgent(ABC, Generic[ASpecificEvaluator, ASpecificEvolvableSubjects]):
 
     def __init__(self, max_loop: int, evolving_strategy: EvolvingStrategy) -> None:
         self.max_loop = max_loop
@@ -28,9 +25,9 @@ class EvoAgent(ABC, Generic[ASpecificEvaluator]):
     @abstractmethod
     def multistep_evolve(
         self,
-        evo: EvolvableSubjects,
+        evo: ASpecificEvolvableSubjects,
         eva: ASpecificEvaluator | Feedback,
-    ) -> Generator[EvolvableSubjects, None, None]:
+    ) -> Generator[ASpecificEvolvableSubjects, None, None]:
         """
         yield EvolvableSubjects for caller for easier process control and logging.
         """
@@ -47,7 +44,7 @@ class RAGEvaluator(Evaluator):
         raise NotImplementedError
 
 
-class RAGEvoAgent(EvoAgent[RAGEvaluator]):
+class RAGEvoAgent(EvoAgent[RAGEvaluator, ASpecificEvolvableSubjects], Generic[ASpecificEvolvableSubjects]):
 
     def __init__(
         self,
@@ -63,7 +60,7 @@ class RAGEvoAgent(EvoAgent[RAGEvaluator]):
     ) -> None:
         super().__init__(max_loop, evolving_strategy)
         self.rag = rag
-        self.evolving_trace: list[EvoStep] = []
+        self.evolving_trace: list[EvoStep[ASpecificEvolvableSubjects]] = []
         self.with_knowledge = with_knowledge
         self.with_feedback = with_feedback
         self.knowledge_self_gen = knowledge_self_gen
@@ -72,9 +69,9 @@ class RAGEvoAgent(EvoAgent[RAGEvaluator]):
 
     def multistep_evolve(
         self,
-        evo: EvolvableSubjects,
+        evo: ASpecificEvolvableSubjects,
         eva: RAGEvaluator | Feedback,
-    ) -> Generator[EvolvableSubjects, None, None]:
+    ) -> Generator[ASpecificEvolvableSubjects, None, None]:
         for evo_loop_id in tqdm(range(self.max_loop), "Implementing"):
             with logger.tag(f"evo_loop_{evo_loop_id}"):
                 # 1. RAG
@@ -91,7 +88,7 @@ class RAGEvoAgent(EvoAgent[RAGEvaluator]):
                 )
 
                 # 3. Pack evolve results
-                es = EvoStep(evo, queried_knowledge)
+                es = EvoStep[ASpecificEvolvableSubjects](evo, queried_knowledge)
 
                 # 4. Evaluation
                 if self.with_feedback:

rdagent/core/evolving_framework.py

@@ -3,7 +3,7 @@ from __future__ import annotations
 import copy
 from abc import ABC, abstractmethod
 from dataclasses import dataclass
-from typing import TYPE_CHECKING, Any
+from typing import TYPE_CHECKING, Any, Generic, TypeVar
 
 from rdagent.core.evaluation import EvaluableObj
 from rdagent.core.knowledge_base import KnowledgeBase
@@ -36,8 +36,11 @@ class EvolvableSubjects(EvaluableObj):
         return copy.deepcopy(self)
 
 
+ASpecificEvolvableSubjects = TypeVar("ASpecificEvolvableSubjects", bound=EvolvableSubjects)
+
+
 @dataclass
-class EvoStep:
+class EvoStep(Generic[ASpecificEvolvableSubjects]):
     """At a specific step,
     based on
     - previous trace
@@ -48,23 +51,24 @@ class EvoStep:
     (optional) After evaluation, we get feedback `feedback`.
     """
 
-    evolvable_subjects: EvolvableSubjects
+    evolvable_subjects: ASpecificEvolvableSubjects
+
     queried_knowledge: QueriedKnowledge | None = None
     feedback: Feedback | None = None
 
 
-class EvolvingStrategy(ABC):
+class EvolvingStrategy(ABC, Generic[ASpecificEvolvableSubjects]):
     def __init__(self, scen: Scenario) -> None:
         self.scen = scen
 
     @abstractmethod
     def evolve(
         self,
-        *evo: EvolvableSubjects,
-        evolving_trace: list[EvoStep] | None = None,
+        *evo: ASpecificEvolvableSubjects,
+        evolving_trace: list[EvoStep[ASpecificEvolvableSubjects]] | None = None,
         queried_knowledge: QueriedKnowledge | None = None,
         **kwargs: Any,
-    ) -> EvolvableSubjects:
+    ) -> ASpecificEvolvableSubjects:
         """The evolving trace is a list of (evolvable_subjects, feedback) ordered
         according to the time.
 
@@ -74,7 +78,7 @@ class EvolvingStrategy(ABC):
         """
 
 
-class RAGStrategy(ABC):
+class RAGStrategy(ABC, Generic[ASpecificEvolvableSubjects]):
     """Retrieval Augmentation Generation Strategy"""
 
     def __init__(self, *args: Any, **kwargs: Any) -> None:
@@ -91,7 +95,7 @@ class RAGStrategy(ABC):
     @abstractmethod
     def query(
         self,
-        evo: EvolvableSubjects,
+        evo: ASpecificEvolvableSubjects,
         evolving_trace: list[EvoStep],
         **kwargs: Any,
     ) -> QueriedKnowledge | None:
@@ -100,7 +104,7 @@ class RAGStrategy(ABC):
     @abstractmethod
     def generate_knowledge(
         self,
-        evolving_trace: list[EvoStep],
+        evolving_trace: list[EvoStep[ASpecificEvolvableSubjects]],
         *,
         return_knowledge: bool = False,
         **kwargs: Any,

rdagent/core/experiment.py

@@ -1,11 +1,13 @@
 from __future__ import annotations
 
+import io
 import os
 import platform
 import re
 import shutil
 import typing
 import uuid
+import zipfile
 from abc import ABC, abstractmethod
 from collections.abc import Sequence
 from copy import deepcopy
@@ -98,6 +100,19 @@ class Workspace(ABC, Generic[ASpecificTask, ASpecificFeedback]):
         Get all the code files in the workspace as a single string.
         """
 
+    # when the workspace is mutable inplace, provide support for creating checkpoints and recovering.
+    @abstractmethod
+    def create_ws_ckp(self) -> None:
+        """
+        Create an in-memory checkpoint of the workspace so it can be restored later.
+        """
+
+    @abstractmethod
+    def recover_ws_ckp(self) -> None:
+        """
+        Restore the workspace from the checkpoint created by :py:meth:`create_ws_ckp`.
+        """
+
 
 ASpecificWS = TypeVar("ASpecificWS", bound=Workspace)
 
@@ -138,6 +153,8 @@ class FBWorkspace(Workspace):
             {}
         )  # The code injected into the folder, store them in the variable to reproduce the former result
         self.workspace_path: Path = RD_AGENT_SETTINGS.workspace_path / uuid.uuid4().hex
+        # In-memory checkpoint data created by ``create_ws_ckp``.
+        self.ws_ckp: bytes | None = None
 
     @staticmethod
     def _format_code_dict(code_dict: dict[str, str]) -> str:
@@ -282,6 +299,58 @@ class FBWorkspace(Workspace):
         )
         return result
 
+    def create_ws_ckp(self) -> None:
+        """
+        Zip the contents of ``workspace_path`` and persist the archive on
+        ``self.ws_ckp`` for later restoration via :py:meth:`recover_ws_ckp`.
+        """
+        buf = io.BytesIO()
+        with zipfile.ZipFile(buf, "w", zipfile.ZIP_DEFLATED) as zf:
+            for file_path in self.workspace_path.rglob("*"):
+                # Only include regular files up to 100 KB so that the checkpoint
+                # remains lightweight. Larger files (for example, datasets) are
+                # expected to be recreated or mounted separately.
+                if file_path.is_symlink():
+                    # Preserve symbolic links within the archive
+                    zi = zipfile.ZipInfo(str(file_path.relative_to(self.workspace_path)))
+                    zi.create_system = 3  # indicates Unix
+                    zi.external_attr = 0o120777 << 16  # symlink file type + 0777 perms
+                    zf.writestr(zi, str(file_path.readlink()))
+                elif file_path.is_file():
+                    size_limit = RD_AGENT_SETTINGS.workspace_ckp_size_limit
+                    if size_limit <= 0 or file_path.stat().st_size <= size_limit:
+                        zf.write(file_path, file_path.relative_to(self.workspace_path))
+        self.ws_ckp = buf.getvalue()
+
+    def recover_ws_ckp(self) -> None:
+        """
+        Restore the workspace directory from the in-memory checkpoint created by
+        :py:meth:`create_ws_ckp`.
+        """
+        if self.ws_ckp is None:
+            msg = "Workspace checkpoint doesn't exist. Call `create_ws_ckp` first."
+            raise RuntimeError(msg)
+        shutil.rmtree(self.workspace_path, ignore_errors=True)
+        self.workspace_path.mkdir(parents=True, exist_ok=True)
+        buf = io.BytesIO(self.ws_ckp)
+        with zipfile.ZipFile(buf, "r") as zf:
+            for info in zf.infolist():
+                dest_path = self.workspace_path / info.filename
+                # File type bits (upper 4) are in high 16 bits of external_attr
+                mode = (info.external_attr >> 16) & 0o170000
+                symlink_mode = 0o120000  # Constant for symlink file type in Unix
+                if mode == symlink_mode:  # Symlink
+                    dest_path.parent.mkdir(parents=True, exist_ok=True)
+                    link_target = zf.read(info).decode()
+                    os.symlink(link_target, dest_path)
+                else:
+                    if info.is_dir():
+                        dest_path.mkdir(parents=True, exist_ok=True)
+                    else:
+                        dest_path.parent.mkdir(parents=True, exist_ok=True)
+                        with dest_path.open("wb") as f:
+                            f.write(zf.read(info))
+
     def __str__(self) -> str:
         return f"Workspace[{self.workspace_path=}" + (
             "]" if self.target_task is None else f",{self.target_task.name=}]"
@@ -355,6 +424,21 @@ class Experiment(
     def result(self, value: object) -> None:
         self.running_info.result = value
 
+    # when the workspace is mutable inplace, provide support for creating checkpoints and recovering.
+    def create_ws_ckp(self) -> None:
+        if self.experiment_workspace is not None:
+            self.experiment_workspace.create_ws_ckp()
+        for ws in self.sub_workspace_list:
+            if ws is not None:
+                ws.create_ws_ckp()
+
+    def recover_ws_ckp(self) -> None:
+        if self.experiment_workspace is not None:
+            self.experiment_workspace.recover_ws_ckp()
+        for ws in self.sub_workspace_list:
+            if ws is not None:
+                ws.recover_ws_ckp()
+
 
 ASpecificExp = TypeVar("ASpecificExp", bound=Experiment)
 ASpecificPlan = TypeVar("ASpecificPlan", bound=ExperimentPlan)