Coverage for src/mockslurm/process_db.py: 88%
82 statements
« prev ^ index » next coverage.py v7.15.3, created at 2026-08-07 13:34 +0000
1# Copyright 2026 CNRS
2# This software is distributed under the terms of the CeCILL-C free software license.
4"""Implements the access to a "slurm database", stored in a HDF5 file.
6HDF5 files are locked during access, which provides a convenient method to
7not have several process modifying the same database easily
8"""
10import copy
11import datetime
12import getpass
13import os
14import time
15from enum import IntEnum
16from pathlib import Path
17from typing import Any
19import h5py
20import numpy as np
23class JobState(IntEnum):
24 COMPLETED = 1
25 RUNNING = 2
26 PENDING = 3
27 CANCELLED = 4
28 FAILED = 5
31class JobReason(IntEnum):
32 NOREASON = 1 # sometimes there are no reason: job is running for instance
33 WaitingForScheduling = 2
34 Dependency = 3
35 DependencyNeverSatisfied = 4
36 NonZeroExitCode = 5
37 JobLaunchFailure = 6
40_DB_DTYPE = np.dtype(
41 [
42 ("PID", np.int64),
43 ("NAME", np.dtype("S128")),
44 ("USER", np.dtype("S128")),
45 ("ACCOUNT", np.dtype("S128")),
46 ("PARTITION", np.dtype("S128")),
47 ("RESERVATION", np.dtype("S128")),
48 ("NODELIST", np.dtype("S16384")),
49 ("TIME", np.int64),
50 ("START_TIME", np.float64),
51 ("CMD", np.dtype("S16384")),
52 ("STATE", np.int64),
53 ("REASON", np.int64),
54 ("EXIT_CODE", np.int64),
55 ]
56)
58DB_DEFAULTS = {
59 "PID": -1,
60 "NAME": "wrap",
61 "USER": getpass.getuser(),
62 "ACCOUNT": getpass.getuser(),
63 "PARTITION": "",
64 "RESERVATION": "",
65 "NODELIST": "mocknode1",
66 "TIME": 0,
67 "START_TIME": datetime.datetime.now().timestamp(),
68 "CMD": "",
69 "STATE": JobState.PENDING,
70 "REASON": JobReason.WaitingForScheduling,
71 "EXIT_CODE": np.iinfo(np.int16).max,
72}
75def find_db_file() -> Path:
76 """Return a path to the location of the mock database hdf5 file.
78 The database file location is searched among various locations typically available.
80 Returns
81 -------
82 Path
83 Path to the mock database hdf5 file.
85 Raises
86 ------
87 FileNotFoundError
88 If a suitable location for the database file could not be found.
89 """
90 possible_locations = [Path("/tmp"), Path("/var/tmp"), Path(os.environ["HOME"]) / ".local", Path()]
91 mock_slurm_db = Path("mock_slurm_db.h5")
93 for p in possible_locations: 93 ↛ 99line 93 didn't jump to line 99 because the loop on line 93 didn't complete
94 # retrieve the user permission for the directory
95 # if it is 7, we can read, write and execute it so we can open the folder and write inside
96 if p.exists() and int(oct(p.stat().st_mode)[-3]) == 7: 96 ↛ 93line 96 didn't jump to line 93 because the condition on line 96 was always true
97 return p / mock_slurm_db
99 raise FileNotFoundError("Could not find a location to write mock slurm DB.")
102def open_file_retry_on_locked(file: Path, mode: str = "a", nb_retries: int = 40, wait_s: float = 0.01) -> h5py.File:
103 """Open `file` as an HDF5 file in `mode`.
105 Since the HDF5 files are locked when another process access them, this function tries to
106 open the file a certain number of time.
108 Parameters
109 ----------
110 file : Path
111 Path to the file to open
112 nb_retries : int, optional
113 Number of times to try to open the file, by default 40
114 wait_s : float, optional
115 Amount of time to wait between 2 attempts to open the file, in seconds, by default 0.1
117 Returns
118 -------
119 h5py.File
120 Open file handle to `file`.
121 """
122 # Failing to open a file on fefs doesn't necessarily mean we won't succeed next time !
123 for _ in range(nb_retries - 1): 123 ↛ 130line 123 didn't jump to line 130 because the loop on line 123 didn't complete
124 try:
125 f = h5py.File(file, mode)
126 return f
127 except BlockingIOError:
128 time.sleep(wait_s)
129 # try 1 last time, let tables error raise if failing again
130 return h5py.File(file, mode)
133def clear_db():
134 """Deletes the database file."""
135 db_file = find_db_file()
136 if db_file.exists():
137 print(f"Deleting db at {db_file}")
138 db_file.unlink()
139 else:
140 print("No file to delete.")
143def get_db_file_handle(db_file: Path) -> h5py.File:
144 """Open the database HDF5 file in append mode.
146 On success, the returned h5py.File contains a dataset with the expected dtypes of the database.
148 Parameters
149 ----------
150 db_file : Path
151 Path to the database file to open
153 Returns
154 -------
155 h5py.File
156 File handle to the database
157 """
158 if not db_file.exists():
159 db = np.empty(
160 dtype=_DB_DTYPE,
161 shape=(0,),
162 )
163 f = open_file_retry_on_locked(db_file)
164 f.create_dataset("SLURM_DB", data=db, maxshape=(None,))
165 return f
167 return open_file_retry_on_locked(db_file)
170def get_db(db_file: h5py.File) -> h5py.Dataset:
171 """Get the Database as a dataset in the HDF5 file
173 Parameters
174 ----------
175 db_file : h5py.File
176 Opened database file handle
178 Returns
179 -------
180 h5py.Dataset
181 Dataset storing the database
182 """
183 return db_file["SLURM_DB"]
186def update_with_default_value(db_dict: dict) -> dict:
187 """Update `db_dict` missing fields with the default value in the database
189 Parameters
190 ----------
191 db_dict : Dict
192 dict containing a database row data, possibly missing some columns to be filled with default values
194 Returns
195 -------
196 Dict
197 dict with all fields expected by the database present, with `dict` values if present, otherwise default values
198 """
199 default_dict = copy.deepcopy(DB_DEFAULTS)
200 default_dict.update(db_dict)
201 return default_dict
204def append_job(db_file: h5py.File, **kwargs) -> int:
205 """Append a job to the database
207 Parameters
208 ----------
209 db_file : h5py.File
210 Opened file handle to the database
212 Returns
213 -------
214 int
215 Index of the job that was appended
216 """
217 dataset = get_db(db_file)
218 dataset.resize(dataset.shape[0] + 1, axis=0)
219 job_data = np.empty(dtype=dataset.dtype, shape=(1,))
220 for k, v in update_with_default_value(kwargs).items():
221 job_data[k] = v
222 dataset[-1] = job_data
223 return dataset.shape[0] - 1
226def update_db_value(db_file: h5py.File, index: int, key: str, value: Any):
227 """Update the `value` of `key` in the database, at `index`
229 Parameters
230 ----------
231 db_file : h5py.File
232 Opened file handle to the database HDF5 file
233 index : int
234 Index of the row to update
235 key : str
236 Field to update
237 value : Any
238 New value for the field
239 """
240 dataset = get_db(db_file)
241 update_value = dataset[index]
242 update_value[key] = value
243 dataset[index] = update_value
246def get_filtered_DB_mask(db_file: h5py.File, fields_values: dict[str, str | list[str]]) -> np.ndarray:
247 """Get a mask selecting the DB rows where the field values are equal to `fields_values` values.
249 Parameters
250 ----------
251 db_file : h5py.File
252 Opened file handle to the database HDF5 file
253 fields_values : Dict[str, str | List[str]]
254 Map from fields to allowed fields values. Rows where the fields value is not equal to
255 one of the field values are not selected.
256 Key: field name, eg "NAME", "USER"
257 values: field value, eg "Robert", ["Robert", "Roberta"]
259 Returns
260 -------
261 np.ndarray
262 Index mask array, True where the row's fields are equal to the `field_values`.
263 """
264 db = get_db(db_file)
265 total_mask = np.ones(shape=(db.shape[0],), dtype=bool)
266 for field, values in fields_values.items():
267 mask = np.zeros(shape=(db.shape[0],), dtype=bool)
268 value_list = values if isinstance(values, list) else [values]
269 for v in value_list:
270 mask |= db[field] == v
271 total_mask &= mask
272 return total_mask