Merge pull request #64 from deepghs/dev/idxcache

dev(narugo): use lru cache for index caches

Author: narugo1992

docs/source/api_doc/index/fetch.rst

@@ -61,3 +61,10 @@ hf_tar_file_download
 
 
 
+hf_tar_cache_reset
+----------------------------------------------
+
+.. autofunction:: hf_tar_cache_reset
+
+
+

docs/source/api_doc/index/local_fetch.rst

@@ -49,3 +49,10 @@ tar_file_download
 
 
 
+tar_cache_reset
+----------------------------------------------
+
+.. autofunction:: tar_cache_reset
+
+
+

hfutils/index/__init__.py

@@ -1,6 +1,8 @@
 from .fetch import hf_tar_list_files, hf_tar_file_download, hf_tar_get_index, hf_tar_file_exists, \
-    ArchiveStandaloneFileIncompleteDownload, ArchiveStandaloneFileHashNotMatch, hf_tar_file_size, hf_tar_file_info
-from .local_fetch import tar_get_index, tar_file_info, tar_file_download, tar_file_size, tar_file_exists, tar_list_files
+    ArchiveStandaloneFileIncompleteDownload, ArchiveStandaloneFileHashNotMatch, hf_tar_file_size, hf_tar_file_info, \
+    hf_tar_cache_reset
+from .local_fetch import tar_get_index, tar_file_info, tar_file_download, tar_file_size, tar_file_exists, \
+    tar_list_files, tar_cache_reset
 from .make import tar_create_index, hf_tar_create_index, tar_get_index_info, hf_tar_create_from_directory, \
     tar_create_index_for_directory
 from .validate import hf_tar_item_validate, hf_tar_validate

hfutils/index/fetch.py

@@ -4,6 +4,7 @@
 from collections import defaultdict
 from typing import Optional, Dict, Union, List
 
+from cachetools import LRUCache
 from huggingface_hub.file_download import http_get, hf_hub_url
 from huggingface_hub.utils import build_hf_headers
 from tqdm import tqdm
@@ -26,7 +27,7 @@ class ArchiveStandaloneFileHashNotMatch(Exception):
 
 
 _HF_TAR_IDX_LOCKS = defaultdict(threading.Lock)
