feat: enhance timeout management and knowledge base handling in CoSTEER components (#1130)

* feat: enhance timeout management and knowledge base handling in CoSTEER components

* fix a little bug

* fix small bug

* fix a small bug

* Update rdagent/scenarios/data_science/loop.py

Co-authored-by: you-n-g <you-n-g@users.noreply.github.com>

* add scale check

* fix a small bug

* fix CI

* use dynamic chat_token_limit & remove repeated lines

* fix CI

* remove useless comment

* fix small bug

* update draft appendix

* fix prompt

* add code correctness as top priority

---------

Co-authored-by: you-n-g <you-n-g@users.noreply.github.com>
This commit is contained in:
Xu Yang
2025-07-31 13:06:28 +08:00
committed by GitHub
parent ce25bafbf7
commit 1aa5cc2535
26 changed files with 337 additions and 131 deletions
@@ -63,11 +63,6 @@ class PipelineMultiProcessEvolvingStrategy(MultiProcessEvolvingStrategy):
data_folder_info = self.scen.processed_data_folder_description
pipeline_task_info = target_task.get_task_information()
queried_similar_successful_knowledge = (
queried_knowledge.task_to_similar_task_successful_knowledge[pipeline_task_info]
if queried_knowledge is not None
else []
)
queried_former_failed_knowledge = (
queried_knowledge.task_to_former_failed_traces[pipeline_task_info] if queried_knowledge is not None else []
)
@@ -82,7 +77,6 @@ class PipelineMultiProcessEvolvingStrategy(MultiProcessEvolvingStrategy):
system_prompt = T(".prompts:pipeline_coder.system").r(
task_desc=pipeline_task_info,
queried_similar_successful_knowledge=queried_similar_successful_knowledge,
queried_former_failed_knowledge=queried_former_failed_knowledge[0],
out_spec=PythonAgentOut.get_spec(),
runtime_environment=self.scen.get_runtime_environment(),
@@ -158,10 +158,17 @@ class PipelineCoSTEEREvaluator(CoSTEEREvaluator):
else:
eda_output = implementation.file_dict.get("EDA.md", None)
queried_similar_successful_knowledge = (
queried_knowledge.task_to_similar_task_successful_knowledge[target_task.get_task_information()]
if queried_knowledge is not None
else []
)
system_prompt = T(".prompts:pipeline_eval.system").r(
is_sub_enabled=test_eval.is_sub_enabled(self.scen.competition),
debug_mode=DS_RD_SETTING.sample_data_by_LLM,
mle_check=(DS_RD_SETTING.sample_data_by_LLM and test_eval.is_sub_enabled(self.scen.competition)),
queried_similar_successful_knowledge=queried_similar_successful_knowledge,
)
user_prompt = T(".prompts:pipeline_eval.user").r(
scenario=self.scen.get_scenario_all_desc(eda_output=eda_output),
@@ -28,19 +28,6 @@ pipeline_coder:
# Specification your code should follow
{% include "scenarios.data_science.share:component_spec.Pipeline" %}
{% if queried_similar_successful_knowledge|length != 0 or queried_former_failed_knowledge|length != 0 %}
# Relevant Information for This Task
{% endif %}
{% if queried_similar_successful_knowledge|length != 0 %}
## Successful Implementations for Similar Models
====={% for similar_successful_knowledge in queried_similar_successful_knowledge %} Model {{ loop.index }}:=====
{{ similar_successful_knowledge.target_task.get_task_information() }}
=====Code:=====
{{ similar_successful_knowledge.implementation.all_codes }}
{% endfor %}
{% endif %}
{% if queried_former_failed_knowledge|length != 0 %}
## Previous Failed Attempts
{% for former_failed_knowledge in queried_former_failed_knowledge %} Attempt {{ loop.index }}:
@@ -102,6 +89,16 @@ pipeline_coder:
```bash
python main.py --debug
```
Please simulate the following code to check whether the code is running in debug mode:
```python
import argparse
parser = argparse.ArgumentParser()
parser.add_argument('--debug', action='store_true', help='Run in debug mode')
args = parser.parse_args()
DEBUG = False
if args.debug:
DEBUG = True
```
In debug mode, you should only sample ten percent of the training data and run the minimum epochs to quickly test the correctness of the code.
In debug mode, you should implement a timer to measure the time taken for your debug configuration and estimate the time required for the full run. Your timer should only measure the time taken for the training part, not the data loading or feature engineering part.
For example:
@@ -115,7 +112,8 @@ pipeline_coder:
```
In debug mode, your code should run faster, so the environment will set a shorter time limit than the standard time limit for your code.
For example, you can sample ten percent of the training data and run for one epoch, then the full run with ten epochs will take one hundred times the time taken for the debug run. The scale is calculated by yourself depending on the data sampling and epoch number you choose. If your full run enables early stopping, the scale should be smaller considering the early stopping will stop the training earlier than the full epochs.
You should sample the data after train valid split. When you split the data after sampling, you might get a class with only one sample which might cause the split strategy to fail.
Be careful about the train-valid split strategy. StratifiedShuffleSplit is highly risk since the data has some categories with only one sample. If you use StratifiedShuffleSplit, you should consider using a try-except block to catch the error and use a different split strategy if the error occurs.
You should sample the data after train valid split. When you split the data after sampling, you might get a class with only one sample which might cause the split strategy to fail.
Your debug code should run exactly the same as the full run, except for the data sampling and epoch number, to ensure the correctness of the code.
You should print total time and estimated time in standard output using print function in the following schema:
=== Start of Debug Information ===
@@ -135,16 +133,17 @@ pipeline_coder:
{% endif %}
## General Guidelines
1. Use the print() function for all output; do not use the logging module.
2. **Avoid all hard-coded values (e.g., fixed dataset sizes)**. Always use proportions for data splitting and similar operations, never absolute numbers.
3. Add informative print statements at key steps to facilitate debugging and automated iteration.
4. For model training, use reasonable epoch numbers. ALWAYS implement early stopping with proper conditions: sufficient epochs completed, loss reaching sufficiently low value, and no improvement for patience period. Save best model checkpoints based on validation performance.
5. Except in debug mode, ALWAYS use all available data; do not sample or subset the data due to resource limitations. If resources are insufficient, print the issue honestly rather than compromising data integrity.
6. Do not use tqdm or similar progress bar tools.
7. **Try-except blocks are ONLY allowed when reading files. If no files are successfully read, it indicates incorrect file paths or reading methods, not a try-except issue. Try-except is PROHIBITED elsewhere in the code. Assert statements are PROHIBITED throughout the entire code.**
8. ATTENTION: ALWAYS use the best saved model (not necessarily final epoch) for predictions. **NEVER create dummy/placeholder submissions (e.g., all 1s, random values)**. If training fails, report failure honestly rather than generating fake submission files.
9. You should ALWAYS generate the complete code rather than partial code.
10. Strictly follow all specifications and general guidelines described above.
1. Code correctness is the top priority. Ensure your code is runnable and produces the expected output even if some task requirements are not fully met because the task itself might contain some errors like the wrong package name or wrong package function names.
2. Use the print() function for all output; do not use the logging module.
3. **Avoid all hard-coded values (e.g., fixed dataset sizes)**. Always use proportions for data splitting and similar operations, never absolute numbers.
4. Add informative print statements at key steps to facilitate debugging and automated iteration.
5. For model training, use reasonable epoch numbers. ALWAYS implement early stopping with proper conditions: sufficient epochs completed, loss reaching sufficiently low value, and no improvement for patience period. Save best model checkpoints based on validation performance.
6. Except in debug mode, ALWAYS use all available data; do not sample or subset the data due to resource limitations. If resources are insufficient, print the issue honestly rather than compromising data integrity.
7. Do not use tqdm or similar progress bar tools.
8. **Try-except blocks are ONLY allowed when reading files. If no files are successfully read, it indicates incorrect file paths or reading methods, not a try-except issue. Try-except is PROHIBITED elsewhere in the code. Assert statements are PROHIBITED throughout the entire code.**
9. ATTENTION: ALWAYS use the best saved model (not necessarily final epoch) for predictions. **NEVER create dummy/placeholder submissions (e.g., all 1s, random values)**. If training fails, report failure honestly rather than generating fake submission files.
10. You should ALWAYS generate the complete code rather than partial code.
11. Strictly follow all specifications and general guidelines described above.
### Output Format
{% if out_spec %}
@@ -269,6 +268,18 @@ pipeline_eval:
- If the submission check returns an error message, you should set the "final_decision" to false and clearly document the issues in the "return_checking" field.
{% endif %}
{% if queried_similar_successful_knowledge|length != 0 %}
### Step 6: Similar Successful Implementations to help Code Improvement
The user has done several similar tasks and get some successful implementations. These code might not be implemented to the same task, but they are similar to your task and they might work well on your dataset.
Please refer to these successful implementation and provide your suggestions in your response on how to correct your current code based on these successful implementations.
## Successful Implementations for Similar Tasks
====={% for similar_successful_knowledge in queried_similar_successful_knowledge %} Similar Task {{ loop.index }}:=====
{{ similar_successful_knowledge.target_task.get_task_information() }}
=====Code:=====
{{ similar_successful_knowledge.implementation.all_codes }}
{% endfor %}
{% endif %}
## Output Format
Please respond with your feedback in the following JSON format without anything else.
```json