feat: add retry mechanism with wait_retry decorator and refactor diff generation (#572)

rdagent/scenarios/data_science/proposal/exp_gen.py

@@ -11,7 +11,8 @@ from rdagent.oai.llm_utils import APIBackend
 from rdagent.scenarios.data_science.experiment.experiment import COMPONENT, DSExperiment
 from rdagent.scenarios.data_science.scen import DataScienceScen
 from rdagent.utils.agent.tpl import T
-from rdagent.utils.repo.diff import generate_diff
+from rdagent.utils.repo.diff import generate_diff_from_dict
+from rdagent.utils.workflow import wait_retry
 
 
 class DSHypothesis(Hypothesis):
@@ -285,10 +286,10 @@ class DSExpGen(ExpGen):
                 exp=sota_exp, heading="Best of previous exploration of the scenario"
             )
             last_exp_diff = "\n".join(
-                generate_diff(
-                    sota_exp.experiment_workspace.workspace_path, last_exp.experiment_workspace.workspace_path
+                generate_diff_from_dict(
+                    sota_exp.experiment_workspace.file_dict, last_exp.experiment_workspace.file_dict
                 )
-            )
+            )  # we use file_dict for hitting the cache when replicate the experiment in another machine.
             exp_and_feedback_desc = T("scenarios.data_science.share:describe.feedback").r(
                 exp_and_feedback=exp_and_feedback
             )
@@ -348,40 +349,51 @@ class DSExpGen(ExpGen):
                     recent_trace_desc="\n".join(recent_trace_desc),
                 )
 
-                resp_dict = json.loads(
-                    APIBackend().build_messages_and_create_chat_completion(
-                        user_prompt=user_prompt, system_prompt=system_prompt, json_mode=True
-                    )
-                )
-                assert "hypothesis_proposal" in resp_dict, "Hypothesis proposal not provided."
-                assert "task_design" in resp_dict, "Task design not provided."
-                task_class = component_info["task_class"]
-                hypothesis_proposal = resp_dict.get("hypothesis_proposal", {})
-                hypothesis = DSHypothesis(
-                    component=component,
-                    hypothesis=hypothesis_proposal.get("hypothesis", ""),
-                    reason=hypothesis_proposal.get("reason", ""),
-                    concise_reason=hypothesis_proposal.get("concise_reason", ""),
-                    concise_observation=hypothesis_proposal.get("concise_observation", ""),
-                    concise_justification=hypothesis_proposal.get("concise_justification", ""),
-                    concise_knowledge=hypothesis_proposal.get("concise_knowledge", ""),
-                )
+                def _append_retry(args: tuple, kwargs: dict) -> tuple[tuple, dict]:
+                    # Only modify the user_prompt on retries (i > 0)
+                    user_prompt = args[0]
+                    user_prompt += "\n\nretrying..."
+                    return (user_prompt,), kwargs
 
