Software veto

.gitignore

@@ -29,6 +29,7 @@ custom_data
 test_input_data
 *.zip
 last_bootstrax_exception.txt
+bootstrax_exceptions
 
 # cProfile output
 *.prof

bin/bootstrax

@@ -15,13 +15,14 @@ How to use
 For more info, see the documentation:
 https://straxen.readthedocs.io/en/latest/bootstrax.html
 """
-__version__ = '2.0.0'
+__version__ = '3.0.0'
 
 import argparse
 import typing
 from datetime import datetime, timedelta, timezone
 import logging
 import multiprocessing
+import multiprocess
 import npshmex
 import os
 import os.path as osp
@@ -39,6 +40,7 @@ import straxen
 import threading
 import pandas as pd
 import typing as ty
+from immutabledict import immutabledict
 import daqnt
 import fnmatch
 from glob import glob
@@ -86,7 +88,12 @@ parser.add_argument(
 parser.add_argument(
     '--max_messages', type=int, default=10,
     help="number of max mailbox messages")
-
+parser.add_argument(
+    '--software_veto_overwrite', default=None,
+    help="Class name of veto plugin to apply (i.e.: RadialVeto). It's overwritten by rundoc value.")
+parser.add_argument(
+    '--test_input_folder', default=None,
+    help="Add to storage")
 
 actions = parser.add_mutually_exclusive_group()
 actions.add_argument(
@@ -109,8 +116,8 @@ args = parser.parse_args()
 
 # The folder that can be used for testing bootstrax (i.e. non production
 # mode). It will be written to:
-test_data_folder = ('/data/test_processed/' if
-                    os.path.exists('/data/test_processed/')
+test_data_folder = ('/data/test_pre_processed/' if
+                    os.path.exists('/data/test_pre_processed/')
                     else './bootstrax/')
 
 # Timeouts in seconds
@@ -270,6 +277,9 @@ if os.access(output_folder, os.W_OK) is not True:
     raise IOError(message)
 
 
+
+
+
 def new_context(cores=args.cores,
                 max_messages=args.max_messages,
                 timeout=500,
@@ -293,8 +303,11 @@ def new_context(cores=args.cores,
         # all other storage frontends except fo the test folder.
         context.storage = [context.storage[0],
                            strax.DataDirectory(output_folder)]
+        if args.test_input_folder is not None:
+            context.storage += [strax.DataDirectory(args.test_input_folder, readonly=True)]
         context.storage[0].readonly = True
         context.storage[0].local_only = True
+
     return context
 
 
@@ -316,6 +329,7 @@ if not args.undying:
 
 
 def main():
+
     if args.cores == -1:
         # Use all of the available cores on this machine
         args.cores = multiprocessing.cpu_count()
@@ -505,7 +519,7 @@ def keep_target(targets, compare_with, n_fails):
     return kept_targets
 
 
-def infer_target(rd: dict) -> dict:
+def infer_target(rd: dict, software_veto_on: bool = False) -> dict:
     """
     Check if the target should be overridden based on the mode of the DAQ for this run
     :param rd: rundoc
@@ -514,6 +528,20 @@ def infer_target(rd: dict) -> dict:
     targets = args.targets.copy()
     post_process = args.post_process.copy()
 
+    if software_veto_on:
+        # TODO study which targets_sv is better to have in targets or post
+        # TODO THIS is bad!!
+        _targets = list(targets)
+        _targets.append('raw_records_sv')
+        _targets.append('peaklets_sv')
+
+        _post = list(post_process)
+        _post.append('event_info_sv')
+
+        targets = tuple(_targets)
+        post_process = tuple(_post)
+
+
     if args.fix_target:
         return {'targets': strax.to_str_tuple(targets),
                 'post_processing': strax.to_str_tuple(post_process)}
@@ -579,6 +607,8 @@ def infer_target(rd: dict) -> dict:
                      f'processing up to {targets} and postprocessing '
                      f'to {post_process}')
 
+
+
     if targets is None or not len(targets):
         targets = 'raw_records'
     if post_process is None or not len(post_process):
