Coverage for python/lsst/daf/butler/datastore/cache_manager.py: 90%
437 statements
« prev ^ index » next coverage.py v7.15.2, created at 2026-07-25 08:37 +0000
1# This file is part of daf_butler.
2#
3# Developed for the LSST Data Management System.
4# This product includes software developed by the LSST Project
5# (http://www.lsst.org).
6# See the COPYRIGHT file at the top-level directory of this distribution
7# for details of code ownership.
8#
9# This software is dual licensed under the GNU General Public License and also
10# under a 3-clause BSD license. Recipients may choose which of these licenses
11# to use; please see the files gpl-3.0.txt and/or bsd_license.txt,
12# respectively. If you choose the GPL option then the following text applies
13# (but note that there is still no warranty even if you opt for BSD instead):
14#
15# This program is free software: you can redistribute it and/or modify
16# it under the terms of the GNU General Public License as published by
17# the Free Software Foundation, either version 3 of the License, or
18# (at your option) any later version.
19#
20# This program is distributed in the hope that it will be useful,
21# but WITHOUT ANY WARRANTY; without even the implied warranty of
22# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
23# GNU General Public License for more details.
24#
25# You should have received a copy of the GNU General Public License
26# along with this program. If not, see <http://www.gnu.org/licenses/>.
28"""Cache management for a datastore."""
30from __future__ import annotations
32__all__ = (
33 "AbstractDatastoreCacheManager",
34 "DatastoreCacheManager",
35 "DatastoreCacheManagerConfig",
36 "DatastoreDisabledCacheManager",
37)
39import atexit
40import contextlib
41import datetime
42import itertools
43import logging
44import os
45import shutil
46import tempfile
47import uuid
48from abc import ABC, abstractmethod
49from collections import defaultdict
50from collections.abc import ItemsView, Iterable, Iterator, KeysView, ValuesView
51from random import Random
52from typing import TYPE_CHECKING, Self
54from pydantic import BaseModel, PrivateAttr
56from lsst.resources import ResourcePath
58from .._config import Config, ConfigSubset
59from .._config_support import processLookupConfigs
60from .._dataset_ref import DatasetId, DatasetRef
62if TYPE_CHECKING:
63 from .._config_support import LookupKey
64 from .._dataset_type import DatasetType
65 from .._storage_class import StorageClass
66 from ..dimensions import DimensionUniverse
68log = logging.getLogger(__name__)
71def remove_cache_directory(directory: str) -> None:
72 """Remove the specified directory and all its contents.
74 Parameters
75 ----------
76 directory : `str`
77 Directory to remove.
78 """
79 log.debug("Removing temporary cache directory %s", directory)
80 shutil.rmtree(directory, ignore_errors=True)
83def _construct_cache_path(root: ResourcePath, ref: DatasetRef, extension: str) -> ResourcePath:
84 """Construct the full path to use for this dataset in the cache.
86 Parameters
87 ----------
88 root : `lsst.resources.ResourcePath`
89 The root of the cache.
90 ref : `DatasetRef`
91 The dataset to look up in or write to the cache.
92 extension : `str`
93 File extension to use for this file. Should include the
94 leading "``.``".
96 Returns
97 -------
98 uri : `lsst.resources.ResourcePath`
99 URI to use for this dataset in the cache.
100 """
101 # Dataset type component is needed in the name if composite
102 # disassembly is happening since the ID is shared for all components.
103 component = ref.datasetType.component()
104 component = f"_{component}" if component else ""
105 return root.join(f"{ref.id}{component}{extension}")
108def _parse_cache_name(cached_location: str) -> tuple[uuid.UUID, str | None, str | None]:
109 """For a given cache name, return its component parts.
111 Changes to ``_construct_cache_path()`` should be reflected here.
113 Parameters
114 ----------
115 cached_location : `str`
116 The name of the file within the cache.
118 Returns
119 -------
120 id : `uuid.UUID`
121 The dataset ID.
122 component : `str` or `None`
123 The name of the component, if present.
124 extension: `str` or `None`
125 The file extension, if present.
126 """
127 # Assume first dot is the extension and so allow .fits.gz
128 root_ext = cached_location.split(".", maxsplit=1)
129 root = root_ext.pop(0)
130 ext = "." + root_ext.pop(0) if root_ext else None
132 parts = root.split("_")
133 try:
134 id_ = uuid.UUID(parts.pop(0))
135 except ValueError as e:
136 raise InvalidCacheFilenameError(f"Invalid or missing ID in cache file name: {cached_location}") from e
137 component = parts.pop(0) if parts else None
138 return id_, component, ext
141class CacheEntry(BaseModel):
142 """Represent an entry in the cache."""
144 name: str
145 """Name of the file."""
147 size: int
148 """Size of the file in bytes."""
150 ctime: datetime.datetime
151 """Creation time of the file."""
153 ref: DatasetId
154 """ID of this dataset."""
156 component: str | None = None
157 """Component for this disassembled composite (optional)."""
159 @classmethod
160 def from_file(cls, file: ResourcePath, root: ResourcePath) -> CacheEntry:
161 """Construct an object from a file name.
163 Parameters
164 ----------
165 file : `lsst.resources.ResourcePath`
166 Path to the file.
167 root : `lsst.resources.ResourcePath`
168 Cache root directory.
169 """
170 file_in_cache = file.relative_to(root)
171 if file_in_cache is None: 171 ↛ 172line 171 didn't jump to line 172 because the condition on line 171 was never true
172 raise ValueError(f"Supplied file {file} is not inside root {root}")
173 id_, component, _ = _parse_cache_name(file_in_cache)
175 stat = os.stat(file.ospath)
176 return cls.model_construct(
177 name=file_in_cache,
178 size=stat.st_size,
179 ref=id_,
180 component=component,
181 ctime=datetime.datetime.fromtimestamp(stat.st_ctime, datetime.UTC),
182 )
185class _MarkerEntry(CacheEntry):
186 pass
189class CacheRegistry(BaseModel):
190 """Collection of cache entries."""
192 _size: int = PrivateAttr(0)
193 """Size of the cache."""
195 _entries: dict[str, CacheEntry] = PrivateAttr({})
196 """Internal collection of cache entries."""
198 _ref_map: dict[DatasetId, list[str]] = PrivateAttr({})
199 """Mapping of DatasetID to corresponding keys in cache registry."""
201 @property
202 def cache_size(self) -> int:
203 return self._size
205 def __getitem__(self, key: str) -> CacheEntry:
206 return self._entries[key]
208 def __setitem__(self, key: str, entry: CacheEntry) -> None:
209 self._size += entry.size
210 self._entries[key] = entry
212 # Update the mapping from ref to path.
213 if entry.ref not in self._ref_map:
214 self._ref_map[entry.ref] = []
215 self._ref_map[entry.ref].append(key)
217 def __delitem__(self, key: str) -> None:
218 entry = self._entries.pop(key)
219 self._decrement(entry)
220 self._ref_map[entry.ref].remove(key)
222 def _decrement(self, entry: CacheEntry | None) -> None:
223 if entry: 223 ↛ exitline 223 didn't return from function '_decrement' because the condition on line 223 was always true
224 self._size -= entry.size
225 if self._size < 0: 225 ↛ 226line 225 didn't jump to line 226 because the condition on line 225 was never true
226 log.warning("Cache size has gone negative. Inconsistent cache records...")
227 self._size = 0
229 def __contains__(self, key: str) -> bool:
230 return key in self._entries
232 def __len__(self) -> int:
233 return len(self._entries)
235 def __iter__(self) -> Iterator[str]: # type: ignore
236 return iter(self._entries)
238 def keys(self) -> KeysView[str]:
239 return self._entries.keys()
241 def values(self) -> ValuesView[CacheEntry]:
242 return self._entries.values()
244 def items(self) -> ItemsView[str, CacheEntry]:
245 return self._entries.items()
247 # An private marker to indicate that pop() should raise if no default
248 # is given.
249 __marker = _MarkerEntry(
250 name="marker",
251 size=0,
252 ref=uuid.UUID("{00000000-0000-0000-0000-000000000000}"),
253 ctime=datetime.datetime.fromtimestamp(0, datetime.UTC),
254 )
256 def pop(self, key: str, default: CacheEntry | None = __marker) -> CacheEntry | None:
257 # The marker for dict.pop is not the same as our marker.
258 if default is self.__marker: 258 ↛ 259line 258 didn't jump to line 259 because the condition on line 258 was never true
259 entry = self._entries.pop(key)
260 else:
261 entry = self._entries.pop(key, self.__marker)
262 # Should not attempt to correct for this entry being removed
263 # if we got the default value.
264 if entry is self.__marker: 264 ↛ 265line 264 didn't jump to line 265 because the condition on line 264 was never true
265 return default
267 self._decrement(entry)
268 # The default entry given to this method may not even be in the cache.
269 if entry and entry.ref in self._ref_map: 269 ↛ 273line 269 didn't jump to line 273 because the condition on line 269 was always true
270 keys = self._ref_map[entry.ref]
271 if key in keys: 271 ↛ 273line 271 didn't jump to line 273 because the condition on line 271 was always true
272 keys.remove(key)
273 return entry
275 def get_dataset_keys(self, dataset_id: DatasetId | None) -> list[str] | None:
276 """Retrieve all keys associated with the given dataset ID.
278 Parameters
279 ----------
280 dataset_id : `DatasetId` or `None`
281 The dataset ID to look up. Returns `None` if the ID is `None`.
283 Returns
284 -------
285 keys : `list` [`str`]
286 Keys associated with this dataset. These keys can be used to lookup
287 the cache entry information in the `CacheRegistry`. Returns
288 `None` if the dataset is not known to the cache.
289 """
290 if dataset_id not in self._ref_map:
291 return None
292 keys = self._ref_map[dataset_id]
293 if not keys:
294 return None
295 return keys
298class DatastoreCacheManagerConfig(ConfigSubset):
299 """Configuration information for `DatastoreCacheManager`."""
301 component = "cached"
302 requiredKeys = ("cacheable",)
305class AbstractDatastoreCacheManager(ABC):
306 """An abstract base class for managing caching in a Datastore.
308 Parameters
309 ----------
310 config : `str` or `DatastoreCacheManagerConfig`
311 Configuration to control caching.
312 universe : `DimensionUniverse`
313 Set of all known dimensions, used to expand and validate any used
314 in lookup keys.
315 """
317 @property
318 def cache_size(self) -> int:
319 """Size of the cache in bytes."""
320 return 0
322 @property
323 def file_count(self) -> int:
324 """Return number of cached files tracked by registry."""
325 return 0
327 @property
328 def download_directory(self) -> ResourcePath | None:
329 """Local directory into which a remote file destined for the cache
330 may be downloaded (`lsst.resources.ResourcePath` or `None`).
332 Notes
333 -----
334 Downloading a soon-to-be-cached remote file into this directory
335 guarantees that the file ends up on the same file system as the cache,
336 so that moving it into the cache is an atomic rename rather than a
337 cross-file-system copy. The base class implementation returns `None`
338 to indicate that no such directory is available and downloads should
339 use the default temporary location.
340 """
341 return None
343 def __init__(self, config: str | DatastoreCacheManagerConfig, universe: DimensionUniverse):
344 if not isinstance(config, DatastoreCacheManagerConfig):
345 config = DatastoreCacheManagerConfig(config)
346 assert isinstance(config, DatastoreCacheManagerConfig)
347 self.config = config
349 @abstractmethod
350 def should_be_cached(self, entity: DatasetRef | DatasetType | StorageClass) -> bool:
351 """Indicate whether the entity should be added to the cache.
353 This is relevant when reading or writing.
355 Parameters
356 ----------
357 entity : `StorageClass` or `DatasetType` or `DatasetRef`
358 Thing to test against the configuration. The ``name`` property
359 is used to determine a match. A `DatasetType` will first check
360 its name, before checking its `StorageClass`. If there are no
361 matches the default will be returned.
363 Returns
364 -------
365 should_cache : `bool`
366 Returns `True` if the dataset should be cached; `False` otherwise.
367 """
368 raise NotImplementedError()
370 @abstractmethod
371 def known_to_cache(self, ref: DatasetRef, extension: str | None = None) -> bool:
372 """Report if the dataset is known to the cache.
374 Parameters
375 ----------
376 ref : `DatasetRef`
377 Dataset to check for in the cache.
378 extension : `str`, optional
379 File extension expected. Should include the leading "``.``".
380 If `None` the extension is ignored and the dataset ID alone is
381 used to check in the cache. The extension must be defined if
382 a specific component is being checked.
384 Returns
385 -------
386 known : `bool`
387 Returns `True` if the dataset is currently known to the cache
388 and `False` otherwise.
390 Notes
391 -----
392 This method can only report if the dataset is known to the cache
393 in this specific instant and does not indicate whether the file
394 can be read from the cache later. `find_in_cache()` should be called
395 if the cached file is to be used.
396 """
397 raise NotImplementedError()
399 @abstractmethod
400 def move_to_cache(self, uri: ResourcePath, ref: DatasetRef) -> ResourcePath | None:
401 """Move a file to the cache.
403 Move the given file into the cache, using the supplied DatasetRef
404 for naming. A call is made to `should_be_cached()` and if the
405 DatasetRef should not be accepted `None` will be returned.
407 Cache expiry can occur during this.
409 Parameters
410 ----------
411 uri : `lsst.resources.ResourcePath`
412 Location of the file to be relocated to the cache. Will be moved.
413 ref : `DatasetRef`
414 Ref associated with this file. Will be used to determine the name
415 of the file within the cache.
417 Returns
418 -------
419 new : `lsst.resources.ResourcePath` or `None`
420 URI to the file within the cache, or `None` if the dataset
421 was not accepted by the cache.
422 """
423 raise NotImplementedError()
425 @abstractmethod
426 @contextlib.contextmanager
427 def find_in_cache(self, ref: DatasetRef, extension: str) -> Iterator[ResourcePath | None]:
428 """Look for a dataset in the cache and return its location.
430 Parameters
431 ----------
432 ref : `DatasetRef`
433 Dataset to locate in the cache.
434 extension : `str`
435 File extension expected. Should include the leading "``.``".
437 Yields
438 ------
439 uri : `lsst.resources.ResourcePath` or `None`
440 The URI to the cached file, or `None` if the file has not been
441 cached.
443 Notes
444 -----
445 Should be used as a context manager in order to prevent this
446 file from being removed from the cache for that context.
447 """
448 raise NotImplementedError()
450 @abstractmethod
451 def remove_from_cache(self, ref: DatasetRef | Iterable[DatasetRef]) -> None:
452 """Remove the specified datasets from the cache.
454 It is not an error for these datasets to be missing from the cache.
456 Parameters
457 ----------
458 ref : `DatasetRef` or iterable of `DatasetRef`
459 The datasets to remove from the cache.
460 """
461 raise NotImplementedError()
463 @abstractmethod
464 def __str__(self) -> str:
465 raise NotImplementedError()
468class DatastoreCacheManager(AbstractDatastoreCacheManager):
469 """A class for managing caching in a Datastore using local files.
471 Parameters
472 ----------
473 config : `str` or `DatastoreCacheManagerConfig`
474 Configuration to control caching.
475 universe : `DimensionUniverse`
476 Set of all known dimensions, used to expand and validate any used
477 in lookup keys.
479 Notes
480 -----
481 Two environment variables can be used to override the cache directory
482 and expiration configuration:
484 * ``$DAF_BUTLER_CACHE_DIRECTORY``
485 * ``$DAF_BUTLER_CACHE_EXPIRATION_MODE``
487 The expiration mode should take the form ``mode=threshold`` so for
488 example to configure expiration to limit the cache directory to 5 datasets
489 the value would be ``datasets=5``.
491 Additionally the ``$DAF_BUTLER_CACHE_DIRECTORY_IF_UNSET`` environment
492 variable can be used to indicate that this directory should be used
493 if no explicit directory has been specified from configuration or from
494 the ``$DAF_BUTLER_CACHE_DIRECTORY`` environment variable.
495 """
497 _temp_exemption_prefix = "exempt/"
498 _tmpdir_prefix = "butler-cache-dir-"
500 def __init__(self, config: str | DatastoreCacheManagerConfig, universe: DimensionUniverse):
501 super().__init__(config, universe)
503 # Expiration mode. Read from config but allow override from
504 # the environment.
505 expiration_mode = self.config.get(("expiry", "mode"))
506 threshold = self.config.get(("expiry", "threshold"))
508 external_mode = os.environ.get("DAF_BUTLER_CACHE_EXPIRATION_MODE")
509 if external_mode:
510 if external_mode == "disabled":
511 expiration_mode = external_mode
512 threshold = 0
513 elif "=" in external_mode:
514 expiration_mode, expiration_threshold = external_mode.split("=", 1)
515 threshold = int(expiration_threshold)
516 else:
517 log.warning(
518 "Unrecognized form (%s) for DAF_BUTLER_CACHE_EXPIRATION_MODE environment variable. "
519 "Ignoring.",
520 external_mode,
521 )
522 if expiration_mode is None:
523 # Force to None to avoid confusion.
524 threshold = None
526 allowed = ("disabled", "datasets", "age", "size", "files")
527 if expiration_mode and expiration_mode not in allowed:
528 raise ValueError(
529 f"Unrecognized value for cache expiration mode. Got {expiration_mode} but expected "
530 + ",".join(allowed)
531 )
533 self._expiration_mode: str | None = expiration_mode
534 self._expiration_threshold: int | None = threshold
535 if self._expiration_threshold is None and self._expiration_mode is not None:
536 raise ValueError(
537 f"Cache expiration threshold must be set for expiration mode {self._expiration_mode}"
538 )
540 # Files in cache, indexed by path within the cache directory.
541 self._cache_entries = CacheRegistry()
543 # No need for additional configuration if the cache has been disabled.
544 if self._expiration_mode == "disabled":
545 log.debug("Cache configured in disabled state.")
546 self._cache_directory: ResourcePath | None = ResourcePath(
547 "datastore_cache_disabled", forceAbsolute=False, forceDirectory=True
548 )
549 return
551 # Set cache directory if it pre-exists, else defer creation until
552 # requested. Allow external override from environment.
553 root = os.environ.get("DAF_BUTLER_CACHE_DIRECTORY") or self.config.get("root")
555 # Allow the execution environment to override the default values
556 # so long as no default value has been set from the line above.
557 if root is None:
558 root = os.environ.get("DAF_BUTLER_CACHE_DIRECTORY_IF_UNSET")
560 self._cache_directory = (
561 ResourcePath(root, forceAbsolute=True, forceDirectory=True) if root is not None else None
562 )
564 if self._cache_directory:
565 if not self._cache_directory.isLocal: 565 ↛ 566line 565 didn't jump to line 566 because the condition on line 565 was never true
566 raise ValueError(
567 f"Cache directory must be on a local file system. Got: {self._cache_directory}"
568 )
569 # Ensure that the cache directory is created. We assume that
570 # someone specifying a permanent cache directory will be expecting
571 # it to always be there. This will also trigger an error
572 # early rather than waiting until the cache is needed.
573 self._cache_directory.mkdir()
575 # Calculate the caching lookup table.
576 self._lut = processLookupConfigs(self.config["cacheable"], universe=universe)
578 # Default decision to for whether a dataset should be cached.
579 self._caching_default = self.config.get("default", False)
581 log.debug(
582 "Cache configuration:\n- root: %s\n- expiration mode: %s",
583 self._cache_directory if self._cache_directory else "tmpdir",
584 f"{self._expiration_mode}={self._expiration_threshold}" if self._expiration_mode else "disabled",
585 )
587 @property
588 def cache_directory(self) -> ResourcePath:
589 if self._cache_directory is None:
590 # Create on demand. Allow the override environment variable
591 # to be used in case it got set after this object was created
592 # but before a cache was used.
593 if cache_dir := os.environ.get("DAF_BUTLER_CACHE_DIRECTORY_IF_UNSET"):
594 # Someone else will clean this up.
595 isTemporary = False
596 msg = "deferred fallback"
597 else:
598 cache_dir = tempfile.mkdtemp(prefix=self._tmpdir_prefix)
599 isTemporary = True
600 msg = "temporary"
602 self._cache_directory = ResourcePath(cache_dir, forceDirectory=True, isTemporary=isTemporary)
603 log.debug("Using %s cache directory at %s", msg, self._cache_directory)
605 # Remove when we no longer need it.
606 if isTemporary:
607 atexit.register(remove_cache_directory, self._cache_directory.ospath)
608 return self._cache_directory
610 @property
611 def _temp_exempt_directory(self) -> ResourcePath:
612 """Return the directory in which to store temporary cache files that
613 should not be expired.
614 """
615 return self.cache_directory.join(self._temp_exemption_prefix)
617 @property
618 def download_directory(self) -> ResourcePath | None:
619 # Docstring inherited.
620 # Stage downloads in the exemption area: it lives on the same file
621 # system as the cache (so the move into the cache is a rename) and is
622 # ignored by the cache scanner, so a partially-downloaded file is never
623 # mistaken for a cache entry.
624 download_dir = self._temp_exempt_directory
625 download_dir.mkdir()
626 return download_dir
628 @property
629 def cache_size(self) -> int:
630 return self._cache_entries.cache_size
632 @property
633 def file_count(self) -> int:
634 return len(self._cache_entries)
636 @classmethod
637 def create_disabled(cls, universe: DimensionUniverse) -> Self:
638 """Create an instance that is disabled by default but can be
639 overridden by the environment.
641 Parameters
642 ----------
643 universe : `DimensionUniverse`
644 The universe to use if the datastore becomes enabled via the
645 environment.
647 Returns
648 -------
649 cache_manager : `DatastoreCacheManager`
650 A new cache manager, that is disabled by default but might be
651 enabled if environment variables are set.
652 """
653 # It is not sufficient to set the mode to disabled, there has to be
654 # enough configuration for the caching to work when enabled. This
655 # means setting the default to true (cache everything), inheriting
656 # inherit the FileDatastore default config (which works for Rubin but
657 # doesn't allow non-Rubin deployments to cache anything, but can in
658 # theory be overridden), or allowing a parameter to be passed in here
659 # defining the cacheable section of the config. For now cache
660 # everything. Supporting a JSON environment variable defining the
661 # cacheable storage classes is also a possibility.
662 config_str = """
663cached:
664 default: true
665 cacheable:
666 irrelevant: false
667 expiry:
668 mode: disabled
669 threshold: 0
670 """
671 config = Config.fromYaml(config_str)
672 return cls(DatastoreCacheManagerConfig(config), universe)
674 @classmethod
675 def set_fallback_cache_directory_if_unset(cls) -> tuple[bool, str]:
676 """Define a fallback cache directory if a fallback not set already.
678 Returns
679 -------
680 defined : `bool`
681 `True` if the fallback directory was newly-defined in this method.
682 `False` if it had already been set.
683 cache_dir : `str`
684 Returns the path to the cache directory that will be used if it's
685 needed. This can allow the caller to run a directory cleanup
686 when it's no longer needed (something that the cache manager
687 can not do because forks should not clean up directories defined
688 by the parent process).
690 Notes
691 -----
692 The fallback directory will not be defined if one has already been
693 defined. This method sets the ``DAF_BUTLER_CACHE_DIRECTORY_IF_UNSET``
694 environment variable only if a value has not previously been stored
695 in that environment variable. Setting the environment variable allows
696 this value to survive into spawned subprocesses. Calling this method
697 will lead to all subsequently created cache managers sharing the same
698 cache.
699 """
700 if cache_dir := os.environ.get("DAF_BUTLER_CACHE_DIRECTORY_IF_UNSET"): 700 ↛ 702line 700 didn't jump to line 702 because the condition on line 700 was never true
701 # A value has already been set.
702 return (False, cache_dir)
704 # As a class method, we do not know at this point whether a cache
705 # directory will be needed so it would be impolite to create a
706 # directory that will never be used.
708 # Construct our own temp name -- 16 characters should have a fairly
709 # low chance of clashing when combined with the process ID.
710 characters = "abcdefghijklmnopqrstuvwxyz0123456789_"
711 rng = Random()
712 tempchars = "".join(rng.choice(characters) for _ in range(16))
714 tempname = f"{cls._tmpdir_prefix}{os.getpid()}-{tempchars}"
716 cache_dir = os.path.join(tempfile.gettempdir(), tempname)
717 os.environ["DAF_BUTLER_CACHE_DIRECTORY_IF_UNSET"] = cache_dir
718 return (True, cache_dir)
720 def should_be_cached(self, entity: DatasetRef | DatasetType | StorageClass) -> bool:
721 # Docstring inherited
722 if self._expiration_mode == "disabled":
723 return False
725 matchName: LookupKey | str = f"{entity} (via default)"
726 should_cache = self._caching_default
728 for key in entity._lookupNames():
729 if key in self._lut:
730 should_cache = bool(self._lut[key])
731 matchName = key
732 break
734 if not isinstance(should_cache, bool): 734 ↛ 735line 734 didn't jump to line 735 because the condition on line 734 was never true
735 raise TypeError(
736 f"Got cache value {should_cache!r} for config entry {matchName!r}; expected bool."
737 )
739 log.debug("%s (match: %s) should%s be cached", entity, matchName, "" if should_cache else " not")
740 return should_cache
742 def _construct_cache_name(self, ref: DatasetRef, extension: str) -> ResourcePath:
743 """Construct the name to use for this dataset in the cache.
745 Parameters
746 ----------
747 ref : `DatasetRef`
748 The dataset to look up in or write to the cache.
749 extension : `str`
750 File extension to use for this file. Should include the
751 leading "``.``".
753 Returns
754 -------
755 uri : `lsst.resources.ResourcePath`
756 URI to use for this dataset in the cache.
757 """
758 return _construct_cache_path(self.cache_directory, ref, extension)
760 def move_to_cache(self, uri: ResourcePath, ref: DatasetRef) -> ResourcePath | None:
761 # Docstring inherited
762 if self._expiration_mode == "disabled":
763 return None
764 if not self.should_be_cached(ref):
765 return None
767 # Write the file using the id of the dataset ref and the file
768 # extension.
769 cached_location = self._construct_cache_name(ref, uri.getExtension())
771 # Run cache expiry to ensure that we have room for this
772 # item.
773 self._expire_cache()
775 # The above reset the in-memory cache status. It's entirely possible
776 # that another process has just cached this file (if multiple
777 # processes are caching on read), so check our in-memory cache
778 # before attempting to cache the dataset.
779 path_in_cache = cached_location.relative_to(self.cache_directory)
780 if path_in_cache and path_in_cache in self._cache_entries: 780 ↛ 781line 780 didn't jump to line 781 because the condition on line 780 was never true
781 return cached_location
783 # Move into the cache. Given that multiple processes might be
784 # sharing a single cache directory, and the file we need might have
785 # been copied in whilst we were checking, allow overwrite without
786 # complaint. Even for a private cache directory it is possible that
787 # a second butler in a subprocess could be writing to it.
788 cached_location.transfer_from(uri, transfer="move", overwrite=True)
789 log.debug("Cached dataset %s to %s", ref, cached_location)
791 self._register_cache_entry(cached_location)
793 return cached_location
795 @contextlib.contextmanager
796 def find_in_cache(self, ref: DatasetRef, extension: str) -> Iterator[ResourcePath | None]:
797 # Docstring inherited
798 if self._expiration_mode == "disabled":
799 yield None
800 return
802 # Short circuit this if the cache directory has not been created yet.
803 if self._cache_directory is None:
804 yield None
805 return
807 cached_location = self._construct_cache_name(ref, extension)
808 if cached_location.exists():
809 log.debug("Found cached file %s for dataset %s.", cached_location, ref)
811 # The cached file could be removed by another process doing
812 # cache expiration so we need to protect against that by making
813 # a copy in a different tree. Use hardlinks to ensure that
814 # we either have the cached file or we don't. This is robust
815 # against race conditions that can be caused by using soft links
816 # and the other end of the link being deleted just after it
817 # is created.
818 path_in_cache = cached_location.relative_to(self.cache_directory)
819 assert path_in_cache is not None, f"Somehow {cached_location} not in cache directory"
821 # Need to use a unique file name for the temporary location to
822 # ensure that two different processes can read the file
823 # simultaneously without one of them deleting it when it's in
824 # use elsewhere. Retain the original filename for easier debugging.
825 random = str(uuid.uuid4())[:8]
826 basename = cached_location.basename()
827 filename = f"{random}-{basename}"
829 temp_location: ResourcePath | None = self._temp_exempt_directory.join(filename)
830 try:
831 if temp_location is not None: 831 ↛ 840line 831 didn't jump to line 840 because the condition on line 831 was always true
832 temp_location.transfer_from(cached_location, transfer="hardlink")
833 except Exception as e:
834 log.debug("Detected error creating hardlink for dataset %s: %s", ref, e)
835 # Any failure will be treated as if the file was not
836 # in the cache. Yielding the original cache location
837 # is too dangerous.
838 temp_location = None
840 try:
841 log.debug("Yielding temporary cache location %s for dataset %s", temp_location, ref)
842 yield temp_location
843 finally:
844 try:
845 if temp_location: 845 ↛ 849line 845 didn't jump to line 849 because the condition on line 845 was always true
846 temp_location.remove()
847 except FileNotFoundError:
848 pass
849 return
851 log.debug("Dataset %s not found in cache.", ref)
852 yield None
853 return
855 def remove_from_cache(self, refs: DatasetRef | Iterable[DatasetRef]) -> None:
856 # Docstring inherited.
858 # Stop early if there are no cache entries anyhow.
859 if len(self._cache_entries) == 0:
860 return
862 if isinstance(refs, DatasetRef):
863 refs = [refs]
865 # Create a set of all the IDs
866 all_ids = {ref.id for ref in refs}
868 keys_to_remove = []
869 for key, entry in self._cache_entries.items():
870 if entry.ref in all_ids:
871 keys_to_remove.append(key)
872 self._remove_from_cache(keys_to_remove)
874 def _register_cache_entry(self, cached_location: ResourcePath, can_exist: bool = False) -> str | None:
875 """Record the file in the cache registry.
877 Parameters
878 ----------
879 cached_location : `lsst.resources.ResourcePath`
880 Location of the file to be registered.
881 can_exist : `bool`, optional
882 If `True` the item being registered can already be listed.
883 This can allow a cache refresh to run without checking the
884 file again. If `False` it is an error for the registry to
885 already know about this file.
887 Returns
888 -------
889 cache_key : `str` or `None`
890 The key used in the registry for this file. `None` if the file
891 no longer exists (it could have been expired by another process).
892 """
893 path_in_cache = cached_location.relative_to(self.cache_directory)
894 if path_in_cache is None: 894 ↛ 895line 894 didn't jump to line 895 because the condition on line 894 was never true
895 raise ValueError(
896 f"Can not register cached file {cached_location} that is not within"
897 f" the cache directory at {self.cache_directory}."
898 )
899 if path_in_cache in self._cache_entries:
900 if can_exist: 900 ↛ 903line 900 didn't jump to line 903 because the condition on line 900 was always true
901 return path_in_cache
902 else:
903 raise ValueError(
904 f"Cached file {cached_location} is already known to the registry"
905 " but this was expected to be a new file."
906 )
907 try:
908 details = CacheEntry.from_file(cached_location, root=self.cache_directory)
909 except FileNotFoundError:
910 return None
911 self._cache_entries[path_in_cache] = details
912 return path_in_cache
914 def scan_cache(self) -> None:
915 """Scan the cache directory and record information about files."""
916 if self._expiration_mode == "disabled": 916 ↛ 917line 916 didn't jump to line 917 because the condition on line 916 was never true
917 return
919 found = set()
920 for file in ResourcePath.findFileResources([self.cache_directory]):
921 assert isinstance(file, ResourcePath), "Unexpectedly did not get ResourcePath from iterator"
923 # Skip any that are found in an exempt part of the hierarchy
924 # since they should not be part of the registry.
925 if file.relative_to(self._temp_exempt_directory) is not None:
926 continue
928 try:
929 path_in_cache = self._register_cache_entry(file, can_exist=True)
930 if path_in_cache: 930 ↛ 920line 930 didn't jump to line 920 because the condition on line 930 was always true
931 found.add(path_in_cache)
932 except InvalidCacheFilenameError:
933 # Skip files that are in the cache directory, but were not
934 # created by us.
935 pass
937 # Find any files that were recorded in the cache but are no longer
938 # on disk. (something else cleared them out?)
939 known_to_cache = set(self._cache_entries)
940 missing = known_to_cache - found
942 if missing: 942 ↛ 943line 942 didn't jump to line 943 because the condition on line 942 was never true
943 log.debug(
944 "Entries no longer on disk but thought to be in cache and so removed: %s", ",".join(missing)
945 )
946 for path_in_cache in missing:
947 self._cache_entries.pop(path_in_cache, None)
949 def known_to_cache(self, ref: DatasetRef, extension: str | None = None) -> bool:
950 """Report if the dataset is known to the cache.
952 Parameters
953 ----------
954 ref : `DatasetRef`
955 Dataset to check for in the cache.
956 extension : `str`, optional
957 File extension expected. Should include the leading "``.``".
958 If `None` the extension is ignored and the dataset ID alone is
959 used to check in the cache. The extension must be defined if
960 a specific component is being checked.
962 Returns
963 -------
964 known : `bool`
965 Returns `True` if the dataset is currently known to the cache
966 and `False` otherwise. If the dataset refers to a component and
967 an extension is given then only that component is checked.
969 Notes
970 -----
971 This method can only report if the dataset is known to the cache
972 in this specific instant and does not indicate whether the file
973 can be read from the cache later. `find_in_cache()` should be called
974 if the cached file is to be used.
976 This method does not force the cache to be re-scanned and so can miss
977 cached datasets that have recently been written by other processes.
978 """
979 if self._expiration_mode == "disabled":
980 return False
981 if self._cache_directory is None:
982 return False
983 if self.file_count == 0:
984 return False
986 if extension is None:
987 # Look solely for matching dataset ref ID and not specific
988 # components.
989 cached_paths = self._cache_entries.get_dataset_keys(ref.id)
990 return bool(cached_paths)
992 else:
993 # Extension is known so we can do an explicit look up for the
994 # cache entry.
995 cached_location = self._construct_cache_name(ref, extension)
996 path_in_cache = cached_location.relative_to(self.cache_directory)
997 assert path_in_cache is not None # For mypy
998 return path_in_cache in self._cache_entries
1000 def _remove_from_cache(self, cache_entries: Iterable[str]) -> None:
1001 """Remove the specified cache entries from cache.
1003 Parameters
1004 ----------
1005 cache_entries : `~collections.abc.Iterable` of `str`
1006 The entries to remove from the cache. The values are the path
1007 within the cache.
1008 """
1009 for entry in cache_entries:
1010 path = self.cache_directory.join(entry)
1012 self._cache_entries.pop(entry, None)
1013 log.debug("Removing file from cache: %s", path)
1014 with contextlib.suppress(FileNotFoundError):
1015 path.remove()
1017 def _expire_cache(self) -> None:
1018 """Expire the files in the cache.
1020 Notes
1021 -----
1022 The expiration modes are defined by the config or can be overridden.
1023 Available options:
1025 * ``files``: Number of files.
1026 * ``datasets``: Number of datasets
1027 * ``size``: Total size of files.
1028 * ``age``: Age of files.
1030 The first three would remove in reverse time order.
1031 Number of files is complicated by the possibility of disassembled
1032 composites where 10 small files can be created for each dataset.
1034 Additionally there is a use case for an external user to explicitly
1035 state the dataset refs that should be cached and then when to
1036 remove them. Overriding any global configuration.
1037 """
1038 if self._expiration_mode is None:
1039 # Expiration has been disabled.
1040 return
1042 # mypy can't be sure we have set a threshold properly
1043 if self._expiration_threshold is None: 1043 ↛ 1044line 1043 didn't jump to line 1044 because the condition on line 1043 was never true
1044 log.warning(
1045 "Requesting cache expiry of mode %s but no threshold set in config.", self._expiration_mode
1046 )
1047 return
1049 # Sync up cache. There is no file locking involved so for a shared
1050 # cache multiple processes may be racing to delete files. Deleting
1051 # a file that no longer exists is not an error.
1052 self.scan_cache()
1054 if self._expiration_mode == "files":
1055 n_files = len(self._cache_entries)
1056 n_over = n_files - self._expiration_threshold
1057 if n_over > 0:
1058 sorted_keys = self._sort_cache()
1059 keys_to_remove = sorted_keys[:n_over]
1060 self._remove_from_cache(keys_to_remove)
1061 return
1063 if self._expiration_mode == "datasets":
1064 # Count the datasets, in ascending timestamp order,
1065 # so that oldest turn up first.
1066 datasets = defaultdict(list)
1067 for key in self._sort_cache():
1068 entry = self._cache_entries[key]
1069 datasets[entry.ref].append(key)
1071 n_datasets = len(datasets)
1072 n_over = n_datasets - self._expiration_threshold
1073 if n_over > 0:
1074 # Keys will be read out in insert order which
1075 # will be date order so oldest ones are removed.
1076 ref_ids = list(datasets.keys())[:n_over]
1077 keys_to_remove = list(itertools.chain.from_iterable(datasets[ref_id] for ref_id in ref_ids))
1078 self._remove_from_cache(keys_to_remove)
1079 return
1081 if self._expiration_mode == "size":
1082 if self.cache_size > self._expiration_threshold:
1083 for key in self._sort_cache(): 1083 ↛ 1087line 1083 didn't jump to line 1087 because the loop on line 1083 didn't complete
1084 self._remove_from_cache([key])
1085 if self.cache_size <= self._expiration_threshold: 1085 ↛ 1083line 1085 didn't jump to line 1083 because the condition on line 1085 was always true
1086 break
1087 return
1089 if self._expiration_mode == "age": 1089 ↛ 1100line 1089 didn't jump to line 1100 because the condition on line 1089 was always true
1090 now = datetime.datetime.now(datetime.UTC)
1091 for key in self._sort_cache():
1092 delta = now - self._cache_entries[key].ctime
1093 if delta.seconds > self._expiration_threshold:
1094 self._remove_from_cache([key])
1095 else:
1096 # We're already in date order.
1097 break
1098 return
1100 raise ValueError(f"Unrecognized cache expiration mode of {self._expiration_mode}")
1102 def _sort_cache(self) -> list[str]:
1103 """Sort the cache entries by time and return the sorted keys.
1105 Returns
1106 -------
1107 sorted : `list` of `str`
1108 Keys into the cache, sorted by time with oldest first.
1109 """
1111 def _sort_by_time(key: str) -> datetime.datetime:
1112 """Sorter key function using cache entry details."""
1113 return self._cache_entries[key].ctime
1115 return sorted(self._cache_entries, key=_sort_by_time)
1117 def __str__(self) -> str:
1118 if self._expiration_mode == "disabled":
1119 return f"""{type(self).__name__}(disabled)"""
1120 cachedir = self._cache_directory if self._cache_directory else "<tempdir>"
1121 return (
1122 f"{type(self).__name__}@{cachedir} ({self._expiration_mode}={self._expiration_threshold},"
1123 f"default={self._caching_default}) "
1124 f"n_files={self.file_count}, n_bytes={self.cache_size}"
1125 )
1128class DatastoreDisabledCacheManager(AbstractDatastoreCacheManager):
1129 """A variant of the datastore cache where no cache is enabled.
1131 Parameters
1132 ----------
1133 config : `str` or `DatastoreCacheManagerConfig`
1134 Configuration to control caching.
1135 universe : `DimensionUniverse`
1136 Set of all known dimensions, used to expand and validate any used
1137 in lookup keys.
1138 """
1140 def __init__(
1141 self,
1142 config: str | DatastoreCacheManagerConfig | None = None,
1143 universe: DimensionUniverse | None = None,
1144 ):
1145 return
1147 def should_be_cached(self, entity: DatasetRef | DatasetType | StorageClass) -> bool:
1148 """Indicate whether the entity should be added to the cache.
1150 Parameters
1151 ----------
1152 entity : `StorageClass` or `DatasetType` or `DatasetRef`
1153 Thing to test against the configuration. The ``name`` property
1154 is used to determine a match. A `DatasetType` will first check
1155 its name, before checking its `StorageClass`. If there are no
1156 matches the default will be returned.
1158 Returns
1159 -------
1160 should_cache : `bool`
1161 Always returns `False`.
1162 """
1163 return False
1165 def move_to_cache(self, uri: ResourcePath, ref: DatasetRef) -> ResourcePath | None:
1166 """Move dataset to cache.
1168 Parameters
1169 ----------
1170 uri : `lsst.resources.ResourcePath`
1171 Location of the file to be relocated to the cache. Will be moved.
1172 ref : `DatasetRef`
1173 Ref associated with this file. Will be used to determine the name
1174 of the file within the cache.
1176 Returns
1177 -------
1178 new : `lsst.resources.ResourcePath` or `None`
1179 Always refuses and returns `None`.
1180 """
1181 return None
1183 @contextlib.contextmanager
1184 def find_in_cache(self, ref: DatasetRef, extension: str) -> Iterator[ResourcePath | None]:
1185 """Look for a dataset in the cache and return its location.
1187 Parameters
1188 ----------
1189 ref : `DatasetRef`
1190 Dataset to locate in the cache.
1191 extension : `str`
1192 File extension expected. Should include the leading "``.``".
1194 Yields
1195 ------
1196 uri : `lsst.resources.ResourcePath` or `None`
1197 Never finds a file. Always returns `None`.
1198 """
1199 yield None
1201 def remove_from_cache(self, ref: DatasetRef | Iterable[DatasetRef]) -> None:
1202 """Remove datasets from cache.
1204 Parameters
1205 ----------
1206 ref : `DatasetRef` or iterable of `DatasetRef`
1207 The datasets to remove from the cache. Always does nothing.
1208 """
1209 return
1211 def known_to_cache(self, ref: DatasetRef, extension: str | None = None) -> bool:
1212 """Report if a dataset is known to the cache.
1214 Parameters
1215 ----------
1216 ref : `DatasetRef`
1217 Dataset to check for in the cache.
1218 extension : `str`, optional
1219 File extension expected. Should include the leading "``.``".
1220 If `None` the extension is ignored and the dataset ID alone is
1221 used to check in the cache. The extension must be defined if
1222 a specific component is being checked.
1224 Returns
1225 -------
1226 known : `bool`
1227 Always returns `False`.
1228 """
1229 return False
1231 def __str__(self) -> str:
1232 return f"{type(self).__name__}()"
1235class InvalidCacheFilenameError(Exception):
1236 """Raised when attempting to register a file in the cache with a name in
1237 the incorrect format.
1238 """