From c7ba8dac1221ea8c593d3570d4caa05eab92a4d1 Mon Sep 17 00:00:00 2001
From: Nick Rossenbach <rossenbach@cs.rwth-aachen.de>
Date: Tue, 12 Dec 2023 14:31:08 +0900
Subject: [PATCH] Add internal cache manager (#13)

* Experimental version of internal cache manager

manages (persistent) file copy to temporary directory, checks for file equivalence,
and also caches internal access history.

* add missing return

* black
---
 CHANGELOG.rst           |  6 ++++++
 returnn/__setup__.py    |  2 +-
 returnn/torch/engine.py |  1 -
 returnn/util/basic.py   | 39 +++++++++++++++++++++++++++++++++------
 4 files changed, 40 insertions(+), 8 deletions(-)

diff --git a/CHANGELOG.rst b/CHANGELOG.rst
index 4f9795f325..b83f30fd44 100644
--- a/CHANGELOG.rst
+++ b/CHANGELOG.rst
@@ -1,6 +1,12 @@
 Changelog
 =========
 
+Current Master (0.3+git)
+------------------------
+
+- Added temporary file cache manager (activated via `use_cache_manager` as for the i6 specific caching)
+
+
 
 Version 0.3
 -----------
diff --git a/returnn/__setup__.py b/returnn/__setup__.py
index 1a7ad5856b..23e6c3c97c 100644
--- a/returnn/__setup__.py
+++ b/returnn/__setup__.py
@@ -8,7 +8,7 @@
 import sys
 
 
-VERSION = "0.3"
+VERSION = "0.3+git"
 
 _my_dir = os.path.dirname(os.path.abspath(__file__))
 # Use realpath to resolve any symlinks. We want the real root-dir, to be able to check the Git revision.
diff --git a/returnn/torch/engine.py b/returnn/torch/engine.py
index b0b064db66..b3be34d774 100644
--- a/returnn/torch/engine.py
+++ b/returnn/torch/engine.py
@@ -13,7 +13,6 @@
 import torch
 import torch.utils.data.datapipes as dp
 from torch import autocast, Tensor
-from torch.cuda import amp
 from torch.utils.data import DataLoader
 from random import random
 import math
diff --git a/returnn/util/basic.py b/returnn/util/basic.py
index f17b062ecf..91e9b9ed23 100644
--- a/returnn/util/basic.py
+++ b/returnn/util/basic.py
@@ -5,6 +5,8 @@
 """
 
 from __future__ import annotations
+
+import shutil
 from typing import Generic, TypeVar
 
 import subprocess
@@ -27,6 +29,7 @@
     import _thread as thread
 import threading
 
+import filecmp
 import typing
 from returnn.log import log
 import builtins
@@ -2407,10 +2410,15 @@ def is_namedtuple(cls):
     return issubclass(cls, tuple) and cls is not tuple
 
 
+# cache for the i6 cache manager
 _cf_cache = {}
 _cf_msg_printed = False
 
 
+# generic tempdir cache
+_tempdir_cache = {}
+
+
 def cf(filename):
     """
     Cache manager. i6 specific.
@@ -2422,18 +2430,37 @@ def cf(filename):
     import os
     from subprocess import check_output
 
+    # first try caching via i6 cf
     if filename in _cf_cache:
         return _cf_cache[filename]
     try:
         cached_fn = check_output(["cf", filename]).strip().decode("utf8")
-    except CalledProcessError:
+        assert os.path.exists(cached_fn)
+        _cf_cache[filename] = cached_fn
+        return cached_fn
+    except (CalledProcessError, FileNotFoundError):
         if not _cf_msg_printed:
-            print("Cache manager: Error occurred, using local file")
+            print("i6 cache manager: Error occurred, using internal cache manager", file=log.v3)
             _cf_msg_printed = True
-        return filename
-    assert os.path.exists(cached_fn)
-    _cf_cache[filename] = cached_fn
-    return cached_fn
+
+    # otherwise do generic caching
+    tmp_root = get_temp_dir()
+    if not os.path.exists(tmp_root):
+        os.mkdir(tmp_root)
+
+    real_filename = os.path.realpath(filename)
+    assert real_filename.startswith("/")
+    if filename in _tempdir_cache:
+        existing_file = _tempdir_cache[filename]
+        if filecmp.cmp(real_filename, existing_file) is True:
+            return _tempdir_cache[filename]
+
+    temp_file = os.path.join(tmp_root, real_filename[1:])  # join without root slash
+    os.makedirs(os.path.dirname(temp_file), exist_ok=True)
+    shutil.copy(real_filename, temp_file)
+    _tempdir_cache[filename] = temp_file
+
+    return temp_file
 
 
 def binary_search_any(cmp, low, high):