From 5f6c2b8875031e8898efa231de0662c68f479828 Mon Sep 17 00:00:00 2001
From: Tim <illking@foxmail.com>
Date: Wed, 22 Jan 2025 14:21:08 +0800
Subject: [PATCH] chore: clean unused scripts (#528)

* clean unused scripts

* reformat
---
 rdagent/scenarios/data_science/dev/runner.py | 25 +++++++++++++++++++-
 1 file changed, 24 insertions(+), 1 deletion(-)

diff --git a/rdagent/scenarios/data_science/dev/runner.py b/rdagent/scenarios/data_science/dev/runner.py
index d403b212..2f34c4fb 100644
--- a/rdagent/scenarios/data_science/dev/runner.py
+++ b/rdagent/scenarios/data_science/dev/runner.py
@@ -1,3 +1,7 @@
+import json
+import os
+from pathlib import Path
+
 import pandas as pd
 
 from rdagent.app.data_science.conf import DS_RD_SETTING
@@ -17,7 +21,7 @@ class DSRunner(Developer[DSExperiment]):
         de = DockerEnv(conf=ds_docker_conf)
 
         # execute workflow
-        stdout = exp.experiment_workspace.execute(env=de, entry="python main.py")
+        stdout = exp.experiment_workspace.execute(env=de, entry="coverage run main.py")
 
         score_fp = exp.experiment_workspace.workspace_path / "scores.csv"
         if not score_fp.exists():
@@ -30,4 +34,23 @@ class DSRunner(Developer[DSExperiment]):
             raise RunnerError(f"Submission file (submission.csv) is not generated, log is:\n{stdout}")
 
         exp.result = pd.read_csv(score_fp, index_col=0)
+
+        # remove unused files
+        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"}
+                logger.info("All used scripts: {}".format(used_files))
+                all_python_files = set(Path(exp.experiment_workspace.workspace_path).rglob("*.py"))
+                unused_files = [
+                    py_file
+                    for py_file in all_python_files
+                    if not (py_file.name in used_files or py_file.name.endswith("test.py"))
+                ]
+                if unused_files:
+                    logger.warning(f"Unused scripts: {unused_files}")
+                    exp.experiment_workspace.inject_files(
+                        {file_path.name: exp.experiment_workspace.DEL_KEY for file_path in unused_files}
+                    )
+            os.remove(exp.experiment_workspace.workspace_path / "coverage.json")
         return exp