perf: accelerate data loading in training

deepmd/env.py

@@ -1,6 +1,7 @@
 # SPDX-License-Identifier: LGPL-3.0-or-later
 import logging
 import os
+import platform
 from configparser import (
     ConfigParser,
 )
@@ -16,6 +17,7 @@
     "GLOBAL_CONFIG",
     "GLOBAL_ENER_FLOAT_PRECISION",
     "GLOBAL_NP_FLOAT_PRECISION",
+    "LRU_CACHE_SIZE",
     "SHARED_LIB_DIR",
     "SHARED_LIB_MODULE",
     "global_float_prec",
@@ -47,6 +49,20 @@
         "DP_INTERFACE_PREC."
     )
 
+# Dynamic calculation of cache size
+_default_lru_cache_size = 512
+LRU_CACHE_SIZE = _default_lru_cache_size
+
+if platform.system() != "Windows":
+    import resource
+
+    soft_limit, hard_limit = resource.getrlimit(resource.RLIMIT_NOFILE)
+    safe_buffer = 128
+    if soft_limit > safe_buffer + _default_lru_cache_size:
+        LRU_CACHE_SIZE = soft_limit - safe_buffer
+    else:
+        LRU_CACHE_SIZE = soft_limit // 2
+
 
 def set_env_if_empty(key: str, value: str, verbose: bool = True) -> None:
     """Set environment variable only if it is empty.