returnn 1.20240621.152144__tar.gz → 1.20240621.194850__tar.gz

{returnn-1.20240621.152144 → returnn-1.20240621.194850}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: returnn
-Version: 1.20240621.152144
+Version: 1.20240621.194850
 Summary: The RWTH extensible training framework for universal recurrent neural networks
 Home-page: https://github.com/rwth-i6/returnn/
 Author: Albert Zeyer

returnn-1.20240621.194850/_setup_info_generated.py

@@ -0,0 +1,2 @@
+version = '1.20240621.194850'
+long_version = '1.20240621.194850+git.ef9cff9'

{returnn-1.20240621.152144 → returnn-1.20240621.194850}/returnn/__main__.py

@@ -29,8 +29,7 @@ from returnn.log import log
 from returnn.config import Config, get_global_config
 from returnn.datasets import Dataset, init_dataset, init_dataset_via_str
 from returnn.datasets.hdf import HDFDataset
-from returnn.util import debug as debug_util
-from returnn.util import basic as util
+from returnn.util import basic as util, debug as debug_util
 from returnn.util.basic import BackendEngine, BehaviorVersion
 
 # These imports are not directly used here, but make them available, as other code imports them from here.

{returnn-1.20240621.152144 → returnn-1.20240621.194850}/returnn/datasets/basic.py

@@ -25,6 +25,7 @@ from returnn.log import log
 from returnn.engine.batch import Batch, BatchSetGenerator
 from returnn.datasets.util.vocabulary import Vocabulary
 from returnn.util.basic import try_run, NumbersDict, OptionalNotImplementedError
+from returnn.util import file_cache
 from returnn.tensor import TensorDict
 
 if TYPE_CHECKING:
@@ -195,6 +196,7 @@ class Dataset(object):
         self.shuffle_frames_of_nseqs = shuffle_frames_of_nseqs
         self.epoch = None
         self.zpad = None
