[BugFix] Fix VD4RL (#1834)

Author: Vincent Moens

torchrl/data/datasets/vd4rl.py

@@ -4,6 +4,8 @@
 # LICENSE file in the root directory of this source tree.
 from __future__ import annotations
 
+import functools
+
 import importlib
 import json
 import logging
@@ -12,14 +14,14 @@
 import shutil
 import tempfile
 from collections import defaultdict
-from concurrent.futures import ThreadPoolExecutor
 from pathlib import Path
 from typing import Callable, List
 
 import numpy as np
 
 import torch
 from tensordict import PersistentTensorDict, TensorDict
+from torch import multiprocessing as mp
 
 from torchrl._utils import KeyDependentDefaultDict
 from torchrl.data.datasets.utils import _get_root_dir
@@ -96,6 +98,8 @@ class VD4RLExperienceReplay(TensorDictReplayBuffer):
             transform that will be appended to the transform list. Supports
             `int` types (square resizing) or a list/tuple of `int` (rectangular
             resizing). Defaults to ``None`` (no resizing).
+        num_workers (int, optional): the number of workers to download the files.
+            Defaults to ``0`` (no multiprocessing).
 
     Attributes:
         available_datasets: a list of accepted entries to be downloaded. These
@@ -173,6 +177,7 @@ def __init__(
         split_trajs: bool = False,
         totensor: bool = True,
         image_size: int | List[int] | None = None,
+        num_workers: int = 0,
         **env_kwargs,
     ):
         if not _has_h5py or not _has_hf_hub:
@@ -191,6 +196,7 @@ def __init__(
         self.root = root
         self.split_trajs = split_trajs
         self.download = download
+        self.num_workers = num_workers
         if self.download == "force" or (self.download and not self._is_downloaded()):
             if self.download == "force":
                 try:
@@ -199,7 +205,9 @@ def __init__(
                         shutil.rmtree(self.data_path)
                 except FileNotFoundError:
                     pass
-            storage = self._download_and_preproc(dataset_id, data_path=self.data_path)
+            storage = self._download_and_preproc(
+                dataset_id, data_path=self.data_path, num_workers=self.num_workers
+            )
         elif self.split_trajs and not os.path.exists(self.data_path):
             storage = self._make_split()
         else:
@@ -251,14 +259,23 @@ def _parse_datasets(cls):
         return sibs
 
     @classmethod
-    def _download_and_preproc(cls, dataset_id, data_path):
+    def _hf_hub_download(cls, subfolder, filename, *, tmpdir):
         from huggingface_hub import hf_hub_download
 
-        files = []
+        return hf_hub_download(
+            "conglu/vd4rl",
+            subfolder=subfolder,
+            filename=filename,
+            repo_type="dataset",
+            cache_dir=str(tmpdir),
+        )
+
+    @classmethod
+    def _download_and_preproc(cls, dataset_id, data_path, num_workers):
+
         tds = []
         with tempfile.TemporaryDirectory() as tmpdir:
             sibs = cls._parse_datasets()
-            # files = []
             total_steps = 0
 
             paths_to_proc = []
@@ -270,19 +287,19 @@ def _download_and_preproc(cls, dataset_id, data_path):
                 for file in sibs[path]:
                     paths_to_proc.append(str(path))
                     files_to_proc.append(str(file.parts[-1]))
-
-            with ThreadPoolExecutor(32) as executor:
-                files = executor.map(
-                    lambda path_file: hf_hub_download(
-                        "conglu/vd4rl",
-                        subfolder=path_file[0],
-                        filename=path_file[1],
-                        repo_type="dataset",
-                        cache_dir=str(tmpdir),
-                    ),
-                    zip(paths_to_proc, files_to_proc),
-                )
-                files = list(files)
+            func = functools.partial(cls._hf_hub_download, tmpdir=tmpdir)
+            if num_workers > 0:
+                with mp.Pool(num_workers) as pool:
+                    files = pool.starmap(
+                        func,
+                        zip(paths_to_proc, files_to_proc),
+                    )
+                    files = list(files)
+            else:
+                files = [
+                    func(subfolder, filename)
+                    for (subfolder, filename) in zip(paths_to_proc, files_to_proc)
+                ]
             logging.info("Downloaded, processing files")
             if _has_tqdm:
                 import tqdm