gavindiaz/NexQuant
1
0
Fork 0
mirror of https://github.com/NicolasBohn/NexQuant.git synced 2026-07-27 23:47:46 +00:00
Code Issues Packages Projects Releases Wiki Activity

feat: offline selector (#1231)

Browse Source 
* offline selector test
* fix score with tensor
* sort sota list

---------

Co-authored-by: Xu Yang <peteryang@vip.qq.com>
This commit is contained in:
Tim
2025-09-09 11:56:42 +08:00
committed by GitHub
parent fc18942339
commit d4c539912a
3 changed files with 793 additions and 145 deletions
Download Patch File Download Diff File Expand all files Collapse all files
rdagent/scenarios/data_science/proposal/exp_gen/select/prompts.yaml
+34 -1
View File
@@ -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 <file>.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 %}
rdagent/scenarios/data_science/proposal/exp_gen/select/submit.py
+758 -143
View File
@@ -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)
rdagent/utils/env.py
+1 -1
View File
@@ -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()