-                task_design = resp_dict.get("task_design", {})
-                task_name = task_design["model_name"] if component == "Model" else component
-                description = task_design.get(
-                    "description", f"{component_info['target_name']} description not provided"
-                )
-                task = task_class(
-                    name=task_name,
-                    description=description,
-                    **{k: task_design.get(k, v) for k, v in component_info.get("extra_params", {}).items()},
-                )
+                @wait_retry(retry_n=5, transform_args_fn=_append_retry)
+                def _f(user_prompt):
+                    resp_dict = json.loads(
+                        APIBackend().build_messages_and_create_chat_completion(
+                            user_prompt=user_prompt, system_prompt=system_prompt, json_mode=True
+                        )
+                    )
+                    assert "hypothesis_proposal" in resp_dict, "Hypothesis proposal not provided."
+                    assert "task_design" in resp_dict, "Task design not provided."
+                    task_class = component_info["task_class"]
+                    hypothesis_proposal = resp_dict.get("hypothesis_proposal", {})
+                    hypothesis = DSHypothesis(
+                        component=component,
+                        hypothesis=hypothesis_proposal.get("hypothesis", ""),
+                        reason=hypothesis_proposal.get("reason", ""),
+                        concise_reason=hypothesis_proposal.get("concise_reason", ""),
+                        concise_observation=hypothesis_proposal.get("concise_observation", ""),
+                        concise_justification=hypothesis_proposal.get("concise_justification", ""),
+                        concise_knowledge=hypothesis_proposal.get("concise_knowledge", ""),
+                    )
+
+                    task_design = resp_dict.get("task_design", {})
+                    task_name = task_design["model_name"] if component == "Model" else component
+                    description = task_design.get(
+                        "description", f"{component_info['target_name']} description not provided"
+                    )
+                    task = task_class(
+                        name=task_name,
+                        description=description,
+                        **{k: task_design.get(k, v) for k, v in component_info.get("extra_params", {}).items()},
+                    )
+                    new_workflow_desc = resp_dict.get("workflow_update", "No update needed")
+                    return hypothesis, task, new_workflow_desc
+
+                hypothesis, task, new_workflow_desc = _f(user_prompt)
 
                 exp = DSExperiment(pending_tasks_list=[[task]], hypothesis=hypothesis)
                 exp.experiment_workspace.inject_code_from_folder(sota_exp.experiment_workspace.workspace_path)
 
