1"""Storage providers backends for Memory caching."""
2
3import collections
4import datetime
5import json
6import operator
7import os
8import os.path
9import shutil
10import threading
11import time
12import types
13import uuid
14import warnings
15from abc import ABCMeta, abstractmethod
16from pickle import PicklingError
17
18from . import numpy_pickle
19from .backports import concurrency_safe_rename
20from .disk import memstr_to_bytes, mkdirp, rm_subdirs
21from .logger import format_time
22
23CacheItemInfo = collections.namedtuple("CacheItemInfo", "path size last_access")
24
25
26class CacheWarning(Warning):
27 """Warning to capture dump failures except for PicklingError."""
28
29 pass
30
31
32class StoreBackendBase(metaclass=ABCMeta):
33 """Helper Abstract Base Class which defines all methods that
34 a StorageBackend must implement."""
35
36 location = None
37
38 @abstractmethod
39 def _open_item(self, f, mode):
40 """Opens an item on the store and return a file-like object.
41
42 This method is private and only used by the StoreBackendMixin object.
43
44 Parameters
45 ----------
46 f: a file-like object
47 The file-like object where an item is stored and retrieved
48 mode: string, optional
49 the mode in which the file-like object is opened allowed valued are
50 'rb', 'wb'
51
52 Returns
53 -------
54 a file-like object
55 """
56
57 @abstractmethod
58 def _item_exists(self, location):
59 """Checks if an item location exists in the store.
60
61 This method is private and only used by the StoreBackendMixin object.
62
63 Parameters
64 ----------
65 location: string
66 The location of an item. On a filesystem, this corresponds to the
67 absolute path, including the filename, of a file.
68
69 Returns
70 -------
71 True if the item exists, False otherwise
72 """
73
74 @abstractmethod
75 def _unlink_item(self, location):
76 """Unlink an item from the store.
77
78 This method is private and only used by the StoreBackendMixin object.
79
80 Parameters
81 ----------
82 location: string
83 The location of an item. On a filesystem, this corresponds to the
84 absolute path, including the filename, of a file.
85 """
86
87 @abstractmethod
88 def _move_item(self, src, dst):
89 """Moves an item from src to dst in the store.
90
91 This method is private and only used by the StoreBackendMixin object.
92
93 Parameters
94 ----------
95 src: string
96 The source location of an item
97 dst: string
98 The destination location of an item
99 """
100
101 @abstractmethod
102 def create_location(self, location):
103 """Creates a location on the store.
104
105 Parameters
106 ----------
107 location: string
108 The location in the store. On a filesystem, this corresponds to a
109 directory.
110 """
111
112 @abstractmethod
113 def clear_location(self, location):
114 """Clears a location on the store.
115
116 Parameters
117 ----------
118 location: string
119 The location in the store. On a filesystem, this corresponds to a
120 directory or a filename absolute path
121 """
122
123 @abstractmethod
124 def get_items(self):
125 """Returns the whole list of items available in the store.
126
127 Returns
128 -------
129 The list of items identified by their ids (e.g filename in a
130 filesystem).
131 """
132
133 @abstractmethod
134 def configure(self, location, verbose=0, backend_options=dict()):
135 """Configures the store.
136
137 Parameters
138 ----------
139 location: string
140 The base location used by the store. On a filesystem, this
141 corresponds to a directory.
142 verbose: int
143 The level of verbosity of the store
144 backend_options: dict
145 Contains a dictionary of named parameters used to configure the
146 store backend.
147 """
148
149
150class StoreBackendMixin(object):
151 """Class providing all logic for managing the store in a generic way.
152
153 The StoreBackend subclass has to implement 3 methods: create_location,
154 clear_location and configure. The StoreBackend also has to provide
155 a private _open_item, _item_exists, _unlink_item and _move_item methods.
156 The _open_item method has to have the same signature as the builtin open
157 and return a file-like object.
158 """
159
160 def load_item(self, call_id, verbose=1, timestamp=None, metadata=None):
161 """Load an item from the store given its id as a list of str."""
162 full_path = os.path.join(self.location, *call_id)
163
164 if verbose > 1:
165 ts_string = (
166 "{: <16}".format(format_time(time.time() - timestamp))
167 if timestamp is not None
168 else ""
169 )
170 signature = os.path.basename(call_id[0])
171 if metadata is not None and "input_args" in metadata:
172 kwargs = ", ".join(
173 "{}={}".format(*item) for item in metadata["input_args"].items()
174 )
175 signature += "({})".format(kwargs)
176 msg = "[Memory]{}: Loading {}".format(ts_string, signature)
177 if verbose < 10:
178 print("{0}...".format(msg))
179 else:
180 print("{0} from {1}".format(msg, full_path))
181
182 mmap_mode = None if not hasattr(self, "mmap_mode") else self.mmap_mode
183
184 filename = os.path.join(full_path, "output.pkl")
185 if not self._item_exists(filename):
186 raise KeyError(
187 "Non-existing item (may have been "
188 "cleared).\nFile %s does not exist" % filename
189 )
190
191 # file-like object cannot be used when mmap_mode is set
192 if mmap_mode is None:
193 with self._open_item(filename, "rb") as f:
194 item = numpy_pickle.load(f)
195 else:
196 item = numpy_pickle.load(filename, mmap_mode=mmap_mode)
197 return item
198
199 def dump_item(self, call_id, item, verbose=1):
200 """Dump an item in the store at the id given as a list of str."""
201 try:
202 item_path = os.path.join(self.location, *call_id)
203 if not self._item_exists(item_path):
204 self.create_location(item_path)
205 filename = os.path.join(item_path, "output.pkl")
206 if verbose > 10:
207 print("Persisting in %s" % item_path)
208
209 def write_func(to_write, dest_filename):
210 with self._open_item(dest_filename, "wb") as f:
211 try:
212 numpy_pickle.dump(to_write, f, compress=self.compress)
213 except PicklingError as e:
214 # TODO(1.5) turn into error
215 warnings.warn(
216 "Unable to cache to disk: failed to pickle "
217 "output. In version 1.5 this will raise an "
218 f"exception. Exception: {e}.",
219 FutureWarning,
220 )
221
222 self._concurrency_safe_write(item, filename, write_func)
223 except Exception as e: # noqa: E722
224 warnings.warn(
225 "Unable to cache to disk. Possibly a race condition in the "
226 f"creation of the directory. Exception: {e}.",
227 CacheWarning,
228 )
229
230 def clear_item(self, call_id):
231 """Clear the item at the id, given as a list of str."""
232 item_path = os.path.join(self.location, *call_id)
233 if self._item_exists(item_path):
234 self.clear_location(item_path)
235
236 def contains_item(self, call_id):
237 """Check if there is an item at the id, given as a list of str."""
238 item_path = os.path.join(self.location, *call_id)
239 filename = os.path.join(item_path, "output.pkl")
240
241 return self._item_exists(filename)
242
243 def get_item_info(self, call_id):
244 """Return information about item."""
245 return {"location": os.path.join(self.location, *call_id)}
246
247 def get_metadata(self, call_id):
248 """Return actual metadata of an item."""
249 try:
250 item_path = os.path.join(self.location, *call_id)
251 filename = os.path.join(item_path, "metadata.json")
252 with self._open_item(filename, "rb") as f:
253 return json.loads(f.read().decode("utf-8"))
254 except: # noqa: E722
255 return {}
256
257 def store_metadata(self, call_id, metadata):
258 """Store metadata of a computation."""
259 try:
260 item_path = os.path.join(self.location, *call_id)
261 self.create_location(item_path)
262 filename = os.path.join(item_path, "metadata.json")
263
264 def write_func(to_write, dest_filename):
265 with self._open_item(dest_filename, "wb") as f:
266 f.write(json.dumps(to_write).encode("utf-8"))
267
268 self._concurrency_safe_write(metadata, filename, write_func)
269 except: # noqa: E722
270 pass
271
272 def contains_path(self, call_id):
273 """Check cached function is available in store."""
274 func_path = os.path.join(self.location, *call_id)
275 return self._item_exists(func_path)
276
277 def clear_path(self, call_id):
278 """Clear all items with a common path in the store."""
279 func_path = os.path.join(self.location, *call_id)
280 if self._item_exists(func_path):
281 self.clear_location(func_path)
282
283 def store_cached_func_code(self, call_id, func_code=None):
284 """Store the code of the cached function."""
285 func_path = os.path.join(self.location, *call_id)
286 if not self._item_exists(func_path):
287 self.create_location(func_path)
288
289 if func_code is not None:
290 filename = os.path.join(func_path, "func_code.py")
291
292 def write_func(to_write, dest_filename):
293 with self._open_item(dest_filename, "wb") as f:
294 f.write(to_write.encode("utf-8"))
295
296 # Written in place, another process could read it half-written and
297 # take that for a changed function, wiping the cache (#1694)
298 self._concurrency_safe_write(func_code, filename, write_func)
299
300 def get_cached_func_code(self, call_id):
301 """Store the code of the cached function."""
302 filename = os.path.join(self.location, *call_id, "func_code.py")
303 try:
304 with self._open_item(filename, "rb") as f:
305 return f.read().decode("utf-8")
306 except: # noqa: E722
307 raise
308
309 def get_cached_func_info(self, call_id):
310 """Return information related to the cached function if it exists."""
311 return {"location": os.path.join(self.location, *call_id)}
312
313 def clear(self):
314 """Clear the whole store content."""
315 self.clear_location(self.location)
316
317 def enforce_store_limits(self, bytes_limit, items_limit=None, age_limit=None):
318 """
319 Remove the store's oldest files to enforce item, byte, and age limits.
320 """
321 items_to_delete = self._get_items_to_delete(bytes_limit, items_limit, age_limit)
322
323 for item in items_to_delete:
324 if self.verbose > 10:
325 print("Deleting item {0}".format(item))
326 try:
327 self.clear_location(item.path)
328 except OSError:
329 # Even with ignore_errors=True shutil.rmtree can raise OSError
330 # with:
331 # [Errno 116] Stale file handle if another process has deleted
332 # the folder already.
333 pass
334
335 def _get_items_to_delete(self, bytes_limit, items_limit=None, age_limit=None):
336 """
337 Get items to delete to keep the store under size, file, & age limits.
338 """
339 if isinstance(bytes_limit, str):
340 bytes_limit = memstr_to_bytes(bytes_limit)
341
342 items = self.get_items()
343 if not items:
344 return []
345
346 size = sum(item.size for item in items)
347
348 if bytes_limit is not None:
349 to_delete_size = size - bytes_limit
350 else:
351 to_delete_size = 0
352
353 if items_limit is not None:
354 to_delete_items = len(items) - items_limit
355 else:
356 to_delete_items = 0
357
358 if age_limit is not None:
359 older_item = min(item.last_access for item in items)
360 if age_limit.total_seconds() < 0:
361 raise ValueError("age_limit has to be a positive timedelta")
362 deadline = datetime.datetime.now() - age_limit
363 else:
364 deadline = None
365
366 if (
367 to_delete_size <= 0
368 and to_delete_items <= 0
369 and (deadline is None or older_item > deadline)
370 ):
371 return []
372
373 # We want to delete first the cache items that were accessed a
374 # long time ago
375 items.sort(key=operator.attrgetter("last_access"))
376
377 items_to_delete = []
378 size_so_far = 0
379 items_so_far = 0
380
381 for item in items:
382 if (
383 (size_so_far >= to_delete_size)
384 and items_so_far >= to_delete_items
385 and (deadline is None or deadline < item.last_access)
386 ):
387 break
388
389 items_to_delete.append(item)
390 size_so_far += item.size
391 items_so_far += 1
392
393 return items_to_delete
394
395 def _concurrency_safe_write(self, to_write, filename, write_func):
396 """Writes an object into a file in a concurrency-safe way."""
397
398 # Temporary name is composed of UUID, process_id and thread_id to avoid
399 # collisions due to concurrent write.
400 # UUID is unique across nodes and time and help avoid collisions, even if
401 # the cache folder is shared by several Python processes with the same pid and
402 # thread id on different nodes of a cluster for instance.
403 thread_id = id(threading.current_thread())
404 temporary_filename = f"{filename}.{uuid.uuid4().hex}-{os.getpid()}-{thread_id}"
405
406 try:
407 write_func(to_write, temporary_filename)
408 except BaseException:
409 if self._item_exists(temporary_filename):
410 self._unlink_item(temporary_filename)
411 raise
412 else:
413 self._move_item(temporary_filename, filename)
414
415 def __repr__(self):
416 """Printable representation of the store location."""
417 return '{class_name}(location="{location}")'.format(
418 class_name=self.__class__.__name__, location=self.location
419 )
420
421
422def _check_hex(s, length):
423 if len(s) != length:
424 return False
425 try:
426 if length % 2 == 1:
427 # bytes.fromhex waits for pair sequences
428 s += "0"
429 bytes.fromhex(s)
430 return True
431 except ValueError:
432 return False
433
434
435def _old_split_id(self, call_id):
436 info_path = os.path.join(self.location, "store_backend_info.json")
437 try:
438 with open(info_path, "rb") as file:
439 info = json.loads(file.read().decode("utf-8"))
440 if info["cache_version"] == 2:
441 self._split_id = types.MethodType(FileSystemStoreBackend._split_id, self)
442 return self._split_id(call_id)
443 except OSError:
444 pass
445 return call_id
446
447
448def _reconstructStoreBackend(cls, location, verbose, compress, mmap_mode):
449 obj = cls.__new__(cls)
450 obj.configure(location, verbose, dict(compress=compress, mmap_mode=mmap_mode))
451 return obj
452
453
454def _replace_cache_dir(old_dir, new_dir):
455 """
456 Move an old cache directory to a newer one depending on
457 the last modification time of output.pkl
458 """
459 old_item = os.path.join(old_dir, "output.pkl")
460 try:
461 old_mtime = os.path.getmtime(old_item)
462 except OSError:
463 # "output.pkl" does not exist, so we don't replace
464 shutil.rmtree(old_dir)
465 return
466 new_item = os.path.join(new_dir, "output.pkl")
467 try:
468 new_mtime = os.path.getmtime(new_item)
469 except OSError:
470 new_mtime = -1
471 if new_mtime < old_mtime:
472 # replace
473 shutil.rmtree(new_dir)
474 os.replace(old_dir, new_dir)
475 else:
476 # don't replace
477 shutil.rmtree(old_dir)
478
479
480def _split_decorator(cls):
481 """Helper to split call_id in all relevant class methods.
482
483 Use a decorator as all methods are not implemented directly
484 in this class.
485 """
486
487 def split_method_decorator(method):
488 def split_method(self, call_id, *args, **kwargs):
489 call_id = self._split_id(call_id)
490 return method(self, call_id, *args, **kwargs)
491
492 return split_method
493
494 for method in [
495 "load_item",
496 "dump_item",
497 "clear_item",
498 "contains_item",
499 "get_item_info",
500 "get_metadata",
501 "store_metadata",
502 ]:
503 setattr(cls, method, split_method_decorator(getattr(StoreBackendMixin, method)))
504 return cls
505
506
507@_split_decorator
508class FileSystemStoreBackend(StoreBackendBase, StoreBackendMixin):
509 """A StoreBackend used with local or network file systems."""
510
511 _open_item = staticmethod(open)
512 _item_exists = staticmethod(os.path.exists)
513 _unlink_item = staticmethod(os.unlink)
514 _move_item = staticmethod(concurrency_safe_rename)
515
516 def clear_location(self, location):
517 """Delete location on store."""
518 if location == self.location:
519 rm_subdirs(location)
520 else:
521 shutil.rmtree(location, ignore_errors=True)
522
523 def create_location(self, location):
524 """Create object location on store"""
525 mkdirp(location)
526
527 def get_items(self):
528 """Returns the whole list of items available in the store."""
529 items = []
530
531 for dirpath, _, filenames in os.walk(self.location):
532 parent, basename = os.path.split(dirpath)
533 if not _check_hex(basename, 32):
534 if not _check_hex(basename, 29):
535 continue
536 if not _check_hex(os.path.basename(parent), 3):
537 continue
538
539 output_filename = os.path.join(dirpath, "output.pkl")
540 try:
541 last_access = os.path.getatime(output_filename)
542 except OSError:
543 try:
544 last_access = os.path.getatime(dirpath)
545 except OSError:
546 # The directory has already been deleted
547 continue
548
549 last_access = datetime.datetime.fromtimestamp(last_access)
550 try:
551 full_filenames = [os.path.join(dirpath, fn) for fn in filenames]
552 dirsize = sum(os.path.getsize(fn) for fn in full_filenames)
553 except OSError:
554 # Either output_filename or one of the files in
555 # dirpath does not exist any more. We assume this
556 # directory is being cleaned by another process already
557 continue
558
559 items.append(CacheItemInfo(dirpath, dirsize, last_access))
560
561 return items
562
563 def configure(self, location, verbose=1, backend_options=None):
564 """Configure the store backend.
565
566 For this backend, valid store options are 'compress' and 'mmap_mode'
567 """
568 if backend_options is None:
569 backend_options = {}
570
571 # setup location directory
572 self.location = location
573 info = None
574 if not os.path.exists(self.location):
575 mkdirp(self.location)
576 info = {"cache_version": 2, "require_update": False}
577
578 # Get info
579 info_path = os.path.join(self.location, "store_backend_info.json")
580 if os.path.exists(info_path):
581 with open(info_path, "rb") as file:
582 info = json.loads(file.read().decode("utf-8"))
583 else:
584 if info is None:
585 # Cache directory without info.
586 # As long as it has subdirectories,
587 # we consider it uses an old cache tree
588 with os.scandir(self.location) as entries:
589 for entry in entries:
590 if entry.is_dir():
591 info = {"cache_version": 1, "require_update": True}
592 break
593 if info is None:
594 info = {"cache_version": 2, "require_update": False}
595 with open(info_path, "wb") as file:
596 file.write(json.dumps(info).encode("utf-8"))
597
598 # Warn if an update is required
599 if info["require_update"]:
600 true_location = (
601 f"'{os.path.dirname(location)}'"
602 if os.path.basename(location) == "joblib"
603 else f"pathlib.Path('{location}')"
604 )
605 warnings.warn(
606 f"The FileSystemStoreBackend at '{self.location}' may contains "
607 "items using old cache storage tree.\n"
608 "The joblib cache tree has recently been updated "
609 "for efficiency reasons.\n"
610 "Starting with joblib 1.9, the old cache tree "
611 "will no longer be supported.\n"
612 f"Please run `joblib.Memory({true_location}).update_cache_tree()` "
613 "to update your cache tree."
614 )
615
616 # Using old split method for older cache version
617 if info["cache_version"] == 1:
618 self._split_id = types.MethodType(_old_split_id, self)
619
620 # Automatically add `.gitignore` file to the cache folder.
621 # XXX: the condition is necessary because in `Memory.__init__`, the user
622 # passed `location` param is modified to be either `{location}` or
623 # `{location}/joblib` depending on input type (`pathlib.Path` vs `str`).
624 # The proper resolution of this inconsistency is tracked in:
625 # https://github.com/joblib/joblib/issues/1684
626 cache_directory = (
627 os.path.dirname(location)
628 if os.path.dirname(location) and os.path.basename(location) == "joblib"
629 else location
630 )
631 gitignore = os.path.join(cache_directory, ".gitignore")
632 if not os.path.exists(gitignore):
633 try:
634 with open(gitignore, "w") as file:
635 file.write("# Created by joblib automatically.\n")
636 file.write("*\n")
637 except OSError as e:
638 warnings.warn(f"Unable to write {gitignore}. Exception: {e}.")
639
640 # item can be stored compressed for faster I/O
641 self.compress = backend_options.get("compress", False)
642
643 # FileSystemStoreBackend can be used with mmap_mode options under
644 # certain conditions.
645 mmap_mode = backend_options.get("mmap_mode")
646 if self.compress and mmap_mode is not None:
647 warnings.warn(
648 "Compressed items cannot be memmapped in a "
649 "filesystem store. Option will be ignored.",
650 stacklevel=2,
651 )
652
653 self.mmap_mode = mmap_mode
654 self.verbose = verbose
655
656 def _split_id(self, call_id):
657 if len(call_id) == 0 or not _check_hex(call_id[-1], 32):
658 return call_id
659 return (*call_id[:-1], call_id[-1][:3], call_id[-1][3:])
660
661 # XXX: To remove in joblib 1.9
662 def update_cache_tree(self):
663 # First info update
664 info_path = os.path.join(self.location, "store_backend_info.json")
665 info = {"cache_version": 2, "require_update": True}
666 with open(info_path, "wb") as file:
667 file.write(json.dumps(info).encode("utf-8"))
668
669 # Find old cache directories
670 old_cache = []
671 for dirpath, _, files in os.walk(self.location):
672 if not _check_hex(os.path.basename(dirpath), 32):
673 continue
674 if not (("output.pkl" in files) or ("metadata.json" in files)):
675 continue
676 old_cache.append(dirpath)
677
678 # Replace old cache directories
679 for dirpath in old_cache:
680 parent, basename = os.path.split(dirpath)
681 newdir_parent = os.path.join(parent, basename[:3])
682 newdir = os.path.join(newdir_parent, basename[3:])
683 if os.path.exists(newdir):
684 # If the new cache directory already exists,
685 # it is replaced depending on the last modification time
686 # of 'output.pkl'
687 _replace_cache_dir(dirpath, newdir)
688 else:
689 os.makedirs(newdir_parent, exist_ok=True)
690 os.replace(dirpath, newdir)
691
692 # Second info update
693 info["require_update"] = False
694 with open(info_path, "wb") as file:
695 file.write(json.dumps(info).encode("utf-8"))
696
697 def __reduce__(self):
698 return (
699 _reconstructStoreBackend,
700 (
701 self.__class__,
702 self.location,
703 self.verbose,
704 self.compress,
705 self.mmap_mode,
706 ),
707 )