From 0db49991cb0b760b1b69ce78d08cb88d2c55c0e4 Mon Sep 17 00:00:00 2001 From: you-n-g Date: Mon, 10 Feb 2025 01:35:54 +0800 Subject: [PATCH] feat: add retry mechanism with wait_retry decorator and refactor diff generation (#572) --- .../data_science/proposal/exp_gen.py | 78 +++++++++++-------- rdagent/utils/repo/diff.py | 77 +++++++++++------- rdagent/utils/workflow.py | 48 +++++++++++- 3 files changed, 141 insertions(+), 62 deletions(-) diff --git a/rdagent/scenarios/data_science/proposal/exp_gen.py b/rdagent/scenarios/data_science/proposal/exp_gen.py index 5d149c04..c77702e6 100644 --- a/rdagent/scenarios/data_science/proposal/exp_gen.py +++ b/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", diff --git a/rdagent/utils/repo/diff.py b/rdagent/utils/repo/diff.py index 4f271fcf..4ecc0dc3 100644 --- a/rdagent/utils/repo/diff.py +++ b/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 diff --git a/rdagent/utils/workflow.py b/rdagent/utils/workflow.py index c8a9156e..407b990f 100644 --- a/rdagent/utils/workflow.py +++ b/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