Coverage for /pythoncovmergedfiles/medio/medio/usr/local/lib/python3.11/site-packages/joblib/_store_backends.py: 24%

Shortcuts on this page

r m x   toggle line displays

j k   next/prev highlighted chunk

0   (zero) top of page

1   (one) first highlighted chunk

347 statements  

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 )