-                new_workflow_desc = resp_dict.get("workflow_update", "No update needed")
                 if new_workflow_desc != "No update needed":
                     workflow_task = WorkflowTask(
                         name="Workflow",

rdagent/utils/repo/diff.py

@@ -1,49 +1,70 @@
 import difflib
+import fnmatch
 from pathlib import Path
-from typing import List
 
 
-def generate_diff(dir1: str, dir2: str) -> List[str]:
+# support two interfaces. add another one like  def generate_diff(file_dict1: dict, file_dict2: dict) -> List[str]:
+# the file_dict1 and file_dict2 are like {file_path: file_content}.
+# the two shuld share code.
+def generate_diff(dir1: str, dir2: str, file_pattern: str = "*.py") -> list[str]:
     """
-    Generate a diff between two directories(from dir1 to dir2), considering only .py files.
-    It is mocking `diff -durN dir1 dir2` in linux.
+    Generate a diff between two directories (from dir1 to dir2) using files that match the specified file pattern.
+    This function mimics the behavior of `diff -durN dir1 dir2` in Linux.
 
     Args:
         dir1 (str): Path to the first directory.
         dir2 (str): Path to the second directory.
+        file_pattern (str, optional): Glob pattern to filter files. Defaults to "*.py".
 
     Returns:
-        List[str]: A list of diffs for .py files that are different between the two directories.
+        list[str]: A list of diffs for files that differ between the two directories.
     """
 
-    diff_files = []
-
-    dir1_files = {f.relative_to(dir1) for f in Path(dir1).rglob("*.py") if f.is_file()}
-    dir2_files = {f.relative_to(dir2) for f in Path(dir2).rglob("*.py") if f.is_file()}
+    dir1_files = {f.relative_to(dir1) for f in Path(dir1).rglob("*") if f.is_file()}
+    dir2_files = {f.relative_to(dir2) for f in Path(dir2).rglob("*") if f.is_file()}
 
     all_files = dir1_files.union(dir2_files)
-
+    file_dict1 = {}
+    file_dict2 = {}
     for file in all_files:
         file1 = Path(dir1) / file
         file2 = Path(dir2) / file
-
-        if file1.exists() and file2.exists():
-            with file1.open() as f1, file2.open() as f2:
-                diff = list(difflib.unified_diff(f1.readlines(), f2.readlines(), fromfile=str(file), tofile=str(file)))
-                if diff:
-                    diff_files.extend(diff)
+        if file1.exists():
+            with file1.open() as f1:
+                file_dict1[str(file)] = f1.read()
         else:
-            if file1.exists():
-                with file1.open() as f1:
-                    diff = list(
-                        difflib.unified_diff(f1.readlines(), [], fromfile=str(file), tofile=str(file) + " (empty file)")
-                    )
-                    diff_files.extend(diff)
-            elif file2.exists():
-                with file2.open() as f2:
-                    diff = list(
-                        difflib.unified_diff([], f2.readlines(), fromfile=str(file) + " (empty file)", tofile=str(file))
-                    )
-                    diff_files.extend(diff)
+            file_dict1[str(file)] = ""
+        if file2.exists():
+            with file2.open() as f2:
+                file_dict2[str(file)] = f2.read()
+        else:
+            file_dict2[str(file)] = ""
+    return generate_diff_from_dict(file_dict1, file_dict2, file_pattern=file_pattern)
 
+
+def generate_diff_from_dict(file_dict1: dict, file_dict2: dict, file_pattern: str = "*.py") -> list[str]:
+    """
+    Generate a diff between two dictionaries of file contents.
+    The dictionaries should be of the format {file_path: file_content}.
+
+    Returns:
+        List[str]: A list of diffs for files that are different between the two dictionaries.
+    """
+    diff_files = []
+    all_files = set(file_dict1.keys()).union(file_dict2.keys())
+    for file in sorted(all_files):
+        if not fnmatch.fnmatch(file, file_pattern):
+            continue
+        content1 = file_dict1.get(file, "")
+        content2 = file_dict2.get(file, "")
+        diff = list(
+            difflib.unified_diff(
+                content1.splitlines(keepends=True),
+                content2.splitlines(keepends=True),
+                fromfile=file if file in file_dict1 else file + " (empty file)",
+                tofile=file if file in file_dict2 else file + " (empty file)",
+            )
+        )
+        if diff:
+            diff_files.extend(diff)
     return diff_files

rdagent/utils/workflow.py

@@ -10,10 +10,11 @@ Postscripts:
 
 import datetime
 import pickle
+import time
 from collections import defaultdict
 from dataclasses import dataclass, field
 from pathlib import Path
-from typing import Callable
+from typing import Any, Callable
 
 from tqdm.auto import tqdm
 
@@ -160,3 +161,48 @@ class LoopBase:
         max_loop = max(session.loop_trace.keys())
         logger.storage.truncate(time=session.loop_trace[max_loop][-1].end)
         return session
+
+
+def wait_retry(
+    retry_n: int = 3, sleep_time: int = 1, transform_args_fn: Callable[[tuple, dict], tuple[tuple, dict]] | None = None
+) -> Callable[[Callable[..., Any]], Callable[..., Any]]:
+    """Decorator to wait and retry the function for retry_n times.
+
+    Example:
+    >>> import time
+    >>> @wait_retry(retry_n=2, sleep_time=1)
+    ... def test_func():
+    ...     global counter
+    ...     counter += 1
+    ...     if counter < 3:
+    ...         raise ValueError("Counter is less than 3")
+    ...     return counter
+    >>> counter = 0
+    >>> try:
+    ...     test_func()
+    ... except ValueError as e:
+    ...     print(f"Caught an exception: {e}")
+    Error: Counter is less than 3
+    Error: Counter is less than 3
+    Caught an exception: Counter is less than 3
+    >>> counter
+    2
+    """
+
+    def decorator(f: Callable[..., Any]) -> Callable[..., Any]:
+        def wrapper(*args: Any, **kwargs: Any) -> Any:
+            for i in range(retry_n):
+                try:
+                    return f(*args, **kwargs)
+                except Exception as e:
+                    print(f"Error: {e}")
+                    time.sleep(sleep_time)
+                    if i == retry_n - 1:
+                        raise e
+                    # Update args and kwargs using the transform function if provided.
+                    if transform_args_fn is not None:
+                        args, kwargs = transform_args_fn(args, kwargs)
+
+        return wrapper
+
+    return decorator