chore: add runtime info (#526)

* add runtime info * Update rdagent/scenarios/data_science/scen/__init__.py Co-authored-by: you-n-g <[email protected]> * reformatted by black --------- Co-authored-by: you-n-g <[email protected]>
microsoft · Jan 21, 2025 · 845b675 · 845b675
1 parent 1236e80
commit 845b675
Show file tree

Hide file tree

Showing 4 changed files with 59 additions and 0 deletions.
diff --git 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 @@ def implement_one_task(
         # 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 @@ def implement_one_task(
         # 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
@@ -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
@@ -6,6 +6,7 @@
 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 @@
     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 @@ def get_scenario_all_desc(self) -> str:
             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
@@ -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)