Add object store multi cache sharding

lib/galaxy/config/sample/object_store_conf.sample.yml

@@ -18,6 +18,20 @@
 #   # cache. By default (true) data is also copied to the cache.
 #   cache_updated_data: true
 #
+# Instead of a single `path`, the cache can be sharded across multiple directories
+# with weighted distribution. All cache entries for the same dataset (including
+# extra_files and metadata) resolve to the same shard deterministically:
+#
+# cache:
+#   dirs:
+#     - path: /fast/cache
+#       weight: 3
+#       size: 500
+#     - path: /slow/cache
+#       weight: 1
+#       size: 200
+#   cache_updated_data: true
+#
 # Most object store types have a `store_by` option which can be set to either `uuid` or `id`. Older Galaxy servers
 # stored datasets by their numeric id (000/dataset_1.dat, 00/dataset_2.dat, ...), whereas newer Galaxy servers store
 # them by UUID (b/5/e/dataset_b5e0301c-4c2e-41ac-b2c1-3c243f91b6ec.dat, ...). Storing by UUID is preferred, storing by

lib/galaxy/objectstore/azure_blob.py

@@ -19,7 +19,9 @@
 from galaxy.util import now
 from ._caching_base import CachingConcreteObjectStore
 from .caching import (
+    CacheShardManager,
     enable_cache_monitor,
+    ObjectId,
     parse_caching_config_dict_from_xml,
 )
 
@@ -139,9 +141,8 @@ def __init__(self, config, config_dict):
                 typed_transfer_dict[key] = int(value)
         self.transfer_dict = typed_transfer_dict
 
-        self.cache_size = cache_dict.get("size") or self.config.object_store_cache_size
-        self.staging_path = cache_dict.get("path") or self.config.object_store_cache_path
         self.cache_updated_data = cache_dict.get("cache_updated_data", True)
+        self._cache_shards = CacheShardManager.from_config(cache_dict, self.config)
 
         self._initialize()
 
@@ -168,11 +169,7 @@ def to_dict(self):
                     "name": self.container_name,
                 },
                 "transfer": self.transfer_dict,
-                "cache": {
-                    "size": self.cache_size,
-                    "path": self.staging_path,
-                    "cache_updated_data": self.cache_updated_data,
-                },
+                "cache": self._cache_config_to_dict(),
             }
         )
         return as_dict
@@ -239,11 +236,11 @@ def _exists_remotely(self, rel_path: str):
     def _blob_client(self, rel_path: str):
         return self.service.get_blob_client(self.container_name, rel_path)
 
-    def _download(self, rel_path):
-        local_destination = self._get_cache_path(rel_path)
+    def _download(self, rel_path, object_id: ObjectId):
+        local_destination = self._get_cache_path(rel_path, object_id)
         try:
             log.debug("Pulling '%s' into cache to %s", rel_path, local_destination)
-            if not self._caching_allowed(rel_path):
+            if not self._caching_allowed(rel_path, object_id=object_id):
                 return False
             else:
                 with self._atomic_download(local_destination) as tmp:

lib/galaxy/objectstore/caching.py

@@ -1,10 +1,12 @@
 """ """
 
+import hashlib
 import logging
 import os
 import threading
 import time
 from math import inf
+from uuid import UUID
 
 from typing_extensions import NamedTuple
 
@@ -22,10 +24,15 @@
 
 FileListT = list[tuple[time.struct_time, str, int]]
 
+#: The type of an object identifier used for cache shard selection.
+#: Can be a numeric id (``store_by="id"``) or a UUID / string
+#: (``store_by="uuid"``).
+ObjectId = int | UUID | str
+
 
 class CacheTarget(NamedTuple):
     path: str
-    size: int  # cache size in gigabytes
+    size: float  # cache size in gigabytes
     limit: float  # cache limit as a percent
 
     def fits_in_cache(self, bytes: int) -> bool:
@@ -43,6 +50,111 @@ def log_description(self) -> str:
         return f"{self.limit} percent of {self.size} gigabytes"
 
 
