Coverage for python/diagnostic_and_simulation_base/diagnostic_and_simulation_base.py: 47%
180 statements
« prev ^ index » next coverage.py v7.16.2, created at 2026-10-14 15:25 +0000
1import json
2import logging
3import multiprocessing
4import sys
5import textwrap
6import traceback
7import typing
8from importlib import resources as resources_py # BUXTON: this is deprecated and needs removing
9from multiprocessing.connection import Connection
10from pathlib import Path
11from time import time as time_py
13import f90nml
15from . import version_storage
16from .nested_dictionary import NestedDict
17from .utility import logging as util_logging
18from .utility.make_settings_json import make_settings_json
21def _create_script_nodes(
22 analysis_name: str,
23 pulseNo_write: int,
24 pulseNo_cal: int | None,
25 run_name: str,
26 run_description: str,
27 workflows: list[str],
28 link_best: bool,
29) -> float:
30 """Creates the MDSplus nodes. Returns how long the creation took [second].
32 A module level function taking nothing but primitives, because it is also the body of the
33 child process started by `start_mds_node_creation`, which must not reach back into the parent.
34 """
36 # Lazy loading of `standard_utility` because it's specific to Tokamak Energy.
37 import standard_utility as util
39 tic = time_py()
40 util.create_script_nodes(
41 script_name=analysis_name,
42 pulseNo_write=pulseNo_write,
43 pulseNo_cal=pulseNo_cal,
44 run_name=run_name,
45 run_info=run_description,
46 workflows=workflows,
47 link_best=link_best,
48 )
49 toc = time_py()
51 return toc - tic
54def _create_script_nodes_in_child(connection: Connection, **keyword_arguments: typing.Any) -> None:
55 """Runs `_create_script_nodes` in the child process, and reports the outcome to the parent.
57 The outcome goes back down `connection` because raising in a child process only prints to
58 stderr and sets an exit code, which does not tell the parent what went wrong. The message is a
59 few hundred bytes at most, so it always fits in the pipe buffer and the child never blocks
60 sending it.
61 """
63 try:
64 connection.send({"seconds": _create_script_nodes(**keyword_arguments)})
65 except BaseException:
66 connection.send({"traceback": traceback.format_exc()})
67 finally:
68 connection.close()
71class DiagnosticAndSimulationBase:
72 """Class for creating Diagnostic Analysis codes, such as: `Gas`, `Efit`, `Ppts`, `Zeff_pi`"""
74 log_records: list[logging.LogRecord] = []
76 @property
77 def log_string(self) -> str:
78 return util_logging.format_logs(self.log_records)
80 def __init__(
81 self,
82 pulseNo: int,
83 run_name: str,
84 run_description: str = "Standard run with default settings",
85 settings_path: str = "default",
86 write_to_mds: bool = True,
87 pulseNo_write: int | None = None,
88 analysis_name: str | None = None,
89 link_run_to_best: bool = False,
90 ) -> None:
91 """Class constructor
93 :param pulseNo: pulse number
94 :param run_name: run_name to save to MDSplus
95 :param run_description: help string for MDSplus Tree
96 :param settings_path: location where code inputs are stored
97 :param write_to_mds: flag to turn on / off writing to MDSplus
98 :param pulseNo_write: pulse number in which data is to be written, if different from the current pulse
99 :param analysis_name: for example "EFITP" is analysis using the "EFIT code"
100 :param link_run_to_best: BEST will be linked to the run
101 """
102 self.analysis_time_start = time_py()
104 # code_name: is the package name from "pip list" in upper case (note, this includes "_")
105 code_name = self.__class__.__module__.split(".")[0].upper() # BUXTON: I don't like this!!
107 # analysis_name
108 if analysis_name is None:
109 analysis_name = code_name
111 # Create a logger with the module name
112 # Get the logging_level
113 logging_level = util_logging.logger.getEffectiveLevel()
115 # Create logger
116 self.logger = logging.getLogger(analysis_name)
117 # Don't propagate messages to the root logger; `mdsthin` uses the root logger and this avoids double logging
118 self.logger.propagate = False
119 # Set the logging level
120 self.logger.setLevel(level=logging_level)
122 # Check if the logger already has handlers to avoid adding duplicates
123 if not self.logger.handlers:
124 # Create and add handler to emit the logs to "stadnard out",
125 # i.e. printing to the terminal (standard behaviour)
126 handler_to_standard_out = logging.StreamHandler(sys.stdout)
127 handler_to_standard_out.setFormatter(util_logging.CustomFormatter())
128 self.logger.addHandler(handler_to_standard_out)
130 # Create and add handler to emit the log "records" into `self.log_records`
131 handler_to_list = util_logging.EmitLogToListHandler(self.log_records)
132 self.logger.addHandler(handler_to_list)
134 # `log_records` is a class-level list we "clear" it when a new instance is created
135 self.log_records.clear()
137 # Store inputs
138 self.pulseNo = pulseNo
139 self.code_name = code_name
140 self.analysis_name = analysis_name
141 self.run_name = run_name.upper()
142 self.run_description = run_description
143 self.write_to_mds = write_to_mds
144 self.link_run_to_best = link_run_to_best
146 # Determine which pulseNo to write to and store in class object
147 if pulseNo_write is None:
148 pulseNo_write = pulseNo
149 self.pulseNo_write = pulseNo_write
151 # Test if "settings_path" is a directory.
152 # If directory doesn't exist, then treat as "relative path"
153 if not (Path(settings_path).is_dir()):
154 python_module = self.code_name.lower()
155 settings_path = f"{resources_py.files(python_module)}/settings/{settings_path}"
157 # Test if settings directory exists
158 if not Path(settings_path).is_dir():
159 raise FileNotFoundError("settings directory not found")
161 # Store the resolved settings path
162 self.settings_path = settings_path
164 # Create empty settings dictionary
165 self.settings = {} # type: dict[str, typing.Any]
167 # State for the background MDSplus node creation, see `start_mds_node_creation`
168 self._mds_node_creation_process: multiprocessing.process.BaseProcess | None = None
169 self._mds_node_creation_connection: Connection | None = None
170 self._mds_nodes_created = False
172 self._load_settings_from_files()
174 # Create results dictionary, with the stuff we already know
175 # Important: results is a 1:1 mapping to MDSplus data-strucutre
176 # TO-DO: Change this to a pre-populated dictionary from the *.csv file
177 self.results = NestedDict()
178 self.results["CODE_VERSION"]["COMPUTER"] = version_storage.__computer__
179 self.results["CODE_VERSION"]["DATETIME"] = version_storage.__datetime__
180 self.results["CODE_VERSION"]["GIT_ID"] = version_storage.__git_short_hash__
181 self.results["CODE_VERSION"]["LIBRARY"] = version_storage.__python_library__
182 self.results["CODE_VERSION"]["PYTHON"] = version_storage.__python__
183 self.results["CODE_VERSION"]["USER"] = version_storage.__user__
184 self.results["CODE_VERSION"]["VERSION"] = version_storage.__version__
186 def __repr__(self) -> str:
187 """Print to screen"""
188 string_output = ""
189 string_output += "╔═════════════════════════════════════════════════════════════════════════════╗\n"
190 string_output += f"║ {f' <{self.__class__.__name__}>':<75} ║\n"
191 string_output += f"║ {f' {version_storage.__version__}':<75} ║\n"
192 string_output += f"║ {' ':<75} ║\n"
193 string_output += f"║ {' pulseNo = ' + f'{self.pulseNo:_}':<75} ║\n"
194 string_output += f"║ {' pulseNo_write = ' + f'{self.pulseNo_write:_}':<75} ║\n"
195 string_output += f"║ {' run_name = ' + str(self.run_name):<75} ║\n"
196 string_output += f"║ {' run_description = ' + str(self.run_description):<75} ║\n"
197 string_output += f"║ {' settings_path = ...':<75} ║\n"
198 return string_output
200 def _load_settings_from_files(self) -> None:
201 """Look in the "settings_path" directory and load all settings files.
202 Will recursively load *.json and *.nml files
204 TODO: add *.csv reader #SUNDAR - is ths even needed?
205 """
207 settings_path = self.settings_path
209 # Load *.json settings, including sub-directories
210 for file in Path(settings_path).glob("**/*.json"):
211 with open(file, "r") as file_id:
212 step_name = f'Loading settings from: "{file.name}"'
213 try:
214 relative_path = str(file.relative_to(settings_path))
215 self.settings[relative_path] = json.load(file_id)
216 self.logger.info(msg=step_name)
217 except Exception as exception_obj:
218 self.logger.exception(msg=step_name)
219 raise exception_obj
221 # Load *.nml (namelist) settings, including sub-directories
222 for file in Path(settings_path).glob("**/*.nml"):
223 with open(file, "r") as file_id:
224 step_name = f"Loading settings from {file.name}"
225 try:
226 relative_path = str(file.relative_to(settings_path))
227 self.settings[relative_path] = f90nml.read(file_id)
228 self.logger.info(msg=step_name)
229 except Exception as exception_obj:
230 self.logger.exception(msg=step_name)
231 raise exception_obj
233 def _get_pulseNo_cal(self) -> int | None:
234 """The calibration pulse number, taken from this code's MDSplus settings file.
236 Returns `None` when the code has no MDSplus settings file, or the file does not name a
237 calibration pulse.
238 """
240 mdsplus_settings_file: str = f"{self.analysis_name}_mdsplus_settings.json"
241 pulseNo_cal: int | None
242 if mdsplus_settings_file in self.settings:
243 pulseNo_cal = self.settings[mdsplus_settings_file].get("calibration", {}).get("pulse", None)
244 else:
245 pulseNo_cal = None
247 return pulseNo_cal
249 def _mds_node_arguments(self, workflows: list[str] | None = None) -> dict[str, typing.Any]:
250 """Everything the node creation needs, taken from the settings as they stand right now.
252 This is deliberately a snapshot. The settings belong to the user right up until the
253 analysis starts - which `MAG` or `PSU2COIL` run to read from, for example - so the node
254 creation has to use the values that were in force when it was started, not whatever they
255 become afterwards. Once the analysis is running they are fixed, which is what makes it
256 safe to hand this to another process and not look at it again.
258 :param workflows: names of the input codes to create the `INPUT.WORKFLOW` structure for.
259 When `None` the names are read out of `self.results`, which means the results must
260 already have been mapped to the database structure.
261 """
263 if workflows is None:
264 workflows = list(self.results["INPUT"]["WORKFLOW"].keys())
266 return {
267 "analysis_name": self.analysis_name,
268 "pulseNo_write": self.pulseNo_write,
269 "pulseNo_cal": self._get_pulseNo_cal(),
270 "run_name": self.run_name,
271 "run_description": self.run_description,
272 "workflows": workflows,
273 "link_best": self.link_run_to_best,
274 }
276 def _create_mds_nodes(self, workflows: list[str] | None = None) -> None:
277 """Creates the MDSplus nodes which `_write_data_to_mds` later writes the data into,
278 here, in this process.
280 Everything this needs comes from the settings, which are read in the constructor, so it
281 does not have to wait for the analysis to finish. See `start_mds_node_creation`.
283 :param workflows: see `_mds_node_arguments`.
284 """
286 seconds = _create_script_nodes(**self._mds_node_arguments(workflows))
287 self._mds_nodes_created = True
289 # Logged so that we can see how much of the node creation was hidden behind the analysis
290 self.logger.info(msg=f"MDSplus nodes created; {seconds * 1e3:,.2f}ms")
292 def _write_data_to_mds(self) -> None:
293 """Writes the data into the MDSplus nodes created by `_create_mds_nodes`.
295 The nodes must already exist, so either `_create_mds_nodes` or `wait_for_mds_node_creation`
296 has to have been called first.
297 """
299 # Lazy loading of `standard_utility` because it's specific to Tokamak Energy.
300 import standard_utility as util
302 # Add settings files to results.
303 # We do this right at the end, as they can be programatially changed, e.g. for scans
304 self.results["INPUT"]["SETTINGS"] = make_settings_json(data=self.settings, json_indent=2)
306 util.write_script_data(
307 script_name=self.analysis_name,
308 pulseNo_write=self.pulseNo_write,
309 data_to_write=self.results.to_dictionary(),
310 pulseNo_cal=self._get_pulseNo_cal(),
311 run_name=self.run_name,
312 run_description=self.run_description,
313 force_write=True,
314 )
316 def _write_to_mds(self) -> None:
317 """Creates the MDSplus nodes and then writes the data into them, one after the other.
319 A code which wants to overlap the node creation with its own computation should instead
320 call `start_mds_node_creation` early on, and then `wait_for_mds_node_creation` followed by
321 `_write_data_to_mds` at the end.
322 """
324 self._create_mds_nodes()
325 self._write_data_to_mds()
327 def start_mds_node_creation(self, workflows: list[str] | None = None) -> None:
328 """Starts creating the MDSplus nodes in a separate process.
330 Creating the nodes is talking to the MDSplus server, and needs none of the results, so it
331 can run while the analysis is still computing. `wait_for_mds_node_creation` must be called
332 before writing, both to make sure the nodes exist and to check that the creation worked.
334 **A separate process, not a thread.** The analysis spends most of its wall time inside
335 `gsfit_rs`, which holds the GIL for the whole call, so a thread makes no progress during
336 exactly the part of the run we want to hide the node creation behind. Measured on a
337 480 time-slice run: as a thread the node creation ran for 279 s of wall clock but only
338 progressed for about 95 s of it, all of it during the database reads.
340 **`fork`, not `spawn`.** `spawn` re-imports the main module in the child, and GSFit's entry
341 scripts run the analysis at module level with no `if __name__ == "__main__"` guard, so the
342 child would start a second reconstruction. `forkserver` re-imports it too, so it is no help
343 either. `fork` does not re-import, and is safe here because nothing has started a thread yet
344 at the point the analysis calls this.
346 **Only where `fork` can be used.** Windows has no `fork` at all, and on macOS forking a
347 process which has already initialised Apple's frameworks can abort the child. On those
348 platforms the nodes are created here and now instead, which costs the overlap with the
349 analysis but is correct everywhere.
351 **Why the test is written the way it is.** The obvious spelling - ask
352 `multiprocessing.get_all_start_methods()` whether `fork` is available, and return early
353 when it is not - is wrong twice over. It lists `fork` on macOS, because the kernel has it,
354 so it does not answer the question we actually care about; and type checkers cannot see
355 through it. Typeshed declares the `fork` half of `multiprocessing` under
356 `if sys.platform != "win32"`, so on a Windows run `get_context("fork")` resolves to plain
357 `BaseContext`, which has no `Process`.
359 Hence a `sys.platform` test, which every checker narrows on. It is written as a positive
360 `if <fork is usable>: ... else: ...` rather than as an early return for the platforms which
361 cannot, because `mypy --warn-unreachable` suppresses its complaint for a block that is
362 unreachable *because of* a platform check, but reports the code following an early return
363 it knows is always taken. The two spellings behave identically at run time; only this one
364 type-checks cleanly on all three platforms, with no suppressions.
366 Note that `sys.platform == "win32" or sys.platform == "darwin"` is deliberate in place of
367 the tidier `sys.platform in ("win32", "darwin")`. `mypy` narrows on `==`, `!=` and
368 `.startswith()` only, not on `in`, so the tuple form silently stops working for it.
370 :param workflows: see `_mds_node_arguments`.
371 """
373 if sys.platform != "win32" and sys.platform != "darwin":
374 context = multiprocessing.get_context("fork")
375 self._mds_node_creation_connection, connection_child = context.Pipe(duplex=False)
377 # `daemon=True`, so that an analysis which fails before it reaches the writing cannot
378 # leave a process behind that nobody is going to wait for
379 self._mds_node_creation_process = context.Process(
380 target=_create_script_nodes_in_child,
381 args=(connection_child,),
382 kwargs=self._mds_node_arguments(workflows),
383 name="mds_node_creation",
384 daemon=True,
385 )
386 self._mds_node_creation_process.start()
388 # The child now holds the only sending end. Closing ours means that if the child dies
389 # without sending anything, the pipe reports end-of-file rather than waiting for a
390 # writer which no longer exists
391 connection_child.close()
392 else:
393 self._create_mds_nodes(workflows)
395 def wait_for_mds_node_creation(self, workflows: list[str] | None = None) -> None:
396 """Waits for the MDSplus node creation to finish, and raises if it did not succeed.
397 The data must not be written unless this returns without raising.
399 If `start_mds_node_creation` was never called then the nodes are created here instead, so
400 this is always safe to call before writing. The same applies when it was called but had to
401 create them synchronously, in which case there is nothing left to wait for.
403 :param workflows: only used when the nodes have to be created here; see
404 `_mds_node_arguments`.
405 """
407 if self._mds_nodes_created:
408 return
410 process = self._mds_node_creation_process
411 connection = self._mds_node_creation_connection
412 self._mds_node_creation_process = None
413 self._mds_node_creation_connection = None
415 if process is None:
416 self._create_mds_nodes(workflows=workflows)
417 return
419 process.join()
421 # The child sends its result before exiting, so by now it is sitting in the pipe. `poll`
422 # rather than a blocking `recv`, because a child which was killed sends nothing at all -
423 # and for that same child `poll` reports readable but `recv` then raises `EOFError`, which
424 # is the end of the pipe rather than an error worth propagating
425 result: dict[str, typing.Any] = {}
426 if connection is not None:
427 try:
428 if connection.poll():
429 result = connection.recv()
430 except EOFError:
431 pass
432 connection.close()
434 if "traceback" in result:
435 self.logger.error(msg=f"Creating the MDSplus nodes failed; not writing\n{result['traceback']}")
436 raise RuntimeError(f"Creating the MDSplus nodes failed:\n{result['traceback']}")
438 if "seconds" not in result:
439 # Killed, or died without managing to say why. Either way the nodes are not trustworthy
440 self.logger.error(msg="Creating the MDSplus nodes failed; not writing")
441 raise RuntimeError(
442 f"Creating the MDSplus nodes failed: the process exited with code {process.exitcode} "
443 "without reporting a result"
444 )
446 # Logged so that we can see how much of the node creation was hidden behind the analysis
447 self.logger.info(msg=f"MDSplus nodes created; {result['seconds'] * 1e3:,.2f}ms")
449 def run_with_log(
450 self,
451 operator: typing.Callable[..., typing.Any],
452 message: str | None,
453 args: dict[typing.Any, typing.Any] | None = None,
454 ) -> typing.Any:
455 if message is None:
456 step_name = ""
457 else:
458 step_name = message
460 try:
461 if args is None:
462 ret_data = operator()
463 else:
464 ret_data = operator(**args)
465 self.logger.info(msg=step_name)
466 return ret_data
467 except Exception as exception_obj:
468 self.logger.exception(msg=step_name)
469 raise exception_obj