+        self._file_cache: Optional[file_cache.FileCache] = None
 
     def __repr__(self):
         return "<%s %r epoch=%s>" % (
@@ -282,6 +284,13 @@ class Dataset(object):
                 return returnn.tf.horovod.get_ctx(config=config).rank() * 16127
         return 0
 
+    def set_file_cache(self, cache: file_cache.FileCache):
+        """
+        Stores the given file cache with the dataset to unregister the files within
+        the cache when the dataset is deinitialized.
+        """
+        self._file_cache = cache
+
     @staticmethod
     def _parse_chunking(chunking):
         """
@@ -1434,7 +1443,13 @@ def init_dataset(
                 kwargs.update(default_kwargs)
             if extra_kwargs:
                 kwargs.update(extra_kwargs)
-            return init_dataset_via_str(config_str=config_str, **kwargs)
+            cache = file_cache.get_instance()
+            kwargs, cached_files = cache.handle_cached_files_in_config(kwargs)
+            obj = init_dataset_via_str(config_str=config_str, **kwargs)
+            assert isinstance(obj, Dataset)
+            if cached_files:
+                obj.set_file_cache(cache)
+            return obj
     assert isinstance(kwargs, dict)
     kwargs = kwargs.copy()
     assert "class" in kwargs
@@ -1448,9 +1463,13 @@ def init_dataset(
             kwargs.setdefault(key, value)
     if extra_kwargs:
         kwargs.update(extra_kwargs)
+    cache = file_cache.get_instance()
+    kwargs, cached_files = cache.handle_cached_files_in_config(kwargs)
     obj = clazz(**kwargs)
     assert isinstance(obj, Dataset)
     obj.initialize()
+    if cached_files:
+        obj.set_file_cache(cache)
     return obj
 
 

{returnn-1.20240621.152144 → returnn-1.20240621.194850}/returnn/datasets/distrib_files.py

@@ -6,7 +6,7 @@ https://github.com/rwth-i6/returnn/issues/1519
 
 from __future__ import annotations
 
-from typing import Optional, Any, Dict, Tuple, List, Sequence, Collection, Callable, Union
+from typing import Optional, Any, Dict, Tuple, List, Sequence, Callable, Union
 import os
 import sys
 import numpy
@@ -136,7 +136,6 @@ class DistributeFilesDataset(CachedDataset2):
         get_sub_epoch_dataset: Callable[[List[FileTree]], Dict[str, Any]],
         preload_next_n_sub_epochs: int = 1,
         buffer_size: int = 1,
-        file_cache_opts: Optional[Dict[str, Any]] = None,
         distrib_shard_files: bool = False,
         _meta_info_cache: Optional[Dict[str, Any]] = None,
         **kwargs,
@@ -157,13 +156,11 @@ class DistributeFilesDataset(CachedDataset2):
         assert preload_next_n_sub_epochs >= 0
         self.preload_next_n_sub_epochs = preload_next_n_sub_epochs
         self.buffer_size = buffer_size
-        self.file_cache_opts = file_cache_opts or {}
         self._file_sizes: Optional[Dict[str, int]] = None  # key -> size. for equal distribution across sub epochs
         self._data_keys: Optional[List[str]] = None
         self._num_seqs: Optional[int] = None
         self._shard_index, self._num_shards = _get_rank_and_size() if distrib_shard_files else (0, 1)
 
-        self._file_cache: Optional[_FileCacheProc] = None
         self._workers: Dict[int, _WorkerProcParent] = {}  # epoch -> worker
         self._files_order_cache: Dict[int, List[List[FileTree]]] = {}  # full epoch (0-indexed) -> files order
 
@@ -205,15 +202,12 @@ class DistributeFilesDataset(CachedDataset2):
             return
         # First, we need to know the num_inputs, num_outputs, total_num_seqs, labels.
         # Init the dataset with the first file.
-        dataset_dict, exit_hook = self._get_sub_dataset_dict(files=[self.files[0]])
-        try:
-            dataset = init_dataset(dataset_dict, extra_kwargs={"seq_ordering": "default"}, parent_dataset=self)
-            self.num_inputs = dataset.num_inputs
-            self.num_outputs = dataset.num_outputs
-            self.labels = dataset.labels
-            self._data_keys = dataset.get_data_keys()
-        finally:
-            exit_hook()
+        dataset_dict = self._get_sub_dataset_dict(files=[self.files[0]])
+        dataset = init_dataset(dataset_dict, extra_kwargs={"seq_ordering": "default"}, parent_dataset=self)
+        self.num_inputs = dataset.num_inputs
+        self.num_outputs = dataset.num_outputs
+        self.labels = dataset.labels
+        self._data_keys = dataset.get_data_keys()
 
     def _lazy_init_file_sizes(self):
         import tree
@@ -224,17 +218,7 @@ class DistributeFilesDataset(CachedDataset2):
             _get_key_for_file_tree(t): sum((os.path.getsize(fn) for fn in tree.flatten(t)), 0) for t in self.files
         }
 
-    def _lazy_init_file_cache_proc(self):
-        if self._file_cache:
-            return
-        self._file_cache = _FileCacheProc(
-            name=f"{self.__class__.__name__} {self.name}", file_cache_opts=self.file_cache_opts
-        )
-
     def __del__(self):
-        if self._file_cache:
-            try_run(self._file_cache.exit, kwargs={"join": False})
-            self._file_cache = None
         for k, worker in self._workers.items():
             try_run(worker.exit, kwargs={"join": False})
         self._workers.clear()
@@ -256,7 +240,6 @@ class DistributeFilesDataset(CachedDataset2):
             return True
 
         self._lazy_init_file_sizes()
-        self._lazy_init_file_cache_proc()
 
         full_epoch_0idx = (epoch - 1) // self.partition_epoch
 
@@ -300,21 +283,20 @@ class DistributeFilesDataset(CachedDataset2):
             files_order: List[List[FileTree]] = self._files_order_cache[full_epoch_0idx_]
             files_for_subep = files_order[(ep_ - 1) % self.partition_epoch]
             print(f"{self}: using files for epoch {ep_}: {files_for_subep}", file=log.v4)
-            dataset_dict, exit_hook = self._get_sub_dataset_dict(files=files_for_subep)
+            dataset_dict = self._get_sub_dataset_dict(files=files_for_subep)
             worker = _WorkerProcParent(
                 name=f"{self.__class__.__name__} {self.name} ep {epoch}",
                 epoch=ep_,
                 full_epoch_0idx=full_epoch_0idx_,
                 dataset_dict=dataset_dict,
                 buffer_size=self.buffer_size,
-                exit_hook=exit_hook,
             )
             self._workers[ep_] = worker
 
         self._num_seqs = self._workers[epoch].get_num_seqs()
         return True
 
-    def _get_sub_dataset_dict(self, files: List[FileTree]) -> Tuple[Dict[str, Any], _FileCacheExitHook]:
+    def _get_sub_dataset_dict(self, files: List[FileTree]) -> Dict[str, Any]:
         dataset_dict = self.get_sub_epoch_dataset(files)
         dataset_dict = extend_dataset_dict_from_parent_dataset(dataset_dict, parent_dataset=self)
         if dataset_dict.get("partition_epoch", 1) != 1:
@@ -324,9 +306,7 @@ class DistributeFilesDataset(CachedDataset2):
                 f"{self}: sub dataset should have explicit seq_ordering "
                 f"(or seq_order_control_dataset for MetaDataset), got: {dataset_dict}"
             )
-        self._lazy_init_file_cache_proc()
-        dataset_dict, exit_hook = self._file_cache.handle_cached_files_in_config(dataset_dict)
-        return dataset_dict, exit_hook
+        return dataset_dict
 
     @staticmethod
     def _distribute_evenly_by_size(
@@ -417,9 +397,6 @@ class DistributeFilesDataset(CachedDataset2):
                 worker.exit()
             self._workers.clear()
             self._files_order_cache.clear()
-            if self._file_cache:
-                self._file_cache.exit()
-                self._file_cache = None
         else:
             if self.epoch in self._workers:
                 worker = self._workers.pop(self.epoch)
@@ -647,95 +624,3 @@ def _worker_proc_loop(
                 raise Exception(f"unknown msg {msg!r}")
     except KeyboardInterrupt:  # when parent dies
         pass
-
-
-class _FileCacheProc:
-    def __init__(self, *, name: str, file_cache_opts: Dict[str, Any]):
-        self.file_cache_opts = file_cache_opts
-
-        parent_conn, child_conn = _mp.Pipe()
-        self.parent_conn: mpConnection = parent_conn
-
-        self.proc = _mp.Process(
-            name=f"{name} file cache",
-            target=_file_cache_proc_loop,
-            args=(file_cache_opts, child_conn),
-            daemon=True,
-        )
-        self.proc.start()
-        # Make sure the child connection is closed here.
-        # It stays open in the child, until the child dies.
-        # When that happens, now any consecutive read on the pipe
-        # should yield an exception -- which is what we want,
-        # otherwise it would just hang.
-        child_conn.close()
-
-    def handle_cached_files_in_config(self, config: Dict[str, Any]) -> Tuple[Dict[str, Any], _FileCacheExitHook]:
-        """:func:`FileCache.handle_cached_files_in_config`"""
-        self.parent_conn.send(("handle_cached_files_in_config", {"config": config}))
-        msg, (res, files) = self.parent_conn.recv()
-        assert msg == "config_and_cached_files" and isinstance(config, dict) and isinstance(files, list)
-        return res, _FileCacheExitHook(self, files)
-
-    def release_files(self, filenames: Collection[str]):
-        """:func:`FileCache.release_files`"""
-        self.parent_conn.send(("release_files", {"filenames": filenames}))
-
-    def exit(self, *, join: bool = True):
-        """exit"""
-        if self.proc.exitcode is not None:  # already stopped?
-            return
-        self.parent_conn.send(("exit", {}))
-        if join:
-            self.proc.join()
-
-    def __del__(self):
-        # noinspection PyBroadException
-        try:
-            self.exit(join=False)
-        except Exception:
-            pass
-        else:
-            try_run(self.proc.join)
-
-
-class _FileCacheExitHook:
-    def __init__(self, file_cache_proc: _FileCacheProc, cached_files: List[str]):
-        self.file_cache_proc = file_cache_proc
-        self.cached_files = cached_files
-
-    def __call__(self):
-        if self.file_cache_proc.proc.exitcode is not None:
-            return
-        self.file_cache_proc.release_files(self.cached_files)
-
-
-def _file_cache_proc_loop(
-    file_cache_opts: Dict[str, Any],
-    parent_conn: mpConnection,
-):
-    if sys.platform == "linux":
-        with open("/proc/self/comm", "w") as f:
-            f.write(f"File cache")
-
-    assert isinstance(file_cache_opts, dict)
-    assert isinstance(parent_conn, mpConnection)
-
-    from returnn.util.file_cache import FileCache
-
-    file_cache = FileCache(**file_cache_opts)
-
-    try:
-        while True:
-            msg, kwargs = parent_conn.recv()
-            if msg == "exit":
-                break
-            elif msg == "handle_cached_files_in_config":
-                res, files = file_cache.handle_cached_files_in_config(**kwargs)
-                parent_conn.send(("config_and_cached_files", (res, files)))
-            elif msg == "release_files":
-                file_cache.release_files(**kwargs)
-            else:
-                raise Exception(f"unknown msg {msg!r}")
-    except KeyboardInterrupt:  # when parent dies
-        pass

{returnn-1.20240621.152144 → returnn-1.20240621.194850}/returnn/util/file_cache.py

@@ -9,9 +9,10 @@ See https://github.com/rwth-i6/returnn/issues/1519 for initial discussion.
 Main class is :class:`FileCache`.
 """
 
-from typing import Any, Collection, List, Tuple
+from typing import Any, Collection, List, Optional, Tuple
 import errno
 import os
+import pathlib
 import time
 import shutil
 from dataclasses import dataclass
@@ -19,9 +20,10 @@ from collections import defaultdict
 from contextlib import contextmanager
 from threading import Thread, Event
 from .basic import expand_env_vars, LockFile, human_bytes_size
+from returnn.config import Config, get_global_config
 
 
-__all__ = ["FileCache", "CachedFile"]
+__all__ = ["FileCache", "CachedFile", "get_instance"]
 
 
 class FileCache:
@@ -78,8 +80,6 @@ class FileCache:
         assert cleanup_disk_usage_wanted_multiplier >= 1.0
         self._cleanup_disk_usage_wanted_multiplier = cleanup_disk_usage_wanted_multiplier
         self._touch_files_thread = _TouchFilesThread(cache_base_dir=self.cache_directory)
-        self._touch_files_thread.start()
-        self._recent_full_cleanup_time = float("-inf")
         assert num_tries > 0
         self._num_tries = num_tries
 
@@ -151,13 +151,21 @@ class FileCache:
             int(self._cleanup_disk_usage_wanted_multiplier * need_at_least_free_space_size),
             int(self._cleanup_disk_usage_wanted_free_ratio * disk_usage.total),
         )
