Coverage for src/sparkle/CLI/help/jobs.py: 100%

1"""File to help with RunRunner jobs.""" 

2 

3import sys 

4import inspect 

5from pathlib import Path 

6 

7from runrunner.base import Status 

8from runrunner.slurm import SlurmRun 

9 

10from sparkle.types import DataFileLock 

11 

12 

13def get_locking_runs( 

14 path: Path, 

15 locks: set[DataFileLock], 

16) -> tuple[list[SlurmRun], list[SlurmRun]]: 

17 """Return running and waiting jobs that lock the given data files. 

18 

19 Args: 

20 path: Path to search for RunRunner job JSON files. 

21 locks: Set of DataFileLock values to check against. 

22 

23 Returns: 

24 A tuple (running, waiting) of SlurmRun lists that overlap with the 

25 given locks. 

26 """ 

27 # Maps job name prefixes to the data file they lock 

28 job_name_locks: dict[str, DataFileLock] = { 

29 "Run Solver ": DataFileLock.PERFORMANCE, 

30 "Run Extractor ": DataFileLock.FEATURE, 

31 } 

32 running = [] 

33 waiting = [] 

34 # Collect only the name prefixes that correspond to the requested lock types, 

35 # so we skip jobs that write to unrelated data files. 

36 relevant_prefixes = { 

37 prefix for prefix, lock in job_name_locks.items() if lock in locks 

38 } 

39 for run in get_runs_from_file(path, filter=[Status.RUNNING, Status.WAITING]): 

40 if any(run.name.startswith(prefix) for prefix in relevant_prefixes): 

41 if run.status == Status.RUNNING: 

42 running.append(run) 

43 else: 

44 waiting.append(run) 

45 return running, waiting 

46 

47 

48def check_running_waiting_jobs( 

49 path: Path, locks: set[DataFileLock] | None = None 

50) -> None: 

51 """Check for running/waiting jobs that lock the calling command's data files. 

52 

53 For standard CLI commands (locks=None): derives the relevant locks by looking 

54 up the calling command's filename in the internal mapping. Raises KeyError if 

55 the caller is not registered. 

56 

57 For cleanup only (locks provided): uses the given lock set directly, since 

58 cleanup determines its locks from CLI flags at runtime. Any other caller 

59 passing explicit locks raises ValueError. 

60 

61 Exits with -1 if running jobs are found. 

62 Asks the user whether to continue if only waiting jobs are found, and 

63 exits with -1 if the user declines. 

64 

65 Args: 

66 path: Path to search for RunRunner job JSON files. 

67 locks: Only cleanup.py may pass this. All other callers must omit it. 

68 """ 

69 # Internal mapping: which data files each CLI command structurally modifies. 

70 # cleanup is intentionally absent — it determines its locks from CLI flags. 

71 cli_command_locks: dict[str, set[DataFileLock]] = { 

72 "add_solver": {DataFileLock.PERFORMANCE}, 

73 "remove_solver": {DataFileLock.PERFORMANCE}, 

74 "add_instances": {DataFileLock.PERFORMANCE, DataFileLock.FEATURE}, 

75 "remove_instances": {DataFileLock.PERFORMANCE, DataFileLock.FEATURE}, 

76 "add_feature_extractor": {DataFileLock.FEATURE}, 

77 "remove_feature_extractor": {DataFileLock.FEATURE}, 

78 "run_portfolio_selector": {DataFileLock.PERFORMANCE, DataFileLock.FEATURE}, 

79 } 

80 

81 def get_locks() -> set[DataFileLock]: 

82 """Validate the caller and return the appropriate lock set.""" 

83 # inspect.stack()[2]: 0=get_locks, 1=check_running_waiting_jobs, 2=caller 

84 caller = Path(inspect.stack()[2].filename).stem 

85 if locks is not None: 

86 # Explicit locks are only permitted for cleanup 

87 if caller != "cleanup": 

88 raise ValueError(f"'{caller}' cannot pass explicit locks. ") 

89 return locks 

90 if caller not in cli_command_locks: 

91 raise KeyError(f"'{caller}' is not registered in the lock commands. ") 

92 return cli_command_locks[caller] 

93 

94 resolved_locks = get_locks() 

95 # Build a human-readable string of the affected data files for the warning 

96 # messages, e.g. "PERFORMANCE_DATA or FEATURE_DATA". 

97 lock_names = " or ".join(lock.value.capitalize() for lock in resolved_locks) 

98 running, waiting = get_locking_runs(path, resolved_locks) 

99 if running: 

100 print( 

101 f"WARNING: There are {len(running)} running job(s) writing to the " 

102 f"{lock_names}. Please cancel them before modifying the platform." 

103 ) 

104 sys.exit(-1) 

105 if waiting: 

106 print( 

107 f"WARNING: There are {len(waiting)} waiting job(s) that will write " 

108 f"to the {lock_names}. These may conflict with this operation. " 

109 "Continue? [y/n]" 

110 ) 

111 if input() != "y": 

112 sys.exit(-1) 

113 

114 

115def get_runs_from_file( 

116 path: Path, print_error: bool = False, filter: list[Status] | None = None 

117) -> list[SlurmRun]: 

118 """Retrieve all run objects from file storage. 

119 

120 Args: 

121 path: Path object where to look recursively for the files. 

122 print_error: Whether to print errors. 

123 filter: If not None, only runs with the given statuses will be 

124 returned. 

125 

126 Returns: 

127 List of all found SlumRun objects. 

128 """ 

129 if not path.exists(): 

130 return [] 

131 runs = [] 

132 for file in path.rglob("*.json"): 

133 # TODO: RunRunner should be adapted to have more general methods for runs 

134 # So this method can work for both local and slurm 

135 try: 

136 run_obj = SlurmRun.from_file(file) 

137 if filter is None or run_obj.status in filter: 

138 runs.append(run_obj) 

139 except Exception as ex: 

140 # Not a (correct) RunRunner JSON file 

141 if print_error: 

142 print(f"[WARNING] Could not load file: {file}. Exception: {ex}") 

143 return runs