-_HF_TAR_IDX_CACHE = {}
+_HF_TAR_IDX_CACHE = LRUCache(maxsize=192)
 
 
 def _hf_tar_get_cache_key(repo_id: str, archive_in_repo: str,
@@ -153,7 +154,7 @@ def hf_tar_get_index(repo_id: str, archive_in_repo: str,
             return idx_data
 
 
-_HF_TAR_IDX_PFILES_CACHE = {}
+_HF_TAR_IDX_PFILES_CACHE = LRUCache(maxsize=192)
 
 
 def _hf_tar_get_processed_files(repo_id: str, archive_in_repo: str,
@@ -718,3 +719,27 @@ def hf_tar_file_download(repo_id: str, archive_in_repo: str, file_in_archive: st
         if os.path.exists(local_file):
             os.remove(local_file)
         raise
+
+
+def hf_tar_cache_reset(maxsize: Optional[int] = None):
+    """
+    Reset the tar archive index caches and optionally resize them.
+
+    :param maxsize: New maximum size for the caches. If None, only clears the caches without resizing.
+    :type maxsize: Optional[int]
+
+    This function performs two operations:
+
+    1. Clears both the index cache and processed files cache
+    2. If maxsize is provided, recreates the caches with the new size
+
+    Example::
+        >>> hf_tar_cache_reset()  # Clear caches
+        >>> hf_tar_cache_reset(maxsize=256)  # Clear and resize caches
+    """
+    global _HF_TAR_IDX_CACHE, _HF_TAR_IDX_PFILES_CACHE
+    _HF_TAR_IDX_CACHE.clear()
+    _HF_TAR_IDX_PFILES_CACHE.clear()
+    if maxsize is not None and _HF_TAR_IDX_CACHE.maxsize != maxsize:
+        _HF_TAR_IDX_CACHE = LRUCache(maxsize=maxsize)
+        _HF_TAR_IDX_PFILES_CACHE = LRUCache(maxsize=maxsize)

hfutils/index/local_fetch.py

@@ -14,7 +14,9 @@
 import os
 from typing import Optional, List
 
-_TAR_IDX_CACHE = {}
+from cachetools import LRUCache
+
+_TAR_IDX_CACHE = LRUCache(maxsize=192)
 
 
 def _tar_get_cache_key(archive_file: str, idx_file: Optional[str] = None):
@@ -77,7 +79,7 @@ def tar_get_index(archive_file: str, idx_file: Optional[str] = None, no_cache: b
         return idx_data
 
 
-_TAR_IDX_PFILES_CACHE = {}
+_TAR_IDX_PFILES_CACHE = LRUCache(maxsize=192)
 
 
 def _tar_get_processed_files(archive_file: str, idx_file: Optional[str] = None, no_cache: bool = False):
@@ -335,3 +337,26 @@ def tar_file_download(archive_file: str, file_in_archive: str, local_file: str,
         if os.path.exists(local_file):
             os.remove(local_file)
         raise
+
+
+def tar_cache_reset(maxsize: Optional[int] = None):
+    """
+    Reset the tar index and processed files caches.
+
+    This function clears both the index cache and processed files cache.
+    Optionally, it can also resize the caches.
+
+    :param maxsize: Optional new maximum size for the caches.
+                   If provided, both caches will be recreated with this size.
+    :type maxsize: Optional[int]
+
+    :example:
+        >>> tar_cache_reset(maxsize=256)  # Reset and resize caches
+        >>> tar_cache_reset()  # Just clear the existing caches
+    """
+    global _TAR_IDX_CACHE, _TAR_IDX_PFILES_CACHE
+    _TAR_IDX_CACHE.clear()
+    _TAR_IDX_PFILES_CACHE.clear()
+    if maxsize is not None and _TAR_IDX_CACHE.maxsize != maxsize:
+        _TAR_IDX_CACHE = LRUCache(maxsize=maxsize)
+        _TAR_IDX_PFILES_CACHE = LRUCache(maxsize=maxsize)

requirements.txt

@@ -9,3 +9,4 @@ urlobject
 fsspec>=2024; python_version > '3.8'
 fsspec>=2024,<=2025.3.0; python_version <= '3.8'
 random_user_agent
+cachetools

test/index/test_fetch.py

@@ -1,11 +1,12 @@
 import os.path
+from unittest.mock import patch, MagicMock
 
 import pytest
 from hbutils.testing import isolated_directory
 from natsort import natsorted
 
 from hfutils.index import hf_tar_list_files, hf_tar_file_exists, hf_tar_file_download, hf_tar_file_info, \
-    hf_tar_file_size
+    hf_tar_file_size, hf_tar_cache_reset
 from test.testings import get_testfile, file_compare
 
 
@@ -201,3 +202,50 @@ def test_hf_tar_file_download_empty(self):
                 local_file='2946001.',
             )
             assert os.path.getsize('2946001.') == 0
+
+
+@pytest.fixture
+def mock_lru_cache():
+    """Fixture to mock the LRUCache class and global variables."""
+    mock_cache1 = MagicMock()
+    mock_cache2 = MagicMock()
+
+    with patch('hfutils.index.fetch._HF_TAR_IDX_CACHE', mock_cache1), \
+            patch('hfutils.index.fetch._HF_TAR_IDX_PFILES_CACHE', mock_cache2), \
+            patch('hfutils.index.fetch.LRUCache') as mock_lru:
+        yield mock_cache1, mock_cache2, mock_lru
+
+
+@pytest.mark.unittest
+class TestHfTarCacheReset:
+    def test_reset_without_maxsize(self, mock_lru_cache):
+        """Test resetting the cache without changing the maxsize."""
+        mock_cache1, mock_cache2, _ = mock_lru_cache
+
+        hf_tar_cache_reset()
+
+        mock_cache1.clear.assert_called_once()
+        mock_cache2.clear.assert_called_once()
+
+    def test_reset_with_same_maxsize(self, mock_lru_cache):
+        """Test resetting the cache with the same maxsize."""
+        mock_cache1, mock_cache2, mock_lru = mock_lru_cache
+        mock_cache1.maxsize = 100
+
+        hf_tar_cache_reset(maxsize=100)
+
+        mock_cache1.clear.assert_called_once()
+        mock_cache2.clear.assert_called_once()
+        mock_lru.assert_not_called()
+
+    def test_reset_with_different_maxsize(self, mock_lru_cache):
+        """Test resetting the cache with a different maxsize."""
+        mock_cache1, mock_cache2, mock_lru = mock_lru_cache
+        mock_cache1.maxsize = 100
+
+        hf_tar_cache_reset(maxsize=200)
+
+        mock_cache1.clear.assert_called_once()
+        mock_cache2.clear.assert_called_once()
+        assert mock_lru.call_count == 2
+        mock_lru.assert_called_with(maxsize=200)

test/index/test_local_fetch.py

@@ -1,11 +1,12 @@
 import os.path
+from unittest.mock import MagicMock, patch
 
 import pytest
 from hbutils.testing import isolated_directory
 from natsort import natsorted
 
 from hfutils.index import tar_list_files, tar_file_exists, tar_file_download, tar_file_info, \
-    tar_file_size
+    tar_file_size, tar_cache_reset
 from test.testings import get_testfile, file_compare
 
 
@@ -178,3 +179,50 @@ def test_tar_file_download_empty(self, local_narugo_test_cos5t_tars):
                 local_file='empty_file',
             )
             assert os.path.getsize('empty_file') == 0
+
+
+@pytest.fixture
+def mock_lru_cache():
+    """Fixture to mock the LRUCache class and global variables."""
+    mock_cache1 = MagicMock()
+    mock_cache2 = MagicMock()
+
+    with patch('hfutils.index.local_fetch._TAR_IDX_CACHE', mock_cache1), \
+            patch('hfutils.index.local_fetch._TAR_IDX_PFILES_CACHE', mock_cache2), \
+            patch('hfutils.index.local_fetch.LRUCache') as mock_lru:
+        yield mock_cache1, mock_cache2, mock_lru
+
+
+@pytest.mark.unittest
+class TestTarCacheReset:
+    def test_reset_without_maxsize(self, mock_lru_cache):
+        """Test resetting the cache without changing the maxsize."""
+        mock_cache1, mock_cache2, _ = mock_lru_cache
+
+        tar_cache_reset()
+
+        mock_cache1.clear.assert_called_once()
+        mock_cache2.clear.assert_called_once()
+
+    def test_reset_with_same_maxsize(self, mock_lru_cache):
+        """Test resetting the cache with the same maxsize."""
+        mock_cache1, mock_cache2, mock_lru = mock_lru_cache
+        mock_cache1.maxsize = 100
+
+        tar_cache_reset(maxsize=100)
+
+        mock_cache1.clear.assert_called_once()
+        mock_cache2.clear.assert_called_once()
+        mock_lru.assert_not_called()
+
+    def test_reset_with_different_maxsize(self, mock_lru_cache):
+        """Test resetting the cache with a different maxsize."""
+        mock_cache1, mock_cache2, mock_lru = mock_lru_cache
+        mock_cache1.maxsize = 100
+
+        tar_cache_reset(maxsize=200)
+
+        mock_cache1.clear.assert_called_once()
+        mock_cache2.clear.assert_called_once()
+        assert mock_lru.call_count == 2
+        mock_lru.assert_called_with(maxsize=200)