-        # If we have enough free space, and we did a full cleanup recently, we don't need to do anything.
-        if want_free_space_size <= disk_usage.free and time.monotonic() - self._recent_full_cleanup_time < 60 * 10:
-            return
-        # Do a full cleanup, i.e. iterate through all files in cache directory and check their mtime.
+        cleanup_timestamp_file = self.cache_directory + "/.recent_full_cleanup"
+        try:
+            last_full_cleanup = os.stat(cleanup_timestamp_file).st_mtime
+        except FileNotFoundError:
+            last_full_cleanup = float("-inf")
         # Get current time now, so that cur_time - mtime is pessimistic,
         # and does not count the time for the cleanup itself.
         cur_time = time.time()
+        # If we have enough free space, and we did a full cleanup recently, we don't need to do anything.
+        if want_free_space_size <= disk_usage.free and cur_time - last_full_cleanup < 60 * 10:
+            return
+        # immediately update the file's timestamp to reduce racyness between worker processes
+        # Path().touch() also creates the file if it doesn't exist yet
+        pathlib.Path(cleanup_timestamp_file).touch(exist_ok=True)
+        # Do a full cleanup, i.e. iterate through all files in cache directory and check their mtime.
         all_files = []  # mtime, neg size (better for sorting), filename
         for root, dirs, files in os.walk(self.cache_directory):
             for rel_fn in files:
