Add new settings to control the package payload caching disk usage (#2023)

Signed-off-by: Ibrahim Sani <[email protected]>

Author: sanikache

docs/source/caching.rst

@@ -180,6 +180,17 @@ you should **not** enable caching on package repositories where packages may get
 overwritten. It is for this reason that caching is disabled for local packages by
 default (see :data:`package_cache_local`).
 
+Control Disk Usage
+------------------
+
+You can control the disk usage of the package cache by using the
+:data:`package_cache_space_buffer` and :data:`package_cache_used_threshold` settings.
+Both have a default value already set to a safe value. You are encouraged to tweak
+them to your liking.
+
+Another way to control the disk usage is to run the :option:`rez-pkg-cache --clean` command
+either manually or as a cron job. See :ref:`caching-cleaning-the-cache`.
+
 Commandline Tool
 ----------------
 
@@ -232,6 +243,8 @@ Caching operations are stored into logfiles within the cache directory. To view:
    rez-pkg-cache 2020-05-23 16:17:45,404 PID-29827 INFO Started caching of variant /home/ajohns/packages/python/3.7.4/package.py[0]...
    rez-pkg-cache 2020-05-23 16:17:46,006 PID-29827 INFO Cached variant to /home/ajohns/package_cache/python/3.7.4/ce1c/a in 0.602037 seconds
 
+.. _caching-cleaning-the-cache:
+
 Cleaning The Cache
 ++++++++++++++++++
 

src/rez/cli/pkg-cache.py

@@ -67,6 +67,8 @@ def setup_parser(parser, completions=False):
 
 
 def add_variant(pkgcache, uri, opts):
+    import shutil
+
     from rez.config import config
     from rez.packages import get_variant_from_uri
     from rez.utils.logging_ import print_info, print_warning
@@ -95,6 +97,14 @@ def add_variant(pkgcache, uri, opts):
         print_info("Already exists: %s", destpath)
     elif status == PackageCache.VARIANT_COPYING:
         print_warning("Another process is currently copying to: %s", destpath)
+    elif status == PackageCache.VARIANT_SKIPPED:
+        free = shutil.disk_usage(config.cache_packages_path).free / 1024**2
+        buffer = config.package_cache_space_buffer / 1024**2
+        print_warning(
+            "Cache no longer accepting new variant due to size limit.\n"
+            f"Remaining cache free space: {free:.2f}MB is near "
+            f"configured buffer {buffer:.2f}MB (config.package_cache_space_buffer)."
+        )
     else:
         print_info("Successfully cached to: %s", destpath)
 

src/rez/config.py

@@ -449,6 +449,8 @@ def _parse_env_var(self, value):
     "shell_error_truncate_cap":                     Int,
     "package_cache_log_days":                       Int,
     "package_cache_max_variant_days":               Int,
+    "package_cache_space_buffer":                   Int,
+    "package_cache_used_threshold":                 Int,
     "package_cache_clean_limit":                    Float,
     "allow_unversioned_packages":                   Bool,
     "package_cache_during_build":                   Bool,

src/rez/package_cache.py

@@ -70,6 +70,7 @@ class PackageCache(object):
     VARIANT_COPY_STALLED = 4  #: Variant payload copy has stalled
     VARIANT_PENDING = 5  #: Variant is pending caching
     VARIANT_REMOVED = 6  #: Variant was deleted
+    VARIANT_SKIPPED = 7  # Variant was not cached due to cache size limit
 
     STATUS_DESCRIPTIONS = {
         VARIANT_NOT_FOUND: "was not found",
@@ -81,6 +82,7 @@ class PackageCache(object):
                               "for more information.",
         VARIANT_PENDING: "is pending caching",
         VARIANT_REMOVED: "was deleted",
+        VARIANT_SKIPPED: "is not being cached due to cache size limit"
     }
 
     _FILELOCK_TIMEOUT = 10
@@ -130,6 +132,80 @@ def get_cached_root(self, variant):
 
         return rootpath
 
+    def get_variant_size(self, free, variant_root):
+        """Get the size of the variant root.
+
+        Args:
+            free: The available free cache space.
+            variant_root: The rez resolved variant root.
+
+        Returns:
+            int: The size in bytes of the variant root (may exceed buffer slightly).
+        """
+        variant_size = 0
+        seen_inodes = set()
+        stack = [variant_root]
+
+        while stack:
+            current = stack.pop()
+            try:
+                with os.scandir(current) as ref:
+                    for entry in ref:
+                        try:
+                            # Since we are following symlinks, track visited inodes to make sure
+                            # we are not double counting files.
+                            st = entry.stat(follow_symlinks=True)
+                            inode = (st.st_dev, st.st_ino)
+                            if inode in seen_inodes:
+                                continue
+                            seen_inodes.add(inode)
+
+                            if stat.S_ISREG(st.st_mode):
+                                variant_size += st.st_size
+                                # Bail out early if variant size will overtake the buffer
+                                # set by config.package_cache_space_buffer.
+                                if (free - variant_size) < config.package_cache_space_buffer:
+                                    return variant_size
+                            elif stat.S_ISDIR(st.st_mode):
+                                stack.append(entry.path)
+                        except OSError:
+                            continue
+            except OSError:
+                continue
+
+        return variant_size
+
+    def cache_near_full(self):
+        """ Get the cache available space
+
+        Returns:
+            bool: True if available cache space is below buffer, otherwise False.
+        """
+        _, _, free = shutil.disk_usage(self.path)
+        return free < config.package_cache_space_buffer
+
+    def variant_meets_space_requirements(self, rez_variant_root):
+        """Check if the cache usage is above config.package_cache_used_threshold.
+        If it is, start throttling the cache by checking each variants size to make sure
+        it's not going to take the cache size below the minimum buffer we set.
+
+        Args:
+            variant_root: The rez resolved variant root.
+
+        Returns:
+            bool:
+                - True if the cache space used is below config.package_cache_used_threshold.
+                - False if (free - variant_size) < config.package_cache_space_buffer, otherwise False.
+        """
+        total, used, free = shutil.disk_usage(self.path)
+        used_percentage = (used / total) * 100 if total else 0.0
+
+        if used_percentage > config.package_cache_used_threshold:
+            variant_size = self.get_variant_size(free, rez_variant_root)
+            return (free - variant_size) > config.package_cache_space_buffer
+
+        return True
+
     def add_variant(self, variant, force=False, wait_for_copying=False, logger=None):
         """Copy a variant's payload into the cache.
 
@@ -168,7 +244,7 @@ def add_variant(self, variant, force=False, wait_for_copying=False, logger=None)
         Returns:
             tuple: 2-tuple:
             - str: Path to cached payload
-            - int: One of VARIANT_FOUND, VARIANT_CREATED, VARIANT_COPYING, VARIANT_COPY_STALLED
+            - int: One of VARIANT_FOUND, VARIANT_CREATED, VARIANT_COPYING, VARIANT_COPY_STALLED, VARIANT_SKIPPED
         """
         from rez.utils.base26 import get_next_base26
 
@@ -259,6 +335,11 @@ def add_variant(self, variant, force=False, wait_for_copying=False, logger=None)
                     )
             return (rootpath, status)
 
