feat: add mlebench submission validitor (#545)

XianBW · web-flow · commit 712d94a7d6f2 · 2025-01-27T00:41:35.000+08:00
* add mlebench submission check

* fix CI

* fix bug
diff --git a/rdagent/components/coder/data_science/workflow/eval.py b/rdagent/components/coder/data_science/workflow/eval.py
@@ -15,7 +15,7 @@
 from rdagent.core.experiment import FBWorkspace, Task
 from rdagent.oai.llm_utils import APIBackend
 from rdagent.utils.agent.tpl import T
-from rdagent.utils.env import DockerEnv, DSDockerConf
+from rdagent.utils.env import DockerEnv, DSDockerConf, MLEBDockerConf
 
 DIRNAME = Path(__file__).absolute().resolve().parent
 
@@ -53,11 +53,22 @@ def evaluate(
                 code="This task has failed too many times, skip implementation.",
                 final_decision=False,
             )
+
+        # DockerEnv for Kaggle Competition
         ds_docker_conf = DSDockerConf()
         ds_docker_conf.extra_volumes = {
             f"{DS_RD_SETTING.local_data_path}/sample/{self.scen.competition}": "/kaggle/input"
         }
         de = DockerEnv(conf=ds_docker_conf)
+
+        # DockerEnv for MLEBench submission validation
+        mle_de_conf = MLEBDockerConf()
+        mle_de_conf.extra_volumes = {
+            f"{DS_RD_SETTING.local_data_path}/zip_files": "/mle/data",
+        }
+        mde = DockerEnv(conf=mle_de_conf)
+        mde.prepare()
+
         # Clean the scores.csv & submission.csv.
         stdout = implementation.execute(env=de, entry=f"rm submission.csv scores.csv")
 
@@ -85,10 +96,21 @@ def evaluate(
         if not submission_fp.exists():
             stdout += "\nSubmission file (submission.csv) is not generated."
         else:
-            check_code = (DIRNAME / "eval_tests" / "submission_check.txt").read_text()
-            implementation.inject_files(**{"submission_check.py": check_code})
+            base_check_code = (DIRNAME / "eval_tests" / "submission_check.txt").read_text()
+            implementation.inject_files(**{"submission_check.py": base_check_code})
+            stdout += "----Submission Check 1-----\n"
             stdout += implementation.execute(env=de, entry="python submission_check.py")
 
+            # MLEBench Check
+            mle_check_code = (
+                (DIRNAME / "eval_tests" / "mle_submission_check.txt")
+                .read_text()
+                .replace("<competition_id>", self.scen.competition)
+            )
+            implementation.inject_files(**{"mle_submission_check.py": mle_check_code})
+            stdout += "----Submission Check 2-----\n"
+            stdout += implementation.execute(env=mde, entry=f"python mle_submission_check.py")
+
         system_prompt = T(".prompts:workflow_eval.system").r(
             scenario=self.scen.get_scenario_all_desc(),
             task_desc=target_task.get_task_information(),
diff --git a/rdagent/components/coder/data_science/workflow/eval_tests/mle_submission_check.txt b/rdagent/components/coder/data_science/workflow/eval_tests/mle_submission_check.txt
@@ -0,0 +1,12 @@
+from pathlib import Path
+
+from mlebench.grade import validate_submission
+from mlebench.registry import registry
+
+COMPETITION_ID = "<competition_id>"
+new_registry = registry.set_data_dir(Path("/mle/data"))
+competition = new_registry.get_competition(COMPETITION_ID)
+
+is_valid, message = validate_submission(Path("submission.csv"), competition)
+
+print(message)
diff --git a/rdagent/scenarios/data_science/dev/runner.py b/rdagent/scenarios/data_science/dev/runner.py
@@ -43,7 +43,7 @@ def develop(self, exp: DSExperiment) -> DSExperiment:
         stdout = exp.experiment_workspace.execute(env=de, entry="coverage json -o coverage.json")
         if Path(exp.experiment_workspace.workspace_path / "coverage.json").exists():
             with open(exp.experiment_workspace.workspace_path / "coverage.json") as f:
-                used_files = set(json.load(f)["files"].keys()) | {"submission_check.py"}
+                used_files = set(json.load(f)["files"].keys()) | {"submission_check.py", "mle_submission_check.py"}
                 logger.info("All used scripts: {}".format(used_files))
                 all_python_files = set(Path(exp.experiment_workspace.workspace_path).rglob("*.py"))
                 unused_files = [
