From 43c8ed2aaff49fc46e19bf6e80e89e669aa02788 Mon Sep 17 00:00:00 2001 From: Tim Date: Fri, 4 Apr 2025 12:08:18 +0800 Subject: [PATCH] feat: joblib cache (#749) * cache function * fix test * bin cache * fix test * fix test * fix test * cache for different source * cache for localenv * remove unnecessary log * reformat * remove unrelated modify --- .../feature/eval_tests/feature_test.txt | 14 +++++++++++++- .../coder/data_science/feature/prompts.yaml | 7 ++++++- .../data_science/model/eval_tests/model_test.txt | 15 ++++++++++++++- .../coder/data_science/model/prompts.yaml | 5 +++++ .../eval_tests/data_loader_test.txt | 12 +++++++++++- .../data_science/raw_data_loader/prompts.yaml | 5 +++++ rdagent/utils/env.py | 15 +++++++++++---- 7 files changed, 65 insertions(+), 8 deletions(-) diff --git a/rdagent/components/coder/data_science/feature/eval_tests/feature_test.txt b/rdagent/components/coder/data_science/feature/eval_tests/feature_test.txt index c7721acf..27393b0f 100644 --- a/rdagent/components/coder/data_science/feature/eval_tests/feature_test.txt +++ b/rdagent/components/coder/data_science/feature/eval_tests/feature_test.txt @@ -28,10 +28,22 @@ X_loaded = deepcopy(X) y_loaded = deepcopy(y) X_test_loaded = deepcopy(X_test) +import sys +import reprlib +from joblib.memory import MemorizedFunc + + +def get_original_code(func): + if isinstance(func, MemorizedFunc): + return func.func.__code__ + return func.__code__ + + def debug_info_print(func): def wrapper(*args, **kwargs): + original_code = get_original_code(func) def local_trace(frame, event, arg): - if event == "return" and frame.f_code == func.__code__: + if event == "return" and frame.f_code == original_code: print("\n" + "="*20 + "Running feat_eng code, local variable values:" + "="*20) for k, v in frame.f_locals.items(): printed = aRepr.repr(v) diff --git a/rdagent/components/coder/data_science/feature/prompts.yaml b/rdagent/components/coder/data_science/feature/prompts.yaml index 56ce2d71..949c0fac 100644 --- a/rdagent/components/coder/data_science/feature/prompts.yaml +++ b/rdagent/components/coder/data_science/feature/prompts.yaml @@ -39,10 +39,15 @@ feature_coder: ```python {{ data_loader_code }} ``` - 3. **Additional Guidance:** + 4. **Additional Guidance:** - If a previous attempt exists, improve upon it without repeating mistakes. - If errors indicate a missing file, find a way to download it or implement an alternative solution. - You should avoid using logging module to output information in your generated code, and instead use the print() function. + 5. You should use the following cache decorator to cache the results of the function: + ```python + from joblib import Memory + memory = Memory(location='/tmp/cache', verbose=0) + @memory.cache``` ## Output Format {% if out_spec %} diff --git a/rdagent/components/coder/data_science/model/eval_tests/model_test.txt b/rdagent/components/coder/data_science/model/eval_tests/model_test.txt index e8fd110c..509f24e2 100644 --- a/rdagent/components/coder/data_science/model/eval_tests/model_test.txt +++ b/rdagent/components/coder/data_science/model/eval_tests/model_test.txt @@ -35,6 +35,17 @@ print(f"test_ids length: {len(test_ids)}") train_X, val_X, train_y, val_y = train_test_split(X, y, test_size=0.8, random_state=42) + +import sys +import reprlib +from joblib.memory import MemorizedFunc + + +def get_original_code(func): + if isinstance(func, MemorizedFunc): + return func.func.__code__ + return func.__code__ + print("train_X:", aRepr.repr(train_X)) print("train_y:", aRepr.repr(train_y)) print("val_X:", aRepr.repr(val_X)) @@ -46,10 +57,12 @@ print(f"val_X.shape: {val_X.shape}" if hasattr(val_X, 'shape') else f"val_X leng print(f"val_y.shape: {val_y.shape}" if hasattr(val_y, 'shape') else f"val_y length: {len(val_y)}") + def debug_info_print(func): def wrapper(*args, **kwargs): + original_code = get_original_code(func) def local_trace(frame, event, arg): - if event == "return" and frame.f_code == func.__code__: + if event == "return" and frame.f_code == original_code: print("\n" + "="*20 + "Running model training code, local variable values:" + "="*20) for k, v in frame.f_locals.items(): printed = aRepr.repr(v) diff --git a/rdagent/components/coder/data_science/model/prompts.yaml b/rdagent/components/coder/data_science/model/prompts.yaml index 0e50e782..39fa2cc4 100644 --- a/rdagent/components/coder/data_science/model/prompts.yaml +++ b/rdagent/components/coder/data_science/model/prompts.yaml @@ -40,6 +40,11 @@ model_coder: {{ feature_code }} 2. You should avoid using logging module to output information in your generated code, and instead use the print() function. 3. If the model can both be implemented by PyTorch and Tensorflow, please use pytorch for broader compatibility. + 4. You should use the following cache decorator to cache the results of the function: + ```python + from joblib import Memory + memory = Memory(location='/tmp/cache', verbose=0) + @memory.cache`` ## Output Format {% if out_spec %} diff --git a/rdagent/components/coder/data_science/raw_data_loader/eval_tests/data_loader_test.txt b/rdagent/components/coder/data_science/raw_data_loader/eval_tests/data_loader_test.txt index 493de773..04e1b509 100644 --- a/rdagent/components/coder/data_science/raw_data_loader/eval_tests/data_loader_test.txt +++ b/rdagent/components/coder/data_science/raw_data_loader/eval_tests/data_loader_test.txt @@ -9,12 +9,22 @@ from load_data import load_data import sys import reprlib +from joblib.memory import MemorizedFunc + + +def get_original_code(func): + if isinstance(func, MemorizedFunc): + return func.func.__code__ + return func.__code__ + + def debug_info_print(func): aRepr = reprlib.Repr() aRepr.maxother=300 def wrapper(*args, **kwargs): + original_code = get_original_code(func) def local_trace(frame, event, arg): - if event == "return" and frame.f_code == func.__code__: + if event == "return" and frame.f_code == original_code: print("\n" + "="*20 + "Running data_load code, local variable values:" + "="*20) for k, v in frame.f_locals.items(): printed = aRepr.repr(v) diff --git a/rdagent/components/coder/data_science/raw_data_loader/prompts.yaml b/rdagent/components/coder/data_science/raw_data_loader/prompts.yaml index 55576285..a45a1e4d 100644 --- a/rdagent/components/coder/data_science/raw_data_loader/prompts.yaml +++ b/rdagent/components/coder/data_science/raw_data_loader/prompts.yaml @@ -334,6 +334,11 @@ data_loader_coder: ## Guidelines 1. Ensure that the dataset is loaded strictly from `/kaggle/input/`, following the exact folder structure described in the **Data Folder Description**, and do not attempt to load data from the current directory (`./`). 2. You should avoid using logging module to output information in your generated code, and instead use the print() function. + 3. You should use the following cache decorator to cache the results of the function: + ```python + from joblib import Memory + memory = Memory(location='/tmp/cache', verbose=0) + @memory.cache``` ## Exploratory Data Analysis (EDA) part(Required): - Before returning the data, you should always add an EDA part describing the data to help the following steps understand the data better. diff --git a/rdagent/utils/env.py b/rdagent/utils/env.py index 6be5ed78..fe0e5c9f 100644 --- a/rdagent/utils/env.py +++ b/rdagent/utils/env.py @@ -1,5 +1,5 @@ """ -The motiviation of the utils is for environment management +The motivation of the utils is for environment management Tries to create uniform environment for the agent to run; - All the code and data is expected included in one folder @@ -210,8 +210,8 @@ class Env(Generic[ASpecificEnvConf]): target_folder = Path(RD_AGENT_SETTINGS.pickle_cache_folder_path_str) / f"utils.env.run" target_folder.mkdir(parents=True, exist_ok=True) - # we must add the information of data (beyound code) into the key. - # Otherwise, all commands operating on data will become invalue (e.g. rm -r submission.csv) + # we must add the information of data (beyond code) into the key. + # Otherwise, all commands operating on data will become invalid (e.g. rm -r submission.csv) # So we recursively walk in the folder and add the sorted relative filename list as part of the key. data_key = [] for path in Path(local_path).rglob("*"): @@ -292,7 +292,7 @@ ASpecificLocalConf = TypeVar("ASpecificLocalConf", bound=LocalConf) class LocalEnv(Env[ASpecificLocalConf]): """ - Sometimes local environment may be more convinient for testing + Sometimes local environment may be more convenient for testing """ def prepare(self) -> None: ... @@ -311,6 +311,9 @@ class LocalEnv(Env[ASpecificLocalConf]): if self.conf.extra_volumes is not None: for lp, rp in self.conf.extra_volumes.items(): volumes[lp] = 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] = "/tmp/cache" for lp, rp in running_extra_volume.items(): volumes[lp] = rp @@ -605,9 +608,13 @@ class DockerEnv(Env[DockerConf]): if local_path is not None: local_path = os.path.abspath(local_path) volumes[local_path] = {"bind": self.conf.mount_path, "mode": "rw"} + if self.conf.extra_volumes is not None: for lp, rp in self.conf.extra_volumes.items(): volumes[lp] = {"bind": rp, "mode": self.conf.extra_volume_mode} + 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] = {"bind": "/tmp/cache", "mode": "rw"} for lp, rp in running_extra_volume.items(): volumes[lp] = {"bind": rp, "mode": self.conf.extra_volume_mode}