+class CacheShard(NamedTuple):
+    path: str
+    weight: int
+    size: float
+
+
+def parse_cache_dirs_from_xml(cache_element) -> list[dict] | None:
+    """Parse <dirs><dir .../></dirs> from a cache XML element.
+
+    Returns a list of dicts with keys ``path``, ``weight``, ``size``
+    (size may be ``None``), or ``None`` if no ``<dirs>`` element is present.
+    """
+    dirs_els = cache_element.findall("dirs")
+    if not dirs_els:
+        return None
+    dir_els = dirs_els[0].findall("dir")
+    if not dir_els:
+        return None
+    return [
+        {
+            "path": d.get("path"),
+            "weight": int(d.get("weight", 1)),
+            "size": float(d.get("size", -1)) if d.get("size") is not None else None,
+        }
+        for d in dir_els
+    ]
+
+
+class CacheShardManager:
+    def __init__(self, shards: list[CacheShard]):
+        if not shards:
+            raise ValueError("CacheShardManager requires at least one shard")
+        self.shards = shards
+        total_weight = sum(s.weight for s in shards)
+        self._weighted_index: list[tuple[int, CacheShard]] = []
+        cumulative = 0
+        for shard in shards:
+            cumulative += shard.weight
+            self._weighted_index.append((cumulative, shard))
+        self._total_weight = total_weight
+
+    def _select_shard(self, object_id: ObjectId) -> CacheShard:
+        digest = hashlib.sha256(str(object_id).encode()).digest()
+        hash_value = int.from_bytes(digest, "big")
+        point = hash_value % self._total_weight
+        for cumulative, shard in self._weighted_index:
+            if point < cumulative:
+                return shard
+        return self.shards[-1]
+
+    def get_cache_path(self, object_id: ObjectId, rel_path: str) -> str:
+        shard = self._select_shard(object_id)
+        return os.path.abspath(os.path.join(shard.path, rel_path))
+
+    def get_cache_target(self, object_id: ObjectId) -> CacheTarget:
+        shard = self._select_shard(object_id)
+        return CacheTarget(shard.path, shard.size, 0.9)
+
+    @property
+    def paths(self) -> list[str]:
+        return [s.path for s in self.shards]
+
+    @property
+    def cache_targets(self) -> list[CacheTarget]:
+        return [CacheTarget(s.path, s.size, 0.9) for s in self.shards]
+
+    def to_config_dict(self) -> dict:
+        """Serialize shard config for backend ``to_dict`` / reconstruction.
+
+        Emits ``dirs`` when there are multiple shards, otherwise the legacy
+        single ``path`` / ``size`` keys.  ``weight`` is intentionally omitted
+        for the single-shard case because it is meaningless with only one shard.
+        """
+        if len(self.shards) == 1:
+            s = self.shards[0]
+            return {"path": s.path, "size": s.size}
+        return {
+            "dirs": [{"path": s.path, "weight": s.weight, "size": s.size} for s in self.shards],
+        }
+
+    @classmethod
+    def from_config(cls, cache_dict: dict, config) -> "CacheShardManager":
+        default_size = cache_dict.get("size") or config.object_store_cache_size
+        default_path = cache_dict.get("path") or config.object_store_cache_path
+
+        dirs = cache_dict.get("dirs")
+        if dirs:
+            shards: list[CacheShard] = []
+            for d in dirs:
+                path = d.get("path")
+                if not path:
+                    continue
+                weight = d.get("weight", 1)
+                if weight <= 0:
+                    continue
+                size = d.get("size")
+                if size is None:
+                    size = default_size
+                shards.append(CacheShard(path=path, weight=weight, size=size))
+            if shards:
+                return cls(shards)
+
+        return cls([CacheShard(path=default_path, weight=1, size=default_size)])
+
+
 def check_caches(targets: list[CacheTarget]):
     for target in targets:
         check_cache(target)
@@ -135,20 +247,15 @@ def parse_caching_config_dict_from_xml(config_xml):
             "monitor": monitor,
             "cache_updated_data": cache_updated_data,
         }
