mirror of
https://github.com/NicolasBohn/NexQuant.git
synced 2026-08-09 04:57:44 +00:00
xuangu-fang
a439d9ef2e
* start to work on multi-trace + async * init ver of async-multi-tarce, to test * add eng-ver log * complete version of async+ mul-trace * debug * fix bug on DS_RD_SETTING.get() * update * fix bug + simplif the usage of async in multi-trace * fix mini bug of arg_name * Move local_selection into class Experiment & clean the code
113 lines
4.9 KiB
Python
113 lines
4.9 KiB
Python
from __future__ import annotations
|
|
|
|
import asyncio
|
|
from datetime import timedelta
|
|
from typing import TYPE_CHECKING
|
|
|
|
from rdagent.app.data_science.conf import DS_RD_SETTING
|
|
from rdagent.core.conf import RD_AGENT_SETTINGS
|
|
from rdagent.core.proposal import ExpGen
|
|
from rdagent.log import rdagent_logger as logger
|
|
from rdagent.log.timer import RD_Agent_TIMER_wrapper, RDAgentTimer
|
|
from rdagent.scenarios.data_science.loop import DataScienceRDLoop
|
|
from rdagent.scenarios.data_science.proposal.exp_gen.merge import ExpGen2Hypothesis
|
|
from rdagent.scenarios.data_science.proposal.exp_gen.trace_scheduler import (
|
|
RoundRobinScheduler,
|
|
TraceScheduler,
|
|
)
|
|
|
|
if TYPE_CHECKING:
|
|
from rdagent.scenarios.data_science.experiment.experiment import DSExperiment
|
|
from rdagent.scenarios.data_science.proposal.exp_gen.base import DSTrace, Experiment
|
|
from rdagent.utils.workflow.loop import LoopBase
|
|
|
|
|
|
class ParallelMultiTraceExpGen(ExpGen):
|
|
"""
|
|
An experiment generation strategy that enables parallel multi-trace exploration.
|
|
|
|
This generator is designed to work with the "Attribute Injection" model.
|
|
It uses a TraceScheduler to determine which parent node to expand, and
|
|
injects this parent context into the experiment object itself.
|
|
"""
|
|
|
|
def __init__(self, *args, **kwargs):
|
|
super().__init__(*args, **kwargs)
|
|
# The underlying generator for creating a single experiment
|
|
self.exp_gen = DataScienceRDLoop._get_exp_gen(
|
|
"rdagent.scenarios.data_science.proposal.exp_gen.DSExpGen", self.scen
|
|
)
|
|
self.merge_exp_gen = ExpGen2Hypothesis(self.scen)
|
|
self.trace_scheduler: TraceScheduler = RoundRobinScheduler()
|
|
self.max_trace_num = DS_RD_SETTING.max_trace_num
|
|
|
|
def gen(self, trace: "DSTrace") -> "Experiment":
|
|
raise NotImplementedError(
|
|
"ParallelMultiTraceExpGen is designed for async usage, please call async_gen instead."
|
|
)
|
|
|
|
async def async_gen(self, trace: DSTrace, loop: LoopBase) -> DSExperiment:
|
|
"""
|
|
Waits for a free execution slot, selects a parent trace using the
|
|
scheduler, generates a new experiment, and injects the parent context
|
|
into it before returning.
|
|
"""
|
|
timer: RDAgentTimer = RD_Agent_TIMER_wrapper.timer
|
|
logger.info(f"Remain time: {timer.remain_time_duration}")
|
|
local_selection: tuple[int, ...] = None
|
|
|
|
while True:
|
|
|
|
if timer.remain_time_duration >= timedelta(hours=DS_RD_SETTING.merge_hours):
|
|
|
|
if DS_RD_SETTING.enable_inject_knowledge_at_root:
|
|
|
|
if len(trace.hist) == 0:
|
|
# set the knowledge base option to True for the first trace
|
|
DS_RD_SETTING.enable_knowledge_base = True
|
|
|
|
else:
|
|
# set the knowledge base option back to False for the other traces
|
|
DS_RD_SETTING.enable_knowledge_base = False
|
|
# step 1: select the parant trace to expand
|
|
# Policy: if we have fewer traces than our target, start a new one.
|
|
if trace.sub_trace_count < self.max_trace_num:
|
|
local_selection = trace.NEW_ROOT
|
|
else:
|
|
# Otherwise, use the scheduler to pick an existing trace to expand.
|
|
local_selection = await self.trace_scheduler.select_trace(trace)
|
|
|
|
if loop.get_unfinished_loop_cnt(loop.loop_idx) < RD_AGENT_SETTINGS.get_max_parallel():
|
|
|
|
# set the local selection as the global current selection for the trace
|
|
trace.set_current_selection(local_selection)
|
|
# step 2: generate the experiment with the local selection
|
|
exp = self.exp_gen.gen(trace)
|
|
|
|
# Inject the local selection to the experiment object
|
|
exp.set_local_selection(local_selection)
|
|
|
|
return exp
|
|
|
|
else:
|
|
# enter the merging stage
|
|
# make sure the all loops are finished
|
|
if loop.get_unfinished_loop_cnt(loop.loop_idx) < 1:
|
|
# disable reset in merging stage
|
|
DS_RD_SETTING.coding_fail_reanalyze_threshold = 100000
|
|
DS_RD_SETTING.consecutive_errors = 100000
|
|
|
|
leaves: list[int] = trace.get_leaves()
|
|
if len(leaves) < 2:
|
|
trace.set_current_selection(selection=(-1,))
|
|
return self.exp_gen.gen(trace)
|
|
else:
|
|
selection = (leaves[0],)
|
|
if trace.sota_exp_to_submit is not None:
|
|
if trace.is_parent(trace.exp2idx(trace.sota_exp_to_submit), leaves[1]):
|
|
selection = (leaves[1],)
|
|
trace.set_current_selection(selection)
|
|
return self.merge_exp_gen.gen(trace)
|
|
|
|
await asyncio.sleep(1)
|