@@ -1169,7 +1199,7 @@ def manual_fail(*, mongo_id=None, number=None, reason=''):
 def run_strax(run_id, input_dir, targets, readout_threads, compressor,
               run_start_time, samples_per_record, cores, max_messages, timeout,
               daq_chunk_duration, daq_overlap_chunk_duration, post_processing,
-              records_compressor, debug=False):
+              records_compressor, software_veto_name=None, debug=False):
     # Check mongo connection
     ping_dbs()
     # Clear the swap memory used by npshmmex
@@ -1193,9 +1223,19 @@ def run_strax(run_id, input_dir, targets, readout_threads, compressor,
                          timeout=timeout,
                          )
 
+        if software_veto_name is not None:
+            # If the software veto is on, let's register the correct extra plugins
+            st = software_veto_register_plugins(st, software_veto_name)
+
         for t in ('raw_records', 'records', 'records_nv', 'hitlets_nv'):
             # Set the (raw)records processor to the inferred one
             st._plugin_class_registry[t].compressor = records_compressor
+            try:
+                st._plugin_class_registry[t+'_sv'].compressor = records_compressor
+            except:
+                pass
+
+
 
         # Make a function for running strax, call the function to process the run
         # This way, it can also be run inside a wrapper to profile strax
@@ -1307,6 +1347,9 @@ def process_run(rd, send_heartbeats=args.production):
             # or use the test-dir:
             if not osp.exists(loc):
                 loc = os.path.join('/live_data/xenonnt_bootstrax_test/', run_id)
+            # or maybe it's in the /data dir:
+            if not osp.exists(loc):
+                loc = os.path.join('/data/xenonnt_bootstrax_test/', run_id)
 
         else:
             for dd in rd['data']:
@@ -1369,10 +1412,14 @@ def process_run(rd, send_heartbeats=args.production):
         except Exception as e:
             fail(f"Could not find start in datetime.datetime object: {str(e)}")
 
-        run_strax_config.update(infer_target(rd))
+        software_veto_on = software_veto_get_name(rd, args)['software_veto_name'] is not None
+        run_strax_config.update(infer_target(rd, software_veto_on))
         run_strax_config.update(infer_mode(rd))
+        run_strax_config.update(software_veto_get_name(rd, args))
         run_strax_config['debug'] = args.debug
-        strax_proc = multiprocessing.Process(
+
+        log.info(f"Processing run!")
+        strax_proc = multiprocess.Process(
             target=run_strax,
             kwargs=run_strax_config)
 
@@ -1642,7 +1689,52 @@ def cleanup_db():
             abandon(mongo_id=rd['_id'])
 
 
+
+
+###########################
+# Software veto functions #
+###########################
+
+
+
+def software_veto_get_name(rd, args=None):
+    """Decide if software veto is on based on arguments or rundoc"""
+    if args is not None:
+        if args.software_veto_overwrite is not None:
+            return {'software_veto_name': args.software_veto_overwrite}
+
+    return {'software_veto_name': rd.get('software_veto', None)}
+
+
+def software_veto_register_plugins(st, veto_name):
+    """Based on the selection that we want to apply, 
+    register the correct veto plugin.
+    Probably we want to pass it as argument and if possible overwrite it from rundoc.
+    Then register a copy of all the plugins. 
+    """
+
+    import straxen.plugins.raw_records_sv.software_veto as software_veto
+    raw_records_sv_plugin = getattr(software_veto, veto_name)
+    st.register(raw_records_sv_plugin)
+
+    import straxen.plugins.raw_records_sv._software_veto_copies as _software_veto_copies
+    st.register_all(_software_veto_copies)
+
+    return st
+
+
+
+
+###########################
+#        __MAIN__         #
+###########################
+
 if __name__ == '__main__':
+
+    # to avoid warnings pymongo fork
+    multiprocessing.set_start_method('spawn')
+    multiprocess.set_start_method('spawn')
+
     if not args.undying:
         main()
     else:
@@ -1653,7 +1745,7 @@ if __name__ == '__main__':
                 raise
             except Exception as fatal_error:
                 log.error(f'Fatal warning:\tran into {fatal_error}. Try '
-                          f'logging error and restart bootstrax')
+                        f'logging error and restart bootstrax')
                 try:
                     log_warning(f'Fatal warning:\tran into {fatal_error}',
                                 priority='error')