+
+        dirs = parse_cache_dirs_from_xml(c_xml)
+        if dirs:
+            cache_dict["dirs"] = dirs
     else:
         cache_dict = {}
     return cache_dict
 
 
-def configured_cache_size(config, config_dict) -> int:
-    cache_config_dict = config_dict.get("cache") or {}
-    cache_size = cache_config_dict.get("size") or config.object_store_cache_size
-    if cache_size != -1:
-        # Convert admin-set GBs to bytes internally for quick comparison
-        cache_size = cache_size * ONE_GIGA_BYTE
-    return cache_size
-
-
 def enable_cache_monitor(config, config_dict) -> tuple[bool, int]:
     cache_config_dict = config_dict.get("cache") or {}
     default_interval = getattr(config, "object_store_cache_monitor_interval", 600)
@@ -170,15 +277,15 @@ def enable_cache_monitor(config, config_dict) -> tuple[bool, int]:
 
 
 class InProcessCacheMonitor:
-    def __init__(self, cache_target: CacheTarget, interval: int = 30, initial_sleep: int | None = 2):
+    def __init__(self, cache_targets: list[CacheTarget], interval: int = 30, initial_sleep: int | None = 2):
         # This Event object is initialized to False
         # It is set to True in shutdown(), causing
         # the cache monitor thread to return/terminate
         self.stop_cache_monitor_event = threading.Event()
         # Helper for interruptable sleep
         self.sleeper = Sleeper()
 
-        self.cache_target = cache_target
+        self.cache_targets = cache_targets
         self.interval = interval
         self.initial_sleep = initial_sleep
 
@@ -194,7 +301,7 @@ def _cache_monitor(self):
                 self.initial_sleep
             )  # startup sleep hack - probably originally implemented to prevent contention at app startup
         while not self.stop_cache_monitor_event.is_set():
-            check_cache(self.cache_target)
+            check_caches(self.cache_targets)
             self.sleeper.sleep(self.interval)
 
     def shutdown(self):

lib/galaxy/objectstore/cloud.py

@@ -8,7 +8,11 @@
 
 from ._caching_base import CachingConcreteObjectStore
 from ._util import UsesAxel
-from .caching import enable_cache_monitor
+from .caching import (
+    CacheShardManager,
+    enable_cache_monitor,
+    ObjectId,
+)
 from .s3 import parse_config_xml
 
 try:
@@ -51,9 +55,8 @@ def __init__(self, config, config_dict):
         self.use_rr = bucket_dict.get("use_reduced_redundancy", False)
         self.max_chunk_size = bucket_dict.get("max_chunk_size", 250)
 
-        self.cache_size = cache_dict.get("size") or self.config.object_store_cache_size
-        self.staging_path = cache_dict.get("path") or self.config.object_store_cache_path
         self.cache_updated_data = cache_dict.get("cache_updated_data", True)
+        self._cache_shards = CacheShardManager.from_config(cache_dict, self.config)
 
         self._initialize()
 
@@ -195,11 +198,7 @@ def _config_to_dict(self):
                 "name": self.bucket_name,
                 "use_reduced_redundancy": self.use_rr,
             },
-            "cache": {
-                "size": self.cache_size,
-                "path": self.staging_path,
-                "cache_updated_data": self.cache_updated_data,
-            },
+            "cache": self._cache_config_to_dict(),
         }
 
     def _get_bucket(self, bucket_name):
@@ -245,13 +244,13 @@ def _exists_remotely(self, rel_path):
             return False
         return exists
 
-    def _download(self, rel_path):
-        local_destination = self._get_cache_path(rel_path)
+    def _download(self, rel_path, object_id: ObjectId):
+        local_destination = self._get_cache_path(rel_path, object_id)
         try:
             log.debug("Pulling key '%s' into cache to %s", rel_path, local_destination)
             key = self.bucket.objects.get(rel_path)
             remote_size = key.size
-            if not self._caching_allowed(rel_path, remote_size):
+            if not self._caching_allowed(rel_path, object_id, remote_size=remote_size):
                 return False
             log.debug("Pulled key '%s' into cache to %s", rel_path, local_destination)
             with self._atomic_download(local_destination) as tmp: