Skip to content

Add option to ignore h5 context manager in preprocessing #1690

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
mmccrackan wants to merge 8 commits into
base: master
Choose a base branch
from
+94 −33
Open

Add option to ignore h5 context manager in preprocessing #1690

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
8 commits
Select commit Hold shift + click to select a range
c7eb357
add option to override h5 ctx temp file creation
mmccrackan Jul 3, 2026
15b4a4c
option to bypass h5 ctx manager for preproc
mmccrackan Jul 3, 2026
af154d2
remove args
mmccrackan Jul 3, 2026
38e805d
fix name
mmccrackan Jul 3, 2026
1072228
fix indent
mmccrackan Jul 3, 2026
632f09a
adding the changes required for make_atomic_filterbin_map.py
chervias Jul 3, 2026
8acefda
update argparse description
mmccrackan Jul 3, 2026
c074056
set ctx arg to default false
mmccrackan Jul 5, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
48 changes: 36 additions & 12 deletions sotodlib/preprocess/preprocess_util.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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.
Expand All @@ -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:
Expand All @@ -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:
Expand Down Expand Up @@ -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
Expand All @@ -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
-------
Expand Down Expand Up @@ -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'])
Expand All @@ -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.
Expand All @@ -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):
Expand All @@ -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:
Expand All @@ -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
Expand Down Expand Up @@ -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
-------
Expand Down Expand Up @@ -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"],
Expand Down Expand Up @@ -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)
Expand All @@ -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``.
Expand Down Expand Up @@ -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:
Expand Down Expand Up @@ -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:
Expand Down
19 changes: 13 additions & 6 deletions sotodlib/site_pipeline/make_atomic_filterbin_map.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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:
Expand Down Expand Up @@ -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)
Expand Down Expand Up @@ -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):
Expand Down Expand Up @@ -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

Expand Down
34 changes: 25 additions & 9 deletions sotodlib/site_pipeline/multilayer_preprocess_tod.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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.
Expand All @@ -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
-------
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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"
Expand Down Expand Up @@ -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 = []

Expand Down Expand Up @@ -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:
Expand Down Expand Up @@ -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


Expand All @@ -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:
Expand All @@ -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__':
Expand Down
Loading