@@ -256,8 +264,6 @@ class FileCache:
             except Exception as exc:
                 print(f"FileCache: Error while removing empty dir {root}: {type(exc).__name__}: {exc}")
 
-        self._recent_full_cleanup_time = time.monotonic()
-
     def handle_cached_files_in_config(self, config: Any) -> Tuple[Any, List[str]]:
         """
         :param config: some config, e.g. dict, or any nested structure
@@ -337,6 +343,17 @@ class FileCache:
         return True
 
 
+def get_instance(config: Optional[Config] = None) -> FileCache:
+    """
+    Returns a file cache instance potentially initialized by the global config.
+
+    Uses defaults if no global config is set.
+    """
+    config = config or get_global_config(return_empty_if_none=True)
+    kwargs = config.typed_value("file_cache_opts") or {}
+    return FileCache(**kwargs)
+
+
 def _copy_with_prealloc(src: str, dst: str):
     """
     Copies the file at `src` to `dst` preallocating the space at `dst` before the
@@ -387,6 +404,7 @@ class _TouchFilesThread(Thread):
         self.files = defaultdict(int)  # usage counter
         self.interval = interval
         self.cache_base_dir = cache_base_dir
+        self._is_started = False  # careful: `_started` is already a member of the base class
 
     def run(self):
         """thread main loop"""
