XAMS-nikhef / amstrax / 15042264012

import argparse

import time

import os

import logging

from datetime import datetime, timedelta

from job_submission import submit_job

from db_utils import query_runs, update_processing_status

logging.basicConfig(level=logging.INFO)

log = logging.getLogger(__name__)

SETUP_FILE = "/data/xenon/xams_v2/setup.sh"

def parse_args():

    parser = argparse.ArgumentParser(description="Autoprocess XAMS live data")

    parser.add_argument(

    parser.add_argument("--output_folder", default="/data/xenon/xams_v2/xams_processed", help="Path where to save processed data")

    parser.add_argument("--timeout", default=20, type=int, help="Time (in seconds) between checks.")

    parser.add_argument("--max_jobs", default=5, type=int, help="Maximum number of jobs to submit simultaneously.")

    parser.add_argument("--run_id", default=None, help="Single run ID to process manually.")

    parser.add_argument("--mem", default=8000, type=int, help="Memory per CPU in MB.")

    parser.add_argument("--logs_path", default="/data/xenon/xams_v2/logs/", help="Path to save job logs.")

    parser.add_argument("--production", action="store_true", help="Run in production mode (will update the rundb).")

    parser.add_argument("--set_config_kwargs", default="{}", help="Dictionary of kwargs to pass to set_config.")

    parser.add_argument("--set_context_kwargs", default="{}", help="Dictionary of kwargs to pass to set_context.")

    parser.add_argument("--amstrax_path", default=None, help="Path to the amstrax directory.")

    parser.add_argument("--corrections_version", default=None, help="Version of corrections to use.")

    parser.add_argument("--queue", default="short", help="Queue to submit jobs to. See Nikhef docs for options.")

    return parser.parse_args()

def main(args):

    log.info("Starting XAMS auto-processing script...")

    amstrax_dir = amstrax.amstrax_dir  # Path to the amstrax directory

    nap_time = int(args.timeout)

    output_folder = args.output_folder

    targets = " ".join(args.target)

    runs_col = amstrax.get_mongo_collection()

    client = amstrax.get_mongo_client()

    processing_db = client["daq"]["processing"]

        auto_processing_on = processing_db.find_one({"name": "auto_processing"})["status"] == "on"

        run_docs_to_do = update_task_list(args, runs_col, auto_processing_on)

        handle_running_jobs(runs_col, production=args.production)

        if not run_docs_to_do:

            log.info(f"No runs to process. Sleeping for {nap_time} seconds.")

            time.sleep(nap_time)

            continue

        submit_new_jobs(args, runs_col, run_docs_to_do, amstrax_dir)

        if args.run_id:

            log.info("Finished processing run.")

            break

        log.info(f"Sleeping for {nap_time} seconds before next check...")

        time.sleep(nap_time)

def update_task_list(args, runs_col, auto_processing_on):

    query = {

    if not auto_processing_on:

        query = {"data": {"$elemMatch": {"type": "live", "host": "stbc"}}, "tags": {"$elemMatch": {"name": "process"}}}

    projection = {"number": 1, "start": 1, "end": 1, "data": 1, "processing_status": 1, "processing_failed": 1}

    sort = [("number", -1)]

    run_docs_to_do = list(runs_col.find(query, projection).sort(sort))

    if args.run_id:

        run_docs_to_do = [runs_col.find_one({"number": int(args.run_id)}, projection)]

    if run_docs_to_do:

        log.info(f'Found {len(run_docs_to_do)} runs to process: {[run_doc["number"] for run_doc in run_docs_to_do]}')

    return run_docs_to_do

def handle_running_jobs(runs_col, production=False):

    query = {"processing_status.status": {"$in": ["submitted", "running"]}}

    projection = {"number": 1, "processing_status": 1}

    run_docs_running = list(runs_col.find(query, projection))

    for run_doc in run_docs_running:

        processing_status = run_doc["processing_status"]

        run_number = run_doc["number"]

        if processing_status["status"] in ["running", "submitted"]:

            if processing_status["time"] < datetime.now() - timedelta(minutes=30):

                new_status = "failed"

                log.info(

                if production:

                    update_processing_status(run_number, new_status, production=production, is_online=True)

                    log.info(f"Would have updated run {run_number} to {new_status}")

def submit_new_jobs(args, runs_col, run_docs_to_do, amstrax_dir):

    query = {"processing_status.status": {"$in": ["submitted", "running"]}}

    projection = {"number": 1, "processing_status": 1}

    run_docs_running = list(runs_col.find(query, projection))

    num_running_jobs = len(run_docs_running)

    log.info(f"Found {num_running_jobs} running jobs.")

    if num_running_jobs >= args.max_jobs:

        log.info(f"Too many jobs running ({num_running_jobs}/{args.max_jobs}).")

        return

    max_jobs_to_submit = args.max_jobs - num_running_jobs

    will_do_run_ids = [int(run_doc["number"]) for run_doc in run_docs_to_do[:max_jobs_to_submit]]

    log.info(f"Submitting jobs for runs: {will_do_run_ids}")

    for run_doc in run_docs_to_do[:max_jobs_to_submit]:

        run_id = f'{int(run_doc["number"]):06}'

        job_name = f"process_{run_id}_online"

        production_flag = "--production" if args.production else ""

        targets = " ".join(args.target)

        arguments = []

        arguments.append(f"--run_id {run_id}")

        arguments.append(f"--targets {targets}")

        arguments.append(f"--output_folder {args.output_folder}")

        if args.corrections_version:

            arguments.append(f"--corrections_version {args.corrections_version}")

        if args.amstrax_path:

            arguments.append(f"--amstrax_path {args.amstrax_path}")

        if args.production:

            arguments.append("--production")

            arguments.append("--allow_raw_records")

            arguments.append("--is_online")

        arguments = " ".join(arguments)

        jobstring = f"""

        if args.production:

            submit_job(

            update_processing_status(

            log.info(f"Would have submitted job for run {run_id}")

if __name__ == "__main__":

    args = parse_args()

    log_name = "auto_processing_online"

    import amstrax

    versions = amstrax.print_versions(

    log = amstrax.get_daq_logger(

    main(args)