From 768229427d034dffead94935a639a1537b079c61 Mon Sep 17 00:00:00 2001 From: Xu Yang Date: Mon, 14 Oct 2024 17:34:09 +0800 Subject: [PATCH] feat: use unified pickle cacher & move llm config into a isolated config (#424) * simplify RDAgent conf * add unified cacher(untested) * fix small bugs * fix a bug * fix a small bug in runner * use hash_key = None to skip cache * fix CI * in factor execution, ignore cache when raise exception * add file locker to avoid mp calling * fix CI * use function __module__ name as folder in cache --- docs/installation_and_configuration.rst | 2 - docs/scens/data_agent_fin.rst | 2 +- docs/scens/data_copilot_fin.rst | 2 +- .../coder/factor_coder/CoSTEER/evaluators.py | 5 +- .../factor_coder/CoSTEER/evolving_strategy.py | 7 +- .../coder/factor_coder/CoSTEER/scheduler.py | 4 +- .../components/coder/factor_coder/config.py | 6 - .../components/coder/factor_coder/factor.py | 50 ++---- .../coder/model_coder/CoSTEER/evaluators.py | 21 ++- .../model_coder/CoSTEER/evolving_strategy.py | 3 +- rdagent/components/coder/model_coder/conf.py | 6 - rdagent/components/coder/model_coder/model.py | 33 ++-- rdagent/components/runner/__init__.py | 21 +-- rdagent/components/runner/conf.py | 16 -- rdagent/core/conf.py | 84 ++-------- rdagent/core/utils.py | 53 +++++++ rdagent/log/ui/web.py | 12 +- rdagent/oai/llm_conf.py | 83 ++++++++++ rdagent/oai/llm_utils.py | 150 +++++++++--------- .../data_mining/developer/model_runner.py | 21 +-- rdagent/scenarios/kaggle/developer/runner.py | 20 +-- .../knowledge_management/vector_base.py | 2 +- .../scenarios/qlib/developer/factor_runner.py | 14 +- .../scenarios/qlib/developer/model_runner.py | 21 +-- .../factor_experiment_loader/pdf_loader.py | 7 +- 25 files changed, 299 insertions(+), 346 deletions(-) delete mode 100644 rdagent/components/runner/conf.py create mode 100644 rdagent/oai/llm_conf.py diff --git a/docs/installation_and_configuration.rst b/docs/installation_and_configuration.rst index 23d50c1c..2ac50d0d 100644 --- a/docs/installation_and_configuration.rst +++ b/docs/installation_and_configuration.rst @@ -138,8 +138,6 @@ Configuration List +------------------------------+--------------------------------------------------+-------------------------+ | prompt_cache_path | Path to prompt cache | ./prompt_cache.db | +------------------------------+--------------------------------------------------+-------------------------+ -| session_cache_folder_location| Path to session cache | ./session_cache_folder | -+------------------------------+--------------------------------------------------+-------------------------+ | max_past_message_include | Maximum number of past messages to include | 10 | +------------------------------+--------------------------------------------------+-------------------------+ diff --git a/docs/scens/data_agent_fin.rst b/docs/scens/data_agent_fin.rst index 30e807aa..e9c42f13 100644 --- a/docs/scens/data_agent_fin.rst +++ b/docs/scens/data_agent_fin.rst @@ -133,6 +133,6 @@ The following environment variables can be set in the `.env` file to customize t .. autopydantic_settings:: rdagent.components.coder.factor_coder.config.FactorImplementSettings :settings-show-field-summary: False - :members: coder_use_cache, data_folder, data_folder_debug, cache_location, enable_execution_cache, file_based_execution_timeout, select_method, select_threshold, max_loop, knowledge_base_path, new_knowledge_base_path + :members: coder_use_cache, data_folder, data_folder_debug, file_based_execution_timeout, select_method, select_threshold, max_loop, knowledge_base_path, new_knowledge_base_path :exclude-members: Config, fail_task_trial_limit, v1_query_former_trace_limit, v1_query_similar_success_limit, v2_query_component_limit, v2_query_error_limit, v2_query_former_trace_limit, v2_error_summary, v2_knowledge_sampler :no-index: diff --git a/docs/scens/data_copilot_fin.rst b/docs/scens/data_copilot_fin.rst index 9c350fdf..89280c61 100644 --- a/docs/scens/data_copilot_fin.rst +++ b/docs/scens/data_copilot_fin.rst @@ -159,6 +159,6 @@ The following environment variables can be set in the `.env` file to customize t .. autopydantic_settings:: rdagent.components.coder.factor_coder.config.FactorImplementSettings :settings-show-field-summary: False - :members: coder_use_cache, data_folder, data_folder_debug, cache_location, enable_execution_cache, file_based_execution_timeout, select_method, select_threshold, max_loop, knowledge_base_path, new_knowledge_base_path + :members: coder_use_cache, data_folder, data_folder_debug, file_based_execution_timeout, select_method, select_threshold, max_loop, knowledge_base_path, new_knowledge_base_path :exclude-members: Config, python_bin, fail_task_trial_limit, v1_query_former_trace_limit, v1_query_similar_success_limit, v2_query_component_limit, v2_query_error_limit, v2_query_former_trace_limit, v2_error_summary, v2_knowledge_sampler :no-index: diff --git a/rdagent/components/coder/factor_coder/CoSTEER/evaluators.py b/rdagent/components/coder/factor_coder/CoSTEER/evaluators.py index 52ae0fd7..ea0e4300 100644 --- a/rdagent/components/coder/factor_coder/CoSTEER/evaluators.py +++ b/rdagent/components/coder/factor_coder/CoSTEER/evaluators.py @@ -19,6 +19,7 @@ from rdagent.core.experiment import Task, Workspace from rdagent.core.prompts import Prompts from rdagent.core.utils import multiprocessing_wrapper from rdagent.log import rdagent_logger as logger +from rdagent.oai.llm_conf import LLM_SETTINGS from rdagent.oai.llm_utils import APIBackend evaluate_prompts = Prompts(file_path=Path(__file__).parent.parent / "prompts.yaml") @@ -118,7 +119,7 @@ class FactorCodeEvaluator(FactorEvaluator): user_prompt=user_prompt, system_prompt=system_prompt, ) - > RD_AGENT_SETTINGS.chat_token_limit + > LLM_SETTINGS.chat_token_limit ): execution_feedback_to_render = execution_feedback_to_render[len(execution_feedback_to_render) // 2 :] else: @@ -521,7 +522,7 @@ class FactorFinalDecisionEvaluator(Evaluator): user_prompt=user_prompt, system_prompt=system_prompt, ) - > RD_AGENT_SETTINGS.chat_token_limit + > LLM_SETTINGS.chat_token_limit ): execution_feedback_to_render = execution_feedback_to_render[len(execution_feedback_to_render) // 2 :] else: diff --git a/rdagent/components/coder/factor_coder/CoSTEER/evolving_strategy.py b/rdagent/components/coder/factor_coder/CoSTEER/evolving_strategy.py index e165a5cd..9e1e9038 100644 --- a/rdagent/components/coder/factor_coder/CoSTEER/evolving_strategy.py +++ b/rdagent/components/coder/factor_coder/CoSTEER/evolving_strategy.py @@ -22,6 +22,7 @@ from rdagent.core.evolving_framework import EvolvingStrategy, QueriedKnowledge from rdagent.core.experiment import Workspace from rdagent.core.prompts import Prompts from rdagent.core.utils import multiprocessing_wrapper +from rdagent.oai.llm_conf import LLM_SETTINGS from rdagent.oai.llm_utils import APIBackend if TYPE_CHECKING: @@ -160,7 +161,7 @@ class FactorEvolvingStrategy(MultiProcessEvolvingStrategy): session.build_chat_completion_message_and_calculate_token( user_prompt, ) - < RD_AGENT_SETTINGS.chat_token_limit + < LLM_SETTINGS.chat_token_limit ): break elif len(queried_former_failed_knowledge_to_render) > 1: @@ -281,7 +282,7 @@ class FactorEvolvingStrategyWithGraph(MultiProcessEvolvingStrategy): ) if ( session_summary.build_chat_completion_message_and_calculate_token(error_summary_user_prompt) - < RD_AGENT_SETTINGS.chat_token_limit + < LLM_SETTINGS.chat_token_limit ): break elif len(queried_similar_error_knowledge_to_render) > 0: @@ -310,7 +311,7 @@ class FactorEvolvingStrategyWithGraph(MultiProcessEvolvingStrategy): session.build_chat_completion_message_and_calculate_token( user_prompt, ) - < RD_AGENT_SETTINGS.chat_token_limit + < LLM_SETTINGS.chat_token_limit ): break elif len(queried_former_failed_knowledge_to_render) > 1: diff --git a/rdagent/components/coder/factor_coder/CoSTEER/scheduler.py b/rdagent/components/coder/factor_coder/CoSTEER/scheduler.py index ecbdf1a4..ffb1bfc4 100644 --- a/rdagent/components/coder/factor_coder/CoSTEER/scheduler.py +++ b/rdagent/components/coder/factor_coder/CoSTEER/scheduler.py @@ -7,10 +7,10 @@ from jinja2 import Environment, StrictUndefined from rdagent.components.coder.factor_coder.CoSTEER.evolvable_subjects import ( FactorEvolvingItem, ) -from rdagent.core.conf import RD_AGENT_SETTINGS from rdagent.core.prompts import Prompts from rdagent.core.scenario import Scenario from rdagent.log import rdagent_logger as logger +from rdagent.oai.llm_conf import LLM_SETTINGS from rdagent.oai.llm_utils import APIBackend scheduler_prompts = Prompts(file_path=Path(__file__).parent.parent / "prompts.yaml") @@ -68,7 +68,7 @@ def LLMSelect( user_prompt=user_prompt, system_prompt=system_prompt, ) - < RD_AGENT_SETTINGS.chat_token_limit + < LLM_SETTINGS.chat_token_limit ): break diff --git a/rdagent/components/coder/factor_coder/config.py b/rdagent/components/coder/factor_coder/config.py index 702b304c..4ef2a980 100644 --- a/rdagent/components/coder/factor_coder/config.py +++ b/rdagent/components/coder/factor_coder/config.py @@ -20,12 +20,6 @@ class FactorImplementSettings(BaseSettings): data_folder_debug: str = "git_ignore_folder/factor_implementation_source_data_debug" """Path to the folder containing partial financial data (for debugging)""" - cache_location: str = "git_ignore_folder/factor_implementation_execution_cache" - """Path to the cache location""" - - enable_execution_cache: bool = True - """Indicates whether to enable the execution cache""" - # TODO: the factor implement specific settings should not appear in this settings # Evolving should have a method specific settings # evolving related config diff --git a/rdagent/components/coder/factor_coder/factor.py b/rdagent/components/coder/factor_coder/factor.py index 72f8b535..723db08a 100644 --- a/rdagent/components/coder/factor_coder/factor.py +++ b/rdagent/components/coder/factor_coder/factor.py @@ -13,6 +13,7 @@ from rdagent.app.kaggle.conf import KAGGLE_IMPLEMENT_SETTING from rdagent.components.coder.factor_coder.config import FACTOR_IMPLEMENT_SETTINGS from rdagent.core.exception import CodeFormatError, CustomRuntimeError, NoOutputError from rdagent.core.experiment import Experiment, FBWorkspace, Task +from rdagent.core.utils import cache_with_pickle from rdagent.log import rdagent_logger as logger from rdagent.oai.llm_utils import md5_hash @@ -80,17 +81,21 @@ class FactorFBWorkspace(FBWorkspace): def __init__( self, *args, - executed_factor_value_dataframe: pd.DataFrame = None, raise_exception: bool = False, **kwargs, ) -> None: super().__init__(*args, **kwargs) - self.executed_factor_value_dataframe = executed_factor_value_dataframe self.raise_exception = raise_exception - def execute( - self, enable_cache: bool = FACTOR_IMPLEMENT_SETTINGS.enable_execution_cache, data_type: str = "Debug" - ) -> Tuple[str, pd.DataFrame]: + def hash_func(self, data_type: str = "Debug") -> str: + return ( + md5_hash(data_type + self.code_dict["factor.py"]) + if ("factor.py" in self.code_dict and not self.raise_exception) + else None + ) + + @cache_with_pickle(hash_func) + def execute(self, data_type: str = "Debug") -> Tuple[str, pd.DataFrame]: """ execute the implementation and get the factor value by the following steps: 1. make the directory in workspace path @@ -108,8 +113,6 @@ class FactorFBWorkspace(FBWorkspace): 1. We will store the function's return value to ensure it behaves as expected. - The cached information will include a tuple with the following: (execution_feedback, executed_factor_value_dataframe, Optional[Exception]) - parameters: - enable_cache: if True, store the factor value in the instance variable, this feature is to be used in the gt implementation to avoid multiple execution on the same gt implementation """ super().execute() if self.code_dict is None or "factor.py" not in self.code_dict: @@ -118,31 +121,6 @@ class FactorFBWorkspace(FBWorkspace): else: return self.FB_CODE_NOT_SET, None with FileLock(self.workspace_path / "execution.lock"): - if FACTOR_IMPLEMENT_SETTINGS.enable_execution_cache: - # NOTE: cache the result for the same code and same data type - target_file_name = md5_hash(data_type + self.code_dict["factor.py"]) - cache_file_path = Path(FACTOR_IMPLEMENT_SETTINGS.cache_location) / f"{target_file_name}.pkl" - Path(FACTOR_IMPLEMENT_SETTINGS.cache_location).mkdir(exist_ok=True, parents=True) - if enable_cache and cache_file_path.exists(): - cached_res = pickle.load(open(cache_file_path, "rb")) - - if len(cached_res) == 2: - # NOTE: this is trying to be compatible with previous results. - # Previously, the exception is not saved. we should not enable the cache mechanism - # othersise we can raise the exception directly. - if self.raise_exception: - pass # pass to disable the cache mechanism - else: - self.executed_factor_value_dataframe = cached_res[1] - return cached_res - else: - # NOTE: (execution_feedback, executed_factor_value_dataframe, Optional[Exception]) - if self.raise_exception and cached_res[-1] is not None: - raise cached_res[-1] - else: - self.executed_factor_value_dataframe = cached_res[1] - return cached_res[:2] - if self.target_task.version == 1: source_data_path = ( Path( @@ -220,14 +198,6 @@ class FactorFBWorkspace(FBWorkspace): else: execution_error = NoOutputError(execution_feedback) - if enable_cache and executed_factor_value_dataframe is not None: - self.executed_factor_value_dataframe = executed_factor_value_dataframe - - if FACTOR_IMPLEMENT_SETTINGS.enable_execution_cache: - pickle.dump( - (execution_feedback, executed_factor_value_dataframe, execution_error), - open(cache_file_path, "wb"), - ) return execution_feedback, executed_factor_value_dataframe def __str__(self) -> str: diff --git a/rdagent/components/coder/model_coder/CoSTEER/evaluators.py b/rdagent/components/coder/model_coder/CoSTEER/evaluators.py index f6381362..903d7430 100644 --- a/rdagent/components/coder/model_coder/CoSTEER/evaluators.py +++ b/rdagent/components/coder/model_coder/CoSTEER/evaluators.py @@ -19,6 +19,7 @@ from rdagent.core.experiment import Task, Workspace from rdagent.core.prompts import Prompts from rdagent.core.utils import multiprocessing_wrapper from rdagent.log import rdagent_logger as logger +from rdagent.oai.llm_conf import LLM_SETTINGS from rdagent.oai.llm_utils import APIBackend evaluate_prompts = Prompts(file_path=Path(__file__).parent.parent / "prompts.yaml") @@ -90,9 +91,11 @@ class ModelCodeEvaluator(Evaluator): Environment(undefined=StrictUndefined) .from_string(evaluate_prompts["evaluator_code_feedback"]["system"]) .render( - scenario=self.scen.get_scenario_all_desc(target_task) - if self.scen is not None - else "No scenario description." + scenario=( + self.scen.get_scenario_all_desc(target_task) + if self.scen is not None + else "No scenario description." + ) ) ) @@ -116,7 +119,7 @@ class ModelCodeEvaluator(Evaluator): user_prompt=user_prompt, system_prompt=system_prompt, ) - > RD_AGENT_SETTINGS.chat_token_limit + > LLM_SETTINGS.chat_token_limit ): execution_feedback_to_render = execution_feedback_to_render[len(execution_feedback_to_render) // 2 :] else: @@ -150,9 +153,11 @@ class ModelFinalEvaluator(Evaluator): Environment(undefined=StrictUndefined) .from_string(evaluate_prompts["evaluator_final_feedback"]["system"]) .render( - scenario=self.scen.get_scenario_all_desc(target_task) - if self.scen is not None - else "No scenario description." + scenario=( + self.scen.get_scenario_all_desc(target_task) + if self.scen is not None + else "No scenario description." + ) ) ) @@ -176,7 +181,7 @@ class ModelFinalEvaluator(Evaluator): user_prompt=user_prompt, system_prompt=system_prompt, ) - > RD_AGENT_SETTINGS.chat_token_limit + > LLM_SETTINGS.chat_token_limit ): execution_feedback_to_render = execution_feedback_to_render[len(execution_feedback_to_render) // 2 :] else: diff --git a/rdagent/components/coder/model_coder/CoSTEER/evolving_strategy.py b/rdagent/components/coder/model_coder/CoSTEER/evolving_strategy.py index 6d3eee5a..154dad10 100644 --- a/rdagent/components/coder/model_coder/CoSTEER/evolving_strategy.py +++ b/rdagent/components/coder/model_coder/CoSTEER/evolving_strategy.py @@ -20,6 +20,7 @@ from rdagent.core.conf import RD_AGENT_SETTINGS from rdagent.core.evolving_framework import EvolvingStrategy from rdagent.core.prompts import Prompts from rdagent.core.utils import multiprocessing_wrapper +from rdagent.oai.llm_conf import LLM_SETTINGS from rdagent.oai.llm_utils import APIBackend from rdagent.scenarios.kaggle.experiment.kaggle_experiment import KG_MODEL_MAPPING @@ -100,7 +101,7 @@ class ModelCoderEvolvingStrategy(EvolvingStrategy): user_prompt=user_prompt, system_prompt=system_prompt, ) - < RD_AGENT_SETTINGS.chat_token_limit + < LLM_SETTINGS.chat_token_limit ): break elif len(queried_former_failed_knowledge_to_render) > 1: diff --git a/rdagent/components/coder/model_coder/conf.py b/rdagent/components/coder/model_coder/conf.py index ea45041b..c600c699 100644 --- a/rdagent/components/coder/model_coder/conf.py +++ b/rdagent/components/coder/model_coder/conf.py @@ -10,10 +10,6 @@ class ModelImplSettings(BaseSettings): coder_use_cache: bool = False - cache_location: str = str( - (Path().cwd() / "git_ignore_folder" / "model_implementation_execution_cache").absolute(), - ) - knowledge_base_path: Union[str, None] = None new_knowledge_base_path: Union[str, None] = None @@ -23,7 +19,5 @@ class ModelImplSettings(BaseSettings): query_similar_success_limit: int = 5 fail_task_trial_limit: int = 20 - enable_execution_cache: bool = True # whether to enable the execution cache - MODEL_IMPL_SETTINGS = ModelImplSettings() diff --git a/rdagent/components/coder/model_coder/model.py b/rdagent/components/coder/model_coder/model.py index a8f700e9..66b4a972 100644 --- a/rdagent/components/coder/model_coder/model.py +++ b/rdagent/components/coder/model_coder/model.py @@ -6,6 +6,7 @@ from typing import Dict, Optional from rdagent.components.coder.model_coder.conf import MODEL_IMPL_SETTINGS from rdagent.core.experiment import Experiment, FBWorkspace, Task +from rdagent.core.utils import cache_with_pickle from rdagent.oai.llm_utils import md5_hash from rdagent.utils.env import KGDockerEnv, QTDockerEnv @@ -71,6 +72,21 @@ class ModelFBWorkspace(FBWorkspace): (version 2) for kaggle we'll make a script to call the fit and predict function in the implementation in file `model.py` after setting the cwd into the directory """ + def hash_func( + self, + batch_size: int = 8, + num_features: int = 10, + num_timesteps: int = 4, + num_edges: int = 20, + input_value: float = 1.0, + param_init_value: float = 1.0, + ) -> str: + target_file_name = f"{batch_size}_{num_features}_{num_timesteps}_{input_value}_{param_init_value}" + for code_file_name in sorted(list(self.code_dict.keys())): + target_file_name = f"{target_file_name}_{self.code_dict[code_file_name]}" + return md5_hash(target_file_name) + + @cache_with_pickle(hash_func) def execute( self, batch_size: int = 8, @@ -82,17 +98,6 @@ class ModelFBWorkspace(FBWorkspace): ): super().execute() try: - if MODEL_IMPL_SETTINGS.enable_execution_cache: - # NOTE: cache the result for the same code - target_file_name = f"{batch_size}_{num_features}_{num_timesteps}_{input_value}_{param_init_value}" - for code_file_name in sorted(list(self.code_dict.keys())): - target_file_name = f"{target_file_name}_{self.code_dict[code_file_name]}" - target_file_name = md5_hash(target_file_name) - cache_file_path = Path(MODEL_IMPL_SETTINGS.cache_location) / f"{target_file_name}.pkl" - Path(MODEL_IMPL_SETTINGS.cache_location).mkdir(exist_ok=True, parents=True) - if cache_file_path.exists(): - return pickle.load(open(cache_file_path, "rb")) - qtde = QTDockerEnv() if self.target_task.version == 1 else KGDockerEnv() qtde.prepare() @@ -121,12 +126,6 @@ PARAM_INIT_VALUE = {param_init_value} raise RuntimeError(f"Error in running the model code: {log}") [execution_feedback_str, execution_model_output] = results - if MODEL_IMPL_SETTINGS.enable_execution_cache: - pickle.dump( - (execution_feedback_str, execution_model_output), - open(cache_file_path, "wb"), - ) - except Exception as e: execution_feedback_str = f"Execution error: {e}\nTraceback: {traceback.format_exc()}" execution_model_output = None diff --git a/rdagent/components/runner/__init__.py b/rdagent/components/runner/__init__.py index 997bcd6a..256b3095 100644 --- a/rdagent/components/runner/__init__.py +++ b/rdagent/components/runner/__init__.py @@ -1,8 +1,7 @@ import pickle from pathlib import Path -from typing import Tuple +from typing import Any, Tuple -from rdagent.components.runner.conf import RUNNER_SETTINGS from rdagent.core.developer import Developer from rdagent.core.experiment import ASpecificExp, Experiment from rdagent.oai.llm_utils import md5_hash @@ -18,16 +17,8 @@ class CachedRunner(Developer[ASpecificExp]): task_info_str = "\n".join(task_info_list) return md5_hash(task_info_str) - def get_cache_result(self, exp: Experiment) -> Tuple[bool, object]: - task_info_key = self.get_cache_key(exp) - Path(RUNNER_SETTINGS.cache_path).mkdir(parents=True, exist_ok=True) - cache_path = Path(RUNNER_SETTINGS.cache_path) / f"{task_info_key}.pkl" - if cache_path.exists(): - return True, pickle.load(open(cache_path, "rb")) - else: - return False, None - - def dump_cache_result(self, exp: Experiment, result: object): - task_info_key = self.get_cache_key(exp) - cache_path = Path(RUNNER_SETTINGS.cache_path) / f"{task_info_key}.pkl" - pickle.dump(result, open(cache_path, "wb")) + def assign_cached_result(self, exp: Experiment, cached_res: Experiment) -> Experiment: + if exp.based_experiments and exp.based_experiments[-1].result is None: + exp.based_experiments[-1].result = cached_res.based_experiments[-1].result + exp.result = cached_res.result + return exp diff --git a/rdagent/components/runner/conf.py b/rdagent/components/runner/conf.py deleted file mode 100644 index 92b777ca..00000000 --- a/rdagent/components/runner/conf.py +++ /dev/null @@ -1,16 +0,0 @@ -from __future__ import annotations - -from pathlib import Path - -from pydantic_settings import BaseSettings - - -class RunnerSettings(BaseSettings): - class Config: - env_prefix = "RUNNER_" # Use RUNNER_ as prefix for environment variables - - cache_result: bool = True # whether to cache the result of the docker execution - cache_path: str = str(Path.cwd() / "runner_cache/") # the path to store the cache - - -RUNNER_SETTINGS = RunnerSettings() diff --git a/rdagent/core/conf.py b/rdagent/core/conf.py index 00f92685..9f57d32d 100644 --- a/rdagent/core/conf.py +++ b/rdagent/core/conf.py @@ -14,84 +14,10 @@ class RDAgentSettings(BaseSettings): # Log configs # TODO: (xiao) think it can be a separate config. log_trace_path: str | None = None - log_llm_chat_content: bool = True - - use_azure: bool = False - use_azure_token_provider: bool = False - managed_identity_client_id: str | None = None - max_retry: int = 10 - retry_wait_seconds: int = 1 - dump_chat_cache: bool = False - use_chat_cache: bool = False - dump_embedding_cache: bool = False - use_embedding_cache: bool = False - prompt_cache_path: str = str(Path.cwd() / "prompt_cache.db") - session_cache_folder_location: str = str(Path.cwd() / "session_cache_folder/") - max_past_message_include: int = 10 - - # Chat configs - openai_api_key: str = "" # TODO: simplify the key design. - chat_openai_api_key: str = "" - chat_azure_api_base: str = "" - chat_azure_api_version: str = "" - chat_model: str = "gpt-4-turbo" - chat_max_tokens: int = 3000 - chat_temperature: float = 0.5 - chat_stream: bool = True - chat_seed: int | None = None - chat_frequency_penalty: float = 0.0 - chat_presence_penalty: float = 0.0 - chat_token_limit: int = ( - 100000 # 100000 is the maximum limit of gpt4, which might increase in the future version of gpt - ) - default_system_prompt: str = "You are an AI assistant who helps to answer user's questions." - - # Embedding configs - embedding_openai_api_key: str = "" - embedding_azure_api_base: str = "" - embedding_azure_api_version: str = "" - embedding_model: str = "" - embedding_max_str_num: int = 50 - - # offline llama2 related config - use_llama2: bool = False - llama2_ckpt_dir: str = "Llama-2-7b-chat" - llama2_tokenizer_path: str = "Llama-2-7b-chat/tokenizer.model" - llams2_max_batch_size: int = 8 # azure document intelligence configs azure_document_intelligence_key: str = "" azure_document_intelligence_endpoint: str = "" - - # server served endpoints - use_gcr_endpoint: bool = False - gcr_endpoint_type: str = "llama2_70b" # or "llama3_70b", "phi2", "phi3_4k", "phi3_128k" - - llama2_70b_endpoint: str = "" - llama2_70b_endpoint_key: str = "" - llama2_70b_endpoint_deployment: str = "" - - llama3_70b_endpoint: str = "" - llama3_70b_endpoint_key: str = "" - llama3_70b_endpoint_deployment: str = "" - - phi2_endpoint: str = "" - phi2_endpoint_key: str = "" - phi2_endpoint_deployment: str = "" - - phi3_4k_endpoint: str = "" - phi3_4k_endpoint_key: str = "" - phi3_4k_endpoint_deployment: str = "" - - phi3_128k_endpoint: str = "" - phi3_128k_endpoint_key: str = "" - phi3_128k_endpoint_deployment: str = "" - - gcr_endpoint_temperature: float = 0.7 - gcr_endpoint_top_p: float = 0.9 - gcr_endpoint_do_sample: bool = False - gcr_endpoint_max_token: int = 100 - # factor extraction conf max_input_duplicate_factor_group: int = 300 max_output_duplicate_factor_group: int = 20 @@ -103,5 +29,15 @@ class RDAgentSettings(BaseSettings): # multi processing conf multi_proc_n: int = 1 + # pickle cache conf + cache_with_pickle: bool = True # whether to use pickle cache + pickle_cache_folder_path_str: str = str( + Path.cwd() / "pickle_cache/", + ) # the path of the folder to store the pickle cache + use_file_lock: bool = ( + True # when calling the function with same parameters, whether to use file lock to avoid + # executing the function multiple times + ) + RD_AGENT_SETTINGS = RDAgentSettings() diff --git a/rdagent/core/utils.py b/rdagent/core/utils.py index 3d0d8a68..86aada12 100644 --- a/rdagent/core/utils.py +++ b/rdagent/core/utils.py @@ -1,14 +1,19 @@ from __future__ import annotations +import functools import importlib import json import multiprocessing as mp import pickle from collections.abc import Callable +from pathlib import Path from typing import Any, ClassVar, NoReturn, cast +from filelock import FileLock from fuzzywuzzy import fuzz # type: ignore[import-untyped] +from rdagent.core.conf import RD_AGENT_SETTINGS + class RDAgentException(Exception): # noqa: N818 pass @@ -103,3 +108,51 @@ def multiprocessing_wrapper(func_calls: list[tuple[Callable, tuple]], n: int) -> with mp.Pool(processes=n) as pool: results = [pool.apply_async(f, args) for f, args in func_calls] return [result.get() for result in results] + + +def cache_with_pickle(hash_func: Callable, post_process_func: Callable | None = None) -> Callable: + """ + This decorator will cache the return value of the function with pickle. + The cache key is generated by the hash_func. The hash function returns a string or None. + If it returns None, the cache will not be used. The cache will be stored in the folder + specified by RD_AGENT_SETTINGS.pickle_cache_folder_path_str with name hash_key.pkl. + The post_process_func will be called with the original arguments and the cached result + to give each caller a chance to process the cached result. The post_process_func should + return the final result. + """ + + def cache_decorator(func: Callable) -> Callable: + @functools.wraps(func) + def cache_wrapper(*args: Any, **kwargs: Any) -> Any: + if RD_AGENT_SETTINGS.cache_with_pickle: + target_folder = Path(RD_AGENT_SETTINGS.pickle_cache_folder_path_str) / func.__module__ + target_folder.mkdir(parents=True, exist_ok=True) + hash_key = hash_func(*args, **kwargs) + if hash_key is not None and (target_folder / (hash_key + ".pkl")).exists(): + with Path.open( + target_folder / (hash_key + ".pkl"), + "rb", + ) as f: + cached_res = pickle.load(f) + return ( + post_process_func(*args, cached_res=cached_res, **kwargs) + if post_process_func is not None + else cached_res + ) + if hash_key is not None and RD_AGENT_SETTINGS.use_file_lock: + with FileLock(target_folder / (hash_key + ".lock")): + result = func(*args, **kwargs) + if hash_key is not None: + with Path.open( + target_folder / (hash_key + ".pkl"), + "wb", + ) as f: + pickle.dump(result, f) + else: + result = func(*args, **kwargs) + + return result + + return cache_wrapper + + return cache_decorator diff --git a/rdagent/log/ui/web.py b/rdagent/log/ui/web.py index b69f7645..c4bbe5ee 100644 --- a/rdagent/log/ui/web.py +++ b/rdagent/log/ui/web.py @@ -291,10 +291,6 @@ class WorkspaceWindow(StWindow): self.container.markdown(f"`{k}`") self.container.code(v, language="python") - # executed_factor_value_dataframe - # if isinstance(ws, FactorFBWorkspace): - # self.container.dataframe(ws.executed_factor_value_dataframe) - class QlibFactorExpWindow(StWindow): def __init__(self, container: DeltaGenerator, show_task_info: bool = False): @@ -622,9 +618,11 @@ class TraceWindow(StWindow): self.hypothesis_decisions[self.hypotheses[-1]] = msg.content.decision self.summary_c.markdown( "\n".join( - f"{id+1}. :green[{self.hypotheses[id].hypothesis}]\n\t>*{self.hypotheses[id].concise_reason}*" - if d - else f"{id+1}. {self.hypotheses[id].hypothesis}\n\t>*{self.hypotheses[id].concise_reason}*" + ( + f"{id+1}. :green[{self.hypotheses[id].hypothesis}]\n\t>*{self.hypotheses[id].concise_reason}*" + if d + else f"{id+1}. {self.hypotheses[id].hypothesis}\n\t>*{self.hypotheses[id].concise_reason}*" + ) for id, (h, d) in enumerate(self.hypothesis_decisions.items()) ) ) diff --git a/rdagent/oai/llm_conf.py b/rdagent/oai/llm_conf.py new file mode 100644 index 00000000..8244a921 --- /dev/null +++ b/rdagent/oai/llm_conf.py @@ -0,0 +1,83 @@ +from __future__ import annotations + +from pathlib import Path + +from pydantic_settings import BaseSettings + + +class LLMSettings(BaseSettings): + log_llm_chat_content: bool = True + + use_azure: bool = False + use_azure_token_provider: bool = False + managed_identity_client_id: str | None = None + max_retry: int = 10 + retry_wait_seconds: int = 1 + dump_chat_cache: bool = False + use_chat_cache: bool = False + dump_embedding_cache: bool = False + use_embedding_cache: bool = False + prompt_cache_path: str = str(Path.cwd() / "prompt_cache.db") + max_past_message_include: int = 10 + + # Chat configs + openai_api_key: str = "" # TODO: simplify the key design. + chat_openai_api_key: str = "" + chat_azure_api_base: str = "" + chat_azure_api_version: str = "" + chat_model: str = "gpt-4-turbo" + chat_max_tokens: int = 3000 + chat_temperature: float = 0.5 + chat_stream: bool = True + chat_seed: int | None = None + chat_frequency_penalty: float = 0.0 + chat_presence_penalty: float = 0.0 + chat_token_limit: int = ( + 100000 # 100000 is the maximum limit of gpt4, which might increase in the future version of gpt + ) + default_system_prompt: str = "You are an AI assistant who helps to answer user's questions." + + # Embedding configs + embedding_openai_api_key: str = "" + embedding_azure_api_base: str = "" + embedding_azure_api_version: str = "" + embedding_model: str = "" + embedding_max_str_num: int = 50 + + # offline llama2 related config + use_llama2: bool = False + llama2_ckpt_dir: str = "Llama-2-7b-chat" + llama2_tokenizer_path: str = "Llama-2-7b-chat/tokenizer.model" + llams2_max_batch_size: int = 8 + + # server served endpoints + use_gcr_endpoint: bool = False + gcr_endpoint_type: str = "llama2_70b" # or "llama3_70b", "phi2", "phi3_4k", "phi3_128k" + + llama2_70b_endpoint: str = "" + llama2_70b_endpoint_key: str = "" + llama2_70b_endpoint_deployment: str = "" + + llama3_70b_endpoint: str = "" + llama3_70b_endpoint_key: str = "" + llama3_70b_endpoint_deployment: str = "" + + phi2_endpoint: str = "" + phi2_endpoint_key: str = "" + phi2_endpoint_deployment: str = "" + + phi3_4k_endpoint: str = "" + phi3_4k_endpoint_key: str = "" + phi3_4k_endpoint_deployment: str = "" + + phi3_128k_endpoint: str = "" + phi3_128k_endpoint_key: str = "" + phi3_128k_endpoint_deployment: str = "" + + gcr_endpoint_temperature: float = 0.7 + gcr_endpoint_top_p: float = 0.9 + gcr_endpoint_do_sample: bool = False + gcr_endpoint_max_token: int = 100 + + +LLM_SETTINGS = LLMSettings() diff --git a/rdagent/oai/llm_utils.py b/rdagent/oai/llm_utils.py index 2f5db594..ef83bd6f 100644 --- a/rdagent/oai/llm_utils.py +++ b/rdagent/oai/llm_utils.py @@ -1,9 +1,7 @@ from __future__ import annotations -import datetime import hashlib import json -import multiprocessing import os import re import sqlite3 @@ -18,10 +16,10 @@ from typing import Any, Optional import numpy as np import tiktoken -from rdagent.core.conf import RD_AGENT_SETTINGS from rdagent.core.utils import SingletonBaseClass from rdagent.log import LogColors from rdagent.log import rdagent_logger as logger +from rdagent.oai.llm_conf import LLM_SETTINGS DEFAULT_QLIB_DOT_PATH = Path("./") @@ -170,7 +168,7 @@ class SQliteLazyCache(SingletonBaseClass): class SessionChatHistoryCache(SingletonBaseClass): def __init__(self) -> None: """load all history conversation json file from self.session_cache_location""" - self.cache = SQliteLazyCache(cache_location=RD_AGENT_SETTINGS.prompt_cache_path) + self.cache = SQliteLazyCache(cache_location=LLM_SETTINGS.prompt_cache_path) def message_get(self, conversation_id: str) -> list[str]: return self.cache.message_get(conversation_id) @@ -182,8 +180,7 @@ class SessionChatHistoryCache(SingletonBaseClass): class ChatSession: def __init__(self, api_backend: Any, conversation_id: str | None = None, system_prompt: str | None = None) -> None: self.conversation_id = str(uuid.uuid4()) if conversation_id is None else conversation_id - self.cfg = RD_AGENT_SETTINGS - self.system_prompt = system_prompt if system_prompt is not None else self.cfg.default_system_prompt + self.system_prompt = system_prompt if system_prompt is not None else LLM_SETTINGS.default_system_prompt self.api_backend = api_backend def build_chat_completion_message(self, user_prompt: str) -> list[dict[str, Any]]: @@ -243,7 +240,7 @@ class APIBackend: """ # FIXME: (xiao) We should avoid using self.xxxx. - # Instead, we can use self.cfg directly. If it's difficult to support different backend settings, we can split them into multiple BaseSettings. + # Instead, we can use LLM_SETTINGS directly. If it's difficult to support different backend settings, we can split them into multiple BaseSettings. def __init__( # noqa: C901, PLR0912, PLR0915 self, *, @@ -260,37 +257,36 @@ class APIBackend: use_embedding_cache: bool | None = None, dump_embedding_cache: bool | None = None, ) -> None: - self.cfg = RD_AGENT_SETTINGS - if self.cfg.use_llama2: + if LLM_SETTINGS.use_llama2: self.generator = Llama.build( - ckpt_dir=self.cfg.llama2_ckpt_dir, - tokenizer_path=self.cfg.llama2_tokenizer_path, - max_seq_len=self.cfg.max_tokens, - max_batch_size=self.cfg.llams2_max_batch_size, + ckpt_dir=LLM_SETTINGS.llama2_ckpt_dir, + tokenizer_path=LLM_SETTINGS.llama2_tokenizer_path, + max_seq_len=LLM_SETTINGS.max_tokens, + max_batch_size=LLM_SETTINGS.llams2_max_batch_size, ) self.encoder = None - elif self.cfg.use_gcr_endpoint: - gcr_endpoint_type = self.cfg.gcr_endpoint_type + elif LLM_SETTINGS.use_gcr_endpoint: + gcr_endpoint_type = LLM_SETTINGS.gcr_endpoint_type if gcr_endpoint_type == "llama2_70b": - self.gcr_endpoint_key = self.cfg.llama2_70b_endpoint_key - self.gcr_endpoint_deployment = self.cfg.llama2_70b_endpoint_deployment - self.gcr_endpoint = self.cfg.llama2_70b_endpoint + self.gcr_endpoint_key = LLM_SETTINGS.llama2_70b_endpoint_key + self.gcr_endpoint_deployment = LLM_SETTINGS.llama2_70b_endpoint_deployment + self.gcr_endpoint = LLM_SETTINGS.llama2_70b_endpoint elif gcr_endpoint_type == "llama3_70b": - self.gcr_endpoint_key = self.cfg.llama3_70b_endpoint_key - self.gcr_endpoint_deployment = self.cfg.llama3_70b_endpoint_deployment - self.gcr_endpoint = self.cfg.llama3_70b_endpoint + self.gcr_endpoint_key = LLM_SETTINGS.llama3_70b_endpoint_key + self.gcr_endpoint_deployment = LLM_SETTINGS.llama3_70b_endpoint_deployment + self.gcr_endpoint = LLM_SETTINGS.llama3_70b_endpoint elif gcr_endpoint_type == "phi2": - self.gcr_endpoint_key = self.cfg.phi2_endpoint_key - self.gcr_endpoint_deployment = self.cfg.phi2_endpoint_deployment - self.gcr_endpoint = self.cfg.phi2_endpoint + self.gcr_endpoint_key = LLM_SETTINGS.phi2_endpoint_key + self.gcr_endpoint_deployment = LLM_SETTINGS.phi2_endpoint_deployment + self.gcr_endpoint = LLM_SETTINGS.phi2_endpoint elif gcr_endpoint_type == "phi3_4k": - self.gcr_endpoint_key = self.cfg.phi3_4k_endpoint_key - self.gcr_endpoint_deployment = self.cfg.phi3_4k_endpoint_deployment - self.gcr_endpoint = self.cfg.phi3_4k_endpoint + self.gcr_endpoint_key = LLM_SETTINGS.phi3_4k_endpoint_key + self.gcr_endpoint_deployment = LLM_SETTINGS.phi3_4k_endpoint_deployment + self.gcr_endpoint = LLM_SETTINGS.phi3_4k_endpoint elif gcr_endpoint_type == "phi3_128k": - self.gcr_endpoint_key = self.cfg.phi3_128k_endpoint_key - self.gcr_endpoint_deployment = self.cfg.phi3_128k_endpoint_deployment - self.gcr_endpoint = self.cfg.phi3_128k_endpoint + self.gcr_endpoint_key = LLM_SETTINGS.phi3_128k_endpoint_key + self.gcr_endpoint_deployment = LLM_SETTINGS.phi3_128k_endpoint_deployment + self.gcr_endpoint = LLM_SETTINGS.phi3_128k_endpoint else: error_message = f"Invalid gcr_endpoint_type: {gcr_endpoint_type}" raise ValueError(error_message) @@ -299,46 +295,48 @@ class APIBackend: "Authorization": ("Bearer " + self.gcr_endpoint_key), "azureml-model-deployment": self.gcr_endpoint_deployment, } - self.gcr_endpoint_temperature = self.cfg.gcr_endpoint_temperature - self.gcr_endpoint_top_p = self.cfg.gcr_endpoint_top_p - self.gcr_endpoint_do_sample = self.cfg.gcr_endpoint_do_sample - self.gcr_endpoint_max_token = self.cfg.gcr_endpoint_max_token + self.gcr_endpoint_temperature = LLM_SETTINGS.gcr_endpoint_temperature + self.gcr_endpoint_top_p = LLM_SETTINGS.gcr_endpoint_top_p + self.gcr_endpoint_do_sample = LLM_SETTINGS.gcr_endpoint_do_sample + self.gcr_endpoint_max_token = LLM_SETTINGS.gcr_endpoint_max_token if not os.environ.get("PYTHONHTTPSVERIFY", "") and hasattr(ssl, "_create_unverified_context"): ssl._create_default_https_context = ssl._create_unverified_context # noqa: SLF001 self.encoder = None else: - self.use_azure = self.cfg.use_azure - self.use_azure_token_provider = self.cfg.use_azure_token_provider - self.managed_identity_client_id = self.cfg.managed_identity_client_id + self.use_azure = LLM_SETTINGS.use_azure + self.use_azure_token_provider = LLM_SETTINGS.use_azure_token_provider + self.managed_identity_client_id = LLM_SETTINGS.managed_identity_client_id # Priority: chat_api_key/embedding_api_key > openai_api_key > os.environ.get("OPENAI_API_KEY") # TODO: Simplify the key design. Consider Pandatic's field alias & priority. self.chat_api_key = ( chat_api_key - or self.cfg.chat_openai_api_key - or self.cfg.openai_api_key + or LLM_SETTINGS.chat_openai_api_key + or LLM_SETTINGS.openai_api_key or os.environ.get("OPENAI_API_KEY") ) self.embedding_api_key = ( embedding_api_key - or self.cfg.embedding_openai_api_key - or self.cfg.openai_api_key + or LLM_SETTINGS.embedding_openai_api_key + or LLM_SETTINGS.openai_api_key or os.environ.get("OPENAI_API_KEY") ) - self.chat_model = self.cfg.chat_model if chat_model is None else chat_model + self.chat_model = LLM_SETTINGS.chat_model if chat_model is None else chat_model self.encoder = tiktoken.encoding_for_model(self.chat_model) - self.chat_api_base = self.cfg.chat_azure_api_base if chat_api_base is None else chat_api_base - self.chat_api_version = self.cfg.chat_azure_api_version if chat_api_version is None else chat_api_version - self.chat_stream = self.cfg.chat_stream - self.chat_seed = self.cfg.chat_seed + self.chat_api_base = LLM_SETTINGS.chat_azure_api_base if chat_api_base is None else chat_api_base + self.chat_api_version = ( + LLM_SETTINGS.chat_azure_api_version if chat_api_version is None else chat_api_version + ) + self.chat_stream = LLM_SETTINGS.chat_stream + self.chat_seed = LLM_SETTINGS.chat_seed - self.embedding_model = self.cfg.embedding_model if embedding_model is None else embedding_model + self.embedding_model = LLM_SETTINGS.embedding_model if embedding_model is None else embedding_model self.embedding_api_base = ( - self.cfg.embedding_azure_api_base if embedding_api_base is None else embedding_api_base + LLM_SETTINGS.embedding_azure_api_base if embedding_api_base is None else embedding_api_base ) self.embedding_api_version = ( - self.cfg.embedding_azure_api_version if embedding_api_version is None else embedding_api_version + LLM_SETTINGS.embedding_azure_api_version if embedding_api_version is None else embedding_api_version ) if self.use_azure: @@ -376,20 +374,22 @@ class APIBackend: self.chat_client = openai.OpenAI(api_key=self.chat_api_key) self.embedding_client = openai.OpenAI(api_key=self.embedding_api_key) - self.dump_chat_cache = self.cfg.dump_chat_cache if dump_chat_cache is None else dump_chat_cache - self.use_chat_cache = self.cfg.use_chat_cache if use_chat_cache is None else use_chat_cache + self.dump_chat_cache = LLM_SETTINGS.dump_chat_cache if dump_chat_cache is None else dump_chat_cache + self.use_chat_cache = LLM_SETTINGS.use_chat_cache if use_chat_cache is None else use_chat_cache self.dump_embedding_cache = ( - self.cfg.dump_embedding_cache if dump_embedding_cache is None else dump_embedding_cache + LLM_SETTINGS.dump_embedding_cache if dump_embedding_cache is None else dump_embedding_cache + ) + self.use_embedding_cache = ( + LLM_SETTINGS.use_embedding_cache if use_embedding_cache is None else use_embedding_cache ) - self.use_embedding_cache = self.cfg.use_embedding_cache if use_embedding_cache is None else use_embedding_cache if self.dump_chat_cache or self.use_chat_cache or self.dump_embedding_cache or self.use_embedding_cache: - self.cache_file_location = self.cfg.prompt_cache_path + self.cache_file_location = LLM_SETTINGS.prompt_cache_path self.cache = SQliteLazyCache(cache_location=self.cache_file_location) # transfer the config to the class if the config is not supposed to change during the runtime - self.use_llama2 = self.cfg.use_llama2 - self.use_gcr_endpoint = self.cfg.use_gcr_endpoint - self.retry_wait_seconds = self.cfg.retry_wait_seconds + self.use_llama2 = LLM_SETTINGS.use_llama2 + self.use_gcr_endpoint = LLM_SETTINGS.use_gcr_endpoint + self.retry_wait_seconds = LLM_SETTINGS.retry_wait_seconds def build_chat_session( self, @@ -423,14 +423,14 @@ class APIBackend: if system_prompt is not None: while "\n\n\n" in system_prompt: system_prompt = system_prompt.replace("\n\n\n", "\n\n") - system_prompt = self.cfg.default_system_prompt if system_prompt is None else system_prompt + system_prompt = LLM_SETTINGS.default_system_prompt if system_prompt is None else system_prompt messages = [ { "role": "system", "content": system_prompt, }, ] - messages.extend(former_messages[-1 * self.cfg.max_past_message_include :]) + messages.extend(former_messages[-1 * LLM_SETTINGS.max_past_message_include :]) messages.append( { "role": "user", @@ -504,7 +504,7 @@ class APIBackend: **kwargs: Any, ) -> Any: assert not (chat_completion and embedding), "chat_completion and embedding cannot be True at the same time" - max_retry = self.cfg.max_retry if self.cfg.max_retry is not None else max_retry + max_retry = LLM_SETTINGS.max_retry if LLM_SETTINGS.max_retry is not None else max_retry for i in range(max_retry): try: if embedding: @@ -544,8 +544,8 @@ class APIBackend: if len(filtered_input_content_list) > 0: for sliced_filtered_input_content_list in [ - filtered_input_content_list[i : i + self.cfg.embedding_max_str_num] - for i in range(0, len(filtered_input_content_list), self.cfg.embedding_max_str_num) + filtered_input_content_list[i : i + LLM_SETTINGS.embedding_max_str_num] + for i in range(0, len(filtered_input_content_list), LLM_SETTINGS.embedding_max_str_num) ]: if self.use_azure: response = self.embedding_client.embeddings.create( @@ -594,8 +594,8 @@ class APIBackend: To make retries useful, we need to enable a seed. This seed is different from `self.chat_seed` for GPT. It is for the local cache mechanism enabled by RD-Agent locally. """ - # TODO: we can add this function back to avoid so much `self.cfg.log_llm_chat_content` - if self.cfg.log_llm_chat_content: + # TODO: we can add this function back to avoid so much `LLM_SETTINGS.log_llm_chat_content` + if LLM_SETTINGS.log_llm_chat_content: logger.info(self._build_log_messages(messages), tag="llm_messages") # TODO: fail to use loguru adaptor due to stream response input_content_json = json.dumps(messages) @@ -605,18 +605,18 @@ class APIBackend: if self.use_chat_cache: cache_result = self.cache.chat_get(input_content_json) if cache_result is not None: - if self.cfg.log_llm_chat_content: + if LLM_SETTINGS.log_llm_chat_content: logger.info(f"{LogColors.CYAN}Response:{cache_result}{LogColors.END}", tag="llm_messages") return cache_result, None if temperature is None: - temperature = self.cfg.chat_temperature + temperature = LLM_SETTINGS.chat_temperature if max_tokens is None: - max_tokens = self.cfg.chat_max_tokens + max_tokens = LLM_SETTINGS.chat_max_tokens if frequency_penalty is None: - frequency_penalty = self.cfg.chat_frequency_penalty + frequency_penalty = LLM_SETTINGS.chat_frequency_penalty if presence_penalty is None: - presence_penalty = self.cfg.chat_presence_penalty + presence_penalty = LLM_SETTINGS.chat_presence_penalty finish_reason = None if self.use_llama2: @@ -626,7 +626,7 @@ class APIBackend: temperature=temperature, ) resp = response[0]["generation"]["content"] - if self.cfg.log_llm_chat_content: + if LLM_SETTINGS.log_llm_chat_content: logger.info(f"{LogColors.CYAN}Response:{resp}{LogColors.END}", tag="llm_messages") elif self.use_gcr_endpoint: body = str.encode( @@ -648,7 +648,7 @@ class APIBackend: req = urllib.request.Request(self.gcr_endpoint, body, self.headers) # noqa: S310 response = urllib.request.urlopen(req) # noqa: S310 resp = json.loads(response.read().decode())["output"] - if self.cfg.log_llm_chat_content: + if LLM_SETTINGS.log_llm_chat_content: logger.info(f"{LogColors.CYAN}Response:{resp}{LogColors.END}", tag="llm_messages") else: kwargs = dict( @@ -673,7 +673,7 @@ class APIBackend: if self.chat_stream: resp = "" # TODO: with logger.config(stream=self.chat_stream): and add a `stream_start` flag to add timestamp for first message. - if self.cfg.log_llm_chat_content: + if LLM_SETTINGS.log_llm_chat_content: logger.info(f"{LogColors.CYAN}Response:{LogColors.END}", tag="llm_messages") for chunk in response: @@ -682,19 +682,19 @@ class APIBackend: if len(chunk.choices) > 0 and chunk.choices[0].delta.content is not None else "" ) - if self.cfg.log_llm_chat_content: + if LLM_SETTINGS.log_llm_chat_content: logger.info(LogColors.CYAN + content + LogColors.END, raw=True, tag="llm_messages") resp += content if len(chunk.choices) > 0 and chunk.choices[0].finish_reason is not None: finish_reason = chunk.choices[0].finish_reason - if self.cfg.log_llm_chat_content: + if LLM_SETTINGS.log_llm_chat_content: logger.info("\n", raw=True, tag="llm_messages") else: resp = response.choices[0].message.content finish_reason = response.choices[0].finish_reason - if self.cfg.log_llm_chat_content: + if LLM_SETTINGS.log_llm_chat_content: logger.info(f"{LogColors.CYAN}Response:{resp}{LogColors.END}", tag="llm_messages") if json_mode: json.loads(resp) diff --git a/rdagent/scenarios/data_mining/developer/model_runner.py b/rdagent/scenarios/data_mining/developer/model_runner.py index a3ef7e3d..2d04149a 100644 --- a/rdagent/scenarios/data_mining/developer/model_runner.py +++ b/rdagent/scenarios/data_mining/developer/model_runner.py @@ -1,27 +1,12 @@ -import shutil -import uuid -from pathlib import Path - -import pandas as pd - -from rdagent.components.coder.model_coder.model import ModelExperiment, ModelFBWorkspace from rdagent.components.runner import CachedRunner -from rdagent.components.runner.conf import RUNNER_SETTINGS -from rdagent.core.developer import Developer from rdagent.core.exception import ModelEmptyError -from rdagent.log import rdagent_logger as logger +from rdagent.core.utils import cache_with_pickle from rdagent.scenarios.data_mining.experiment.model_experiment import DMModelExperiment -from rdagent.utils.env import DMDockerEnv class DMModelRunner(CachedRunner[DMModelExperiment]): + @cache_with_pickle(CachedRunner.get_cache_key, CachedRunner.assign_cached_result) def develop(self, exp: DMModelExperiment) -> DMModelExperiment: - if RUNNER_SETTINGS.cache_result: - cache_hit, result = self.get_cache_result(exp) - if cache_hit: - exp.result = result - return exp - if exp.sub_workspace_list[0].code_dict.get("model.py") is None: raise ModelEmptyError("model.py is empty") # to replace & inject code @@ -32,7 +17,5 @@ class DMModelRunner(CachedRunner[DMModelExperiment]): result = exp.experiment_workspace.execute(run_env=env_to_use) exp.result = result - if RUNNER_SETTINGS.cache_result: - self.dump_cache_result(exp, result) return exp diff --git a/rdagent/scenarios/kaggle/developer/runner.py b/rdagent/scenarios/kaggle/developer/runner.py index 74740a42..489c19df 100644 --- a/rdagent/scenarios/kaggle/developer/runner.py +++ b/rdagent/scenarios/kaggle/developer/runner.py @@ -4,10 +4,10 @@ import shutil from pathlib import Path from rdagent.components.runner import CachedRunner -from rdagent.components.runner.conf import RUNNER_SETTINGS from rdagent.core.exception import CoderError, FactorEmptyError, ModelEmptyError from rdagent.core.experiment import ASpecificExp from rdagent.core.prompts import Prompts +from rdagent.core.utils import cache_with_pickle from rdagent.oai.llm_utils import md5_hash from rdagent.scenarios.kaggle.experiment.kaggle_experiment import ( KGFactorExperiment, @@ -27,15 +27,11 @@ class KGCachedRunner(CachedRunner[ASpecificExp]): codes = "\n".join(codes) return md5_hash(codes) + @cache_with_pickle(get_cache_key, CachedRunner.assign_cached_result) def init_develop(self, exp: KGFactorExperiment | KGModelExperiment) -> KGFactorExperiment | KGModelExperiment: """ For the initial development, the experiment serves as a benchmark for feature engineering. """ - if RUNNER_SETTINGS.cache_result: - cache_hit, result = self.get_cache_result(exp) - if cache_hit: - exp.result = result - return exp env_to_use = {"PYTHONPATH": "./"} @@ -43,13 +39,11 @@ class KGCachedRunner(CachedRunner[ASpecificExp]): exp.result = result - if RUNNER_SETTINGS.cache_result: - self.dump_cache_result(exp, result) - return exp class KGModelRunner(KGCachedRunner[KGModelExperiment]): + @cache_with_pickle(KGCachedRunner.get_cache_key, KGCachedRunner.assign_cached_result) def develop(self, exp: KGModelExperiment) -> KGModelExperiment: if exp.based_experiments and exp.based_experiments[-1].result is None: exp.based_experiments[-1] = self.init_develop(exp.based_experiments[-1]) @@ -65,12 +59,6 @@ class KGModelRunner(KGCachedRunner[KGModelExperiment]): model_file_name = f"model/model_{model_type.lower()}.py" exp.experiment_workspace.inject_code(**{model_file_name: sub_ws.code_dict["model.py"]}) - if RUNNER_SETTINGS.cache_result: - cache_hit, result = self.get_cache_result(exp) - if cache_hit: - exp.result = result - return exp - env_to_use = {"PYTHONPATH": "./"} result = exp.experiment_workspace.execute(run_env=env_to_use) @@ -79,8 +67,6 @@ class KGModelRunner(KGCachedRunner[KGModelExperiment]): raise CoderError("No result is returned from the experiment workspace") exp.result = result - if RUNNER_SETTINGS.cache_result: - self.dump_cache_result(exp, result) return exp diff --git a/rdagent/scenarios/kaggle/knowledge_management/vector_base.py b/rdagent/scenarios/kaggle/knowledge_management/vector_base.py index 417178de..a07fd37a 100644 --- a/rdagent/scenarios/kaggle/knowledge_management/vector_base.py +++ b/rdagent/scenarios/kaggle/knowledge_management/vector_base.py @@ -1,9 +1,9 @@ +import json from datetime import datetime, timezone from pathlib import Path from typing import List, Union import pandas as pd -from _pytest.cacheprovider import json from jinja2 import Environment, StrictUndefined from rdagent.components.knowledge_management.vector_base import Document, PDVectorBase diff --git a/rdagent/scenarios/qlib/developer/factor_runner.py b/rdagent/scenarios/qlib/developer/factor_runner.py index 9b893514..503bd373 100644 --- a/rdagent/scenarios/qlib/developer/factor_runner.py +++ b/rdagent/scenarios/qlib/developer/factor_runner.py @@ -6,12 +6,11 @@ import pandas as pd from pandarallel import pandarallel from rdagent.core.conf import RD_AGENT_SETTINGS -from rdagent.core.utils import multiprocessing_wrapper +from rdagent.core.utils import cache_with_pickle, multiprocessing_wrapper pandarallel.initialize(verbose=1) from rdagent.components.runner import CachedRunner -from rdagent.components.runner.conf import RUNNER_SETTINGS from rdagent.core.exception import FactorEmptyError from rdagent.log import rdagent_logger as logger from rdagent.scenarios.qlib.experiment.factor_experiment import QlibFactorExperiment @@ -71,6 +70,7 @@ class QlibFactorRunner(CachedRunner[QlibFactorExperiment]): IC_max = IC_max.unstack().max(axis=0) return new_feature.iloc[:, IC_max[IC_max < 0.99].index] + @cache_with_pickle(CachedRunner.get_cache_key, CachedRunner.assign_cached_result) def develop(self, exp: QlibFactorExperiment) -> QlibFactorExperiment: """ Generate the experiment by processing and combining factor data, @@ -79,12 +79,6 @@ class QlibFactorRunner(CachedRunner[QlibFactorExperiment]): if exp.based_experiments and exp.based_experiments[-1].result is None: exp.based_experiments[-1] = self.develop(exp.based_experiments[-1]) - if RUNNER_SETTINGS.cache_result: - cache_hit, result = self.get_cache_result(exp) - if cache_hit: - exp.result = result - return exp - if exp.based_experiments: SOTA_factor = None if len(exp.based_experiments) > 1: @@ -120,8 +114,6 @@ class QlibFactorRunner(CachedRunner[QlibFactorExperiment]): ) exp.result = result - if RUNNER_SETTINGS.cache_result: - self.dump_cache_result(exp, result) return exp @@ -143,7 +135,7 @@ class QlibFactorRunner(CachedRunner[QlibFactorExperiment]): for exp in exp_or_list: # Iterate over sub-implementations and execute them to get each factor data message_and_df_list = multiprocessing_wrapper( - [(implementation.execute, (False, "All")) for implementation in exp.sub_workspace_list], + [(implementation.execute, ("All",)) for implementation in exp.sub_workspace_list], n=RD_AGENT_SETTINGS.multi_proc_n, ) for message, df in message_and_df_list: diff --git a/rdagent/scenarios/qlib/developer/model_runner.py b/rdagent/scenarios/qlib/developer/model_runner.py index fec16c9a..ed83ac2f 100644 --- a/rdagent/scenarios/qlib/developer/model_runner.py +++ b/rdagent/scenarios/qlib/developer/model_runner.py @@ -1,17 +1,7 @@ -import shutil -import uuid -from pathlib import Path - -import pandas as pd - -from rdagent.components.coder.model_coder.model import ModelExperiment, ModelFBWorkspace from rdagent.components.runner import CachedRunner -from rdagent.components.runner.conf import RUNNER_SETTINGS -from rdagent.core.developer import Developer from rdagent.core.exception import ModelEmptyError -from rdagent.log import rdagent_logger as logger +from rdagent.core.utils import cache_with_pickle from rdagent.scenarios.qlib.experiment.model_experiment import QlibModelExperiment -from rdagent.utils.env import QTDockerEnv class QlibModelRunner(CachedRunner[QlibModelExperiment]): @@ -27,13 +17,8 @@ class QlibModelRunner(CachedRunner[QlibModelExperiment]): - let LLM modify model.py """ + @cache_with_pickle(CachedRunner.get_cache_key, CachedRunner.assign_cached_result) def develop(self, exp: QlibModelExperiment) -> QlibModelExperiment: - if RUNNER_SETTINGS.cache_result: - cache_hit, result = self.get_cache_result(exp) - if cache_hit: - exp.result = result - return exp - if exp.sub_workspace_list[0].code_dict.get("model.py") is None: raise ModelEmptyError("model.py is empty") # to replace & inject code @@ -49,7 +34,5 @@ class QlibModelRunner(CachedRunner[QlibModelExperiment]): result = exp.experiment_workspace.execute(qlib_config_name="conf.yaml", run_env=env_to_use) exp.result = result - if RUNNER_SETTINGS.cache_result: - self.dump_cache_result(exp, result) return exp diff --git a/rdagent/scenarios/qlib/factor_experiment_loader/pdf_loader.py b/rdagent/scenarios/qlib/factor_experiment_loader/pdf_loader.py index 74e86cd6..a9e61c92 100644 --- a/rdagent/scenarios/qlib/factor_experiment_loader/pdf_loader.py +++ b/rdagent/scenarios/qlib/factor_experiment_loader/pdf_loader.py @@ -22,6 +22,7 @@ from rdagent.core.conf import RD_AGENT_SETTINGS from rdagent.core.prompts import Prompts from rdagent.core.utils import multiprocessing_wrapper from rdagent.log import rdagent_logger as logger +from rdagent.oai.llm_conf import LLM_SETTINGS from rdagent.oai.llm_utils import APIBackend from rdagent.scenarios.qlib.factor_experiment_loader.json_loader import ( FactorExperimentLoaderFromDict, @@ -86,9 +87,9 @@ def classify_report_from_dict( user_prompt=content, system_prompt=classify_prompt, ) - > RD_AGENT_SETTINGS.chat_token_limit + > LLM_SETTINGS.chat_token_limit ): - content = content[: -(RD_AGENT_SETTINGS.chat_token_limit // 100)] + content = content[: -(LLM_SETTINGS.chat_token_limit // 100)] vote_list = [] for _ in range(vote_time): @@ -374,7 +375,7 @@ def __check_factor_duplication_simulate_json_mode( APIBackend().build_messages_and_calculate_token( user_prompt=current_df.to_string(), system_prompt=document_process_prompts["factor_duplicate_system"] ) - > RD_AGENT_SETTINGS.chat_token_limit + > LLM_SETTINGS.chat_token_limit ): working_list.append(current_df.iloc[: current_df.shape[0] // 2, :]) working_list.append(current_df.iloc[current_df.shape[0] // 2 :, :])