Coverage for gpaw/test/cachedfilehandler.py: 78%

1from abc import ABC, abstractmethod 

2from contextlib import contextmanager 

3from pathlib import Path 

4 

5from gpaw.mpi import world 

6 

7 

8class CachedFilesHandler(ABC): 

9 """Base class for objects which handle the writing 

10 and caching of session-scoped pytest fixtures""" 

11 

12 def __init__(self, folder: Path, outformat: str): 

13 self.folder = folder 

14 self.outformat = outformat 

15 

16 self.cached_files = {} 

17 for file in folder.glob('*' + outformat): 

18 self.cached_files[file.name[:-len(outformat)]] = file 

19 

20 def __getitem__(self, name: str) -> Path: 

21 if name in self.cached_files: 

22 return self.cached_files[name] 

23 

24 filepath = self.folder / (name + self.outformat) 

25 

26 lockfile = self.folder / f'{name}.lock' 

27 

28 for _attempt in range(60): # ~60s timeout 

29 file_exist = 0 

30 if world.rank == 0: 

31 file_exist = int(filepath.exists()) 

32 file_exist = world.sum_scalar(file_exist) 

33 

34 if file_exist: 

35 self.cached_files[name] = filepath 

36 return self.cached_files[name] 

37 

38 try: 

39 with world_temporary_lock(lockfile): 

40 work_path = filepath.with_suffix('.tmp' + self.outformat) 

41 self._calculate_and_write(name, work_path) 

42 

43 # By now files should exist *and* be fully written, by us. 

44 # Rename them to the final intended paths: 

45 if world.rank == 0: 

46 work_path.rename(filepath) 

47 

48 except Locked: 

49 import time 

50 time.sleep(1) 

51 

52 raise RuntimeError(f'{self.__class__.__name__} fixture generation ' 

53 f'takes too long: {name}. Consider using pytest ' 

54 '--cache-clear if there are stale lockfiles, ' 

55 'else write faster tests.') 

56 

57 @abstractmethod 

58 def _calculate_and_write(self, name, work_path): 

59 pass 

60 

61 

62class Locked(FileExistsError): 

63 pass 

64 

65 

66@contextmanager 

67def world_temporary_lock(path): 

68 if world.rank == 0: 

69 try: 

70 with temporary_lock(path): 

71 world.sum_scalar(1) 

72 yield 

73 except Locked: 

74 world.sum_scalar(0) 

75 raise 

76 else: 

77 status = world.sum_scalar(0) 

78 if status: 

79 yield 

80 else: 

81 raise Locked 

82 

83 

84@contextmanager 

85def temporary_lock(path): 

86 fd = None 

87 try: 

88 with path.open('x') as fd: 

89 yield 

90 except FileExistsError: 

91 raise Locked() 

92 finally: 

93 if fd is not None: 

94 path.unlink()