Simpler pip info recording to logs

browsergym/experiments/src/browsergym/experiments/loop.py

@@ -1,9 +1,9 @@
 import gzip
+import importlib.metadata
 import json
 import logging
 import os
 import pickle
-import subprocess
 import sys
 import time
 import traceback
@@ -174,12 +174,26 @@ def _make_dir(self, exp_root):
     def run(self):
         """Run the experiment and save the results"""
 
-        # save pip info
-        _save_pip_info(self.exp_dir)
-
         # start writing logs to run logfile
         self._set_logger()
 
+        # log python environment info
+        # from https://stackoverflow.com/a/78160009/1264211
+        python_dists = "\n".join(
+            sorted(
+                [
+                    f'{dist.metadata["Name"]}=={dist.metadata["Version"]}'
+                    for dist in importlib.metadata.distributions()
+                ]
+            )
+        )
+        logger.info(
+            f"""\
+Python version: {sys.version}
+Python installed distributions:
+{python_dists}
+"""
+        )
         episode_info = []
         env, step_info, err_msg, stack_trace = None, None, None, None
         try:
@@ -455,19 +469,6 @@ def _aggregate_episode_stats(episode_info: list[StepInfo]):
     return aggregated_stats
 
 
-def _save_pip_info(
-    exp_dir,
-):
-    # from https://pip.pypa.io/en/stable/user_guide/#using-pip-from-your-program
-    pip_freeze = subprocess.check_output([sys.executable, "-m", "pip", "freeze"], encoding="utf-8")
-    pip_list = subprocess.check_output([sys.executable, "-m", "pip", "list"], encoding="utf-8")
-
-    with open(exp_dir / "pip_freeze.txt", "w") as f:
-        f.write(pip_freeze)
-    with open(exp_dir / "pip_list.txt", "w") as f:
-        f.write(pip_list)
-
-
 def _save_summary_info(
     episode_info: list[StepInfo],
     exp_dir,
@@ -525,8 +526,6 @@ class ExpResult:
         task_video_path: Path, the path to the task video. (if record_video=True)
         combined_video_path: Path, the path to the combined video. (if video was
             combined)
-        pip_freeze: str, a `pip freeze` of the python environment
-        pip_list: str, a `pip list` of the python environment
     """
 
     def __init__(self, exp_dir) -> None:
@@ -537,8 +536,6 @@ def __init__(self, exp_dir) -> None:
         self._screenshots = {}
         self._flat_exp_args = None
         self._logs = None
-        self._pip_freeze = None
-        self._pip_list = None
 
     @property
     def exp_args(self) -> ExpArgs:
@@ -645,18 +642,6 @@ def logs(self):
             self._logs = (self.exp_dir / "experiment.log").read_text()
         return self._logs
 
-    @property
-    def pip_freeze(self):
-        if self._pip_freeze is None:
-            self._pip_freeze = (self.exp_dir / "pip_freeze.txt").read_text()
-        return self._pip_freeze
-
-    @property
-    def pip_list(self):
-        if self._pip_list is None:
-            self._pip_list = (self.exp_dir / "pip_list.txt").read_text()
-        return self._pip_list
-
 
 EXP_RESULT_CACHE = {}