+        # Block adding new variant to cache from rez-pkg-cache --add-variants
+        # if the cache size is almost full.
+        if self.cache_near_full() or not self.variant_meets_space_requirements(variant_root):
+            return (rootpath, self.VARIANT_SKIPPED)
+
         # 1.
         path = self._get_hash_path(variant)
         os.makedirs(path, exist_ok=True)
@@ -767,6 +848,14 @@ def _lock(self):
     def _run_caching_step(self, state, wait_for_copying=False):
         logger = state["logger"]
 
+        # Keep the cache daemon alive until the cache size reaches its min threshold.
+        if self.cache_near_full():
+            logger.info(
+                "Cache storage has reached the configured threshold of "
+                f"{config.package_cache_space_buffer / 1024**2:.2f}MB, caching will now stop."
+            )
+            return False
+
         # pick a random pending variant to copy
         pending_filenames = set(os.listdir(self._pending_dir))
         if not wait_for_copying:
@@ -787,6 +876,13 @@ def _run_caching_step(self, state, wait_for_copying=False):
             raise
 
         variant = get_variant(variant_handle_dict)
+        variant_root = getattr(variant, "root")
+
+        if not self.variant_meets_space_requirements(variant_root):
+            # variant cannot be cached due to its size, so remove as a pending variant.
+            logger.info(f"Variant {variant_root} is too big to be cached due to remaining cache space.")
+            safe_remove(filepath)
+            return True
 
         # copy the variant and log activity
         logger.info("Started caching of variant %s...", variant.uri)

