From 9bb214fff2109cacad0bf5cd25e1763abb1ddf4d Mon Sep 17 00:00:00 2001
From: Xu Yang <peteryang@vip.qq.com>
Date: Mon, 23 Sep 2024 20:27:34 +0800
Subject: [PATCH] fix: update code to fix a small bug in model cache md5 hash
 (#303)

* update code to fix a small bug in model cache md5 hash

* fix another bug dumping the wrong name to costeer model

* fix a black CI
---
 .../model_coder/CoSTEER/evolving_strategy.py     | 16 ++--------------
 rdagent/components/coder/model_coder/model.py    |  7 ++++---
 2 files changed, 6 insertions(+), 17 deletions(-)

diff --git a/rdagent/components/coder/model_coder/CoSTEER/evolving_strategy.py b/rdagent/components/coder/model_coder/CoSTEER/evolving_strategy.py
index b67f2cc6..2b9880d7 100644
--- a/rdagent/components/coder/model_coder/CoSTEER/evolving_strategy.py
+++ b/rdagent/components/coder/model_coder/CoSTEER/evolving_strategy.py
@@ -146,21 +146,9 @@ class ModelCoderEvolvingStrategy(EvolvingStrategy):
             n=RD_AGENT_SETTINGS.multi_proc_n,
         )
 
-        model_file_mapping = {
-            "XGBoost": "model_xgb.py",
-            "RandomForest": "model_rf.py",
-            "LightGBM": "model_lgb.py",
-            "NN": "model_nn.py",
-        }
-
         for index, target_index in enumerate(to_be_finished_task_index):
-            if evo.sub_workspace_list[target_index] is None:
-                evo.sub_workspace_list[target_index] = ModelFBWorkspace(target_task=evo.sub_tasks[target_index])
-            model_type = evo.sub_tasks[target_index].model_type
-            if model_type in model_file_mapping:
-                evo.sub_workspace_list[target_index].inject_code(**{model_file_mapping[model_type]: result[index]})
-            else:
-                raise ValueError(f"Unsupported model type: {model_type}")
+            evo.sub_workspace_list[target_index] = ModelFBWorkspace(target_task=evo.sub_tasks[target_index])
+            evo.sub_workspace_list[target_index].inject_code(**{"model.py": result[index]})
 
         evo.corresponding_selection = to_be_finished_task_index
 
diff --git a/rdagent/components/coder/model_coder/model.py b/rdagent/components/coder/model_coder/model.py
index 7e66ba95..eca3b432 100644
--- a/rdagent/components/coder/model_coder/model.py
+++ b/rdagent/components/coder/model_coder/model.py
@@ -84,9 +84,10 @@ class ModelFBWorkspace(FBWorkspace):
         try:
             if MODEL_IMPL_SETTINGS.enable_execution_cache:
                 # NOTE: cache the result for the same code
-                target_file_name = md5_hash(
-                    f"{batch_size}_{num_features}_{num_timesteps}_{input_value}_{param_init_value}_{self.code_dict['model.py']}"
-                )
+                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():