From b60f5fc506f2ae3ca001f06ff112e99a05d6b393 Mon Sep 17 00:00:00 2001 From: you-n-g Date: Mon, 23 Jun 2025 10:38:27 +0800 Subject: [PATCH] fix: refine details (#979) * feat: add parquet preview and extract common DataFrame preview logic * refactor: improve error messages, prompts, regex, and session loading * lint --- .../coder/data_science/pipeline/eval.py | 10 ++--- .../data_science/dev/runner/__init__.py | 6 ++- .../scenarios/data_science/dev/runner/eval.py | 8 ++-- .../data_science/dev/runner/prompts.yaml | 2 +- rdagent/scenarios/data_science/scen/utils.py | 40 ++++++++++++++----- rdagent/utils/agent/ret.py | 3 +- rdagent/utils/env.py | 1 + rdagent/utils/workflow/loop.py | 18 ++++++++- 8 files changed, 62 insertions(+), 26 deletions(-) diff --git a/rdagent/components/coder/data_science/pipeline/eval.py b/rdagent/components/coder/data_science/pipeline/eval.py index e441c458..9f8859df 100644 --- a/rdagent/components/coder/data_science/pipeline/eval.py +++ b/rdagent/components/coder/data_science/pipeline/eval.py @@ -74,13 +74,13 @@ class PipelineCoSTEEREvaluator(CoSTEEREvaluator): # Check model names (index) if not score_df.index.is_unique: - score_check_text += "\n[Error] The score dataframe contains duplicate model names." + score_check_text += "\n[Error] The file 'scores.csv' contains duplicate model names." score_ret_code = 1 if "ensemble" not in model_set_in_scores: - score_check_text += "\n[Error] The score dataframe doesn't contain the ensemble model." + score_check_text += "\n[Error] The file 'scores.csv' doesn't contain the ensemble model." score_ret_code = 1 if score_ret_code != 0: - score_check_text += f"The score_df is:\n{score_df}" + score_check_text += f"The dataframe in file 'scores.csv' is:\n{score_df}" # Check metric name (columns) if score_df.columns.tolist() != [self.scen.metric_name]: @@ -125,10 +125,6 @@ class PipelineCoSTEEREvaluator(CoSTEEREvaluator): ) stdout += "\n" + submission_check_out - eda_output = implementation.file_dict.get("EDA.md", None) - - eda_output = implementation.file_dict.get("EDA.md", None) - if not isinstance(implementation, FBWorkspace): eda_output = None else: diff --git a/rdagent/scenarios/data_science/dev/runner/__init__.py b/rdagent/scenarios/data_science/dev/runner/__init__.py index 48ff7ae7..aaa1c455 100644 --- a/rdagent/scenarios/data_science/dev/runner/__init__.py +++ b/rdagent/scenarios/data_science/dev/runner/__init__.py @@ -114,8 +114,10 @@ class DSCoSTEERRunner(CoSTEER): exp.sub_tasks = [ CoSTEERTask( name="Debug running solution", - description=f"The whole workflow of the solution has finished with some execution error, please check the error message and debug the whole code repo.\nCurrent code repo md5: {md5_hash(exp.experiment_workspace.all_codes)}", - ) + description=f"You'll be provided with the source code and the running and testing stdout. " + "Please check the error messages and debug the source code if any errors occur.\n" + f"Current code repo md5: {md5_hash(exp.experiment_workspace.all_codes)}", + ), ] exp = super().develop(exp) # run strategy(code implementation & evaluation loops) exp.sub_tasks = bak_sub_tasks diff --git a/rdagent/scenarios/data_science/dev/runner/eval.py b/rdagent/scenarios/data_science/dev/runner/eval.py index ba39b933..dc94eb18 100644 --- a/rdagent/scenarios/data_science/dev/runner/eval.py +++ b/rdagent/scenarios/data_science/dev/runner/eval.py @@ -60,7 +60,7 @@ class DSCoSTEERCoSTEEREvaluator(CoSTEEREvaluator): eda_output = "No EDA output." implementation.inject_files(**{"EDA.md": eda_output}) stdout = remove_eda_part(stdout) - stdout += f"The code executed {'successfully' if execute_ret_code == 0 else 'failed'}. {'EDA output is emmitted. ' if eda_output else ''}" + stdout += f"The code executed {'successfully' if execute_ret_code == 0 else 'failed'}. {'The EDA output is removed from the stdout. ' if eda_output else ''}" # Check score file score_fp = implementation.workspace_path / "scores.csv" @@ -82,13 +82,13 @@ class DSCoSTEERCoSTEEREvaluator(CoSTEEREvaluator): # in Pipeline task, we only check ensemble in scores.csv if DS_RD_SETTING.coder_on_whole_pipeline: if not score_df.index.is_unique: - score_check_text += "\n[Error] The score dataframe contains duplicate model names." + score_check_text += "\n[Error] The file 'scores.csv' contains duplicate model names." score_ret_code = 1 if "ensemble" not in model_set_in_scores: - score_check_text += "\n[Error] The score dataframe doesn't contain the ensemble model." + score_check_text += "\n[Error] The file 'scores.csv' doesn't contain the ensemble model." score_ret_code = 1 if score_ret_code != 0: - score_check_text += f"The score_df is:\n{score_df}" + score_check_text += f"The dataframe in file 'scores.csv' is:\n{score_df}" else: if model_set_in_scores != model_set_in_folder.union({"ensemble"}): score_check_text += f"\n[Error] The scores dataframe does not contain the correct model names as index.\ncorrect model names are: {model_set_in_folder.union({'ensemble'})}\nscore_df is:\n{score_df}" diff --git a/rdagent/scenarios/data_science/dev/runner/prompts.yaml b/rdagent/scenarios/data_science/dev/runner/prompts.yaml index 4fdd6e9a..eb70dd11 100644 --- a/rdagent/scenarios/data_science/dev/runner/prompts.yaml +++ b/rdagent/scenarios/data_science/dev/runner/prompts.yaml @@ -49,7 +49,7 @@ DSCoSTEER_eval: user: |- --------- code base --------- {{ code }} - --------- test stdout --------- + --------- the stdout of code execution and testing --------- {{ stdout }} DSCoSTEER_debugger: diff --git a/rdagent/scenarios/data_science/scen/utils.py b/rdagent/scenarios/data_science/scen/utils.py index 3113dcdb..6f112016 100644 --- a/rdagent/scenarios/data_science/scen/utils.py +++ b/rdagent/scenarios/data_science/scen/utils.py @@ -306,19 +306,17 @@ def _walk(path: Path): yield p -def preview_csv(p: Path, file_name: str, simple=True, show_nan_columns=False) -> str: - """Generate a textual preview of a csv file +def preview_df(df: pd.DataFrame, file_name: str, simple=True, show_nan_columns=False) -> str: + """Generate a textual preview of a dataframe Args: - p (Path): the path to the csv file + df (pd.DataFrame): the dataframe to preview file_name (str): the file name to use in the preview simple (bool, optional): whether to use a simplified version of the preview. Defaults to True. Returns: str: the textual preview """ - df = pd.read_csv(p) - out = [] out.append(f"-> {file_name} has {df.shape[0]} rows and {df.shape[1]} columns.") @@ -358,6 +356,18 @@ def preview_csv(p: Path, file_name: str, simple=True, show_nan_columns=False) -> return "\n".join(out) +def preview_csv(p: Path, file_name: str, simple=True, show_nan_columns=False) -> str: + """Generate a textual preview of a csv file""" + df = pd.read_csv(p) + return preview_df(df, file_name, simple=simple, show_nan_columns=show_nan_columns) + + +def preview_parquet(p: Path, file_name: str, simple=True, show_nan_columns=False) -> str: + """Generate a textual preview of a parquet file""" + df = pd.read_parquet(p) + return preview_df(df, file_name, simple=simple, show_nan_columns=show_nan_columns) + + def preview_json(p: Path, file_name: str): """Generate a textual preview of a json file using a generated json schema""" builder = SchemaBuilder() @@ -388,7 +398,9 @@ def preview_json(p: Path, file_name: str): return f"-> {file_name} has auto-generated json schema:\n" + builder.to_json(indent=2) -def describe_data_folder_v2(base_path, include_file_details=True, simple=False, show_nan_columns=False): +def describe_data_folder_v2( + base_path, include_file_details=True, simple=False, show_nan_columns=False, max_length: int = 6_000 +): """ Generate a textual preview of a directory, including an overview of the directory structure and previews of individual files @@ -404,6 +416,8 @@ def describe_data_folder_v2(base_path, include_file_details=True, simple=False, out.append(preview_csv(fn, file_name, simple=simple, show_nan_columns=show_nan_columns)) elif fn.suffix == ".json": out.append(preview_json(fn, file_name)) + elif fn.suffix == ".parquet": + out.append(preview_parquet(fn, file_name, simple=simple, show_nan_columns=show_nan_columns)) elif fn.suffix in plaintext_files: if get_file_len_size(fn)[0] < 30: with open(fn) as f: @@ -411,16 +425,22 @@ def describe_data_folder_v2(base_path, include_file_details=True, simple=False, if fn.suffix in code_files: content = f"```\n{content}\n```" out.append(f"-> {file_name} has content:\n\n{content}") + if len("\n\n".join(out)) > max_length: + break result = "\n\n".join(out) # if the result is very long we generate a simpler version - if len(result) > 6_000 and not simple: + if len(result) > max_length and not simple: return describe_data_folder_v2( - base_path, include_file_details=include_file_details, simple=True, show_nan_columns=show_nan_columns + base_path, + include_file_details=include_file_details, + simple=True, + show_nan_columns=show_nan_columns, + max_length=max_length, ) # if still too long, we truncate - if len(result) > 6_000 and simple: - return result[:6_000] + "\n... (truncated)" + if len(result) > max_length and simple: + return result[:max_length] + "\n... (truncated)" return result diff --git a/rdagent/utils/agent/ret.py b/rdagent/utils/agent/ret.py index c8be2802..820bd644 100644 --- a/rdagent/utils/agent/ret.py +++ b/rdagent/utils/agent/ret.py @@ -31,7 +31,8 @@ class PythonAgentOut(AgentOut): @classmethod def extract_output(cls, resp: str): - match = re.search(r".*```[Pp]ython\n(.*)\n```.*", resp, re.DOTALL) + # We use lazy mode (.*?) to only extract the first code block in the response. + match = re.search(r".*```[Pp]ython\n(.*?)\n```.*", resp, re.DOTALL) if match: code = match.group(1) code = re.sub(r"", "", code, flags=re.IGNORECASE) diff --git a/rdagent/utils/env.py b/rdagent/utils/env.py index 26f7e96f..f76abe1a 100644 --- a/rdagent/utils/env.py +++ b/rdagent/utils/env.py @@ -116,6 +116,7 @@ def pull_image_with_progress(image: str) -> None: class EnvConf(ExtendedBaseSettings): + # TODO: add prefix .... default_entry: str extra_volumes: dict = {} running_timeout_period: int = 3600 # 10 minutes diff --git a/rdagent/utils/workflow/loop.py b/rdagent/utils/workflow/loop.py index 36b32345..d8055ae6 100644 --- a/rdagent/utils/workflow/loop.py +++ b/rdagent/utils/workflow/loop.py @@ -250,7 +250,11 @@ class LoopBase: current_step = self.step_idx[li] self.pbar.n = current_step next_step = self.step_idx[li] % len(self.steps) - self.pbar.set_postfix(loop_index=li, step_index=next_step, step_name=self.steps[next_step]) + self.pbar.set_postfix( + loop_index=li + next_step_idx // len(self.steps), + step_index=next_step, + step_name=self.steps[next_step], + ) self._check_exit_conditions_on_step() else: logger.warning(f"Step forward {si} of loop {li} is skipped.") @@ -411,6 +415,18 @@ class LoopBase: An instance of LoopBase with the loaded session. """ path = Path(path) + # if the path is a directory, load the latest session + if path.is_dir(): + if path.name != "__session__": + path = path / "__session__" + + if not path.exists(): + raise FileNotFoundError(f"No session file found in {path}") + + # iterate the dump steps in increasing order + files = sorted(path.glob("*/*_*"), key=lambda f: (int(f.parent.name), int(f.name.split("_")[0]))) + path = files[-1] + logger.info(f"Loading latest session from {path}") with path.open("rb") as f: session = cast(LoopBase, pickle.load(f))