src/rez/rezconfig.py

@@ -302,6 +302,29 @@
 # Logs are written to :file:`{pkg-cache-root}/.sys/log/{filename}.log`
 package_cache_log_days = 7
 
+# Define a default minimum of 100MB of free space buffer for the cache in bytes.
+# This is required to avoid writing to a full cache and for cleaning the cache
+# when running :option:`rez-pkg-cache --clean`.
+# Note: Reported disk usage may vary across different file systems due to differences
+# in block size, allocation strategies and metadata overhead.
+# 100MB = 100 * 1024 * 1024 = 104857600.
+#
+# .. note::
+#    Reported disk usage may vary across different file systems due to differences
+#    in block size, allocation strategies and metadata overhead.
+package_cache_space_buffer = 104857600
+
+# The last variant being cached can take the cache size below the minimum buffer threshold we set.
+# To guard against this, we define a maximum cache usage threshold of 80%. We start throttling the cache
+# at this point by checking the size of each variant against the :data:`package_cache_space_buffer`.
+# If the pending variant about to be cached will take the cache size below the :data:`package_cache_space_buffer`,
+# don't cache it. When setting this value, subtract from your total disk space the fraction of disk space that
+# will be consumed by the largest variant you support and add the :data:`package_cache_space_buffer`.
+#
+# .. note::
+#    Reported disk usage may vary across different file systems due to differences
+#    in block size, allocation strategies and metadata overhead.
+package_cache_used_threshold = 80
 
 ###############################################################################
 # Package Resolution

src/rez/tests/test_package_cache.py

@@ -10,6 +10,8 @@
 import os.path
 import time
 import subprocess
+import tempfile
+from unittest.mock import patch
 
 from rez.tests.util import TestBase, TempdirMixin, restore_os_environ, \
     install_dependent
@@ -19,6 +21,9 @@
 from rez.exceptions import PackageCacheError
 from rez.utils.filesystem import canonical_path
 
+# Simulated total disk size (1 GiB) for disk space tests.
+VIRTUAL_GIGABYTE = 1024 * 1024 * 1024
+
 
 class TestPackageCache(TestBase, TempdirMixin):
     @classmethod
@@ -239,3 +244,99 @@ def test_caching_on_resolve_synchronous(self):
                 "Reference %r should resolve to %s, but resolves to %s"
                 % (ref, cached_root, root)
             )
