diff --git a/amstrax/auto_processing_new/db_utils.py b/amstrax/auto_processing_new/db_utils.py index 0c6fbc0d..ff870c1d 100644 --- a/amstrax/auto_processing_new/db_utils.py +++ b/amstrax/auto_processing_new/db_utils.py @@ -120,16 +120,17 @@ def add_data_entry( else: log.info(f"Would add data entry to run {run_id} (dry run).") -def query_runs(query): + +def get_run_doc(run_id): """ - Query the rundb to retrieve runs based on a specific filter. + Get the document for a specific run from the rundb. - :param query: MongoDB query object. - :return: List of runs matching the query. + :param run_id: ID of the run to retrieve. + :return: Document for the run. """ runsdb = amstrax.get_mongo_collection() - results = list(runsdb.find(query)) - log.info(f"Found {len(results)} runs matching query.") - return results + run_doc = runsdb.find_one({'number': int(run_id)}) + log.info(f"Found run {run_id} in the database.") + return run_doc diff --git a/amstrax/auto_processing_new/offline_processing.py b/amstrax/auto_processing_new/offline_processing.py index a496cf57..d045ab64 100644 --- a/amstrax/auto_processing_new/offline_processing.py +++ b/amstrax/auto_processing_new/offline_processing.py @@ -23,6 +23,7 @@ def parse_args(): run_selection.add_argument( "--run_file", type=str, help="File with run IDs to process. It should contain one run ID per line." ) + run_selection.add_argument("--run_range", type=str, help="Range of run IDs to process (e.g., 1000-2000).") # Arguments for processing parser.add_argument( @@ -99,27 +100,30 @@ def check_for_production(args): raise ValueError("Output folder not specified.") -def main(args): - """ - Main function for offline job submission of selected runs. - """ +def get_run_ids_from_args(args): - # Check the run selection method + run_ids = [] if args.run_id: - run_docs = [{"number": int(run_id)} for run_id in args.run_id] + run_ids = args.run_id elif args.run_file: - # It should contain a list of run numbers, one per line with open(args.run_file, "r") as f: - run_numbers = f.readlines() - run_docs = [{"number": int(run_number)} for run_number in run_numbers] - else: - log.error("Either --run_id or --run_file must be provided.") - return + run_ids = f.readlines() + elif args.run_range: + run_range = args.run_range.split("-") + run_ids = list(range(int(run_range[0]), int(run_range[1]) + 1)) - # Limit the number of runs to process - run_docs = run_docs[args.start_from : args.start_from + args.max_runs] + return run_ids +def main(args): + """ + Main function for offline job submission of selected runs. + """ + + # Get the run IDs from the arguments + run_docs = get_run_ids_from_args(args) + # Limit the number of runs to process + run_docs = run_docs[args.start_from : args.start_from + args.max_runs] log.info(f"We are about to process {len(run_docs)} runs.") # Submit jobs for each run diff --git a/amstrax/auto_processing_new/process.py b/amstrax/auto_processing_new/process.py index b2095b95..3bef63a5 100644 --- a/amstrax/auto_processing_new/process.py +++ b/amstrax/auto_processing_new/process.py @@ -38,6 +38,9 @@ def __init__(self, args): self.setup_amstrax() self.setup_production() + self.run_doc = self.db_utils.get_run_doc(self.run_id) + self.infer_special_modes() + def setup_amstrax(self): @@ -52,6 +55,39 @@ def setup_amstrax(self): self.db_utils = self.amstrax.db_utils + def get_run_doc_info(self): + + # let's just print out the main info from the rundoc + # like mode, start, end, duration, user, comments, tags + # we'll need to format the date and time properly, and calculate the duration + # use datetime to convert the date to a readable format + log.info(f"Run document for run {self.run_id}:") + log.info(f" *** Mode: {self.run_doc.get('mode')}") + log.info(f" *** Start: {datetime.datetime.fromtimestamp(self.run_doc.get('start') / 1e9)}") + log.info(f" *** End: {datetime.datetime.fromtimestamp(self.run_doc.get('end') / 1e9)}") + log.info(f" *** Duration: {(self.run_doc.get('end') - self.run_doc.get('start'))/1e9:.2f} seconds") + log.info(f" *** User: {self.run_doc.get('user')}") + log.info(f" *** Comments: {self.run_doc.get('comments')}") + log.info(f" *** Tags: {self.run_doc.get('tags')}") + + def infer_special_modes(self): + + # Check if there is led in the run_doc + if "ledcalibration" in self.run_doc.get('mode'): + # Add the LEDCalibration plugin to the context + log.info("Detected LED calibration run.") + log.info("Adding LEDCalibration plugin to the context.") + ax = self.amstrax + self.st.register([ + ax.DAQReader, + ax.RecordsLED, + ax.LEDCalibration + ]) + + # override the targets to process only the LEDCalibration + self.targets = ["raw_records", "records_led", "led_calibration"] + + def setup_production(self): if self.production: