From 10abb201175f26e6b2d47357fe8d8312ab4f7e48 Mon Sep 17 00:00:00 2001
From: XianBW <36835909+XianBW@users.noreply.github.com>
Date: Mon, 27 Jan 2025 00:41:35 +0800
Subject: [PATCH] feat: add mlebench submission validitor (#545)

* add mlebench submission check

* fix CI

* fix bug
---
 .../coder/data_science/workflow/eval.py       | 28 +++++++++++++++++--
 .../eval_tests/mle_submission_check.txt       | 12 ++++++++
 rdagent/scenarios/data_science/dev/runner.py  |  2 +-
 3 files changed, 38 insertions(+), 4 deletions(-)
 create mode 100644 rdagent/components/coder/data_science/workflow/eval_tests/mle_submission_check.txt

diff --git a/rdagent/components/coder/data_science/workflow/eval.py b/rdagent/components/coder/data_science/workflow/eval.py
index 7f2e9d8e..7867f6e5 100644
--- 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.evolving_framework import QueriedKnowledge
 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 @@ class WorkflowGeneralCaseSpecEvaluator(CoSTEEREvaluator):
                 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 @@ class WorkflowGeneralCaseSpecEvaluator(CoSTEEREvaluator):
         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
new file mode 100644
index 00000000..7930fa06
--- /dev/null
+++ 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)
\ No newline at end of file
diff --git a/rdagent/scenarios/data_science/dev/runner.py b/rdagent/scenarios/data_science/dev/runner.py
index dd7713f5..f221d8b6 100644
--- a/rdagent/scenarios/data_science/dev/runner.py
+++ b/rdagent/scenarios/data_science/dev/runner.py
@@ -43,7 +43,7 @@ class DSRunner(Developer[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 = [