From d4c539912abdb60e9d8950e7ea1186fd32bfeef3 Mon Sep 17 00:00:00 2001 From: Tim Date: Tue, 9 Sep 2025 11:56:42 +0800 Subject: [PATCH] feat: offline selector (#1231) * offline selector test * fix score with tensor * sort sota list --------- Co-authored-by: Xu Yang --- .../proposal/exp_gen/select/prompts.yaml | 35 +- .../proposal/exp_gen/select/submit.py | 901 +++++++++++++++--- rdagent/utils/env.py | 2 +- 3 files changed, 793 insertions(+), 145 deletions(-) diff --git a/rdagent/scenarios/data_science/proposal/exp_gen/select/prompts.yaml b/rdagent/scenarios/data_science/proposal/exp_gen/select/prompts.yaml index 8aa6baef..55d373cd 100644 --- a/rdagent/scenarios/data_science/proposal/exp_gen/select/prompts.yaml +++ b/rdagent/scenarios/data_science/proposal/exp_gen/select/prompts.yaml @@ -76,5 +76,38 @@ auto_sota_selector: # SOTA Experiments and Feedback {{ historical_sota_exp_with_desc_and_scores }} +sample_data: + system: |- + You are a senior machine learning engineer. + Generate a single, self-contained Python script that strictly follows the user's instructions. + Requirements: + - The script MUST be runnable via `python .py` without extra arguments unless specified. + - Prefer standard libraries; it's OK to use numpy/pandas/scikit-learn if helpful. + - Use robust error handling and clear messages. + - Use relative paths only and create missing directories when needed. + - Keep the script concise and well-commented. + user: |- + Full runnable model code: + {{ reference_code }} + Write a separate script based on this code to sample 80% of the data (while maintaining the class proportions as much as possible) as the new train set, + and 20% as the new test set. Save the new train and test in the `{{ input_folder }}` folder. + Save test label with id to `{{ input_folder }}/label.csv`, which is to be used for grading. + Load source data from path `./source` directory. + Please make sure the new test set has the same columns as the original test set. + Please make sure all files used in the original code and exists in source folder are also available in the `{{ input_folder }}` folder. + Ignore all files that do not exist in the read only source folder. + {% if error %} + {{ error }} + {% endif %} - +grade: + user: |- + Metric method according to {{ reference_code }} + `{{ input_folder }}/label.csv` generated by {{ sample_code }} + Write a Python script named `grade.py` to evaluate `submission.csv` produced by a model. + Input files (relative to current working directory): + - `{{ input_folder }}/label.csv` and `submission.csv` + Output format: `{'score': float, 'metric': str}` + {% if error %} + {{ error }} + {% endif %} diff --git a/rdagent/scenarios/data_science/proposal/exp_gen/select/submit.py b/rdagent/scenarios/data_science/proposal/exp_gen/select/submit.py index 93e7ff9e..8187d6fd 100644 --- a/rdagent/scenarios/data_science/proposal/exp_gen/select/submit.py +++ b/rdagent/scenarios/data_science/proposal/exp_gen/select/submit.py @@ -1,192 +1,807 @@ import json -import random -from typing import Dict, Tuple +import os +import pickle +import re +import shutil +import time +from pathlib import Path +from typing import Any, Dict, List, Optional, Tuple +import fire import numpy as np import pandas as pd +from loguru import logger from rdagent.app.data_science.conf import DS_RD_SETTING +from rdagent.components.coder.data_science.conf import get_ds_env +from rdagent.core.experiment import FBWorkspace from rdagent.core.proposal import ExperimentFeedback, SOTAexpSelector, Trace -from rdagent.log import rdagent_logger as logger +from rdagent.core.utils import multiprocessing_wrapper +from rdagent.log.storage import FileStorage +from rdagent.log.utils import extract_json from rdagent.oai.llm_conf import LLM_SETTINGS -from rdagent.oai.llm_utils import APIBackend, md5_hash +from rdagent.oai.llm_utils import APIBackend from rdagent.scenarios.data_science.experiment.experiment import DSExperiment -from rdagent.scenarios.data_science.proposal.exp_gen.base import DSHypothesis, DSTrace +from rdagent.utils.agent.ret import PythonAgentOut from rdagent.utils.agent.tpl import T +from rdagent.utils.fmt import shrink_text from rdagent.utils.workflow import wait_retry +# --- Configuration Constants --- +MAX_API_RETRIES = int(os.getenv("MAX_API_RETRIES", 5)) +DEFAULT_NUM_WORKERS = int(os.getenv("DEFAULT_NUM_WORKERS", 2)) +MAX_SOTA_CANDIDATES = int(os.getenv("MAX_SOTA_CANDIDATES", 6)) + +logger.add("selector.log") +# ============================================================================== +# ## SOTA Selector Implementations +# ============================================================================== + class GlobalSOTASelector(SOTAexpSelector): """ - return the latest SOTA experiment from the trace to submit + Selects the single best State-Of-The-Art (SOTA) experiment from the entire trace history. """ - def __init__( - self, - ): - print(f"Using global SOTA policy by default") - - def get_sota_exp_to_submit(self, trace: Trace) -> DSExperiment | None: + def __init__(self): + logger.info("Using selector policy: GlobalSOTASelector") + def get_sota_exp_to_submit(self, trace: Trace, **kwargs) -> DSExperiment | None: + """ + Returns the single best experiment from all historical runs. + """ return trace.sota_experiment(search_type="all") class AutoSOTAexpSelector(SOTAexpSelector): """ - retrieve a list of SOTA experiments from the trace, then call the LLM to select the best one + Uses an LLM to select the best SOTA experiment from a list of candidates. + Candidates are retrieved from the leaves of the experiment trace tree. + """ + + def __init__(self): + logger.info("Using selector policy: AutoSOTAexpSelector") + + @wait_retry(retry_n=MAX_API_RETRIES) + def get_sota_exp_to_submit(self, trace: Trace, **kwargs) -> DSExperiment | None: + """ + Retrieves SOTA experiments, then uses an LLM to choose the most promising one. + """ + sota_exp_fb_list = self.collect_sota_candidates(trace) + + if not sota_exp_fb_list: + logger.info("AutoSOTASelector: No SOTA experiments found in trace.") + return None + + if len(sota_exp_fb_list) == 1: + logger.info("AutoSOTASelector: Only one SOTA candidate found, selecting it.") + return sota_exp_fb_list[0][0] + + logger.info(f"AutoSOTASelector: {len(sota_exp_fb_list)} SOTA candidates found. Querying LLM for selection.") + + # Build prompt for LLM + sota_prompt_text = "Historical SOTA experiments:\n\n" + system_prompt = T(".prompts:auto_sota_selector.system").r(scenario=trace.scen.get_scenario_all_desc()) + for i, (exp, _) in enumerate(sota_exp_fb_list): + if exp and exp.result is not None: + current_final_score = pd.DataFrame(exp.result).loc["ensemble"].iloc[0] + desc = T("scenarios.data_science.share:describe.exp").r(exp=exp) + new_experiment_content = f"""SOTA experiment No. {i+1}: + Description: {desc} + Final score: {current_final_score}\n\n""" + + temp_user_prompt = T(".prompts:auto_sota_selector.user").r( + historical_sota_exp_with_desc_and_scores=sota_prompt_text + new_experiment_content, + ) + + token_size = APIBackend().build_messages_and_calculate_token( + user_prompt=temp_user_prompt, + system_prompt=system_prompt, + ) + if token_size >= LLM_SETTINGS.chat_token_limit: + logger.warning(f"Token limit reached at experiment {i+1}. Stopping.") + break + + sota_prompt_text += new_experiment_content + + # Query LLM + user_prompt = T(".prompts:auto_sota_selector.user").r(historical_sota_exp_with_desc_and_scores=sota_prompt_text) + + response = APIBackend().build_messages_and_create_chat_completion( + user_prompt=user_prompt, + system_prompt=system_prompt, + json_mode=True, + json_target_type=Dict[str, Any], + ) + response_dict = json.loads(response) + selected_idx = response_dict.get("selected_SOTA_idx") + + # Process LLM response + if selected_idx and isinstance(selected_idx, int) and 0 < selected_idx <= len(sota_exp_fb_list): + sota_submit = sota_exp_fb_list[selected_idx - 1][0] + logger.info(f"AutoSOTASelector: LLM selected experiment No. {selected_idx}.") + return sota_submit + + logger.warning("AutoSOTASelector: LLM selection was invalid. Falling back to the latest SOTA experiment.") + return sota_exp_fb_list[-1][0] if sota_exp_fb_list else None + + def collect_sota_candidates(self, trace: Trace) -> list: + """Helper to gather SOTA experiments from trace leaves.""" + leaves = trace.get_leaves() + if len(leaves) < 2: + return trace.experiment_and_feedback_list_after_init( + return_type="sota", search_type="all", max_retrieve_num=DS_RD_SETTING.max_sota_retrieved_num + ) + + logger.info(f"AutoSOTASelector: {len(leaves)} branches found, collecting SOTA from each.") + all_sota_candidates = [] + num_per_trace = max(DS_RD_SETTING.max_sota_retrieved_num // len(leaves), 2) + + for leaf in leaves: + sota_from_branch = trace.experiment_and_feedback_list_after_init( + return_type="sota", search_type="ancestors", selection=(leaf,), max_retrieve_num=num_per_trace + ) + all_sota_candidates.extend(sota_from_branch) + + # Remove duplicates and limit total number of candidates + unique_sota_list = list(set(all_sota_candidates)) + is_higher_better = trace.scen.metric_direction + unique_sota_list.sort( + key=lambda exp_fb: pd.DataFrame(exp_fb[0].result).loc["ensemble"].iloc[0], + reverse=is_higher_better, + ) + return unique_sota_list[: DS_RD_SETTING.max_sota_retrieved_num] + + +class BestValidSelector(SOTAexpSelector): + """ + Selects the top N experiments based on their performance score. + Can operate across the entire trace or on a per-branch basis. + """ + + def __init__(self, num_candidates: int = 1, use_decision: bool = True, each_trace: bool = False): + """ + Args: + num_candidates (int): The number of top experiments to return. + use_decision (bool): If True, filters out experiments marked with a negative decision. + each_trace (bool): If True, selects top candidates from each branch instead of globally. + """ + logger.info( + f"Using selector policy: BestValidSelector (num_candidates={num_candidates}, each_trace={each_trace})" + ) + self.num_candidates = num_candidates + self.use_decision = use_decision + self.each_trace = each_trace + + def get_sota_exp_to_submit(self, trace: Trace, **kwargs) -> DSExperiment | None: + """ + Sorts all valid experiments by score and returns the top N. + """ + top_experiments = self.collect_sota_candidates(trace) + if top_experiments: + return top_experiments[0] + return None + + def collect_sota_candidates(self, trace: Trace) -> list[DSExperiment] | None: + """Helper to gather SOTA experiments from trace leaves.""" + """ + Sorts all valid experiments by score and returns the top N. + """ + direction_sign = 1 if trace.scen.metric_direction else -1 + + def get_sort_key(exp_fb: Tuple[DSExperiment, ExperimentFeedback]) -> Tuple[bool, float]: + exp, feedback = exp_fb + score = -np.inf + if exp.result is not None: + try: + score = pd.DataFrame(exp.result).loc["ensemble"].iloc[0] + if isinstance(score, str): + score = float(score.strip("tensor()")) + score = direction_sign * score + except: + logger.warning(f"Failed to extract score from result {exp.result}") + + # Sort key prioritizes decision (True > False), then score + return (feedback.decision, score) if self.use_decision else score + + # Collect candidates + if self.each_trace: + candidate_list = [] + leaves = trace.get_leaves() + num_per_leaf = max(self.num_candidates // len(leaves), 1) + for leaf in leaves: + branch_experiments = trace.experiment_and_feedback_list_after_init( + return_type="all", search_type="ancestors", selection=(leaf,) + ) + if branch_experiments: + branch_experiments.sort(key=get_sort_key, reverse=True) + candidate_list.extend(branch_experiments[:num_per_leaf]) + # Remove duplicates + candidate_list = list(set(candidate_list)) + else: + candidate_list = trace.experiment_and_feedback_list_after_init(return_type="all", search_type="all") + + if not candidate_list: + logger.info("BestValidSelector: No experiments found in trace.") + return None + + # Sort and select the top N + candidate_list.sort(key=get_sort_key, reverse=True) + + top_experiments = [exp for exp, _ in candidate_list[: self.num_candidates]] + logger.info(f"BestValidSelector: Selected {len(top_experiments)} experiments.") + return top_experiments + + +class ValidationSelector(SOTAexpSelector): + """ + A meta-selector that re-validates candidates from a base selector. + + It then generates a consistent validation dataset and grading script, + re-runs all candidates on this new data, and returns the best performer. """ def __init__( self, + candidate: List[Tuple[DSExperiment, str]], + direction_sign: int, + competition: str, + only_sample: bool, + sample_code_path: str, + sample_rate: float = 0.8, ): - print(f"Using auto SOTA policy") + self.candidate = candidate + self.direction_sign = direction_sign + self.competition = competition + self.only_sample = only_sample + self.sample_code_path = Path(sample_code_path) + self.sample_rate = sample_rate + self.hypothesis_loop_id = {exp.hypothesis.hypothesis: loop_id for exp, loop_id in self.candidate} + self.hypothesis_exp = {exp.hypothesis.hypothesis: exp for exp, loop_id in self.candidate} - @wait_retry(retry_n=5) def get_sota_exp_to_submit(self, trace: Trace) -> DSExperiment | None: - # retrieve all SOTA experiments from the trace + """Helper to gather SOTA experiments from trace leaves.""" + """ + Sorts all valid experiments by score and returns the top N. + """ - sota_exp_fb_list = trace.experiment_and_feedback_list_after_init( - return_type="sota", search_type="all", max_retrieve_num=DS_RD_SETTING.max_sota_retrieved_num + mock_folder = f"/tmp/mock/{self.competition}" + + try: + data_py_code, grade_py_code = self._prepare_validation_scripts( + reference_exp=self.candidate[0][0], competition=self.competition, mock_folder=mock_folder + ) + except RuntimeError as e: + logger.error(f"ValidationSelector: Failed to prepare validation environment. {e}") + shutil.rmtree(mock_folder, ignore_errors=True) + return None + + validation_tasks = [ + (process_experiment, (exp, self.competition, mock_folder, grade_py_code, loop_id)) + for exp, loop_id in self.candidate + ] + results = multiprocessing_wrapper(validation_tasks, n=min(DEFAULT_NUM_WORKERS, (len(self.candidate) + 1) // 2)) + + if not results: + logger.warning("ValidationSelector: Validation run produced no results.") + return None + + # 4. Process results and select the best one + valid_results = [ + ( + self.hypothesis_exp.get(exp.hypothesis.hypothesis), + self.hypothesis_loop_id.get(exp.hypothesis.hypothesis), + valid_score, + test_score, + ) + for exp, valid_score, test_score in results + if test_score is not None + ] + if not valid_results: + logger.warning("ValidationSelector: No candidates scored successfully during validation.") + return None + + valid_results.sort(key=lambda x: (x[3]) * self.direction_sign, reverse=True) + best_exp, best_loop_id = valid_results[0][0], valid_results[0][1] + + for loop_id, valid_score, test_score in [(i[1], i[2], i[3]) for i in valid_results]: + logger.info(f"ValidationSelector: Loop_id={loop_id} ->valid score={valid_score}, test score={test_score}") + logger.info( + f"ValidationSelector: Best experiment from validation is loop_id={best_loop_id} with valid score={valid_results[0][2]}, test score={valid_results[0][3]}" ) - logger.info(f"Auto SOTA selector: Found {len(sota_exp_fb_list)} SOTA experiments") - if len(sota_exp_fb_list) == 0: - logger.info("Auto SOTA selector: No SOTA in trace yet") + if len(valid_results) <= 1 or valid_results[0][3] == valid_results[-1][3]: + logger.warning(f"ValidationSelector: There aren't enough scores to compare, current: {len(valid_results)}.") return None - elif len(sota_exp_fb_list) == 1: - sota_idx_in_trace = trace.hist.index(sota_exp_fb_list[0]) - logger.info( - f"Auto SOTA selector: Only one SOTA in trace, using it, which is the No. {sota_idx_in_trace + 1} in the trace" - ) - return sota_exp_fb_list[0][0] + return best_exp - else: - logger.info( - f"Auto SOTA selector: Multiple SOTA in trace, calling LLM to select the best one in {DS_RD_SETTING.max_sota_retrieved_num} SOTA experiments" - ) + def print_code(self, data_py_code: str, grade_py_code: str): + logger.info("Successfully ran data.py.") + print("======== data.py ========") + print(data_py_code) + print("======== grade.py ========") + print(grade_py_code) + print("======== code end ========") - SOTA_exp_with_desc_and_scores = "Historical SOTA experiments:\n\n" + def _prepare_validation_scripts( + self, reference_exp: DSExperiment, competition: str, mock_folder: str + ) -> Tuple[str, str]: + """Generates and verifies data.py and grade.py using an LLM.""" + input_folder = T("scenarios.data_science.share:scen.input_path").r() + mock_input_path = Path(mock_folder) / input_folder + mock_input_path.mkdir(parents=True, exist_ok=True) - leaves: list[int] = trace.get_leaves() + data_py_path = Path(mock_folder) / "data.py" + grade_py_path = Path(mock_folder) / "grade.py" + label_path = Path(mock_folder) / "workspace_input/label.csv" + reference_code = reference_exp.experiment_workspace.file_dict.get("main.py", "") + if not reference_code: + raise RuntimeError("ValidationSelector: No code found in the reference experiment.") - if len(leaves) >= 2: - - logger.info( - f"Auto SOTA selector: {len(leaves)} traces found, collecting SOTA experiments from each trace" + if (self.sample_code_path / competition / "data.py").exists(): + shutil.copy(self.sample_code_path / competition / "data.py", data_py_path) + shutil.copy(self.sample_code_path / competition / "grade.py", grade_py_path) + data_py_code = data_py_path.read_text() + grade_py_code = grade_py_path.read_text() + if not label_path.exists(): + ws = FBWorkspace() + if self.sample_rate != 0.8: + data_py_code = data_py_code.replace("0.8", str(self.sample_rate)).replace( + "0.2", str(round(1 - self.sample_rate, 2)) + ) + ws.inject_code_from_file_dict(reference_exp.experiment_workspace) + ws.inject_files(**{f"data.py": data_py_code}) + env = get_ds_env( + extra_volumes={ + str(Path(mock_folder) / input_folder): {"bind": input_folder, "mode": "rw"}, + f"{DS_RD_SETTING.local_data_path}/{competition}": "./source", + }, + running_timeout_period=DS_RD_SETTING.full_timeout, ) - # multiple trace case, collect the latest SOTA experiments from each trace - new_sota_exp_fb_list: list[tuple[DSExperiment, ExperimentFeedback]] = [] - # calculate the number of SOTA experiments to retrieve from each trace, prevent it from becoming zero - max_sota_retrieved_num_per_trace = max(DS_RD_SETTING.max_sota_retrieved_num // len(leaves), 5) - # recall, due to the integer division, the final number of SOTA experiments to retrieve may be different - for leaf in leaves: - sota_exp_fb_list_per_trace = trace.experiment_and_feedback_list_after_init( - return_type="sota", - search_type="ancestors", - selection=(leaf,), - max_retrieve_num=max_sota_retrieved_num_per_trace, + result = ws.run( + env=env, entry=f"python data.py --cache-buster={time.time()}" + ) # Do not cache the result + if result.exit_code == 0: + self.print_code(data_py_code, grade_py_code) + return data_py_code, grade_py_code + + # --- Generate data.py if needed --- + if not data_py_path.exists() or not label_path.exists(): + logger.info(f"Generating synthetic data script: {data_py_path}") + data_py_code = self._generate_and_run_script( + script_type="data", + prompt_template_key="sample_data", + reference_exp=reference_exp, + competition=competition, + mock_folder=mock_folder, + prompt_kwargs={"reference_code": reference_code, "input_folder": input_folder}, + ) + data_py_path.write_text(data_py_code) + + data_py_code = data_py_path.read_text() + + # --- Generate grade.py if needed --- + if not grade_py_path.exists(): + logger.info(f"Generating grading script: {grade_py_path}") + grade_py_code = self._generate_and_run_script( + script_type="grade", + prompt_template_key="grade", + reference_exp=reference_exp, + competition=competition, + mock_folder=mock_folder, + prompt_kwargs={ + "reference_code": reference_code, + "sample_code": data_py_code, + "input_folder": input_folder, + }, + ) + grade_py_path.write_text(grade_py_code) + self.print_code(data_py_code, grade_py_code) + return data_py_code, grade_py_path.read_text() + + def _generate_and_run_script( + self, + script_type: str, + prompt_template_key: str, + reference_exp: DSExperiment, + competition: str, + mock_folder: str, + prompt_kwargs: dict, + ) -> str: + """A helper to generate, run, and validate a script (data.py or grade.py).""" + system_prompt = T(".prompts:sample_data.system").r() # Generic system prompt for both + input_folder = T("scenarios.data_science.share:scen.input_path").r() + + err_msg = "" + for _ in range(MAX_API_RETRIES): + user_prompt = T(f".prompts:{prompt_template_key}.user").r(error=err_msg, **prompt_kwargs) + + generated_code = PythonAgentOut.extract_output( + APIBackend().build_messages_and_create_chat_completion( + user_prompt=user_prompt, system_prompt=system_prompt + ) + ) + + # Create a temporary workspace to test the generated script + ws = FBWorkspace() + ws.inject_code_from_file_dict(reference_exp.experiment_workspace) + ws.inject_files(**{f"{script_type}.py": generated_code}) + + if script_type == "data": + # For data.py, we need the original data to sample from + env = get_ds_env( + extra_volumes={ + str(Path(mock_folder) / input_folder): {"bind": input_folder, "mode": "rw"}, + f"{DS_RD_SETTING.local_data_path}/{competition}": "./source", + }, + running_timeout_period=DS_RD_SETTING.full_timeout, + ) + else: # For grade.py, we only need the generated data + shutil.copy( + str(Path(mock_folder) / "submission.csv"), + str(ws.workspace_path / "submission.csv"), + ) + env = get_ds_env( + extra_volumes={str(Path(mock_folder) / input_folder): {"bind": input_folder, "mode": "rw"}} + ) + + result = ws.run( + env=env, entry=f"python {script_type}.py --cache-buster={time.time()}" + ) # Do not cache the result + stdout = re.sub(r"^chmod:.*\n?", "", result.stdout, flags=re.MULTILINE) + + if result.exit_code == 0: + logger.info(f"Successfully generated and ran {script_type}.py.") + if script_type == "data": + env = get_ds_env( + extra_volumes={str(Path(mock_folder) / input_folder): {"bind": input_folder, "mode": "rw"}}, + running_timeout_period=DS_RD_SETTING.full_timeout, ) - logger.info( - f"Auto SOTA selector: Collected {len(sota_exp_fb_list_per_trace)} SOTA experiments from trace with leaf #. {leaf}" - ) - - new_sota_exp_fb_list.extend(sota_exp_fb_list_per_trace) - - sota_exp_fb_list = list(set(new_sota_exp_fb_list)) - - if len(sota_exp_fb_list) == 0: - logger.info("Auto SOTA selector: No SOTA in trace yet") - return None - - elif len(sota_exp_fb_list) == 1: - logger.info("Auto SOTA selector: Only one SOTA in trace, using it") - return sota_exp_fb_list[0][0] + result = ws.run(env=env, entry=f"python main.py --cache-buster={time.time()}") + stdout = re.sub(r"^chmod:.*\n?", "", result.stdout, flags=re.MULTILINE) + if result.exit_code == 0: + # move submission.csv to mock_folder + if Path(ws.workspace_path / "submission.csv").exists(): + shutil.copy( + str(ws.workspace_path / "submission.csv"), + str(Path(mock_folder) / "submission.csv"), + ) + return generated_code + else: + err_msg = "No submission.csv found in workspace after running main.py with generated data." + else: + err_msg = f"Error in main.py with generated data: {shrink_text(stdout, context_lines=20, line_len=500)}" else: - logger.info( - f"Auto SOTA selector: select {len(sota_exp_fb_list)} of {len(new_sota_exp_fb_list)} SOTA experiments found in all traces, calling LLM to select the best one" - ) - if len(sota_exp_fb_list) > DS_RD_SETTING.max_sota_retrieved_num: - sota_exp_fb_list = sorted( - sota_exp_fb_list, - key=lambda exp_fb: pd.DataFrame(exp_fb[0].result).loc["ensemble"].iloc[0], - reverse=not trace.scen.metric_direction, - )[-DS_RD_SETTING.max_sota_retrieved_num :] - - system_prompt = T(".prompts:auto_sota_selector.system").r(scenario=trace.scen.get_scenario_all_desc()) - for i, (exp, ef) in enumerate(sota_exp_fb_list): - if exp: - current_final_score = pd.DataFrame(exp.result).loc["ensemble"].iloc[0] - desc = T("scenarios.data_science.share:describe.exp").r( - exp=exp, heading="SOTA of previous exploration of the scenario" - ) - new_experiment_content = f"""SOTA experiment No. {i+1}: - Description: {desc} - Final score: {current_final_score}\n\n""" - - temp_user_prompt = T(".prompts:auto_sota_selector.user").r( - historical_sota_exp_with_desc_and_scores=SOTA_exp_with_desc_and_scores + new_experiment_content, - ) - - token_size = APIBackend().build_messages_and_calculate_token( - user_prompt=temp_user_prompt, - system_prompt=system_prompt, - ) - if token_size >= LLM_SETTINGS.chat_token_limit: - logger.warning(f"Token limit reached at experiment {i+1}. Stopping.") - break - - SOTA_exp_with_desc_and_scores += new_experiment_content - - user_prompt = T(".prompts:auto_sota_selector.user").r( - historical_sota_exp_with_desc_and_scores=SOTA_exp_with_desc_and_scores - ) - - response = APIBackend().build_messages_and_create_chat_completion( - user_prompt=user_prompt, - system_prompt=system_prompt, - json_mode=True, - json_target_type=Dict[str, str | int], - ) - - response_dict = json.loads(response) - - sota_submit_idx = response_dict.get("selected_SOTA_idx", None) - - if sota_submit_idx and int(sota_submit_idx) - 1 < len(sota_exp_fb_list): - sota_submit = sota_exp_fb_list[int(sota_submit_idx) - 1] - sota_idx_in_trace = trace.hist.index(sota_submit) - logger.info( - f"Auto SOTA selector: selected SOTA experiment No. {sota_submit_idx} to submit, which is the No. {sota_idx_in_trace + 1} in the trace" - ) - return sota_submit[0] + score = _parsing_score(stdout) + if score is not None: + return generated_code + else: + err_msg = f"No score found in stdout: {stdout}." else: - # no SOTA experiment to submit, using the latest SOTA experiment - if len(sota_exp_fb_list) > 0: - logger.info("Auto SOTA selector: No SOTA experiment to submit, using the latest SOTA experiment") - return sota_exp_fb_list[-1][0] - else: - logger.info("Auto SOTA selector: No SOTA experiment in trace yet") - return None + err_msg = f"Error in {script_type}.py: {shrink_text(stdout, context_lines=20, line_len=500)}" + + logger.warning(f"Attempt to generate {script_type}.py failed. Retrying... Error: {err_msg}") + raise RuntimeError(f"Failed to generate a working {script_type}.py after {MAX_API_RETRIES} attempts.") -class BestValidSelector(SOTAexpSelector): - def get_sota_exp_to_submit(self, trace: Trace) -> DSExperiment | None: - sota_exp_fb_list = trace.experiment_and_feedback_list_after_init(return_type="all", search_type="all") - direction_sign = 1 if trace.scen.metric_direction else -1 +# ============================================================================== +# ## Worker and Utility Functions +# ============================================================================== - def get_sort_key(exp_fb: tuple[DSExperiment, ExperimentFeedback]) -> tuple[bool, float]: - score = -np.inf - result: pd.DataFrame | None = exp_fb[0].result - if result is not None: - score = result.loc["ensemble"].iloc[0] - return (exp_fb[1].decision, direction_sign * score) - if len(sota_exp_fb_list) == 0: - logger.info("Best Valid SOTA selector: No SOTA in trace yet") - return None +def process_experiment( + exp: DSExperiment, competition: str, folder: str, grade_py_code: str, loop_id: str +) -> Tuple[DSExperiment, Optional[float], Optional[float]]: + """ + Worker function to process a single experiment in an isolated directory. + This function is designed to be called by a multiprocessing pool. + """ + if loop_id is None: + logger.error("Could not find loop_id for a given experiment.") + loop_id = "unknown" + + input_folder = T("scenarios.data_science.share:scen.input_path").r() + valid_score = None + + try: + ws = FBWorkspace() + logger.info(f"Experiment files: {exp.experiment_workspace.file_dict.keys()}") + ws.inject_code_from_file_dict(exp.experiment_workspace) + + # Run main script + env = get_ds_env( + extra_volumes={f"/tmp/mock/{competition}/{input_folder}": input_folder}, + running_timeout_period=DS_RD_SETTING.full_timeout, + ) + result = ws.run(env=env, entry="python main.py") + execute_ret_code = result.exit_code + logger.info(f"Ran {competition}/{loop_id}/main.py; exit_code: {execute_ret_code}") + + # Run grading script if main script succeeded + grade_stdout = "" + if execute_ret_code == 0: + score_fp = ws.workspace_path / "scores.csv" + if score_fp.exists(): + try: + valid_score = pd.read_csv(score_fp, index_col=0).loc["ensemble"].iloc[0] + except Exception as e: + logger.error(f"Error parsing valid score from {score_fp}: {e}") + ws.inject_files(**{"grade.py": grade_py_code}) + env.conf.running_timeout_period = DS_RD_SETTING.debug_timeout + result = ws.run(env=env, entry="python grade.py") + if result.exit_code == 0: + grade_stdout = re.sub(r"^chmod:.*\n?", "", result.stdout, flags=re.MULTILINE) + logger.info(f"Ran grade.py for {competition}/{loop_id}; exit_code: {result.exit_code}") else: - sota_exp_fb_list = sorted(sota_exp_fb_list, key=get_sort_key, reverse=True) - return sota_exp_fb_list[0][0] + logger.warning(f"Skipping grading for {competition}/{loop_id} due to main.py execution failure.") + + except Exception as e: + logger.error(f"CRITICAL ERROR while processing experiment {competition}/{loop_id}: {e}") + return exp, None, None + + # Score parsing + return exp, valid_score, _parsing_score(grade_stdout) -# TODO: more advanced sota exp selector (e.g. LLM-based, merge exp with multiple sub-trace) +def _parsing_score(grade_stdout: str) -> Optional[float]: + for line in grade_stdout.splitlines(): + line = line.strip() + if "score" not in line: + continue + m = re.search(r"\{[^{}]*(?:\{[^{}]*\}[^{}]*)*\}", line) + if not m: + continue + json_str = m.group(0) + try: + # Priority 1: JSON parsing + return float(json.loads(json_str)["score"]) + except: + pass + try: + # Priority 2: Eval dict + return float(eval(json_str)["score"]) + except: + pass + try: + # Priority 3: Regex for the last number in the string + return float(re.findall(r"[-+]?\d*\.\d+|\d+", json_str)[-1]) + except: + pass + return None + + +def check_hit(selected_exp: DSExperiment, trace: Trace, sota_result: Dict[str, Any]) -> bool: + """Checks if any of the selected experiments are considered medal-winning.""" + if not selected_exp: + return False + + index = trace.exp2idx(selected_exp) + # Check by loop_id if available + if hasattr(trace, "idx2loop_id"): + loop_id = trace.idx2loop_id.get(index) + if loop_id and loop_id in sota_result.get("medal_loops", []): + return True + # Fallback to checking by index + if index in sota_result.get("medal_loops_index", []): + return True + return False + + +def try_get_loop_id(trace: Trace, exp: DSExperiment): + index = trace.exp2idx(exp) + if hasattr(trace, "idx2loop_id"): + return trace.idx2loop_id.get(index) + return index + + +# ============================================================================== +# ## Main Orchestration Logic +# ============================================================================== + + +def evaluate_one_trace( + selector_name: str, + trace: Trace, + debug: bool, + only_sample: bool, + sample_code_path: str, + sota_result: dict[str, Any] = {}, + experiment: str = "validation", + log_path: Path | None = None, + sample_rate: float = 0.8, +) -> Tuple[str, bool]: + """ + Loads a single trace, uses the specified selector to pick an experiment, + and checks if the selection was a "hit" (a known SOTA solution). + """ + competition = trace.scen.competition + hit = False + + # Example of scenario-specific adjustment + if competition == "detecting-insults-in-social-commentary": + trace.scen.metric_direction = 1 + direction_sign = 1 if trace.scen.metric_direction else -1 + + # --- Selector Instantiation --- + # The core logic is now encapsulated in these selectors. + if selector_name == "global": + selector = GlobalSOTASelector() + elif selector_name == "auto": + selector = AutoSOTAexpSelector() + elif selector_name == "best_valid": + # These params can be configured or passed via CLI + selector = BestValidSelector(num_candidates=1, use_decision=True, each_trace=False) + + if selector_name == "validation": + if not Path(f"{DS_RD_SETTING.local_data_path}/{competition}").exists(): + logger.warning(f"Competition {DS_RD_SETTING.local_data_path}/{competition} does not exist, skipping.") + return competition, False + # The ValidationSelector is used to select the best re-test score. + quick_selector = BestValidSelector(num_candidates=1, use_decision=True, each_trace=False) + quick_selected_exps = quick_selector.get_sota_exp_to_submit(trace) + if debug: + quick_hit = check_hit(quick_selected_exps, trace, sota_result) + logger.info(f"BestvalidSelector for {experiment} - {competition}: {'HIT' if quick_hit else 'MISS'}") + + base_selector = BestValidSelector(num_candidates=MAX_SOTA_CANDIDATES, use_decision=True, each_trace=True) + candidate_exps = base_selector.collect_sota_candidates(trace) + if not candidate_exps: + logger.info("ValidationSelector: Base selector returned no candidates.") + return competition, False + + logger.info(f"ValidationSelector: Received {len(candidate_exps)} candidates for validation.") + if debug: + pool_hit = any(check_hit(candidate_exp, trace, sota_result) for candidate_exp in candidate_exps) + if not pool_hit: + logger.info("ValidationSelector: Base selector's candidates did not hit any SOTA. Skipping validation.") + return competition, False + + selector = ValidationSelector( + candidate=[(exp, try_get_loop_id(trace, exp)) for exp in candidate_exps], + direction_sign=direction_sign, + competition=competition, + only_sample=only_sample, + sample_code_path=sample_code_path, + sample_rate=sample_rate, + ) + + selected_sota_exps = selector.get_sota_exp_to_submit(trace) + if selector_name == "validation" and selected_sota_exps is None: + selected_sota_exps = quick_selected_exps + + # --- Run Selection and Check for Hit --- + logger.info(f"Running selector '{selector_name}' on trace for competition '{competition}'...") + if debug: + hit = check_hit(selected_sota_exps, trace, sota_result) + logger.info(f"Result for {experiment} - {competition}: {'HIT' if hit else 'MISS'}") + elif selector_name == "validation": + loop_id = selector.hypothesis_loop_id.get(selected_sota_exps.hypothesis.hypothesis) + logger.info(f"Selected loop for {experiment} - {competition}: {loop_id=}") + sota_mle_score_paths = [i for i in log_path.rglob(f"Loop_{loop_id}/running/mle_score/**/*.pkl")] + if len(sota_mle_score_paths): + with sota_mle_score_paths[0].open("rb") as f: + sota_mle_score = extract_json(pickle.load(f)) + hit = sota_mle_score.get("any_medal", False) + return competition, hit + + +def select_on_existing_trace( + selector_name: str, + trace_root: str = "", + experiment: str | None = None, + competition: str | None = None, + debug: bool = False, + only_sample: bool = False, + sample_code_path: str = "", + sample_rate: float = 0.8, +): + """ + Offline evaluation of a SOTA experiment selector on existing traces. + + Args: + selector_name (str): Name of the selector to use. Options: 'global', 'auto', 'best_valid', 'validation'. + trace_root (str): Path to the root directory containing trace folders. + experiment (str | None): Name of the experiment to evaluate, e.g., "devoted-burro+massive-perch". + competition (str | None): Name of the competition to evaluate, e.g., "detecting-insults-in-social-commentary". + debug (bool): If True, debug mode. + only_sample (bool): If True, only generates the sample code. + sample_code_path (str): Path to the sample code. + """ + result_dict = {} + trace_root_path = Path(trace_root) + + # Prepare list of tasks for multiprocessing + tasks = [] + if debug: + for trace_folder in trace_root_path.iterdir(): + if not trace_folder.is_dir(): + continue + if experiment is not None: + if trace_folder.name not in experiment: + continue + for trace_pkl_path in trace_folder.glob("*.pkl"): + if competition is not None and not competition in str(trace_pkl_path): + continue + sota_result = {} + trace = pickle.load(trace_pkl_path.open("rb")) + try: + sota_loops_file = trace_folder / f"{trace_pkl_path.stem.split('_')[0]}_loops.json" + with open(sota_loops_file, "r") as f: + sota_result = json.load(f) + except FileNotFoundError: + logger.warning(f"Could not find SOTA loops file for {trace.scen.competition}, skipping.") + continue + + if not sota_result.get("medal_loops"): + logger.info(f"No Medal loops defined for {trace.scen.competition}, skipping.") + continue + + tasks.append( + ( + evaluate_one_trace, + ( + selector_name, + trace, + debug, + only_sample, + sample_code_path, + sota_result, + trace_pkl_path.parent.name, + None, + sample_rate, + ), + ) + ) + else: + log_path = next( + d for d in Path("log").iterdir() if d.is_dir() and d.name != "pickle_cache" and not d.name.startswith("20") + ) + logger.info(f"Loading trace from {log_path}") + log_storage = FileStorage(log_path) + all_traces = list(log_storage.iter_msg(tag="trace")) + if not all_traces: + logger.error("No valid trace found in log directory.") + return + + trace = all_traces[-1].content + tasks.append( + ( + evaluate_one_trace, + (selector_name, trace, debug, only_sample, sample_code_path, {}, "validation", log_path, sample_rate), + ) + ) + + if not tasks: + logger.error(f"No .pkl trace files found in subdirectories of {trace_root}") + return + + # Run evaluation in parallel + hit_list = multiprocessing_wrapper(tasks, n=1) # n=1 for sequential debugging, increase for parallel runs + + # Aggregate and report results + hit_count = sum(hit for _, hit in hit_list if hit is not None) + total_valid_traces = len(hit_list) + + print("\n" + "=" * 50) + print(f"Evaluation Summary for Selector: '{selector_name}'") + print(f"Total Traces Processed: {total_valid_traces}") + print(f"Total Hits: {hit_count}") + if total_valid_traces > 0: + hit_rate = (hit_count / total_valid_traces) * 100 + print(f"Hit Rate: {hit_rate:.2f}%") + print("=" * 50 + "\n") + + result_dict["summary"] = { + "hit": hit_count, + "total": total_valid_traces, + "hit_rate": hit_rate if total_valid_traces > 0 else 0, + } + result_dict["details"] = [{comp: hit} for comp, hit in hit_list] + + with open(f"result_{selector_name}.json", "w") as f: + json.dump(result_dict, f, indent=4) + logger.info(f"Results saved to result_{selector_name}.json") + + +if __name__ == "__main__": + fire.Fire(select_on_existing_trace) diff --git a/rdagent/utils/env.py b/rdagent/utils/env.py index 2be5ac41..bc1512d1 100644 --- a/rdagent/utils/env.py +++ b/rdagent/utils/env.py @@ -476,7 +476,7 @@ class LocalEnv(Env[ASpecificLocalConf]): volumes = {} if self.conf.extra_volumes is not None: for lp, rp in self.conf.extra_volumes.items(): - volumes[lp] = rp + volumes[lp] = rp["bind"] if isinstance(rp, dict) else rp cache_path = "/tmp/sample" if "/sample/" in "".join(self.conf.extra_volumes.keys()) else "/tmp/full" Path(cache_path).mkdir(parents=True, exist_ok=True) volumes[cache_path] = T("scenarios.data_science.share:scen.cache_path").r()