diff --git a/sotodlib/preprocess/preprocess_util.py b/sotodlib/preprocess/preprocess_util.py index 88127ad3c..bea9e930f 100644 --- a/sotodlib/preprocess/preprocess_util.py +++ b/sotodlib/preprocess/preprocess_util.py @@ -1019,7 +1019,7 @@ def find_db(obs_id, configs, dets, context=None, logger=None): return dbexist -def cleanup_archive(configs, logger=None): +def cleanup_archive(configs, logger=None, use_h5_ctx=False): """This function finds the final preprocess archive file and deletes any datasets that are not found in the preprocess database. This helps avoid cases where the database writing was interrupted in a previous run. @@ -1030,6 +1030,9 @@ def cleanup_archive(configs, logger=None): Filepath or dictionary containing the preprocess configuration file. logger : PythonLogger Optional. Logger object or None will generate a new one. + use_h5_ctx : bool + Optional. Use h5 context manager to avoid metadata corruption and locked sqlite + dbs. For use on site-computing. """ if type(configs) == str: @@ -1056,7 +1059,9 @@ def cleanup_archive(configs, logger=None): key=lambda t: t[0])[1] db_datasets = [d['dataset'] for d in db.inspect()] - with H5ContextManager(latest_file, mode="r+") as f: + + h5_file_func = H5ContextManager if use_h5_ctx else h5py.File + with h5_file_func(latest_file, mode="r+") as f: keys = list(f.keys()) for key in keys: if key not in db_datasets: @@ -1126,7 +1131,7 @@ def get_preproc_group_out_dict(obs_id, configs, dets, context=None, subdir='temp def save_group_and_cleanup(obs_id, configs, context=None, subdir='temp', - logger=None, remove=False): + logger=None, remove=False, use_h5_ctx=False): """This function checks if any temporary files exist from a preprocessing run and will either add them to the config policy file and create an entry in the manifest db by calling ``cleanup_mandb``. If the file exists but @@ -1150,6 +1155,9 @@ def save_group_and_cleanup(obs_id, configs, context=None, subdir='temp', remove : bool Optional. Default is False. Whether to remove a file if found. Used when ``overwrite`` is True in driving functions. + use_h5_ctx : bool + Optional. Use h5 context manager to avoid metadata corruption and locked sqlite + dbs. For use on site-computing. Returns ------- @@ -1184,7 +1192,7 @@ def save_group_and_cleanup(obs_id, configs, context=None, subdir='temp', try: if not remove: cleanup_mandb(outputs_grp, (obs_id, g), - (None, None, None), configs, logger) + (None, None, None), configs, logger, use_h5_ctx=use_h5_ctx) else: # if we're overwriting, remove file so it will re-run os.remove(outputs_grp['temp_file']) @@ -1208,7 +1216,7 @@ def save_group_and_cleanup(obs_id, configs, context=None, subdir='temp', def cleanup_obs(obs_id, policy_dir, errlog, configs, context=None, - subdir='temp', remove=False): + subdir='temp', remove=False, use_h5_ctx=False): """For a given obs id, this function will search the policy_dir directory if it exists for any files with that obsnum in their filename. If any are found, it will run save_group_and_cleanup for that obs id. @@ -1230,6 +1238,9 @@ def cleanup_obs(obs_id, policy_dir, errlog, configs, context=None, remove : bool Optional. Default is False. Whether to remove a file if found. Used when ``overwrite`` is True in driving functions. + use_h5_ctx : bool + Optional. Use h5 context manager to avoid metadata corruption and locked sqlite + dbs. For use on site-computing. """ if os.path.exists(policy_dir): @@ -1241,7 +1252,8 @@ def cleanup_obs(obs_id, policy_dir, errlog, configs, context=None, if found: errors = save_group_and_cleanup(obs_id, configs, context, - subdir=subdir, remove=remove) + subdir=subdir, remove=remove, + use_h5_ctx=use_h5_ctx) if errors[0] is not None: with open(errlog, 'a') as f: @@ -1252,7 +1264,8 @@ def cleanup_obs(obs_id, policy_dir, errlog, configs, context=None, def preproc_or_load_group(obs_id, configs_init, dets, configs_proc=None, logger=None, overwrite=False, save_archive=False, save_proc_aman=True, compress=False, - skip_missing=False, ignore_cfg_check=False): + skip_missing=False, ignore_cfg_check=False, + use_h5_ctx=False): """ This function is expected to receive a single obs_id, and dets dictionary. The dets dictionary must match the grouping specified in the preprocess @@ -1302,6 +1315,9 @@ def preproc_or_load_group(obs_id, configs_init, dets, configs_proc=None, ignore_cfg_check : bool If True, do not attempt to validate that configs_init is the same as the config used to create the existing init db when running ``multilayer_load_and_preprocess``. + use_h5_ctx : bool + Optional. Use h5 context manager to avoid metadata corruption and locked sqlite + dbs. For use on site-computing. Returns ------- @@ -1470,7 +1486,8 @@ def preproc_or_load_group(obs_id, configs_init, dets, configs_proc=None, if save_archive: logger.info(f"Adding result to init db for {obs_id}: {group}") cleanup_mandb(out_dict_init, (obs_id, group), (None, None, None), - configs_init, logger=logger, overwrite=overwrite) + configs_init, logger=logger, overwrite=overwrite, + use_h5_ctx=use_h5_ctx) # Make init plots if make_lmsi_init: new_plots = os.path.join(configs_init["plot_dir"], @@ -1539,7 +1556,8 @@ def preproc_or_load_group(obs_id, configs_init, dets, configs_proc=None, if save_archive: logger.info(f"Adding result to proc db for {obs_id}: {group}") cleanup_mandb(out_dict_proc, (obs_id, group), (None, None, None), - configs_proc, logger=logger, overwrite=overwrite) + configs_proc, logger=logger, overwrite=overwrite, + use_h5_ctx=use_h5_ctx) if 'valid_data' in aman.preprocess: aman.preprocess.move('valid_data', None) aman.preprocess.merge(proc_aman) @@ -1562,7 +1580,8 @@ def preproc_or_load_group(obs_id, configs_init, dets, configs_proc=None, return aman, out_dict_init, out_dict_proc, (None, None, None) -def cleanup_mandb(out_dict, out_meta, errors, configs, logger=None, overwrite=False, db_manager=None): +def cleanup_mandb(out_dict, out_meta, errors, configs, logger=None, overwrite=False, + db_manager=None, use_h5_ctx=False): """Function to update the manifest db when data is collected from the ``preproc_or_load_group`` function. If used in an mpi framework this function is expected to be run from rank 0 after a ``comm.gather``. @@ -1601,6 +1620,9 @@ def cleanup_mandb(out_dict, out_meta, errors, configs, logger=None, overwrite=Fa db_manager : DbBatchManager, optional External database batch manager for optimized operations. If provided, uses the manager instead of creating individual connections. + use_h5_ctx : bool + Optional. Use h5 context manager to avoid metadata corruption and locked sqlite + dbs. For use on site-computing. """ if logger is None: @@ -1628,8 +1650,10 @@ def cleanup_mandb(out_dict, out_meta, errors, configs, logger=None, overwrite=Fa src_file = out_dict['temp_file'] - with H5ContextManager(dest_file, mode='a') as f_dest: - with H5ContextManager(src_file, mode='r') as f_src: + h5_file_func = H5ContextManager if use_h5_ctx else h5py.File + + with h5_file_func(dest_file, mode='a') as f_dest: + with h5_file_func(src_file, mode='r') as f_src: for dts in f_src.keys(): # If the dataset or group already exists, delete it to overwrite if overwrite and dts in f_dest: diff --git a/sotodlib/site_pipeline/make_atomic_filterbin_map.py b/sotodlib/site_pipeline/make_atomic_filterbin_map.py index 2ee8e007c..969000d17 100644 --- a/sotodlib/site_pipeline/make_atomic_filterbin_map.py +++ b/sotodlib/site_pipeline/make_atomic_filterbin_map.py @@ -110,6 +110,9 @@ class Cfg: Data type for timestreams dtype_map: str Data type for maps + use_h5_ctx : bool + Use h5 context manager to avoid metadata corruption and locked sqlite + dbs. For use on site-computing. """ def __init__( self, @@ -151,6 +154,7 @@ def __init__( wn_label: str = 'preprocess.noiseQ_mapmaking.std', apply_wobble: bool = True, compress: bool = True, + use_h5_ctx: bool = False, ) -> None: self.context = context self.preprocess_config = preprocess_config @@ -190,6 +194,7 @@ def __init__( self.wn_label = wn_label self.apply_wobble = apply_wobble self.compress = compress + self.use_h5_ctx = use_h5_ctx @classmethod def from_yaml(cls, path) -> "Cfg": with open(path, "r") as f: @@ -270,6 +275,8 @@ def main( verbose = args.verbose - args.quiet + use_h5_ctx = args.use_h5_ctx + recenter = None if args.center_at: recenter = mapmaking.parse_recentering(args.center_at) @@ -415,15 +422,15 @@ def main( for obs in obslists_arr: obs_id = obs[0][0] if len(preprocess_config)==1: - preprocess_util.cleanup_obs(obs_id, policy_dir_init, errlog[0], preprocess_config[0], subdir='temp', remove=False) + preprocess_util.cleanup_obs(obs_id, policy_dir_init, errlog[0], preprocess_config[0], subdir='temp', remove=False, use_h5_ctx=use_h5_ctx) else: - preprocess_util.cleanup_obs(obs_id, policy_dir_init, errlog[0], preprocess_config[0], subdir='temp', remove=False) - preprocess_util.cleanup_obs(obs_id, policy_dir_proc, errlog[1], preprocess_config[1], subdir='temp_proc', remove=False) + preprocess_util.cleanup_obs(obs_id, policy_dir_init, errlog[0], preprocess_config[0], subdir='temp', remove=False, use_h5_ctx=use_h5_ctx) + preprocess_util.cleanup_obs(obs_id, policy_dir_proc, errlog[1], preprocess_config[1], subdir='temp_proc', remove=False, use_h5_ctx=use_h5_ctx) # remove datasets from final archive file not found in db - preprocess_util.cleanup_archive(preprocess_config[0], L) + preprocess_util.cleanup_archive(preprocess_config[0], L, use_h5_ctx) if len(preprocess_config) > 1: - preprocess_util.cleanup_archive(preprocess_config[1], L) + preprocess_util.cleanup_archive(preprocess_config[1], L, use_h5_ctx) run_list = [] for oi, ol in enumerate(obslists_arr): @@ -509,7 +516,7 @@ def main( oid = outputs[ii][idx_prepoc]['db_data']['obs:obs_id'] group = [v for k, v in outputs[ii][idx_prepoc]['db_data'].items() if 'dets' in k] preprocess_util.cleanup_mandb(outputs[ii][idx_prepoc], (oid, group), (errors[ii], None, None), - preprocess_config[idx_prepoc], L) + preprocess_config[idx_prepoc], L, use_h5_ctx=use_h5_ctx) L.info("Done") return True diff --git a/sotodlib/site_pipeline/multilayer_preprocess_tod.py b/sotodlib/site_pipeline/multilayer_preprocess_tod.py index ff82a2e86..f12778393 100644 --- a/sotodlib/site_pipeline/multilayer_preprocess_tod.py +++ b/sotodlib/site_pipeline/multilayer_preprocess_tod.py @@ -32,7 +32,8 @@ def multilayer_preprocess_tod(obs_id: str, group: list, verbosity: int = 0, compress: bool = False, - overwrite: bool = False): + overwrite: bool = False, + use_h5_ctx: bool = False): """Meant to be run as part of a batched script, this function calls the preprocessing pipeline a specific Observation ID and group combination and saves the results in the ManifestDb specified in the configs. @@ -55,6 +56,9 @@ def multilayer_preprocess_tod(obs_id: str, Whether or not to compress the preprocessing h5 files. overwrite : bool If True, overwrite contents of temporary h5 files. + use_h5_ctx : bool + Use h5 context manager to avoid metadata corruption and locked sqlite + dbs. For use on site-computing. Returns ------- @@ -85,6 +89,7 @@ def multilayer_preprocess_tod(obs_id: str, overwrite=overwrite, save_archive=False, compress=compress, + use_h5_ctx=use_h5_ctx, ) return out_dict_init, out_dict_proc, errors @@ -191,7 +196,8 @@ def _main(executor: Union["MPICommExecutor", "ProcessPoolExecutor"], compress: bool = False, run_from_jobdb: bool = False, raise_error: bool = False, - pb_path: Optional[str] = None): + pb_path: Optional[str] = None, + use_h5_ctx: Optional[bool] = False): init_temp_subdir = "temp" proc_temp_subdir = "temp_proc" @@ -259,13 +265,15 @@ def _main(executor: Union["MPICommExecutor", "ProcessPoolExecutor"], for obs_id in obs_list: pp_util.cleanup_obs(obs_id, init_policy_dir, errlog, configs_init, - context_init, subdir=init_temp_subdir, remove=overwrite) + context_init, subdir=init_temp_subdir, remove=overwrite, + use_h5_ctx=use_h5_ctx) pp_util.cleanup_obs(obs_id, proc_policy_dir, errlog, configs_proc, - context_proc, subdir=proc_temp_subdir, remove=overwrite) + context_proc, subdir=proc_temp_subdir, remove=overwrite, + use_h5_ctx=use_h5_ctx) # remove datasets from final archive file not found in init db for config in (configs_init, configs_proc): - pp_util.cleanup_archive(config, logger) + pp_util.cleanup_archive(config, logger, use_h5_ctx) run_list = [] @@ -435,11 +443,11 @@ def _main(executor: Union["MPICommExecutor", "ProcessPoolExecutor"], logger.info(f"Adding future result to init db for {obs_id}: {group}") pp_util.cleanup_mandb(out_dict_init, out_meta, errors, configs_init, logger, overwrite, - db_manager=db_mgr_init) + db_manager=db_mgr_init, use_h5_ctx=use_h5_ctx) logger.info(f"Adding future result to proc db for {obs_id}: {group}") pp_util.cleanup_mandb(out_dict_proc, out_meta, errors, configs_proc, logger, overwrite, - db_manager=db_mgr_proc) + db_manager=db_mgr_proc, use_h5_ctx=use_h5_ctx) # update jobdb if jobdb_path is not None: @@ -559,6 +567,12 @@ def get_parser(parser=None): type=str, default=None ) + parser.add_argument( + '--use-h5-ctx', + help="Use h5 context manager to prevent file corruption and db lock errrors (for site-computing).", + type=bool, + default=False + ) return parser @@ -577,7 +591,8 @@ def main(configs_init: str, nproc: int = 4, run_from_jobdb: bool = False, raise_error: bool = False, - pb_path: Optional[str] = None): + pb_path: Optional[str] = None, + use_h5_ctx: Optional[bool] = False): rank, executor, as_completed_callable = get_exec_env(nproc) if rank == 0: @@ -598,7 +613,8 @@ def main(configs_init: str, compress=compress, run_from_jobdb=run_from_jobdb, raise_error=raise_error, - pb_path=pb_path) + pb_path=pb_path, + use_h5_ctx=use_h5_ctx) if __name__ == '__main__': diff --git a/sotodlib/site_pipeline/preprocess_tod.py b/sotodlib/site_pipeline/preprocess_tod.py index 111342894..1db35c575 100644 --- a/sotodlib/site_pipeline/preprocess_tod.py +++ b/sotodlib/site_pipeline/preprocess_tod.py @@ -95,6 +95,7 @@ def preprocess_tod(configs: Union[str, dict], verbosity: int = 0, compress: bool = False, overwrite: bool = False, + use_h5_ctx: bool = False ): """Meant to be run as part of a batched script, this function calls the preprocessing pipeline a specific Observation ID and group combination @@ -116,6 +117,9 @@ def preprocess_tod(configs: Union[str, dict], Whether or not to compress the preprocessing h5 files. overwrite : bool If True, overwrite contents of temporary h5 files. + use_h5_ctx : bool + Use h5 context manager to avoid metadata corruption and locked sqlite + dbs. For use on site-computing. Returns ------- @@ -141,6 +145,7 @@ def preprocess_tod(configs: Union[str, dict], overwrite=overwrite, save_archive=False, compress=compress, + use_h5_ctx=use_h5_ctx, ) return out_dict, errors @@ -162,7 +167,8 @@ def _main(executor: Union["MPICommExecutor", "ProcessPoolExecutor"], compress: bool = False, run_from_jobdb: bool = False, raise_error: bool = False, - pb_path: Optional[str] = None): + pb_path: Optional[str] = None, + use_h5_ctx: Optional[bool] = False): temp_subdir = "temp" @@ -209,10 +215,10 @@ def _main(executor: Union["MPICommExecutor", "ProcessPoolExecutor"], ) for obs_id in obs_list: pp_util.cleanup_obs(obs_id, policy_dir, errlog, configs, context, - subdir=temp_subdir, remove=overwrite) + subdir=temp_subdir, remove=overwrite, use_h5_ctx=use_h5_ctx) # remove datasets from final archive file not found in db - pp_util.cleanup_archive(configs, logger) + pp_util.cleanup_archive(configs, logger, use_h5_ctx) run_list = [] @@ -340,7 +346,7 @@ def _main(executor: Union["MPICommExecutor", "ProcessPoolExecutor"], futures.remove(future) pp_util.cleanup_mandb(out_dict, out_meta, errors, configs, - logger, overwrite, db_manager=db_manager) + logger, overwrite, db_manager=db_manager, use_h5_ctx=use_h5_ctx) # update jobdb if jobdb_path is not None: @@ -456,6 +462,12 @@ def get_parser(parser=None): type=str, default=None ) + parser.add_argument( + '--use-h5-ctx', + help="Use h5 context manager to prevent file corruption and db lock errors (for site-computing).", + type=bool, + default=False + ) return parser @@ -473,7 +485,8 @@ def main(configs: str, compress: bool = False, run_from_jobdb: bool = False, raise_error: bool = False, - pb_path: Optional[str] = None): + pb_path: Optional[str] = None, + use_h5_ctx: Optional[bool] = False): rank, executor, as_completed_callable = get_exec_env(nproc) if rank == 0: @@ -493,7 +506,8 @@ def main(configs: str, compress=compress, run_from_jobdb=run_from_jobdb, raise_error=raise_error, - pb_path=pb_path) + pb_path=pb_path, + use_h5_ctx=use_h5_ctx) if __name__ == '__main__': main_launcher(main, get_parser)