Fixing load_dataset to be more memoery efficient

Author: MarcCote

.github/actions/test-if-changes/action.yml

@@ -45,6 +45,7 @@ runs:
         DEBUG_GYM_DEBUG: 1
       shell: bash
       run: |
+        free -h
         pytest ${{ inputs.test-files }} -vv -n 16 --timeout=600 --cov=debug_gym --cov-report=term-missing
     - name: Store coverage report
       uses: actions/upload-artifact@v4

debug_gym/gym/envs/r2egym.py

@@ -17,8 +17,6 @@
 from debug_gym.gym.utils import filter_problems
 from debug_gym.logger import DebugGymLogger
 
-main_logger = logging.getLogger(__name__)
-
 
 def decolor_dict_keys(key):
     """Remove ANSI escape codes"""
@@ -265,10 +263,6 @@ def load_dataset(
         prepull_images: bool = False,
         logger: DebugGymLogger | None = None,
     ) -> dict:
-        main_logger.info(
-            f"Loading R2E-Gym dataset `{dataset_id}` (rev: {dataset_revision})..."
-        )
-
         logger = logger or DebugGymLogger("debug_gym")
         data_path = Path(dataset_id)
 
@@ -285,7 +279,6 @@ def load_dataset(
             # Loading from HuggingFace or a folder.
             ds = load_dataset(dataset_id, revision=dataset_revision)
 
-        main_logger.info("Dataset loaded.")
         # Select the split.
         ds = ds[split]
 
@@ -297,18 +290,18 @@ def load_dataset(
         def extract_instance_id(docker_image: str) -> str:
             return docker_image.split("/", 1)[-1]
 
-        dataset = {
+        id2idx = {
             extract_instance_id(docker_image): i
             for i, docker_image in enumerate(ds["docker_image"])
         }
-        problems = filter_problems(dataset, problems, custom_splits, excluded_ids)
-        dataset = {problem: ds[dataset[problem]] for problem in problems}
+        problems = filter_problems(id2idx, problems, custom_splits, excluded_ids)
+        dataset = {problem: ds[id2idx[problem]] for problem in problems}
 
         # add instance id to each example (name of the image)
-        for instance_id in dataset:
-            dataset[instance_id]["instance_id"] = instance_id
+        for instance_id, task_data in dataset.items():
+            task_data["instance_id"] = instance_id
 
-        image_names = set(example["docker_image"] for example in dataset.values())
+        image_names = set(task_data["docker_image"] for task_data in dataset.values())
         logger.debug(
             f"Loaded {len(dataset)} tasks across {len(image_names)} Docker images from {dataset_id}."
         )

debug_gym/gym/envs/swe_bench.py

@@ -185,12 +185,13 @@ def load_dataset(
     ) -> dict:
         ds = datasets.load_dataset(dataset_id, revision=dataset_revision)[split]
 
-        dataset = {problem["instance_id"]: problem for problem in ds}
-        problems = filter_problems(dataset, problems)
-        dataset = {id: i for id, i in dataset.items() if id in problems}
+        # Memory efficient filtering of problems.
+        id2idx = {id: i for i, id in enumerate(ds["instance_id"])}
+        problems = filter_problems(id2idx, problems)
+        dataset = {problem: ds[id2idx[problem]] for problem in problems}
 
         image_names = set(
-            f"sweb.eval.x86_64.{id.replace('__', '_1776_')}" for id in problems
+            f"sweb.eval.x86_64.{id.replace('__', '_1776_')}" for id in dataset
         )
 
         if prepull_images:

debug_gym/gym/envs/swe_smith.py

@@ -177,11 +177,12 @@ def load_dataset(
             custom_splits = yaml.safe_load(f)
             excluded_ids = custom_splits.get("excluded", [])
 
-        dataset = {d["instance_id"]: d for d in ds}
-        problems = filter_problems(dataset, problems, custom_splits, excluded_ids)
-        dataset = {pid: dataset[pid] for pid in problems}
+        # Memory efficient filtering of problems.
+        id2idx = {id: i for i, id in enumerate(ds["instance_id"])}
+        problems = filter_problems(id2idx, problems, custom_splits, excluded_ids)
+        dataset = {problem: ds[id2idx[problem]] for problem in problems}
 
-        image_names = set([problem["image_name"] for problem in dataset.values()])
+        image_names = set(task_data["image_name"] for task_data in dataset.values())
         logger.debug(
             f"Loaded {len(dataset)} tasks across {len(image_names)} Docker images from {dataset_id}."
         )

tests/gym/envs/conftest.py

@@ -25,8 +25,6 @@
     },
 }
 
-logger = logging.getLogger(__name__)
-
 
 def make_env_factory(env_name, worker_id, tmp_path_factory):
     """Build the `env_name`'s docker image only once."""
@@ -35,27 +33,22 @@ def make_env_factory(env_name, worker_id, tmp_path_factory):
     env_class = kwargs.pop("env_class")
 
     def _make_env():
-        logger.info("\n**Calling load_dataset.**\n")
         dataset = env_class.load_dataset(
             problems=kwargs["problems"], prepull_images=True
         )
         task_data = next(iter(dataset.values()))
-        logger.info(f"\n**Creating env.** {env_class}\n")
         env = env_class(task_data=task_data)
-        logger.info("\n**Done.**\n")
         return env
 
     if worker_id == "master":
         # Not running with pytest-xdist or we are in the master process
-        logger.warning("Environment initialized in master process.")
         _make_env()
     else:
         # When running with pytest-xdist, synchronize between workers using a lock
         root_tmp_dir = tmp_path_factory.getbasetemp().parent
         lock_file = root_tmp_dir / f"{env_class.__name__}_init.lock"
         with FileLock(str(lock_file)):
             # Only the first worker to acquire the lock will initialize the environment
-            logger.warning(f"Environment running in worker {worker_id}.")
             _make_env()
 
     return _make_env