From bdfdc1689f60bf72c3fc5a06a9aad1cade8e7af6 Mon Sep 17 00:00:00 2001
From: Tim <illking@foxmail.com>
Date: Tue, 21 Jan 2025 16:15:45 +0800
Subject: [PATCH] chore: add runtime info (#526)

* add runtime info

* Update rdagent/scenarios/data_science/scen/__init__.py

Co-authored-by: you-n-g <you-n-g@users.noreply.github.com>

* reformatted by black

---------

Co-authored-by: you-n-g <you-n-g@users.noreply.github.com>
---
 .../data_science/raw_data_loader/__init__.py  |  2 +
 .../data_science/raw_data_loader/prompts.yaml |  3 ++
 .../scenarios/data_science/scen/__init__.py   | 14 +++++++
 .../data_science/scen/runtime_info.py         | 40 +++++++++++++++++++
 4 files changed, 59 insertions(+)
 create mode 100644 rdagent/scenarios/data_science/scen/runtime_info.py

diff --git a/rdagent/components/coder/data_science/raw_data_loader/__init__.py b/rdagent/components/coder/data_science/raw_data_loader/__init__.py
index 73a7a349..fbd26a19 100644
--- a/rdagent/components/coder/data_science/raw_data_loader/__init__.py
+++ b/rdagent/components/coder/data_science/raw_data_loader/__init__.py
@@ -55,6 +55,7 @@ class DataLoaderMultiProcessEvolvingStrategy(MultiProcessEvolvingStrategy):
         # return a workspace with "load_data.py", "spec/load_data.md" inside
         # assign the implemented code to the new workspace.
         competition_info = self.scen.get_scenario_all_desc()
+        runtime_environment = self.scen.get_runtime_environment()
         data_folder_info = self.scen.processed_data_folder_description
         data_loader_task_info = target_task.get_task_information()
 
@@ -88,6 +89,7 @@ class DataLoaderMultiProcessEvolvingStrategy(MultiProcessEvolvingStrategy):
         # TODO: We may move spec into a separated COSTEER task
         if "spec/data_loader.md" not in workspace.file_dict:  # Only generate the spec once
             system_prompt = T(".prompts:spec.system").r(
+                runtime_environment=runtime_environment,
                 task_desc=data_loader_task_info,
                 competition_info=competition_info,
                 folder_spec=data_folder_info,
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 b77e6402..baa939c9 100644
--- a/rdagent/components/coder/data_science/raw_data_loader/prompts.yaml
+++ b/rdagent/components/coder/data_science/raw_data_loader/prompts.yaml
@@ -6,6 +6,9 @@ spec:
     Currently, you are working on a Kaggle competition project. 
     This project involves analyzing data and building models to beat other competitors, with the code being generated by large language models.
 
+    The runtime environment you are working in includes the following libraries and their respective versions:
+    {{ runtime_environment }}
+
     Your overall task is provided below:
     {{ task_desc }}
     
diff --git a/rdagent/scenarios/data_science/scen/__init__.py b/rdagent/scenarios/data_science/scen/__init__.py
index 9c6d7cbf..b40e5750 100644
--- a/rdagent/scenarios/data_science/scen/__init__.py
+++ b/rdagent/scenarios/data_science/scen/__init__.py
@@ -6,6 +6,7 @@ import pandas as pd
 from PIL import Image, TiffTags
 
 from rdagent.app.data_science.conf import DS_RD_SETTING
+from rdagent.core.experiment import FBWorkspace
 from rdagent.core.scenario import Scenario
 from rdagent.log import rdagent_logger as logger
 from rdagent.oai.llm_utils import APIBackend
@@ -14,6 +15,7 @@ from rdagent.scenarios.kaggle.kaggle_crawler import (
     leaderboard_scores,
 )
 from rdagent.utils.agent.tpl import T
+from rdagent.utils.env import DockerEnv, DSDockerConf
 
 
 def read_csv_head(file_path, indent=0, lines=5, max_col_width=100):
@@ -304,6 +306,18 @@ class DataScienceScen(Scenario):
             metric_direction=self.metric_direction,
         )
 
+    def get_runtime_environment(self) -> str:
+        # TODO:  add it into base class.  Environment should(i.e. `DSDockerConf`) should be part of the scenario class.
+        ds_docker_conf = DSDockerConf()
+        de = DockerEnv(conf=ds_docker_conf)
+        implementation = FBWorkspace()
+        fname = "temp.py"
+        implementation.inject_files(
+            **{fname: (Path(__file__).absolute().resolve().parent / "runtime_info.py").read_text()}
+        )
+        stdout = implementation.execute(env=de, entry=f"python {fname}")
+        return stdout
+
     def _get_data_folder_description(self) -> str:
         return describe_data_folder(Path(DS_RD_SETTING.local_data_path) / self.competition)
 
diff --git a/rdagent/scenarios/data_science/scen/runtime_info.py b/rdagent/scenarios/data_science/scen/runtime_info.py
new file mode 100644
index 00000000..5634cee6
--- /dev/null
+++ b/rdagent/scenarios/data_science/scen/runtime_info.py
@@ -0,0 +1,40 @@
+import platform
+import sys
+from importlib.metadata import distributions
+
+
+def print_runtime_info():
+    print(f"Python {sys.version} on {platform.system()} {platform.release()}")
+
+
+def get_installed_packages():
+    return {dist.metadata["Name"].lower(): dist.version for dist in distributions()}
+
+
+def print_filtered_packages(installed_packages, filtered_packages):
+    for package_name in filtered_packages:
+        version = installed_packages.get(package_name.lower())
+        if version:
+            print(f"{package_name}=={version}")
+
+
+if __name__ == "__main__":
+    print_runtime_info()
+    filtered_packages = [
+        "transformers",
+        "accelerate",
+        "torch",
+        "tensorflow",
+        "pandas",
+        "numpy",
+        "scikit-learn",
+        "scipy",
+        "lightgbm",
+        "vtk",
+        "opencv-python",
+        "keras",
+        "matplotlib",
+        "pydicom",
+    ]
+    installed_packages = get_installed_packages()
+    print_filtered_packages(installed_packages, filtered_packages)