From 641d09164c52265dd29a31232d57b22616a93f26 Mon Sep 17 00:00:00 2001
From: Glenn Jocher <glenn.jocher@ultralytics.com>
Date: Thu, 19 Sep 2024 21:06:57 +0200
Subject: [PATCH] New `PERSISTENT_CACHE` (#16373)

Signed-off-by: UltralyticsAssistant <web@ultralytics.com>
Co-authored-by: UltralyticsAssistant <web@ultralytics.com>
---
 docs/en/reference/utils/__init__.md |  4 ++
 ultralytics/utils/__init__.py       | 58 +++++++++++++++++++++++++++++
 ultralytics/utils/torch_utils.py    | 16 +++++---
 3 files changed, 72 insertions(+), 6 deletions(-)

diff --git a/docs/en/reference/utils/__init__.md b/docs/en/reference/utils/__init__.md
index b8147dac..9c5f2b44 100644
--- a/docs/en/reference/utils/__init__.md
+++ b/docs/en/reference/utils/__init__.md
@@ -39,6 +39,10 @@ keywords: Ultralytics, utils, TQDM, Python, ML, Machine Learning utilities, YOLO
 
 <br><br><hr><br>
 
+## ::: ultralytics.utils.PersistentCacheDict
+
+<br><br><hr><br>
+
 ## ::: ultralytics.utils.plt_settings
 
 <br><br><hr><br>
diff --git a/ultralytics/utils/__init__.py b/ultralytics/utils/__init__.py
index 7328a5fd..5ba58d08 100644
--- a/ultralytics/utils/__init__.py
+++ b/ultralytics/utils/__init__.py
@@ -3,6 +3,7 @@
 import contextlib
 import importlib.metadata
 import inspect
+import json
 import logging.config
 import os
 import platform
@@ -14,6 +15,7 @@ import time
 import urllib
 import uuid
 from pathlib import Path
+from threading import Lock
 from types import SimpleNamespace
 from typing import Union
 
@@ -1136,6 +1138,61 @@ class SettingsManager(dict):
         self.save()
 
 
+class PersistentCacheDict(dict):
+    """A thread-safe dictionary that persists data to a JSON file for caching purposes."""
+
+    def __init__(self, file_path=USER_CONFIG_DIR / "persistent_cache.json"):
+        """Initializes a thread-safe persistent cache dictionary with a specified file path for storage."""
+        super().__init__()
+        self.file_path = Path(file_path)
+        self.lock = Lock()
+        self._load()
+
+    def _load(self):
+        """Load the persistent cache from a JSON file into the dictionary, handling errors gracefully."""
+        try:
+            if self.file_path.exists():
+                with open(self.file_path) as f:
+                    self.update(json.load(f))
+        except json.JSONDecodeError:
+            print(f"Error decoding JSON from {self.file_path}. Starting with an empty cache.")
+        except Exception as e:
+            print(f"Error reading from {self.file_path}: {e}")
+
+    def _save(self):
+        """Save the current state of the cache dictionary to a JSON file, ensuring thread safety."""
+        try:
+            self.file_path.parent.mkdir(parents=True, exist_ok=True)
+            with open(self.file_path, "w") as f:
+                json.dump(dict(self), f, indent=2)
+        except Exception as e:
+            print(f"Error writing to {self.file_path}: {e}")
+
+    def __setitem__(self, key, value):
+        """Store a key-value pair in the cache and persist the updated cache to disk."""
+        with self.lock:
+            super().__setitem__(key, value)
+            self._save()
+
+    def __delitem__(self, key):
+        """Remove an item from the PersistentCacheDict and update the persistent storage."""
+        with self.lock:
+            super().__delitem__(key)
+            self._save()
+
+    def update(self, *args, **kwargs):
+        """Update the dictionary with key-value pairs from other mappings or iterables, ensuring thread safety."""
+        with self.lock:
+            super().update(*args, **kwargs)
+            self._save()
+
+    def clear(self):
+        """Clears all entries from the persistent cache dictionary, ensuring thread safety."""
+        with self.lock:
+            super().clear()
+            self._save()
+
+
 def deprecation_warn(arg, new_arg):
     """Issue a deprecation warning when a deprecated argument is used, suggesting an updated argument."""
     LOGGER.warning(
@@ -1171,6 +1228,7 @@ def vscode_msg(ext="ultralytics.ultralytics-snippets") -> str:
 # Check first-install steps
 PREFIX = colorstr("Ultralytics: ")
 SETTINGS = SettingsManager()  # initialize settings
+PERSISTENT_CACHE = PersistentCacheDict()  # initialize persistent cache
 DATASETS_DIR = Path(SETTINGS["datasets_dir"])  # global datasets directory
 WEIGHTS_DIR = Path(SETTINGS["weights_dir"])  # global weights directory
 RUNS_DIR = Path(SETTINGS["runs_dir"])  # global runs directory
diff --git a/ultralytics/utils/torch_utils.py b/ultralytics/utils/torch_utils.py
index 758a4e11..75a90b76 100644
--- a/ultralytics/utils/torch_utils.py
+++ b/ultralytics/utils/torch_utils.py
@@ -110,13 +110,17 @@ def autocast(enabled: bool, device: str = "cuda"):
 
 def get_cpu_info():
     """Return a string with system CPU information, i.e. 'Apple M2'."""
-    with contextlib.suppress(Exception):
-        import cpuinfo  # pip install py-cpuinfo
+    from ultralytics.utils import PERSISTENT_CACHE  # avoid circular import error
 
-        k = "brand_raw", "hardware_raw", "arch_string_raw"  # keys sorted by preference (not all keys always available)
-        info = cpuinfo.get_cpu_info()  # info dict
-        string = info.get(k[0] if k[0] in info else k[1] if k[1] in info else k[2], "unknown")
-        return string.replace("(R)", "").replace("CPU ", "").replace("@ ", "")
+    if "cpu_info" not in PERSISTENT_CACHE:
+        with contextlib.suppress(Exception):
+            import cpuinfo  # pip install py-cpuinfo
+
+            k = "brand_raw", "hardware_raw", "arch_string_raw"  # keys sorted by preference
+            info = cpuinfo.get_cpu_info()  # info dict
+            string = info.get(k[0] if k[0] in info else k[1] if k[1] in info else k[2], "unknown")
+            PERSISTENT_CACHE["cpu_info"] = string.replace("(R)", "").replace("CPU ", "").replace("@ ", "")
+    return PERSISTENT_CACHE.get("cpu_info", "unknown")
 
     return "unknown"