Coverage for src / lstautorta / Auto_RTA.py: 23%
178 statements
« prev ^ index » next coverage.py v7.13.5, created at 2026-08-10 11:56 +0000
1# Copyright 2026 CNRS
2# This software is distributed under the terms of the CeCILL-C free software license.
4"""Automatically start/stop RTA reconstruction pipeline for new runs during a observation night.
6New runs are found querying the TCU pymongo database regularly.
7This script:
8- copies the conda environment and reconstruction model to the RAM of the slurm nodes
9- queries the TCU pymongo database for new runs regularly and for each new run:
10 - stops the r0->dl1 daemons for the previous runs
11 - starts the r0->dl1 daemons for the new run, using a static configuration (CDB configuration)
12 from disk, and writing the dynamic configuration to disk as well.
13 - starts the "engineering gui" scripts allowing to monitor data processing for the run.
14- stops at a fixed hour according to its configuration file
15- cleans the RAM of slurm worker nodes
16"""
18import argparse
19import datetime
20import json
21import logging
22import shlex
23import signal
24import subprocess as sp
25import time
26from pathlib import Path
27from subprocess import CalledProcessError, CompletedProcess
28from threading import Thread
29from typing import Annotated, NamedTuple
31from annotated_types import Gt, Le
32from pymongo.errors import PyMongoError
34from lstautorta.config.configuration import (
35 AutoRTAConfiguration,
36 DataStreamConnectionConfiguration,
37 ObservationParameters,
38)
39from lstautorta.observation_data import ObsInfo, get_current_run_info
40from lstautorta.paths import RecoPathStructure
41from lstautorta.shared_observation import DEFAULT_SHARED_OBS_PATH, write_shared_obs
42from lstautorta.utils.logging import LOGGING_LEVELS_DICT, init_logging
43from lstautorta.utils.slurm import (
44 job_statistics_from_squeue_output,
45 parse_slurm_job_ID,
46 subprocess_run_and_raise_exception_on_error,
47)
50class ConnectionJobInfo(NamedTuple):
51 """Store a slurm computing daemon's information."""
53 tel_id: Annotated[int, Gt(0)]
54 hostname: str
55 port: Annotated[int, Gt(0), Le(65535)]
56 slurm_reservation: str
57 slurm_node: str
60def assign_worker_to_data_connection(
61 slurm_nodes: dict[str, list[str]],
62 tel_ids_to_data_servers: dict[Annotated[int, Gt(0)], list[DataStreamConnectionConfiguration]],
63) -> list[ConnectionJobInfo]:
64 """Assign a worker node from `slurm_nodes` to each data server connection in `tel_ids_to_data_servers`.
66 Parameters
67 ----------
68 slurm_nodes : Dict[str, List[str]]
69 Mapping from slurm reservation to slurm nodes, see AutoRTAConfiguration.slurm_nodes field.
70 tel_ids_to_data_servers : Dict[Annotated[int, Gt(0)], List[DataStreamConnectionConfiguration]]
71 Mapping from telescope ID to data servers connections, see AutoRTAConfiguration.tel_ids_to_data_servers
73 Returns
74 -------
75 Dict[Tuple[int, str, str], Tuple[str, str]]
76 Map from (tel_id, hostname, port) to (slurm_reservation, slurm_nodename)
77 """
78 # "Flatten" the slurm nodes information into a list of tuple (reservation, node)
79 node_list = [(reservation, node) for reservation, nodes in slurm_nodes.items() for node in nodes]
80 tel_to_node_map = []
81 node_idx = 0
82 # Now make the tuples from telescope connections and available nodes.
83 for tel_id, tel_data_server_connections in tel_ids_to_data_servers.items():
84 for connection_idx, connection in enumerate(tel_data_server_connections):
85 tel_to_node_map.append(
86 ConnectionJobInfo(
87 tel_id=tel_id,
88 hostname=connection.hostname,
89 port=connection.port,
90 slurm_reservation=node_list[node_idx + connection_idx][0],
91 slurm_node=node_list[node_idx + connection_idx][1],
92 )
93 )
95 node_idx += len(tel_data_server_connections)
96 return tel_to_node_map
99def srun_cmd_worker_nodes(
100 connection_jobs_info: list[ConnectionJobInfo],
101 cmd: str,
102 additional_slurm_params: list[str] | None = None,
103 error_level=LOGGING_LEVELS_DICT["CRITICAL"],
104) -> list[CompletedProcess]:
105 """Submit a slurm command with srun on all worker nodes in `connection_jobs_info`.
107 Parameters
108 ----------
109 connection_jobs_info : List[ConnectionJobInfo]
110 List of connection job information: `cmd` will be executed for each node entry in this list.
111 cmd : str
112 Command to run with srun.
113 additional_slurm_params : List[str] or None
114 List of additional slurm jobs parameters, separated between args and values, for instance ["--mem", "20G"]
115 Optionnal, default is None.
117 Returns
118 -------
119 List[subprocess.CompletedProcess]
120 List of completed processes.
121 """
122 completed_processes = []
123 for job_info in connection_jobs_info:
124 srun_cmd = " ".join(
125 [
126 "srun",
127 *(additional_slurm_params if additional_slurm_params is not None else ""),
128 f"--reservation={job_info.slurm_reservation}",
129 f"--nodelist={job_info.slurm_node}",
130 cmd,
131 ]
132 )
133 logging.info(f"Running {srun_cmd}")
134 completed_processes.append(
135 subprocess_run_and_raise_exception_on_error(
136 shlex.split(srun_cmd),
137 success_log_string=f"Success on node {job_info.slurm_node}",
138 failure_log_string=f"Failure on node {job_info.slurm_node} with {srun_cmd}",
139 error_level=error_level,
140 log_level=LOGGING_LEVELS_DICT["DEBUG"],
141 )
142 )
143 return completed_processes
146def scancel_jobs(
147 job_ids: list[int], signal: signal.Signals, delay_s: float | None = None, ignore_error: bool = False
148) -> None:
149 """Run `scancel on `job_ids`, sending `signal` after `delay_s` seconds.
151 Use the --quiet argument of scancel to not raise error if the jobs are already stopped.
153 Parameters
154 ----------
155 job_ids : List[str]
156 List of job ids to scancel
157 signal : signal.Signals
158 Signal to send with scancel
159 delay_s : float | None
160 Amount of time in second to wait before performing the scancel
161 ignore_error : bool, optional
162 If True, the scancel is run directly with subprocess.srun, and any error happening in the subprocess
163 is simply ignored. This is usefull when running scancel -s KILL after a scancel -s INT: if the SIGINT
164 stopped the job already, the SIGKILL would have exit code 1 even with --quiet, but we want to ignore
165 the error in this case.
166 """
167 # note: --full or -f is required for r0_dl1 daemons to receive signal
168 scancel_cmd = " ".join(["scancel", "--full", "-s", str(signal), *[str(job_id) for job_id in job_ids]])
169 if delay_s is not None:
170 time.sleep(delay_s)
171 if ignore_error:
172 try:
173 sp.run(shlex.split(scancel_cmd), capture_output=True, text=True, check=True)
174 logging.debug(f"Stopping jobs with {scancel_cmd}")
175 except CalledProcessError as error:
176 logging.info(
177 f"Ignoring error of {scancel_cmd} caused by jobs already been stopped. Error info:\nstdout: {error.stdout}\nstderr: {error.stderr}"
178 )
179 else:
180 subprocess_run_and_raise_exception_on_error(
181 shlex.split(scancel_cmd),
182 success_log_string=f"Stopping job with {scancel_cmd}",
183 failure_log_string=f"FAILURE to stop job with {scancel_cmd}",
184 error_level=LOGGING_LEVELS_DICT["ERROR"],
185 log_level=LOGGING_LEVELS_DICT["DEBUG"],
186 )
189def stop_rta(slurm_reservations: list[str], slurm_account: str, r0dl1_job_name: str) -> Thread | None:
190 """Stop the r0dl1 daemons.
192 The r0dl1 jobs are immediately send a SIGINT signal, which should tell them to
193 gracefully shut down. A SIGKILL is also scheduled to run 10 seconds later to
194 ensure the jobs are indeed stopped.
196 Parameters
197 ----------
198 slurm_reservations : List[str]
199 List of slurm reservation to search for r0dl1 daemons.
200 slurm_account : str
201 Slurm account to use when searching for the r0dl1 daemons.
202 r0dl1_job_name : str
203 Name of the r0dl1 jobs in the CDB configuration
205 Returns
206 -------
207 stop_thread : Thread | None
208 Started thread that will SIGKILL the r0dl1 jobs after 10 secs.
209 None if there were no jobs to stop.
210 """
211 job_ids = parse_slurm_job_ID(slurm_reservations, slurm_account, r0dl1_job_name)
212 logging.info(f"Found r0dl1 job ids {job_ids} to stop.")
213 if job_ids:
214 # immediately send the SIGINT
215 scancel_jobs(job_ids, signal.SIGINT, None)
216 # start a detached thread to send the SIGKILL in 10 seconds.
217 # this allows autorta to continue and start new run immediately, while jobs are shuting down.
218 stop_thread = Thread(target=scancel_jobs, args=(job_ids, signal.SIGKILL, 10.0, True))
219 stop_thread.start()
220 return stop_thread
221 return None
224def nuke_rta(slurm_account: str):
225 """Hard stop of All RTA jobs: scancels all jobs of `slurm_account`.
227 Parameters
228 ----------
229 slurm_account : str
230 slurm account
231 """
232 nuke_rta_cmd = " ".join(["scancel", "-u", slurm_account])
233 subprocess_run_and_raise_exception_on_error(
234 shlex.split(nuke_rta_cmd),
235 f"Stopped RTA with {nuke_rta_cmd}",
236 f"Could not stop RTA with {nuke_rta_cmd}",
237 error_level=LOGGING_LEVELS_DICT["CRITICAL"],
238 log_level=LOGGING_LEVELS_DICT["WARNING"],
239 )
242def write_reco_manager_observation_config(
243 obs_info: ObsInfo,
244 obs_dir: Path,
245 night_path_structure: RecoPathStructure,
246 auto_rta_config: AutoRTAConfiguration,
247 output_path: Path,
248) -> None:
249 """Write the observation configuration for the reco-manager.
251 Parameters
252 ----------
253 obs_info : ObsInfo
254 Observation parameters from observation DB
255 obs_dir : Path
256 Path to the observation data directory
257 night_path_structure : RecoPathStructure
258 Path structure of the night
259 auto_rta_config : AutoRTAConfiguration
260 Configuration of Auto RTA
261 output_path: Path
262 Path where to write the hiperta_stream_start configuration
263 """
264 hiperta_obs_config = ObservationParameters.model_validate(
265 {
266 "sb_id": 1, # no scheduling block in LST
267 "obs_id": obs_info.obs_id,
268 "tel_id": 1, # only 1 tel
269 "RA_pointing": obs_info.RA,
270 "DEC_pointing": obs_info.DEC,
271 "dl1_dir": str(night_path_structure.dl1_dir(obs_dir)),
272 "dl2_dir": str(night_path_structure.dl2_dir(obs_dir)),
273 "dl3_dir": str(night_path_structure.dl3_dir(obs_dir)),
274 "log_dir": str(night_path_structure.log_dir(obs_dir)),
275 "reco_manager_log_file": str(night_path_structure.log_dir(obs_dir) / "hiperta_stream_start.log"),
276 "data_stream_connections": auto_rta_config.tel_ids_to_data_servers[1], # only do tel ID 1 for LST 1
277 "slurm_nodelists": auto_rta_config.slurm_nodes,
278 }
279 )
280 with output_path.open("w") as obs_config_f:
281 obs_config_f.write(hiperta_obs_config.model_dump_json(indent=4))
284def main() -> None:
285 """Entrypoint of Auto_RTA.
287 - parse the available slurm nodes
288 - check that the slurm nodes are in "connected" network mode
289 - copy conda environment to slurm nodes (usually in /dev/shm/ (RAM))
290 - query database for current RUN and for each run:
291 - stops previous RTA reconstruction slurm jobs
292 - starts new R0-DL1 jobs for the new RUN with static configuration (CDB config) and
293 a newly written dynamic config
294 - starts engineering GUI plotting scripts
295 - stops at a fixed hour set in configuration, after cleaning slurm nodes memory.
296 """
297 start_time = datetime.datetime.now(datetime.UTC)
299 # initially write log where the script is called (home if cron job)
300 # so that errors can be logged if we can't parse the config
301 init_logging(log_level="DEBUG", log_filename="lstautorta.log")
303 # Load configuration
304 parser = argparse.ArgumentParser(
305 description="Automatic Starting of the RTA Reconstruction for an observation night",
306 formatter_class=argparse.ArgumentDefaultsHelpFormatter,
307 )
308 parser.add_argument(
309 "-c", "--config", dest="config", type=str, required=True, help="LST auto RTA configuration file."
310 )
311 args = parser.parse_args()
312 with open(args.config) as config_file:
313 config = AutoRTAConfiguration.model_validate_json(config_file.read())
315 # Create night data directory
316 path_structure = RecoPathStructure(config.data_dir, start_time)
317 path_structure.create_night_data_dir()
318 # now update logging to write log in night's directory
319 init_logging(
320 log_level=config.log_level,
321 log_filename=path_structure.night_data_dir / "lstautorta.log",
322 )
324 # get the stop time during today
325 stop_time = start_time.replace(
326 hour=config.stop_time_UTC_hours, minute=config.stop_time_UTC_minutes, second=0, microsecond=0
327 )
328 # If we are passed it, it is actually tomorrow
329 if stop_time < start_time:
330 stop_time += datetime.timedelta(days=1)
332 logging.info("Start RTA at " + str(start_time))
333 logging.info(f"Found worker nodes : {config.slurm_nodes}")
335 connection_jobs_info = assign_worker_to_data_connection(config.slurm_nodes, config.tel_ids_to_data_servers)
336 for job_info in connection_jobs_info:
337 logging.info(
338 f'Will run tel {job_info.tel_id} connection {{"hostname": {job_info.hostname}, "port": {job_info.port}}} r0->dl1 job on {job_info.slurm_node} (reservation {job_info.slurm_reservation})'
339 )
341 logging.info("Reading R0-DL1 jobnames from CDB configuration")
342 # load CDB configuration, checking it exists, and parsing r0_dl1_job_name
343 with open(config.hiperta_CDB_config_file) as CDB_f:
344 CDB_config = json.load(CDB_f)
345 r0_dl1_job_name = CDB_config["r0_dl1_params"]["r0_dl1_job_name"]
346 logging.info(f"Found job name: {r0_dl1_job_name}")
348 if config.copy_env:
349 # For all srun commands here: error level is critical because RTA can not run if environment can not be copied
350 logging.info("Loading the environment on worker nodes for the night")
351 logging.info(f"Cleaning previous content in {config.env_archive_extraction_path}")
352 srun_cmd_worker_nodes(
353 connection_jobs_info,
354 f"rm -rf {config.env_archive_extraction_path}",
355 error_level=LOGGING_LEVELS_DICT["CRITICAL"],
356 )
357 logging.info(f"Cleaning previous content in {config.models_archive_copy_path}")
358 srun_cmd_worker_nodes(
359 connection_jobs_info,
360 f"rm -rf {config.models_archive_copy_path}",
361 error_level=LOGGING_LEVELS_DICT["CRITICAL"],
362 )
363 logging.info(f"Extracting environment from {config.env_archive} to {config.env_archive_extraction_path}")
364 srun_cmd_worker_nodes(
365 connection_jobs_info,
366 f"bash -c 'mkdir {config.env_archive_extraction_path} && tar -xzf {config.env_archive} -C {config.env_archive_extraction_path}'",
367 ["--mem", "20G"],
368 error_level=LOGGING_LEVELS_DICT["CRITICAL"],
369 )
370 logging.info(
371 f"Copying reconstruction models from {config.models_archive_path} to {config.models_archive_copy_path}"
372 )
373 srun_cmd_worker_nodes(
374 connection_jobs_info,
375 f"cp -rf {config.models_archive_path} {config.models_archive_copy_path}",
376 ["--mem", "20G"],
377 error_level=LOGGING_LEVELS_DICT["CRITICAL"],
378 )
380 if config.check_node_connection:
381 comp_proc_ib0_cat = srun_cmd_worker_nodes(connection_jobs_info, "cat /sys/class/net/ib0/mode")
382 # If a node is not connected, raise error
383 if not all("connected" in process.stdout.strip() for process in comp_proc_ib0_cat):
384 raise RuntimeError(
385 "Not all RTA slurm nodes are connected to ib0! Found {}".format(
386 ", ".join(
387 [
388 f"{job_info.slurm_node}: {process.stdout.strip()}"
389 for job_info, process in zip(connection_jobs_info, comp_proc_ib0_cat)
390 ]
391 )
392 )
393 )
395 logging.info("RTA ready for the night !")
397 current_obs_info = ObsInfo(None, None, None, None, None, None)
398 while not (datetime.datetime.now(datetime.UTC) > stop_time):
399 # query TCU DB for observation
400 try:
401 obs_info = get_current_run_info(config.db_hostname, 10)
402 except PyMongoError:
403 logging.error("Error retrieving observation information from DB, ignoring...", exc_info=True)
404 obs_info = ObsInfo(None, None, None, None, None, None)
406 # Check observation data
407 obs_info_is_none = obs_info.RA is None or obs_info.DEC is None
408 obs_recent_enough = True # default value if we couldn't get an obs_info
409 if obs_info_is_none:
410 logging.warning("Queried observation information had no RA DEC.")
411 elif config.ignore_old_observation:
412 obs_info_tstart_datetime = datetime.datetime.fromtimestamp(obs_info.time_start_camera, datetime.UTC)
413 obs_info_time_delta = datetime.datetime.now(datetime.UTC) - obs_info_tstart_datetime
414 obs_recent_enough = obs_info_time_delta < datetime.timedelta(hours=4)
415 if not obs_recent_enough:
416 logging.info(
417 f"Queried observation {obs_info.obs_id} has start time {obs_info_tstart_datetime}, timedelta wrt now: {obs_info_time_delta}. Too old to start RTA"
418 )
420 if (not obs_info_is_none) and obs_recent_enough:
421 try:
422 write_shared_obs(DEFAULT_SHARED_OBS_PATH, obs_info)
423 except Exception:
424 logging.error("Failed to write shared observation file at %s", DEFAULT_SHARED_OBS_PATH, exc_info=True)
426 # Start RTA
427 if current_obs_info.obs_id != obs_info.obs_id and (not obs_info_is_none) and obs_recent_enough:
428 # We got a new observation!
429 logging.info(
430 f"Got new observation! ID: {obs_info.obs_id} - RA: {obs_info.RA} - DEC: {obs_info.DEC} - SOURCE.RA: {obs_info.source_RA} - SOURCE.DEC: {obs_info.source_DEC}"
431 )
433 try:
434 logging.info("Stopping RTA")
435 stop_rta(config.slurm_reservations, config.slurm_account, r0_dl1_job_name)
436 except Exception:
437 logging.exception("Could not stop RTA. NEXT RUN DATA MAY BE ACQUIRED BY PREVIOUS RUN DAEMONS !")
438 # Note: we could nuke_rta, but it would also kill DQ, SCI jobs, etc...
440 obs_dir = path_structure.create_observation_data_dirs(str(obs_info.obs_id), True)
441 logging.info(f"Created directories for obs {obs_info.obs_id} at {obs_dir}")
443 reco_manager_obs_config_path = (
444 path_structure.log_dir(obs_dir) / "hiperta_stream_start_observation_config.json"
445 )
446 logging.info(f"Writing reco-manager observation configuration at {reco_manager_obs_config_path}")
447 write_reco_manager_observation_config(
448 obs_info,
449 obs_dir,
450 path_structure,
451 config,
452 reco_manager_obs_config_path,
453 )
455 # Note: hiperta_stream_start has to be started on a worker node as well, because it
456 # will read the training r0dl1 configuration, which is in the model's archive copied to /dev/shm
457 # (it reads the r0dl1 configuration from the path set in the CDB config, that must point to /dev/shm)
458 # Note2: to run several commands with srun, need to wrap with bash -c '...'
459 hiperta_stream_start_cmd = " ".join(
460 [
461 f"bash -c 'export PATH={config.env_archive_extraction_path}/bin/:$PATH ; {config.env_archive_extraction_path}/bin/hiperta_stream_start",
462 "-c",
463 config.hiperta_CDB_config_file,
464 "-d",
465 str(reco_manager_obs_config_path),
466 "'",
467 ]
468 )
469 try:
470 # use 1st node to start hiperta_stream
471 # Note: hiperta_stream must run on a worker node as well, because it needs access to
472 # the r0dl1 training configuration, which is found form the path in the CDB configuration,
473 # which points to /dev/shm/model_archive/...
474 hiperta_stream_job_info = connection_jobs_info[0]
475 hiperta_stream_srun_cmd = " ".join(
476 [
477 "srun",
478 f"--reservation={hiperta_stream_job_info.slurm_reservation}",
479 f"--nodelist={hiperta_stream_job_info.slurm_node}",
480 hiperta_stream_start_cmd,
481 ]
482 )
483 logging.info(f"Starting hiperta_stream_start with {hiperta_stream_srun_cmd}")
484 subprocess_run_and_raise_exception_on_error(
485 shlex.split(hiperta_stream_srun_cmd),
486 success_log_string=f"hiperta_stream started with {hiperta_stream_job_info.slurm_node}",
487 failure_log_string=f"Failed to start hiperta_stream on {hiperta_stream_job_info.slurm_node} with {hiperta_stream_start_cmd}",
488 error_level=LOGGING_LEVELS_DICT[
489 "ERROR"
490 ], # we might want to continue even if we can't start reco-manager
491 log_level=LOGGING_LEVELS_DICT["DEBUG"],
492 )
493 # Only if we could start RTA we update the current obs_info
494 # Otherwise, we will loop, see a new observation again, stop RTA and re-start, etc ...
495 current_obs_info = obs_info
496 except sp.SubprocessError:
497 logging.exception(f"Failed to start reco-manager with {hiperta_stream_start_cmd}")
498 # continue anyway, obs_info is not updated so we will try again to start
500 # If not starting RTA: Query squeue for some statistics on the running jobs
501 else:
502 try:
503 running_jobs_info = subprocess_run_and_raise_exception_on_error(
504 shlex.split(f'squeue -u {config.slurm_account} --format="%T,%R"'),
505 failure_log_string="Could not parse slurm info while waiting for next observation",
506 error_level=LOGGING_LEVELS_DICT["ERROR"],
507 log_level=LOGGING_LEVELS_DICT["DEBUG"],
508 ).stdout
509 n_jobs, n_running, n_pending, n_request_node_not_available = job_statistics_from_squeue_output(
510 running_jobs_info
511 )
512 logging.info(
513 f"{config.slurm_account} job statistics: nb_jobs: {n_jobs} - nb_running: {n_running} - nb_pending: {n_pending} - nb_node_not_available: {n_request_node_not_available}"
514 )
515 except sp.SubprocessError:
516 pass # continue loop anyway
518 # sleep until we get an observation
519 logging.info(f"RTA waiting for next observation, current observation is {obs_info.obs_id}")
520 time.sleep(1.0)
522 logging.info("End of the night: Stop the RTA")
523 try:
524 logging.info("Stopping RTA")
525 sigkill_thread = stop_rta(config.slurm_reservations, config.slurm_account, r0_dl1_job_name)
526 if sigkill_thread is not None: # it is None if there are no jobs to kill
527 sigkill_thread.join(timeout=180.0) # 3 minutes timeout but job should end right after delay of 10 sec
528 else:
529 logging.warning("Found no r0_dl1 jobs to stop !")
530 except Exception:
531 logging.exception(f"Could not stop RTA normally, nuking all {config.slurm_account} jobs")
532 nuke_rta(config.slurm_account)
534 # Clean up node environment
535 logging.info("Cleaning nodes copied files")
536 logging.info(f"Cleaning content in {config.env_archive_extraction_path}")
537 srun_cmd_worker_nodes(
538 connection_jobs_info,
539 f"rm -rf {config.env_archive_extraction_path}",
540 error_level=LOGGING_LEVELS_DICT["CRITICAL"],
541 )
542 logging.info(f"Cleaning content in {config.models_archive_copy_path}")
543 srun_cmd_worker_nodes(
544 connection_jobs_info,
545 f"rm -rf {config.models_archive_copy_path}",
546 error_level=LOGGING_LEVELS_DICT["CRITICAL"],
547 )
548 logging.info("Done cleaning nodes.")
550 logging.info("RTA Done for the night. Good day !")
553if __name__ == "__main__": 553 ↛ 554line 553 didn't jump to line 554 because the condition on line 553 was never true
554 main()