remove use of tmp folder

benchmarl/algorithms/common.py

@@ -3,12 +3,12 @@
 #  This source code is licensed under the license found in the
 #  LICENSE file in the root directory of this source tree.
 #
-
+import os
+import uuid
 import pathlib
 from abc import ABC, abstractmethod
 from dataclasses import dataclass
 from enum import Enum
-import tempfile
 from typing import Any, Callable, Dict, Iterable, List, Optional, Tuple, Type
 
 from tensordict import TensorDictBase
@@ -74,7 +74,6 @@ def __init__(self, experiment):
         self._losses_and_updaters = {}
         self._policies_for_loss = {}
         self._policies_for_collection = {}
-        self._memmap_dir = None
 
         self._check_specs()
 
@@ -151,11 +150,10 @@ def get_storage(self, memory_size: int, physical_storage: PhysicalStorage) -> La
                 device=self.device if self.on_policy else self.buffer_device,
             )
         else:
-            self._memmap_dir = tempfile.TemporaryDirectory()
             storage = LazyMemmapStorage(
                 memory_size,
                 device=self.device if self.on_policy else self.buffer_device,
-                scratch_dir=self._memmap_dir.name
+                scratch_dir=f".memmap-{uuid.uuid4().hex}",
             )
 
         return storage
@@ -359,10 +357,6 @@ def process_loss_vals(
         """
         return loss_vals
 
-    def __del__(self) -> None:
-        if self._memmap_dir is not None:
-            self._memmap_dir.cleanup()
-
 
 @dataclass
 class AlgorithmConfig:

benchmarl/experiment/experiment.py

@@ -8,8 +8,9 @@
 
 import copy
 import importlib
-
 import os
+import shutil
+
 import time
 from collections import deque, OrderedDict
 from dataclasses import dataclass, MISSING
@@ -21,6 +22,7 @@
 from tensordict import TensorDictBase
 from tensordict.nn import TensorDictSequential
 from torchrl.collectors import SyncDataCollector
+from torchrl.data import LazyMemmapStorage
 
 from torchrl.envs import ParallelEnv, SerialEnv, TransformedEnv
 from torchrl.envs.transforms import Compose
@@ -770,6 +772,10 @@ def close(self):
         self.test_env.close()
         self.logger.finish()
 
+        for group in self.replay_buffers:
+            if isinstance(self.replay_buffers[group].storage, LazyMemmapStorage):
+                shutil.rmtree(self.replay_buffers[group].storage.scratch_dir, ignore_errors=True)
+
     def _get_excluded_keys(self, group: str):
         excluded_keys = []
         for other_group in self.group_map.keys():