+
+    @patch('rez.package_cache.shutil.disk_usage')
+    def test_cache_near_full_true(self, mock_du):
+        """cache_near_full returns True when free < buffer."""
+        mock_du.return_value = (VIRTUAL_GIGABYTE, 500_000_000, 10_000_000)  # 10MB free
+        pkgcache = self._pkgcache()
+        self.update_settings({'package_cache_space_buffer': 50_000_000})  # 50MB buffer
+        self.assertTrue(pkgcache.cache_near_full())
+
+    @patch('rez.package_cache.shutil.disk_usage')
+    def test_cache_near_full_false(self, mock_du):
+        """cache_near_full returns False when free >= buffer."""
+        mock_du.return_value = (VIRTUAL_GIGABYTE, 400_000_000, 300_000_000)  # 300MB free
+        pkgcache = self._pkgcache()
+        self.update_settings({'package_cache_space_buffer': 50_000_000})
+        self.assertFalse(pkgcache.cache_near_full())
+
+    @patch('rez.package_cache.shutil.disk_usage')
+    def test_variant_meets_space_requirements_below_threshold(self, mock_du):
+        """Below used threshold always True."""
+        mock_du.return_value = (VIRTUAL_GIGABYTE, 300_000_000, 700_000_000)  # 30% used
+        pkgcache = self._pkgcache()
+        self.update_settings({'package_cache_used_threshold': 80})
+        with tempfile.TemporaryDirectory() as d:
+            self.assertTrue(pkgcache.variant_meets_space_requirements(d))
+
+    @patch('rez.package_cache.shutil.disk_usage')
+    def test_variant_meets_space_requirements_above_threshold_sufficient(self, mock_du):
+        """Above threshold but variant fits inside buffer => True."""
+        mock_du.return_value = (VIRTUAL_GIGABYTE, 900_000_000, 150_000_000)  # ~83.8% used (>80%), 150MB free
+        pkgcache = self._pkgcache()
+        with tempfile.TemporaryDirectory() as d:
+            # create small file (1MB)
+            with open(os.path.join(d, 'f.bin'), 'wb') as f:
+                f.write(b'0' * 1_000_000)
+            self.update_settings(
+                {
+                    'package_cache_used_threshold': 80,
+                    'package_cache_space_buffer': 100_000_000,
+                }
+            )
+            self.assertTrue(pkgcache.variant_meets_space_requirements(d))  # 150MB - 1MB > 100MB
+
+    @patch('rez.package_cache.shutil.disk_usage')
+    def test_variant_meets_space_requirements_above_threshold_insufficient(self, mock_du):
+        """
+        Above threshold and variant would breach buffer => False.
+
+        Uses a mocked get_variant_size to avoid filesystem/storage variations
+        that caused the real size to appear smaller on some platforms (making
+        the predicate unexpectedly True). The logic we need to exercise is:
+            used% > threshold AND (free - variant_size) <= buffer -> False
+        """
+        mock_du.return_value = (VIRTUAL_GIGABYTE, 900_000_000, 150_000_000)  # ~83.8% used (>80%), 150MB free
+        pkgcache = self._pkgcache()
+        with tempfile.TemporaryDirectory() as d:
+            self.update_settings(
+                {
+                    'package_cache_used_threshold': 80,
+                    'package_cache_space_buffer': 100_000_000,
+                }
+            )
+            with patch.object(pkgcache, 'get_variant_size', return_value=70_000_000):  # 70MB
+                # free (150MB) - variant_size (70MB) = 80MB <= buffer (100MB) => False
+                self.assertFalse(pkgcache.variant_meets_space_requirements(d))  # expected False
+
+    def test_variant_meets_space_requirements_invalid_path(self):
+        """Invalid path returns True (early exit) per implementation."""
+        pkgcache = self._pkgcache()
+        self.assertTrue(pkgcache.variant_meets_space_requirements(None))
+        self.assertTrue(pkgcache.variant_meets_space_requirements(""))
+        self.assertTrue(
+            pkgcache.variant_meets_space_requirements(
+                os.path.join(self.root, "path", "does", "not", "exist")
+            )
+        )
+
+    def test_add_variant_skipped_cache_near_full(self):
+        """add_variant returns VARIANT_SKIPPED when cache_near_full True."""
+        pkgcache = self._pkgcache()
+        package = get_package("versioned", "3.0")
+        variant = next(package.iter_variants())
+        with patch.object(pkgcache, 'cache_near_full', return_value=True), \
+             patch.object(pkgcache, 'variant_meets_space_requirements', return_value=True):
+            _, status = pkgcache.add_variant(variant)
+            self.assertEqual(status, PackageCache.VARIANT_SKIPPED)
+
+    def test_add_variant_skipped_variant_too_large(self):
+        """add_variant returns VARIANT_SKIPPED when variant fails space requirements."""
+        pkgcache = self._pkgcache()
+        package = get_package("versioned", "3.0")
+        variant = next(package.iter_variants())
+        with patch.object(pkgcache, 'cache_near_full', return_value=False), \
+             patch.object(pkgcache, 'variant_meets_space_requirements', return_value=False):
+            _, status = pkgcache.add_variant(variant)
+            self.assertEqual(status, PackageCache.VARIANT_SKIPPED)