@@ -400,9 +418,17 @@ class _TouchFilesThread(Thread):
             if self.stop.wait(self.interval):
                 return
 
+    def start_once(self):
+        """reentrant variant of start() that can safely be called multiple times"""
+        if self._is_started:
+            return
+        self._is_started = True
+        self.start()
+
     def files_extend(self, files: Collection[str]):
         """append"""
         assert isinstance(files, (list, set, tuple))
+        self.start_once()
         for file in files:
             self.files[file] += 1
 

{returnn-1.20240621.152144 → returnn-1.20240621.194850}/returnn.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: returnn
-Version: 1.20240621.152144
+Version: 1.20240621.194850
 Summary: The RWTH extensible training framework for universal recurrent neural networks
 Home-page: https://github.com/rwth-i6/returnn/
 Author: Albert Zeyer

{returnn-1.20240621.152144 → returnn-1.20240621.194850}/tests/test_Util.py

@@ -870,7 +870,12 @@ def test_file_cache():
     # Check cleanup mechanism.
     cache._cleanup_files_always_older_than_days = 4 / 60 / 60 / 24
     print("Cache older than secs:", cache._cleanup_files_always_older_than_days * 24 * 60 * 60)
-    print("Recent cleanup is ago:", time.monotonic() - cache._recent_full_cleanup_time, "secs")
+    cleanup_timestamp_file = cache_dir + "/.recent_full_cleanup"
+    try:
+        last_full_cleanup = os.stat(cleanup_timestamp_file).st_mtime
+    except FileNotFoundError:
+        last_full_cleanup = float("-inf")
+    print("Recent cleanup is ago:", time.time() - last_full_cleanup, "secs")
     # Note: After the first cleanup, there are still dirs maybe left over, because dirs are only cleaned
     # when their mtime is a older than some threshold.
     # However, when we just deleted some file inside, that would have updated the mtime of the dir.
@@ -879,14 +884,17 @@ def test_file_cache():
     prev_count_dirs = None
     while True:
         time.sleep(0.1)
-        cache._recent_full_cleanup_time = float("-inf")  # reset
+        os.unlink(cleanup_timestamp_file)  # reset
+        time.sleep(0.1)
         cache.cleanup()
         count_files = 0
         count_dirs = 0
         for root, dirs, files in os.walk(cache_dir):
+            # ignore cleanup timestamp file
             for fn in dirs + files:
                 print(f"{root}/{fn}", "age:", time.time() - os.stat(f"{root}/{fn}").st_mtime, "sec")
-            count_files += len(files)
+            leftover_files = [f for f in files if "recent_full_cleanup" not in f]
+            count_files += len(leftover_files)
             count_dirs += len(dirs)
         print("count files:", count_files, "count dirs:", count_dirs)
         assert count_files == 0

returnn-1.20240621.152144/_setup_info_generated.py

@@ -1,2 +0,0 @@
-version = '1.20240621.152144'
-long_version = '1.20240621.152144+git.6db5121'