diff --git a/README.md b/README.md new file mode 100644 index 0000000..0001a6f --- /dev/null +++ b/README.md @@ -0,0 +1,69 @@ +# End-to-end processing of ILT HBA data with flocs + +This package aims to provide relatively simple end-to-end automatic processing of ILT HBA data. Where `flocs-runners` provides the interface to running pipelines, `flocs-processing` is the scaffolding to tie it together. Data reduction is coordinated via a dedicated SQLite database that holds information on which observations to process, which pipelines to run for them and all of the related statuses. Orchestration of all the pipelines is handled via Airflow through a DAG. + +The autoPILOT package (https://github.com/LOFAR-VLBI/autoPILOT) needs to be on PYTHONPATH to enable the automatic calibrator assessment. + +## Folder setup +Flocs-processing requires three folders to be setup: + +* A processing folder -- this is where data is stored while processing +* A data folder -- this is where the input data is found +* An output folder -- this is where finished pipeline outputs are copied to, and searched for in steps that depend on it. + +The expected naming directory structure for input data is `//{calibrator,target}`. Inside the calibrator and target folders, the observations should follow the usual `LXXXXXX` naming scheme. These **must** match the SAS IDs in the database for flocs to be able to find them. + +## Database setup +A database for processing is created via `flocs-processing create-database`. This will create an empty database with the necessary columns. Datasets to process can be added via `flocs-processing add-field`. + +## Processing data +To start processing data, Airflow needs to be running. This will be delegated to `flocs-processing process-from-database` in the future, but for now requires running Airflow manually. For setup do the following: + +1. Install airflow: `uv pip install apache-airflow` +2. Set up a folder that wil contain all of Airflow's own stuff and assign it to the `AIRFLOW_HOME` environment variable. +3. Run `airflow config list --defaults > "${AIRFLOW_HOME}/airflow.cfg"` +4. Define `AIRFLOW__CORE__DAGS_FOLDER` as `${AIRFLOW_HOME}/dags` and create the folder. Copy the DAGs inside `flocs_processing/dags` to this folder. +5. Define `AIRFLOW__CORE__LOAD_EXAMPLES` as `False` + +Finally, define the following airflow variables: + +``` +export AIRFLOW_HOME=/path/to/some/folder/for/airflow +export AIRFLOW__CORE__DAGS_FOLDER=$AIRFLOW_HOME/dags +export AIRFLOW__CORE__LOAD_EXAMPLES=False +export AIRFLOW__CORE__PARALLELISM=32 +export AIRFLOW__LOGGING__DAG_PROCESSOR_CHILD_PROCESS_LOG_DIRECTORY=$AIRFLOW_HOME/logs/dag_processor +export AIRFLOW__CORE__PLUGINS_FOLDER=$AIRFLOW_HOME/plugins +export AIRFLOW__DATABASE__SQL_ALCHEMY_CONN="sqlite:///$AIRFLOW_HOME/airflow.db" +export AIRFLOW__LOGGING__BASE_LOG_FOLDER=$AIRFLOW_HOME/logs +``` + +For a small test, you can run `airflow standalone` to start the Airflow instance for a small test. For proper deployment, it is recommended by the Airflow docs to not use `standalone`. First we'll set up a persistent JWT secret for authentication purpose. + +``` +mkdir -p "$HOME/.config/airflow" +chmod 700 "$HOME/.config/airflow" +openssl rand -hex 32 > "$HOME/.config/airflow/jwt_secret" +chmod 600 "$HOME/.config/airflow/jwt_secret" +export AIRFLOW__API_AUTH__JWT_SECRET="$(cat "$HOME/.config/airflow/jwt_secret")" +``` + +Next, initialise Airflow's own database with + +``` +airflow db migrate +``` + +Finally, to start the necessary Airflow services, execute them like follows: + +``` +tmux new-session -d -s airflow-api-server "bash -c 'source $HOME/source_airflow.sh && airflow api-server; exec bash'" +tmux new-session -d -s airflow-triggerer "bash -c 'source $HOME/source_airflow.sh && airflow triggerer; exec bash'" +tmux new-session -d -s airflow-dag-processor "bash -c 'source $HOME/source_airflow.sh && airflow dag-processor; exec bash'" +tmux new-session -d -s airflow-scheduler "bash -c 'source $HOME/source_airflow.sh && airflow scheduler; exec bash'" +``` + +This should start four tmux sessions with these services running in the background. The credentials to log into e.g. the web interface will be stored in `${AIRFLOW_HOME}/simple_auth_manager_passwords.json.generated`. The Airflow instance will start on port 8080. You can access it via `localhost:8080` in your browser. If it is running on a remote cluster, you can set up a tunnel via e.g. `ssh -N -L 8080:localhost:8080 ` to forward it to your local machine. + +Once `flocs-processing` is complete the processing loop will be automatic, but for now the user must trigger the DAG manually. On the "Dags" tab you should now see the flocs DAGs available. To manually trigger one, click on the name and on the subsequent page use the "Trigger" button in the top right. + diff --git a/flocs_processing/dags/pilot_single_target.py b/flocs_processing/dags/pilot_single_target.py new file mode 100644 index 0000000..a3e045f --- /dev/null +++ b/flocs_processing/dags/pilot_single_target.py @@ -0,0 +1,588 @@ +from enum import Enum +import datetime +import functools +import os +import pathlib +import re +import sqlite3 +import subprocess +import time + +from airflow.exceptions import AirflowFailException, AirflowSkipException +from airflow.sdk import dag, get_current_context, task +from airflow.task.trigger_rule import TriggerRule +from flocs_lta.lta_search import ObservationStager +from stager_access import get_surls_requested, get_surls_online + +# Need to replace this with a config file +TABLE_NAME = "" +DATABASE = "" +SLURM_ACCOUNT = "" +SLURM_QUEUE = "" +DATA_DIR = "" +OUTPUT_DIR = "" +PROCESSING_DIR = "" +NN_MODEL_CACHE = "" + + +@functools.total_ordering +class PIPELINE_STATUS(Enum): + nothing = 0 + downloaded = 1 + finished = 2 + running = 3 + processing = 98 + error = 99 + + def __eq__(self, other): + if other.__class__ is int: + return self.value == other + elif other.__class__ is self.__class__: + return self.value == other.value + else: + raise NotImplementedError + + def __lt__(self, other): + if self.__class__ is not other.__class__: + raise NotImplementedError + return self.value < other.value + + +def get_db_columns(): + with sqlite3.connect(DATABASE) as db: + db.row_factory = sqlite3.Row + cursor = db.cursor() + columns = "target_name,priority,finished,downloaded,sas_id_calibrator1,sas_id_calibrator2,sas_id_calibrator_final,sas_id_target,status_calibrator1,status_calibrator2,status_target,status_vlbi_delay,status_vlbi_dd" + field = cursor.execute( + f"select {columns} from {TABLE_NAME} where finished==0 order by priority desc" + ).fetchall() + print(field) + return field + + +def set_status_processing(name, identifier, target): + with sqlite3.connect(DATABASE) as db: + cursor = db.cursor() + cursor.execute( + f"update {TABLE_NAME} set status_{identifier}={PIPELINE_STATUS.processing.value} where target_name=='{name}' and sas_id_target=='{target}'" + ) + + +def set_status_finished(name, identifier, target): + with sqlite3.connect(DATABASE) as db: + cursor = db.cursor() + cursor.execute( + f"update {TABLE_NAME} set status_{identifier}={PIPELINE_STATUS.finished.value} where target_name=='{name}' and sas_id_target=='{target}'" + ) + + +def set_status_downloaded(name, target): + with sqlite3.connect(DATABASE) as db: + cursor = db.cursor() + cursor.execute( + f"update {TABLE_NAME} set downloaded=1 where target_name=='{name}' and sas_id_target=='{target}'" + ) + + +def set_field_finished(name, target): + query = f"update {TABLE_NAME} set finished=1 where target_name=='{name}' and sas_id_target=='{target}'" + with sqlite3.connect(DATABASE) as db: + cursor = db.cursor() + cursor.execute(query) + + +def set_final_calibrator(name, target, final_cal): + with sqlite3.connect(DATABASE) as db: + cursor = db.cursor() + cursor.execute( + f"update {TABLE_NAME} set sas_id_calibrator_final={final_cal} where target_name=='{name}' and sas_id_target=='{target}'" + ) + + +def get_most_recent_run(searchpath: str, sas_id: str, pipeline: str) -> pathlib.Path: + rundirs = pathlib.Path(searchpath) + rundirs_sorted = sorted(rundirs.iterdir()) + if pipeline: + rundirs_sorted_filtered = [ + d + for d in rundirs_sorted + if ((sas_id in d.parts[-1]) and (pipeline in d.parts[-1])) and d.is_dir() + ] + else: + rundirs_sorted_filtered = [d for d in rundirs_sorted if sas_id in d.parts[-1]] + rundir_final = rundirs_sorted_filtered[-1].absolute() + return rundir_final + + +@dag(max_active_runs=1) +def pilot_single_target(): + @task + def get_unprocessed_target(): + field = dict(get_db_columns()[0]) + print(field["target_name"]) + return field + + @task.short_circuit + def check_fields(): + fields = get_db_columns() + return bool(fields) + + @task + def download_field(field): + if field["downloaded"]: + return field + else: + has_cal1 = False + stage_calibrators = False + if field["sas_id_calibrator1"]: + ms_folder = f"L{field['sas_id_calibrator1']}" + cal1_full_path = os.path.join( + DATA_DIR, field["target_name"], "calibrator", ms_folder + ) + if os.path.exists(cal1_full_path): + has_cal1 = True + else: + stage_calibrators = True + + has_cal2 = False + if field["sas_id_calibrator2"]: + ms_folder = f"L{field['sas_id_calibrator2']}" + cal2_full_path = os.path.join( + DATA_DIR, field["target_name"], "calibrator", ms_folder + ) + if os.path.exists(cal2_full_path): + has_cal2 = True + else: + stage_calibrators = True + if field["sas_id_target"]: + ms_folder = f"L{field['sas_id_target']}" + target_full_path = os.path.join( + DATA_DIR, field["target_name"], "target", ms_folder + ) + if os.path.exists(target_full_path): + stage_target = False + else: + stage_target = True + else: + raise AirflowFailException( + f"No target SAS ID in database for field {field['target_name']}" + ) + + if stage_calibrators or stage_target: + print(f"Field {field['sas_id_target']} is not downloaded.") + stager = ObservationStager(get_surls=True) + stager.find_observation_by_sasid( + "ALL", + field["sas_id_target"], + None, + 120e6, + 168e6, + ) + if stage_calibrators: + stager.find_nearest_calibrators(2, 120e6, 168e6) + stage_id_calibrators = stager.stage_calibrators() + if stage_target: + stage_id_target = stager.stage_target() + else: + return field + + calibrator_staged = False + target_staged = False + calibrator_downloaded = has_cal1 or has_cal2 + target_downloaded = not stage_target + while True: + if len(get_surls_online(stage_id_calibrators)) == len( + get_surls_requested(stage_id_calibrators) + ): + calibrator_staged = True + if calibrator_staged and not calibrator_downloaded: + dl_path = os.path.join(DATA_DIR, field["target_name"], "calibrator") + cmd = ( + f"flocs-lta download --outdir {dl_path} {stage_id_calibrators}" + ) + with open( + f"log_download_calibrators_{field['target_name']}.txt", + "w+", + ) as f_out, open( + f"log_download_calibrators_{field['target_name']}.txt", + "w+", + ) as f_err: + proc = subprocess.run( + cmd, shell=True, text=True, stdout=f_out, stderr=f_err + ) + if not proc.returncode: + calibrator_downloaded = True + else: + raise RuntimeError + + if len(get_surls_online(stage_id_target)) == len( + get_surls_requested(stage_id_target) + ): + calibrator_staged = True + if target_staged and not target_downloaded: + dl_path = os.path.join(DATA_DIR, field["target_name"], "target") + cmd = f"flocs-lta download --outdir {dl_path} {stage_id_target}" + with open( + f"log_download_calibrators_{field['target_name']}.txt", + "w+", + ) as f_out, open( + f"log_download_calibrators_{field['target_name']}.txt", + "w+", + ) as f_err: + proc = subprocess.run( + cmd, shell=True, text=True, stdout=f_out, stderr=f_err + ) + if not proc.returncode: + set_status_downloaded( + field["target_name"], + field["sas_id_target"], + ) + target_downloaded = True + else: + raise RuntimeError + if calibrator_downloaded and target_downloaded: + break + time.sleep(60) + + @task + def run_linc_calibrator1(field): + if (field["status_calibrator1"] == PIPELINE_STATUS.finished) or ( + field["status_calibrator1"] == PIPELINE_STATUS.running + ): + print( + f"Flux density calibrator {field['sas_id_calibrator1']} for observation {field['target_name']} {field['sas_id_target']} already processed." + ) + return field + else: + print( + f"Processing flux density calibrator {field['sas_id_calibrator1']} for observation {field['target_name']} {field['sas_id_target']}" + ) + ms_folder = f"L{field['sas_id_calibrator1']}" + set_status_processing( + field["target_name"], "calibrator1", field["sas_id_target"] + ) + outdir = os.path.join(OUTPUT_DIR, field["target_name"]) + cmd = f"flocs-run linc calibrator --runner toil --scheduler slurm --slurm-account {SLURM_ACCOUNT} --slurm-queue {SLURM_QUEUE} --rundir {PROCESSING_DIR} --outdir {outdir} {os.path.join(DATA_DIR, field['target_name'], 'calibrator', ms_folder)}" + if not os.path.isdir(outdir): + os.mkdir(outdir) + print(cmd) + with open( + f"log_LINC_calibrator_{field['target_name']}_{field['sas_id_calibrator1']}.txt", + "w+", + ) as f_out, open( + f"log_LINC_calibrator_{field['target_name']}_{field['sas_id_calibrator1']}_err.txt", + "w+", + ) as f_err: + proc = subprocess.run( + cmd, shell=True, text=True, stdout=f_out, stderr=f_err + ) + success = False + pattern = re.compile(r"Workflow.* stopped. Success: True") + if not proc.returncode: + f_err.seek(0) + if pattern.search(f_err.read()): + success = True + if success: + set_status_finished( + field["target_name"], "calibrator1", field["sas_id_target"] + ) + else: + raise RuntimeError + return field + + @task + def run_linc_calibrator2(field): + if (field["status_calibrator2"] == PIPELINE_STATUS.finished) or ( + field["status_calibrator2"] == PIPELINE_STATUS.running + ): + print( + f"Flux density calibrator {field['sas_id_calibrator2']} for observation {field['target_name']} {field['sas_id_target']} already processed." + ) + return field + else: + print( + f"Processing flux density calibrator {field['sas_id_calibrator2']} for observation {field['target_name']} {field['sas_id_target']}" + ) + ms_folder = f"L{field['sas_id_calibrator2']}" + set_status_processing( + field["target_name"], "calibrator2", field["sas_id_target"] + ) + outdir = os.path.join(OUTPUT_DIR, field["target_name"]) + cmd = f"flocs-run linc calibrator --runner toil --scheduler slurm --slurm-account {SLURM_ACCOUNT} --slurm-queue {SLURM_QUEUE} --rundir {PROCESSING_DIR} --outdir {outdir} {os.path.join(DATA_DIR, field['target_name'], 'calibrator', ms_folder)}" + if not os.path.isdir(outdir): + os.mkdir(outdir) + print(cmd) + with open( + f"log_LINC_calibrator_{field['target_name']}_{field['sas_id_calibrator2']}.txt", + "w+", + ) as f_out, open( + f"log_LINC_calibrator_{field['target_name']}_{field['sas_id_calibrator2']}_err.txt", + "w+", + ) as f_err: + proc = subprocess.run( + cmd, shell=True, text=True, stdout=f_out, stderr=f_err + ) + success = False + pattern = re.compile(r"Workflow.* stopped. Success: True") + if not proc.returncode: + f_err.seek(0) + if pattern.search(f_err.read()): + success = True + if success: + set_status_finished( + field["target_name"], "calibrator2", field["sas_id_target"] + ) + else: + raise RuntimeError + return field + + @task(trigger_rule=TriggerRule.ONE_DONE) + def select_best_calibrator(result1, result2): + if result1["sas_id_calibrator_final"]: + return result1 + elif result2["sas_id_calibrator_final"]: + return result2 + elif result1 and result2: + print("Selecting between cal1 and cal2") + # Need actual selection logic here + set_final_calibrator( + result1["target_name"], + result1["sas_id_target"], + result1["sas_id_calibrator1"], + ) + return result1 + elif result1 and (not result2): + print("Only cal 1 succeeded, continuing with that") + set_final_calibrator( + result1["target_name"], + result1["sas_id_target"], + result1["sas_id_calibrator1"], + ) + return result1 + elif (not result1) and result2: + print("Only cal 2 succeeded, continuing with that") + set_final_calibrator( + result2["target_name"], + result2["sas_id_target"], + result2["sas_id_calibrator2"], + ) + return result2 + else: + raise AirflowFailException("No calibrators succeeded; stopping processing.") + + @task + def run_linc_target(field): + if (field["status_target"] == PIPELINE_STATUS.finished) or ( + field["status_target"] == PIPELINE_STATUS.running + ): + return field + else: + print( + f"Processing target observation {field['target_name']} {field['sas_id_target']} with calibrator {field['sas_id_calibrator_final']}" + ) + ms_folder = f"L{field['sas_id_target']}" + outdir = os.path.join(OUTPUT_DIR, field["target_name"]) + calibrator_path = get_most_recent_run( + outdir, field["sas_id_calibrator_final"], "LINC_calibrator" + ) + calibrator_solutions = ( + calibrator_path / "results_LINC_calibrator" / "cal_solutions.h5" + ) + set_status_processing( + field["target_name"], "target", field["sas_id_target"] + ) + cmd = f"flocs-run linc target --runner toil --scheduler slurm --slurm-account {SLURM_ACCOUNT} --slurm-queue {SLURM_QUEUE} --rundir {PROCESSING_DIR} --outdir {outdir} --cal-solutions {calibrator_solutions} {os.path.join(DATA_DIR, field['target_name'], 'target', ms_folder)}" + if not os.path.isdir(outdir): + os.mkdir(outdir) + print(cmd) + with open( + f"log_LINC_target_{field['target_name']}_{field['sas_id_target']}.txt", + "w+", + ) as f_out, open( + f"log_LINC_target_{field['target_name']}_{field['sas_id_target']}_err.txt", + "w+", + ) as f_err: + proc = subprocess.run( + cmd, shell=True, text=True, stdout=f_out, stderr=f_err + ) + success = False + pattern = re.compile(r"Workflow.* stopped. Success: True") + if not proc.returncode: + f_err.seek(0) + if pattern.search(f_err.read()): + success = True + if success: + set_status_finished( + field["target_name"], "target", field["sas_id_target"] + ) + else: + raise RuntimeError + return field + + @task + def validate_linc_target(field): + return field + + @task(retries=0, retry_delay=datetime.timedelta(seconds=5)) + def run_vlbi_delay(field): + if (field["status_vlbi_delay"] == PIPELINE_STATUS.finished) or ( + field["status_vlbi_delay"] == PIPELINE_STATUS.running + ): + return field + else: + print( + f"Processing delay calibration for {field['target_name']} {field['sas_id_target']}" + ) + outdir = os.path.join(OUTPUT_DIR, field["target_name"]) + target_path = get_most_recent_run( + outdir, field["sas_id_target"], "LINC_target" + ) + target_ms_path = target_path / "results_LINC_target" / "results" + set_status_processing( + field["target_name"], "vlbi_delay", field["sas_id_target"] + ) + + delay_cat = os.path.join(outdir, "delay_calibrators.csv") + image_cat = os.path.join(outdir, "image_catalogue.csv") + + proc = subprocess.run( + "detect_bad_slurm_nodes.sh", + shell=True, + text=True, + stdout=subprocess.PIPE, + ) + bad_nodes = proc.stdout.strip() + if bad_nodes: + print(f"Excluding the following bad nodes from scheduling: {bad_nodes}") + os.environ["TOIL_SLURM_ARGS"] = f"--exclude={bad_nodes}" + + context = get_current_context() + if context["ti"].try_number == 1: + cmd = f"flocs-run vlbi delay-calibration --runner toil --scheduler slurm --slurm-account {SLURM_ACCOUNT} --slurm-queue {SLURM_QUEUE} --rundir {PROCESSING_DIR} --outdir {outdir} --ms-suffix dp3concat --delay-calibrator {delay_cat} --image-catalogue {image_cat} {target_ms_path}" + else: + # Extract the previous working directory + flocs_workdir = "" + print( + f"Scanning log_VLBI_delay-calibration_{field['target_name']}_{field['sas_id_target']}.txt for workdir." + ) + with open( + f"log_VLBI_delay-calibration_{field['target_name']}_{field['sas_id_target']}.txt" + ) as f_out: + for line in f_out.readlines(): + print(line) + if "Running workflow with" in line: + flocs_workdir = line.split(" ")[-1].strip() + break + if not flocs_workdir: + raise RuntimeError( + "Could not retrieve PILOT workdir. Flocs probably crashed before launching." + ) + print(f"Resuming failed PILOT run in {flocs_workdir}") + cmd = f"flocs-run vlbi delay-calibration --runner toil --scheduler slurm --slurm-account {SLURM_ACCOUNT} --slurm-queue {SLURM_QUEUE} --rundir {flocs_workdir} --restart --outdir {outdir} --ms-suffix dp3concat --delay-calibrator {delay_cat} --image-catalogue {image_cat} {target_ms_path}" + if not os.path.isdir(outdir): + os.mkdir(outdir) + print(cmd) + with open( + f"log_VLBI_delay-calibration_{field['target_name']}_{field['sas_id_target']}.txt", + "w+", + ) as f_out, open( + f"log_VLBI_delay-calibration_{field['target_name']}_{field['sas_id_target']}_err.txt", + "w+", + ) as f_err: + proc = subprocess.run( + cmd, shell=True, text=True, stdout=f_out, stderr=f_err + ) + success = False + pattern = re.compile(r"Workflow.* stopped. Success: True") + if not proc.returncode: + f_err.seek(0) + if pattern.search(f_err.read()): + success = True + + if success: + set_status_finished( + field["target_name"], "vlbi_delay", field["sas_id_target"] + ) + else: + raise RuntimeError + return field + + @task + def run_ddf_subtract(field): + return field + + @task + def run_vlbi_ddcal(field): + if (field["status_vlbi_dd"] == PIPELINE_STATUS.finished) or ( + field["status_vlbi_dd"] == PIPELINE_STATUS.running + ): + return field + else: + print( + f"Processing ILT dd calibration for {field['target_name']} {field['sas_id_target']}" + ) + outdir = os.path.join(OUTPUT_DIR, field["target_name"]) + target_path = get_most_recent_run( + outdir, field["sas_id_target"], "LINC_target" + ) + target_ms_path = target_path / "results_LINC_target" / "results" + print(f"Using LINC target run: {target_path}") + + sols_path = get_most_recent_run( + outdir, field["sas_id_target"], "VLBI_delay" + ) + sols_path = sols_path / "results_VLBI_delay-calibration" + sols = list(sols_path.glob("merged*selfcalcycle???_linearfulljones*.h5"))[0] + print(f"Using PILOT delay calibration solutions: {sols}") + + source_cat = os.path.join(DATA_DIR, field["target_name"], "vlbi_target.csv") + if not os.path.isfile(source_cat): + raise AirflowFailException(f"{source_cat} not found.") + + set_status_processing( + field["target_name"], "vlbi_dd", field["sas_id_target"] + ) + cmd = f"flocs-run vlbi dd-calibration --runner toil --scheduler slurm --slurm-time 24:00:00 --slurm-account {SLURM_ACCOUNT} --slurm-queue {SLURM_QUEUE} --rundir {PROCESSING_DIR} --outdir {outdir} --delay-solset {sols} --phasediff-score 10.0 --source-catalogue {source_cat} --model-cache {NN_MODEL_CACHE} --ms-suffix .dp3concat {target_ms_path}" + if not os.path.isdir(outdir): + os.mkdir(outdir) + print(cmd) + with open( + f"log_VLBI_dd-calibration_{field['target_name']}_{field['sas_id_target']}.txt", + "w+", + ) as f_out, open( + f"log_VLBI_dd-calibration_{field['target_name']}_{field['sas_id_target']}_err.txt", + "w+", + ) as f_err: + proc = subprocess.run( + cmd, shell=True, text=True, stdout=f_out, stderr=f_err + ) + success = False + pattern = re.compile(r"Workflow.* stopped. Success: True") + if not proc.returncode: + f_err.seek(0) + if pattern.search(f_err.read()): + success = True + if success: + set_status_finished( + field["target_name"], "vlbi_dd", field["sas_id_target"] + ) + set_field_finished(field["target_name"], field["sas_id_target"]) + else: + raise RuntimeError + return field + + proceed = check_fields() + get_field = get_unprocessed_target() + field = download_field(get_field) + result_cal1 = run_linc_calibrator1(field) + result_cal2 = run_linc_calibrator2(field) + best_cal = select_best_calibrator(result_cal1, result_cal2) + result_targ = run_linc_target(best_cal) + linc_is_valid = validate_linc_target(result_targ) + result_vlbi_delay = run_vlbi_delay(linc_is_valid) + result_vlbi_dd = run_vlbi_ddcal(result_vlbi_delay) + + proceed >> get_field + + +pilot_single_target() diff --git a/flocs_processing/dags/pilot_widefield.py b/flocs_processing/dags/pilot_widefield.py new file mode 100644 index 0000000..75a4961 --- /dev/null +++ b/flocs_processing/dags/pilot_widefield.py @@ -0,0 +1,614 @@ +from flocs_processing.db_utils import PIPELINE_STATUS, FlocsDB +from flocs_processing.pipeline_runners import ( + get_most_recent_run, + run_linc_calibrator_cwltool, + run_linc_calibrator_toil, + run_linc_target_cwltool, + run_linc_target_toil, + run_pilot_delay_cwltool, + run_pilot_delay_toil, + run_pilot_ddcal_cwltool, + run_pilot_ddcal_toil, + run_pilot_facet_imaging_toil, + run_pilot_facet_subtract_toil, + run_pilot_intermediate_image_toil, + run_pilot_process_ddf_toil, + run_prepare_ddf, + run_prepare_ddf_subtract, + launch_ddf_pipeline, +) + +import configparser +import datetime +import os +import pathlib +import random +import re +import sqlite3 +import subprocess +import time + +from airflow.exceptions import AirflowFailException +from airflow.sdk import dag, get_current_context, task +from airflow.providers.standard.sensors.python import PythonSensor +from airflow.sdk.exceptions import AirflowSkipException +from airflow.task.trigger_rule import TriggerRule +from flocs_lta.lta_search import ObservationStager +from ilotss.assess_calibrators import assess_and_compare +from stager_access import get_surls_requested, get_surls_online + +if "FLOCS_AIRFLOW_CONFIG" not in os.environ: + raise RuntimeError( + "FLOCS_AIRFLOW_CONFIG environment variable not set. Please point this to a valid configuration file." + ) + +CONFIG_FILE: str = os.getenv("FLOCS_AIRFLOW_CONFIG") or "" +if not os.path.isfile(CONFIG_FILE): + raise RuntimeError(f"{CONFIG_FILE} is not a valid file") + +parser = configparser.ConfigParser() +parser.optionxform = str # ty: ignore[invalid-assignment] +with open(CONFIG_FILE, "r") as config: + parser.read_string("[DEFAULT]\n" + config.read()) + +print("Config summary:") +for k, v in parser["DEFAULT"].items(): + print(f"{k}: {v}") + +TABLE_NAME = parser["DEFAULT"]["TABLE_NAME"] +DATABASE = parser["DEFAULT"]["DATABASE"] +SLURM_ACCOUNT = parser["DEFAULT"]["SLURM_ACCOUNT"] +SLURM_QUEUE = parser["DEFAULT"]["SLURM_QUEUE"] +DATA_DIR = parser["DEFAULT"]["DATA_DIR"] +OUTPUT_DIR = parser["DEFAULT"]["OUTPUT_DIR"] +PROCESSING_DIR = parser["DEFAULT"]["PROCESSING_DIR"] +NN_MODEL_CACHE = parser["DEFAULT"]["NN_MODEL_CACHE"] +DDF_CONFIG = parser["DEFAULT"]["DDF_CONFIG"] +FLUX_CALIBRATOR_TEMPLATE = parser["DEFAULT"]["FLUX_CALIBRATOR_TEMPLATE"] +NEEDS_MANUAL_APPROVAL_DELAY = parser.getboolean( + "DEFAULT", "NEEDS_MANUAL_APPROVAL_DELAY" +) + +CWL_RUNNER_LINC_CALIBRATOR = "cwltool" +CWL_RUNNER_LINC_TARGET = "toil" +CWL_RUNNER_PILOT_DELAY = "toil" +CWL_RUNNER_PILOT_DDCAL = "toil" + +CURRENT_DB = FlocsDB(DATABASE, TABLE_NAME) + + +def get_approval(field, identifier, needs_approval): + if not needs_approval: + return field + with sqlite3.connect(DATABASE) as db: + db.row_factory = sqlite3.Row + cursor = db.cursor() + columns = f"sas_id_target,status_{identifier}" + field = cursor.execute( + f"select {columns} from {TABLE_NAME} where sas_id_target=='{field['sas_id_target']}'" + ).fetchall() + status = field[0][f"status_{identifier}"] + if status == PIPELINE_STATUS.finished.value: + return field + + +@dag(max_active_runs=1) +def pilot_widefield(): + @task + def get_unprocessed_target(): + field = None + for dbrow in CURRENT_DB.get_db_columns(): + is_processing = False + row = dict(dbrow) + status_keys = filter(lambda x: x.startswith("status_"), row.keys()) + for key in status_keys: + if row[key] == PIPELINE_STATUS.processing.value: + is_processing = True + break + if not is_processing: + # Only select a field if nothing is processing it. + field = row + break + if not field: + raise AirflowSkipException("No unprocessed fields found.") + print(field["target_name"]) + return field + + @task.short_circuit + def check_fields(): + fields = CURRENT_DB.get_db_columns() + return bool(fields) + + @task + def download_field(field): + if field["downloaded"]: + return field + else: + stage_calibrators = False + num_downloaded_calib1 = 0 + num_downloaded_calib2 = 0 + num_staged_calib = 0 + num_staged_targ = 0 + if os.path.exists(f"srms_{field['sas_id_target']}_calibrators.txt"): + print("Found srm file; counting calibrator SRMs.") + out = subprocess.check_output( + f"wc -l srms_{field['sas_id_target']}_calibrators.txt | cut -f 1 -d ' '", + text=True, + shell=True, + ) + num_staged_calib = int(out.strip()) + + if os.path.exists(f"srms_{field['sas_id_target']}.txt"): + print("Found srm file; counting target SRMs.") + out = subprocess.check_output( + f"wc -l srms_{field['sas_id_target']}.txt | cut -f 1 -d ' '", + text=True, + shell=True, + ) + num_staged_targ = int(out.strip()) + + if field["sas_id_calibrator1"]: + ms_folder = f"L{field['sas_id_calibrator1']}" + cal1_full_path = os.path.join( + DATA_DIR, field["target_name"], "calibrator", ms_folder + ) + if os.path.exists(cal1_full_path): + num_downloaded_calib1 = len( + list(pathlib.Path(cal1_full_path).glob("*.MS")) + ) + else: + stage_calibrators = True + + if field["sas_id_calibrator2"]: + ms_folder = f"L{field['sas_id_calibrator2']}" + cal2_full_path = os.path.join( + DATA_DIR, field["target_name"], "calibrator", ms_folder + ) + if os.path.exists(cal2_full_path): + num_downloaded_calib2 = len( + list(pathlib.Path(cal2_full_path).glob("*.MS")) + ) + else: + stage_calibrators = True + + num_downloaded_calib = num_downloaded_calib1 + num_downloaded_calib2 + if num_downloaded_calib == num_staged_calib: + print( + f"Number of staged calibrator MSes ({num_staged_calib}) equals number of downloaded MSes ({num_downloaded_calib}); not staging calibrators again." + ) + stage_calibrators = False + else: + print( + f"Number of staged calibrator MSes ({num_staged_calib}) does NOT equal number of downloaded MSes ({num_downloaded_calib}); restaging calibrators and resuming download." + ) + stage_calibrators = True + + stage_target = False + if field["sas_id_target"]: + ms_folder = f"L{field['sas_id_target']}" + target_full_path = os.path.join( + DATA_DIR, field["target_name"], "target", ms_folder + ) + if os.path.exists(target_full_path): + num_downloaded_targ = len( + list(pathlib.Path(target_full_path).glob("*.MS")) + ) + if num_downloaded_targ == num_staged_targ: + print( + f"Number of staged target MSes ({num_staged_targ}) equals number of downloaded MSes ({num_downloaded_targ}); not staging target again." + ) + stage_target = False + else: + print( + f"Number of staged target MSes ({num_staged_targ}) does NOT equal number of downloaded MSes ({num_downloaded_targ}); staging target again and resuming download." + ) + stage_target = True + else: + raise AirflowFailException( + f"No target SAS ID in database for field {field['target_name']}" + ) + + if stage_calibrators or stage_target: + print(f"Field {field['sas_id_target']} is not downloaded.") + stager = ObservationStager(get_surls=True) + stager.find_observation_by_sasid( + "ALL", + field["sas_id_target"], + None, + 120, + 168, + ) + if stage_calibrators: + stager.find_nearest_calibrators(2, 120, 168) + stage_id_calibrators = stager.stage_calibrators() + if stage_target: + stage_id_target = stager.stage_target() + else: + return field + + calibrator_staged = False + target_staged = False + calibrator_downloaded = not stage_calibrators + target_downloaded = not stage_target + while True: + if not calibrator_downloaded: + if len(get_surls_online(stage_id_calibrators)) == len( + get_surls_requested(stage_id_calibrators) + ): + calibrator_staged = True + if calibrator_staged and not calibrator_downloaded: + dl_path = os.path.join( + DATA_DIR, field["target_name"], "calibrator" + ) + cmd = f"flocs-lta download --outdir {dl_path} {stage_id_calibrators}" + with ( + open( + f"log_download_calibrators_{field['target_name']}.txt", + "w+", + ) as f_out, + open( + f"log_download_calibrators_{field['target_name']}_err.txt", + "w+", + ) as f_err, + ): + proc = subprocess.run( + cmd, shell=True, text=True, stdout=f_out, stderr=f_err + ) + if not proc.returncode: + calibrator_downloaded = True + else: + raise RuntimeError + + if not target_downloaded: + if len(get_surls_online(stage_id_target)) == len( + get_surls_requested(stage_id_target) + ): + target_staged = True + if target_staged and not target_downloaded: + dl_path = os.path.join(DATA_DIR, field["target_name"], "target") + cmd = f"flocs-lta download --outdir {dl_path} {stage_id_target}" + with ( + open( + f"log_download_target_{field['target_name']}.txt", + "w+", + ) as f_out, + open( + f"log_download_target_{field['target_name']}_err.txt", + "w+", + ) as f_err, + ): + proc = subprocess.run( + cmd, shell=True, text=True, stdout=f_out, stderr=f_err + ) + if not proc.returncode: + CURRENT_DB.set_status_downloaded( + field["target_name"], + field["sas_id_target"], + ) + target_downloaded = True + else: + raise RuntimeError + if calibrator_downloaded and target_downloaded: + break + time.sleep(60) + + @task + def run_linc_calibrator1(field): + field = dict(CURRENT_DB.get_db_columns(field["sas_id_target"])[0]) + if not field["sas_id_calibrator1"]: + raise AirflowSkipException("Calibrator 1 does not exist, skipping.") + if field["status_calibrator1"] == PIPELINE_STATUS.finished: + print( + f"Flux density calibrator {field['sas_id_calibrator1']} for observation {field['target_name']} {field['sas_id_target']} already processed." + ) + return field + else: + if CWL_RUNNER_LINC_CALIBRATOR == "cwltool": + run_linc_calibrator_cwltool(field, calibrator_field=1, db=CURRENT_DB) + elif CWL_RUNNER_LINC_CALIBRATOR == "toil": + run_linc_calibrator_toil(field, calibrator_field=1, db=CURRENT_DB) + else: + raise RuntimeError("Invalid CWL runner specified.") + return field + + @task + def run_linc_calibrator2(field): + field = dict(CURRENT_DB.get_db_columns(field["sas_id_target"])[0]) + if not field["sas_id_calibrator2"]: + raise AirflowSkipException("Calibrator 2 does not exist, skipping.") + if field["status_calibrator2"] == PIPELINE_STATUS.finished: + print( + f"Flux density calibrator {field['sas_id_calibrator2']} for observation {field['target_name']} {field['sas_id_target']} already processed." + ) + return field + else: + if CWL_RUNNER_LINC_CALIBRATOR == "cwltool": + run_linc_calibrator_cwltool(field, calibrator_field=2, db=CURRENT_DB) + elif CWL_RUNNER_LINC_CALIBRATOR == "toil": + run_linc_calibrator_toil(field, calibrator_field=2, db=CURRENT_DB) + else: + raise RuntimeError("Invalid CWL runner specified.") + return field + + @task(trigger_rule=TriggerRule.ALL_DONE) + def select_best_calibrator(result1, result2): + if result1["sas_id_calibrator_final"]: + return result1 + elif result2["sas_id_calibrator_final"]: + return result2 + elif result1 and result2: + cal_template = pathlib.Path(FLUX_CALIBRATOR_TEMPLATE) + if not cal_template.is_file(): + cal = random.randint(1, 2) + print( + f"No flux density calibrator template found. Randomly selected calibrator{cal}" + ) + if cal == 1: + CURRENT_DB.set_final_calibrator( + result1["target_name"], + result1["sas_id_target"], + result1["sas_id_calibrator1"], + ) + return result1 + elif cal == 2: + CURRENT_DB.set_final_calibrator( + result2["target_name"], + result2["sas_id_target"], + result2["sas_id_calibrator2"], + ) + return result2 + else: + outdir = os.path.join(OUTPUT_DIR, result1["target_name"]) + calibrator1_path = get_most_recent_run( + outdir, result1["sas_id_calibrator1"], "LINC_calibrator" + ) + calibrator1_solutions = ( + calibrator1_path / "results_LINC_calibrator" / "cal_solutions.h5" + ) + + calibrator2_path = get_most_recent_run( + outdir, result2["sas_id_calibrator2"], "LINC_calibrator" + ) + calibrator2_solutions = ( + calibrator2_path / "results_LINC_calibrator" / "cal_solutions.h5" + ) + assess_cal1 = assess_and_compare( + FLUX_CALIBRATOR_TEMPLATE, + [calibrator1_solutions], + ) + assess_cal2 = assess_and_compare( + FLUX_CALIBRATOR_TEMPLATE, + [calibrator2_solutions], + ) + score1 = assess_cal1[0]["score"] + score2 = assess_cal2[0]["score"] + print(f"Calibrator 1 score: {score1}") + print(f"Calibrator 2 score: {score2}") + match score1 <= score2: + case True: + print("Best score for calibrator1") + CURRENT_DB.set_final_calibrator( + result1["target_name"], + result1["sas_id_target"], + result1["sas_id_calibrator1"], + ) + return result1 + case False: + print("Best score for calibrator2") + CURRENT_DB.set_final_calibrator( + result2["target_name"], + result2["sas_id_target"], + result2["sas_id_calibrator2"], + ) + return result2 + elif result1 and (not result2): + print("Only cal 1 succeeded, continuing with that") + CURRENT_DB.set_final_calibrator( + result1["target_name"], + result1["sas_id_target"], + result1["sas_id_calibrator1"], + ) + return result1 + elif (not result1) and result2: + print("Only cal 2 succeeded, continuing with that") + CURRENT_DB.set_final_calibrator( + result2["target_name"], + result2["sas_id_target"], + result2["sas_id_calibrator2"], + ) + return result2 + else: + raise AirflowFailException("No calibrators succeeded; stopping processing.") + + @task + def run_linc_target(field): + if (field["status_target"] == PIPELINE_STATUS.finished) or ( + field["status_target"] == PIPELINE_STATUS.processing + ): + return field + else: + if CWL_RUNNER_LINC_TARGET == "cwltool": + run_linc_target_cwltool(field, CURRENT_DB) + elif CWL_RUNNER_LINC_TARGET == "toil": + run_linc_target_toil(field, CURRENT_DB) + else: + raise RuntimeError("Invalid CWL runner specified.") + return field + + @task + def validate_linc_target(field): + return field + + @task(retries=0, retry_delay=datetime.timedelta(seconds=5)) + def run_vlbi_delay(field): + if field["status_vlbi_delay"] == PIPELINE_STATUS.finished: + return field + else: + if CWL_RUNNER_PILOT_DELAY == "cwltool": + run_pilot_delay_cwltool(field, CURRENT_DB) + elif CWL_RUNNER_PILOT_DELAY == "toil": + run_pilot_delay_toil(field, CURRENT_DB) + else: + raise RuntimeError("Invalid CWL runner specified.") + return field + + @task + def run_ddf_pipeline(field): + field = dict(CURRENT_DB.get_db_columns(field["sas_id_target"])[0]) + logsdir = os.path.join(OUTPUT_DIR, field["target_name"], "logs") + if field["status_ddf"] == PIPELINE_STATUS.processing: + print( + f"ddf-pipeline for {field['target_name']} {field['sas_id_target']} should be running, attempting to resume polling..." + ) + with open( + os.path.join( + logsdir, + f"log_DDF-pipeline_{field['target_name']}_{field['sas_id_target']}.txt", + ), + "r", + ) as f_out: + jobid = None + for line in f_out.readlines(): + if "Submitted batch job" in line: + jobid = line.strip().split()[-1] + break + else: + raise AirflowFailException("Failed to recover job id from log.") + + while True: + print(f"Polling DDF-pipeine job {jobid}") + poll_cmd = f"sacct -X -j {jobid} --format=State --noheader" + status = subprocess.run( + poll_cmd, shell=True, text=True, capture_output=True + ).stdout.strip() + if (status == "RUNNING") or (status == "PENDING"): + time.sleep(60) + elif status == "COMPLETED": + CURRENT_DB.set_status_finished( + field["target_name"], + "ddf", + field["sas_id_target"], + ) + return field + elif ( + (status == "FAILED") + or ("TIMEOUT" in status) + or ("CANCELELD" in status) + or ("OUT_OF_MEM" in status) + ): + raise RuntimeError( + f"DDF-pipeline for {field['target_name']} {field['sas_id_target']} failed." + ) + if field["status_ddf"] == PIPELINE_STATUS.finished: + return field + else: + launch_ddf_pipeline(field, CURRENT_DB) + return field + + @task + def run_ddf_subtract(field): + field = dict(CURRENT_DB.get_db_columns(field["sas_id_target"])[0]) + if field["status_vlbi_ddf_subtract"] == PIPELINE_STATUS.finished: + return field + else: + run_pilot_process_ddf_toil(field, CURRENT_DB) + return field + + @task + def run_vlbi_ddcal(field): + field = dict(CURRENT_DB.get_db_columns(field["sas_id_target"])[0]) + if (field["status_vlbi_dd"] == PIPELINE_STATUS.finished) or ( + field["status_vlbi_dd"] == PIPELINE_STATUS.processing + ): + return field + else: + if CWL_RUNNER_PILOT_DDCAL == "cwltool": + run_pilot_ddcal_cwltool(field, CURRENT_DB) + elif CWL_RUNNER_PILOT_DDCAL == "toil": + run_pilot_ddcal_toil(field, CURRENT_DB) + else: + raise RuntimeError("Invalid CWL runner specified.") + return field + + @task + def prepare_ddf(field): + mses_averaged = run_prepare_ddf(field) + if mses_averaged: + return field + else: + raise RuntimeError("No averaged MSes for ddf-pipeline found.") + + @task + def prepare_ddf_subtract(field): + mses_delay_corrected = run_prepare_ddf_subtract(field) + if mses_delay_corrected: + return field + else: + raise RuntimeError("No delay-corrected MSes for ddf subtract found.") + + @task + def run_vlbi_image_intermediate(field): + field = dict(CURRENT_DB.get_db_columns(field["sas_id_target"])[0]) + if field["status_vlbi_intermediate_img"] == PIPELINE_STATUS.finished: + return field + else: + run_pilot_intermediate_image_toil(field, CURRENT_DB) + return field + + @task + def run_vlbi_facet_subtract(field): + field = dict(CURRENT_DB.get_db_columns(field["sas_id_target"])[0]) + if field["status_vlbi_facet_subtract"] == PIPELINE_STATUS.finished: + return field + else: + run_pilot_facet_subtract_toil(field, CURRENT_DB) + return field + + @task + def run_vlbi_facet_imaging(field): + field = dict(CURRENT_DB.get_db_columns(field["sas_id_target"])[0]) + if field["status_vlbi_facet_imaging"] == PIPELINE_STATUS.finished: + return field + else: + run_pilot_facet_imaging_toil(field, CURRENT_DB) + return field + + proceed = check_fields() + get_field = get_unprocessed_target() + field = download_field(get_field) + result_cal1 = run_linc_calibrator1(field) + result_cal2 = run_linc_calibrator2(field) + best_cal = select_best_calibrator(result_cal1, result_cal2) + result_targ = run_linc_target(best_cal) + linc_is_valid = validate_linc_target(result_targ) + result_vlbi_delay = run_vlbi_delay(linc_is_valid) + + proceed >> get_field + await_approval_delay = PythonSensor( + task_id="approve_delay", + python_callable=get_approval, + poke_interval=60, + timeout=86400 * 7, + mode="poke", + op_args=[result_vlbi_delay, "vlbi_delay", NEEDS_MANUAL_APPROVAL_DELAY], + ) + result_prepare_ddf = prepare_ddf(result_vlbi_delay) + result_ddf = run_ddf_pipeline(result_prepare_ddf) + result_prepare_ddf_subtract = prepare_ddf_subtract(result_ddf) + result_ddf_subtract = run_ddf_subtract(result_prepare_ddf_subtract) + result_vlbi_dd = run_vlbi_ddcal(result_ddf_subtract) + result_vlbi_interm_img = run_vlbi_image_intermediate(result_vlbi_dd) + result_vlbi_facet_subtract = run_vlbi_facet_subtract(result_vlbi_interm_img) + _result_vlbi_facet_img = run_vlbi_facet_imaging(result_vlbi_facet_subtract) + + ( + await_approval_delay + >> result_prepare_ddf + >> result_ddf + >> result_prepare_ddf_subtract + >> result_ddf_subtract + >> result_vlbi_dd + ) + + +pilot_widefield() diff --git a/flocs_processing/db_utils.py b/flocs_processing/db_utils.py new file mode 100644 index 0000000..af85aa7 --- /dev/null +++ b/flocs_processing/db_utils.py @@ -0,0 +1,103 @@ +from enum import Enum +import functools +import sqlite3 + + +@functools.total_ordering +class PIPELINE_STATUS(Enum): + nothing = 0 + downloaded = 1 + finished = 2 + await_approval = 3 + processing = 98 + error = 99 + + def __eq__(self, other): + if other.__class__ is int: + return self.value == other + elif other.__class__ is self.__class__: + return self.value == other.value + else: + raise NotImplementedError + + def __lt__(self, other): + if self.__class__ is not other.__class__: + raise NotImplementedError + return self.value < other.value + + +class FlocsDB: + def __init__(self, dbname: str, db_table: str): + self.DATABASE = dbname + self.TABLE_NAME = db_table + + def get_db_columns(self, obsid: str = None): + with sqlite3.connect(self.DATABASE) as db: + db.row_factory = sqlite3.Row + cursor = db.cursor() + columns = "target_name,priority,finished,downloaded,sas_id_calibrator1,sas_id_calibrator2,sas_id_calibrator_final,sas_id_target,status_calibrator1,status_calibrator2,status_target,status_vlbi_delay,status_vlbi_dd,status_ddf,status_vlbi_ddf_subtract,status_vlbi_intermediate_img,status_vlbi_facet_subtract,status_vlbi_facet_img" + if obsid: + field = cursor.execute( + f"select {columns} from {self.TABLE_NAME} where sas_id_target=='{obsid}' and finished==0 order by priority desc" + ).fetchall() + else: + field = cursor.execute( + f"select {columns} from {self.TABLE_NAME} where finished==0 order by priority desc" + ).fetchall() + print(field) + return field + + def set_status_nothing(self, name, identifier, target): + with sqlite3.connect(self.DATABASE) as db: + cursor = db.cursor() + cursor.execute( + f"update {self.TABLE_NAME} set status_{identifier}={PIPELINE_STATUS.nothing.value} where target_name=='{name}' and sas_id_target=='{target}'" + ) + + def set_status_failed(self, name, identifier, target): + with sqlite3.connect(self.DATABASE) as db: + cursor = db.cursor() + cursor.execute( + f"update {self.TABLE_NAME} set status_{identifier}={PIPELINE_STATUS.error.value} where target_name=='{name}' and sas_id_target=='{target}'" + ) + + def set_status_processing(self, name, identifier, target): + with sqlite3.connect(self.DATABASE) as db: + cursor = db.cursor() + cursor.execute( + f"update {self.TABLE_NAME} set status_{identifier}={PIPELINE_STATUS.processing.value} where target_name=='{name}' and sas_id_target=='{target}'" + ) + + def set_status_await_approval(self, name, identifier, target): + with sqlite3.connect(self.DATABASE) as db: + cursor = db.cursor() + cursor.execute( + f"update {self.TABLE_NAME} set status_{identifier}={PIPELINE_STATUS.await_approval.value} where target_name=='{name}' and sas_id_target=='{target}'" + ) + + def set_status_finished(self, name, identifier, target): + with sqlite3.connect(self.DATABASE) as db: + cursor = db.cursor() + cursor.execute( + f"update {self.TABLE_NAME} set status_{identifier}={PIPELINE_STATUS.finished.value} where target_name=='{name}' and sas_id_target=='{target}'" + ) + + def set_status_downloaded(self, name, target): + with sqlite3.connect(self.DATABASE) as db: + cursor = db.cursor() + cursor.execute( + f"update {self.TABLE_NAME} set downloaded=1 where target_name=='{name}' and sas_id_target=='{target}'" + ) + + def set_field_finished(self, name, target): + query = f"update {self.TABLE_NAME} set finished=1 where target_name=='{name}' and sas_id_target=='{target}'" + with sqlite3.connect(self.DATABASE) as db: + cursor = db.cursor() + cursor.execute(query) + + def set_final_calibrator(self, name, target, final_cal): + with sqlite3.connect(self.DATABASE) as db: + cursor = db.cursor() + cursor.execute( + f"update {self.TABLE_NAME} set sas_id_calibrator_final={final_cal} where target_name=='{name}' and sas_id_target=='{target}'" + ) diff --git a/flocs_processing/flocs_processing.py b/flocs_processing/flocs_processing.py index f23ef2e..db9becb 100644 --- a/flocs_processing/flocs_processing.py +++ b/flocs_processing/flocs_processing.py @@ -1,51 +1,32 @@ #!/usr/bin/env python -from astropy.table import Table -from concurrent.futures import ProcessPoolExecutor +from .db_utils import FlocsDB +from .processors import FlocsAirflowProcessor from cyclopts import Parameter from enum import Enum -from rich.console import Console -from typing import Annotated import cyclopts import functools -import glob -import os -import pathlib -import re import subprocess -import sqlite3 -import threading -import time +import structlog +from typing import Annotated, Literal, Optional, get_args app = cyclopts.App() +logger = structlog.getLogger() -@functools.total_ordering -class PIPELINE(Enum): - download = 0 - linc_calibrator = 1 - linc_target = 2 - vlbi_delay = 3 - - def __eq__(self, other): - if self.__class__ is not other.__class__: - raise NotImplementedError - return self.value == other.value - - def __lt__(self, other): - if self.__class__ is not other.__class__: - raise NotImplementedError - return self.value < other.value - - def __hash__(self): - return hash(self.value) - -PIPELINE_NAMES: dict[PIPELINE, str] = { - PIPELINE.download: "not downloaded", - PIPELINE.linc_calibrator: "LINC Calibrator", - PIPELINE.linc_target: "LINC Target", - PIPELINE.vlbi_delay: "PILOT delay calibration", -} +PIPELINES = Literal[ + "all", + "calibrator1", + "calibrator2", + "target", + "vlbi_delay", + "ddf", + "vlbi_ddf_subtract", + "vlbi_dd", + "vlbi_intermediate_img", + "vlbi_facet_subtract", + "vlbi_facet_img", +] @functools.total_ordering @@ -68,24 +49,6 @@ def __lt__(self, other): return self.value < other.value -@functools.total_ordering -class STAGING_STATUS(Enum): - error = -1 - not_staged = 0 - in_progress = 1 - finished = 2 - - def __eq__(self, other): - if self.__class__ is not other.__class__: - raise NotImplementedError - return self.value == other.value - - def __lt__(self, other): - if self.__class__ is not other.__class__: - raise NotImplementedError - return self.value < other.value - - @app.command() def create_database( dbname: Annotated[ @@ -98,18 +61,24 @@ def create_database( "linc" ], ): + pipeline_str = ",".join(pipelines) pipelines = list(map(str.lower, pipelines)) - dbstr = f"create table {table_name}(source_name text default NULL" + dbstr = f"create table {table_name}(target_name text default NULL, pipelines text default '{pipeline_str}', priority int default 0, finished bit default 0, downloaded bit default 0" if "linc" in pipelines: - dbstr += ", sas_id_calibrator1 text default NULL, sas_id_calibrator2 text default NULL, sas_id_calibrator_final text default NULL, sas_id_target text primary key default NULL, status_calibrator1 smallint default 0, status_calibrator2 smallint default 0, status_target smallint default 0" + dbstr += f", sas_id_calibrator1 text default NULL, sas_id_calibrator2 text default NULL, sas_id_calibrator_final text default NULL, sas_id_target text primary key default NULL, status_calibrator1 smallint default {PIPELINE_STATUS.nothing.value}, status_calibrator2 smallint default {PIPELINE_STATUS.nothing.value}, status_target smallint default {PIPELINE_STATUS.nothing.value}" if "ddf-pipeline" in pipelines: - dbstr += ", status_ddf smallint default 0" + dbstr += f", status_ddf smallint default {PIPELINE_STATUS.nothing.value}" if "vlbi-delay-widefield" in pipelines: - dbstr += ", status_ddf smallint default 0" - dbstr += ", status_delay smallint default 0" + dbstr += f", status_vlbi_delay smallint default {PIPELINE_STATUS.nothing.value}" + dbstr += f", status_vlbi_dd smallint default {PIPELINE_STATUS.nothing.value}" + dbstr += f", status_vlbi_ddf_subtract smallint default {PIPELINE_STATUS.nothing.value}" + dbstr += f", status_vlbi_intermediate_img smallint default {PIPELINE_STATUS.nothing.value}" + dbstr += f", status_vlbi_facet_subtract smallint default {PIPELINE_STATUS.nothing.value}" + dbstr += f", status_vlbi_facet_imaging smallint default {PIPELINE_STATUS.nothing.value}" if "vlbi-delay-single-target" in pipelines: - dbstr += ", status_delay smallint default 0" + dbstr += f", status_vlbi_delay smallint default {PIPELINE_STATUS.nothing.value}" + dbstr += f", status_vlbi_dd smallint default {PIPELINE_STATUS.nothing.value}" dbstr += ");" cmd = ["sqlite3", dbname, dbstr] @@ -137,18 +106,21 @@ def add_field( table_name: Annotated[ str, Parameter(help="Database table that will be processed.") ] = "processing_flocs", + pipelines: Annotated[ + str, Parameter(help="Pipelines this field needs to be processed with.") + ] = "", ): - dbstr = f"insert into {table_name} (source_name" + dbstr = f"insert into {table_name} (target_name" if len(sas_id_calibrators) == 1: dbstr += ", sas_id_calibrator1" if len(sas_id_calibrators) == 2: dbstr += ", sas_id_calibrator1, sas_id_calibrator2" dbstr += f", sas_id_target) values ('{field_name}', " if len(sas_id_calibrators) == 1: - dbstr += f"{sas_id_calibrators[0]}" + dbstr += f"'{sas_id_calibrators[0]}', " if len(sas_id_calibrators) == 2: - dbstr += f"{sas_id_calibrators[0]}, {sas_id_calibrators[1]}, " - dbstr += f"{sas_id_target})" + dbstr += f"'{sas_id_calibrators[0]}', '{sas_id_calibrators[1]}', " + dbstr += f"'{sas_id_target}')" cmd = ["sqlite3", dbname, dbstr] print(f"Adding field {field_name} to {table_name} via: {" ".join(cmd)}") @@ -159,598 +131,77 @@ def add_field( @app.command() -def process_from_database( +def update_field( + field_name: Annotated[str, Parameter(help="Name of the source/field to add.")], + sas_id_target: Annotated[ + str, + Parameter(help="SAS ID of the target to add.", consume_multiple=True), + ], dbname: Annotated[ str, Parameter(help="Sqlite3 database from which processing will be done.") ], - rundir: Annotated[ - str, - Parameter( - help="Directory where data is located and processing will take place." - ), + pipeline: Annotated[ + PIPELINES, + Parameter(help="Pipeline of which to update the status"), ], - slurm_queues: Annotated[ - list[str], Parameter(help="Slurm queues that jobs can be submitted to.") + set_status: Annotated[ + Literal["nothing", "downloaded", "processing", "failed", "success"], + Parameter(help="Set the status of the given pipeline."), ], - slurm_account: Annotated[str, Parameter(help="Slurm account to submit under.")], table_name: Annotated[ str, Parameter(help="Database table that will be processed.") ] = "processing_flocs", ): - fp = FlocsSlurmProcessor( - database=dbname, - slurm_queues=slurm_queues, - slurm_account=slurm_account, - table_name=table_name, - rundir=rundir, - ) - fp.start_processing_loop() - - -class FlocsSlurmProcessor: - def __init__( - self, - database: str, - slurm_queues: list, - slurm_account: str, - rundir: str, - table_name: Annotated[ - str, Parameter(help="Database table to start processing in.") - ] = "flocs_processing", - ): - self.DATABASE = database - self.SLURM_QUEUES = ",".join(slurm_queues) - self.SLURM_ACCOUNT = slurm_account - self.TABLE_NAME = table_name - self.RUNDIR = rundir - - def launch_calibrator(self, field_name, sas_id, restart: bool = False): - if not restart: - try: - cmd = f"flocs-run linc calibrator --record-toil-stats --scheduler slurm --rundir {self.RUNDIR}/{field_name}/rundir/ --outdir {self.RUNDIR}/{field_name} --slurm-queue {self.SLURM_QUEUES} --slurm-time 24:00:00 --slurm-account {self.SLURM_ACCOUNT} --runner toil --save-raw-solutions {self.RUNDIR}/{field_name}/calibrator/L{sas_id}" - print(cmd) - with open( - f"{field_name}/log_LINC_calibrator_{field_name}_{sas_id}.txt", "a" - ) as f_out, open( - f"{field_name}/log_LINC_calibrator_{field_name}_{sas_id}_err.txt", - "a", - ) as f_err: - proc = subprocess.run( - cmd, shell=True, text=True, stdout=f_out, stderr=f_err - ) - if not proc.returncode: - return True - else: - return False - except subprocess.CalledProcessError: - print("something went wrong") - else: - rundirs = pathlib.Path(f"{self.RUNDIR}/{field_name}/rundir") - rundirs_sorted = sorted(rundirs.iterdir(), key=os.path.getctime) - rundirs_sorted_filtered = [ - d for d in rundirs_sorted if sas_id in d.parts[-1] - ] - # Last directory touched for this source - rundir_final = rundirs_sorted_filtered[-1].parts[-1] - try: - cmd = f"flocs-run linc calibrator --record-toil-stats --scheduler slurm --rundir {self.RUNDIR}/{field_name}/rundir/{rundir_final} --outdir {self.RUNDIR}/{field_name} --restart --slurm-queue {self.SLURM_QUEUES} --slurm-time 24:00:00 --slurm-account {self.SLURM_ACCOUNT} --runner toil --save-raw-solutions {self.RUNDIR}/{field_name}/calibrator/L{sas_id}" - print(cmd) - with open( - f"{field_name}/log_LINC_calibrator_{field_name}_{sas_id}.txt", "a" - ) as f_out, open( - f"{field_name}/log_LINC_calibrator_{field_name}_{sas_id}_err.txt", - "a", - ) as f_err: - proc = subprocess.run( - cmd, shell=True, text=True, stdout=f_out, stderr=f_err - ) - if not proc.returncode: - return True - else: - return False - except subprocess.CalledProcessError: - print("something went wrong") - return False - - def launch_target(self, field_name, sas_id, sas_id_cal, restart: bool = False): - if not restart: - try: - cal_sol_path = glob.glob( - f"{self.RUNDIR}/{field_name}/LINC_calibrator_L{sas_id_cal}*/results_LINC_calibrator/cal_solutions.h5" - )[0] - cmd = f"flocs-run linc target --record-toil-stats --scheduler slurm --rundir {self.RUNDIR}/{field_name}/rundir/ --outdir {self.RUNDIR}/{field_name} --slurm-queue {self.SLURM_QUEUES} --slurm-time 48:00:00 --slurm-account {self.SLURM_ACCOUNT} --runner toil --output-fullres-data --min-unflagged-fraction 0.05 --cal-solutions {cal_sol_path} {self.RUNDIR}/{field_name}/target/L{sas_id}/" - print(cmd) - with open( - f"{field_name}/log_LINC_target_{field_name}_{sas_id}.txt", "w" - ) as f_out, open( - f"{field_name}/log_LINC_target_{field_name}_{sas_id}_err.txt", "w" - ) as f_err: - proc = subprocess.run( - cmd, shell=True, text=True, stdout=f_out, stderr=f_err - ) - if not proc.returncode: - pattern = re.compile(r"Workflow.* stopped. Success: False") - if pattern.search(proc.stderr): - return False - return True - else: - return False - except subprocess.CalledProcessError: - print("something went wrong") - else: - rundirs = pathlib.Path(f"{self.RUNDIR}/{field_name}/rundir") - rundirs_sorted = sorted(rundirs.iterdir(), key=os.path.getctime) - rundirs_sorted_filtered = [ - d for d in rundirs_sorted if sas_id in d.parts[-1] - ] - # Last directory touched for this source - rundir_final = rundirs_sorted_filtered[-1].parts[-1] - try: - cal_sol_path = glob.glob( - f"{self.RUNDIR}/{field_name}/LINC_calibrator_L{sas_id_cal}*/results_LINC_calibrator/cal_solutions.h5" - )[0] - cmd = f"flocs-run linc target --record-toil-stats --scheduler slurm --rundir {self.RUNDIR}/{field_name}/rundir/{rundir_final} --restart --outdir {self.RUNDIR}/{field_name} --slurm-queue {self.SLURM_QUEUES} --slurm-time 48:00:00 --slurm-account {self.SLURM_ACCOUNT} --runner toil --output-fullres-data --min-unflagged-fraction 0.05 --cal-solutions {cal_sol_path} {self.RUNDIR}/{field_name}/target/L{sas_id}/" - print(cmd) - with open( - f"{field_name}/log_LINC_target_{field_name}_{sas_id}.txt", "a" - ) as f_out, open( - f"{field_name}/log_LINC_target_{field_name}_{sas_id}_err.txt", "a" - ) as f_err: - proc = subprocess.run( - cmd, shell=True, text=True, stdout=f_out, stderr=f_err - ) - if not proc.returncode: - return True - else: - return False - return True - except subprocess.CalledProcessError: - print("something went wrong") - return False - - def launch_vlbi_delay(self, field_name, sas_id, restart: bool = False): - if not restart: - print(f"Generating input catalogue(s) for {field_name}") - rundirs = pathlib.Path(f"{self.RUNDIR}/{field_name}/") - rundirs_sorted = sorted(rundirs.iterdir(), key=os.path.getctime) - rundirs_sorted_filtered = [ - d - for d in rundirs_sorted - if ((sas_id in d.parts[-1]) and (d.parts[-1].startswith(("LINC_target")))) - ] - # Last LINC target reduction for this source - linc_target_dir = rundirs_sorted_filtered[-1] - first_ms = glob.glob( - f"{linc_target_dir}/results_LINC_target/results/*.dp3concat" - )[0] - - with open( - f"{rundirs}/log_VLBI_delay-calibration_plot_field_{field_name}_{sas_id}.txt", - "w", - ) as f_out, open( - f"{rundirs}/log_VLBI_delay-calibration_plot_field_{field_name}_{sas_id}_err.txt", - "w", - ) as f_err: - cmd = f"lofar-vlbi-plot --output_dir {rundirs} --MS {first_ms} --continue_no_lotss --vlass" - proc = subprocess.run( - cmd, shell=True, text=True, stdout=subprocess.PIPE - ) - if proc.returncode: - return False - delay_csv = rundirs / "delay_calibrators.csv" - if not os.path.isfile(delay_csv): - print(f"Failed to find delay_calibrators.csv for {field_name}") - return False - dc = Table.read(delay_csv) - model_image = rundirs / f"{dc[0]['Observation']}_vlass.fits" - try: - cmd = f"flocs-run vlbi delay-calibration --record-toil-stats --scheduler slurm --rundir {rundirs/'rundir'} --outdir {rundirs} --slurm-queue {self.SLURM_QUEUES} --slurm-time 48:00:00 --slurm-account {self.SLURM_ACCOUNT} --runner toil --delay-calibrator {delay_csv} --model-image {model_image} --ms-suffix dp3concat {linc_target_dir/'results_LINC_target'/'results'}" - print(cmd) - os.chdir(rundirs) - with open( - f"{rundirs}/log_VLBI_delay-calibration_{field_name}_{sas_id}.txt", - "w", - ) as f_out, open( - f"{rundirs}/log_VLBI_delay-calibration_{field_name}_{sas_id}_err.txt", - "w", - ) as f_err: - proc = subprocess.run( - cmd, shell=True, text=True, stdout=f_out, stderr=f_err - ) - if not proc.returncode: - pattern = re.compile(r"Workflow.* stopped. Success: False") - if pattern.search(proc.stderr): - return False - return True - else: - return False - except subprocess.CalledProcessError: - print("something went wrong") - return False - else: - rundirs = pathlib.Path(f"{self.RUNDIR}/{field_name}/") - rundirs_sorted = sorted(rundirs.iterdir(), key=os.path.getctime) - rundirs_sorted_filtered = [ - d - for d in rundirs_sorted - if ((sas_id in d.parts[-1]) and (d.parts[-1].startswith(("LINC_target")))) - ] - # Last LINC target reduction for this source - linc_target_dir = rundirs_sorted_filtered[-1] - print(f"{linc_target_dir=}") - - vlbi_rundirs = pathlib.Path(f"{self.RUNDIR}/{field_name}/rundir") - vlbi_rundirs_sorted = sorted(vlbi_rundirs.iterdir(), key=os.path.getctime) - # vlbi_rundirs_sorted_filtered = [d for d in vlbi_rundirs_sorted if ((sas_id in d.parts[-1]) and ("delay" in d.parts[-1]))] - vlbi_rundirs_sorted_filtered = [ - d for d in vlbi_rundirs_sorted if ("delay" in d.parts[-1]) - ] - vlbi_dir = vlbi_rundirs_sorted_filtered[-1] - print(f"{vlbi_dir=}") - - delay_csv = rundirs / "delay_calibrators.csv" - try: - cmd = f"flocs-run vlbi delay-calibration --record-toil-stats --scheduler slurm --rundir {vlbi_dir} --restart --outdir {self.RUNDIR}/{field_name} --slurm-queue {self.SLURM_QUEUES} --slurm-time 48:00:00 --slurm-account {self.SLURM_ACCOUNT} --runner toil --delay-calibrator {delay_csv} --ms-suffix dp3concat {linc_target_dir/'results_LINC_target/results'}" - print(cmd) - os.chdir(rundirs) - with open( - f"{rundirs}/log_VLBI_delay-calibration_{field_name}_{sas_id}.txt", - "w", - ) as f_out, open( - f"{rundirs}/log_VLBI_delay-calibration_{field_name}_{sas_id}_err.txt", - "w", - ) as f_err: - proc = subprocess.run( - cmd, shell=True, text=True, stdout=f_out, stderr=f_err - ) - if not proc.returncode: - pattern = re.compile(r"Workflow.* stopped. Success: False") - if pattern.search(proc.stderr): - return False - return True - else: - return False - except subprocess.CalledProcessError: - print("something went wrong") - return False - - def summarise_status(self): - console = Console(highlight=False) - console.print(f"General statistics for {self.DATABASE}", style="bold") - with sqlite3.connect(self.DATABASE) as db: - cursor = db.cursor() - not_started = cursor.execute( - f"select count(source_name) from {self.TABLE_NAME} where (status_calibrator1=={PIPELINE_STATUS.nothing.value} or status_calibrator2=={PIPELINE_STATUS.nothing.value})" - ).fetchall()[0][0] - downloaded = cursor.execute( - f"select count(source_name) from {self.TABLE_NAME} where (status_calibrator1=={PIPELINE_STATUS.downloaded.value} or status_calibrator2=={PIPELINE_STATUS.downloaded.value})" - ).fetchall()[0][0] - processing = cursor.execute( - f"select count(source_name) from {self.TABLE_NAME} where (status_calibrator1=={PIPELINE_STATUS.processing.value} or status_calibrator2=={PIPELINE_STATUS.processing.value})" - ).fetchall()[0][0] - finished = cursor.execute( - f"select count(source_name) from {self.TABLE_NAME} where (status_calibrator1=={PIPELINE_STATUS.finished.value} or status_calibrator2=={PIPELINE_STATUS.finished.value})" - ).fetchall()[0][0] - error = cursor.execute( - f"select count(source_name) from {self.TABLE_NAME} where (status_calibrator1=={PIPELINE_STATUS.error.value} or status_calibrator2=={PIPELINE_STATUS.error.value})" - ).fetchall()[0][0] - console.print("Flux density calibrators", style="bold") - console.print(f"= {not_started} calibrators not yet downloaded") - console.print(f"= {downloaded} calibrators downloaded", style="yellow") - console.print(f"= {processing} calibrators processing", style="cyan") - console.print(f"= {finished} calibrators finished", style="green") - console.print(f"= {error} calibrators failed", style="red") - - not_started = cursor.execute( - f"select count(source_name) from {self.TABLE_NAME} where status_target=={PIPELINE_STATUS.nothing.value}" - ).fetchall()[0][0] - downloaded = cursor.execute( - f"select count(source_name) from {self.TABLE_NAME} where status_target=={PIPELINE_STATUS.downloaded.value}" - ).fetchall()[0][0] - finished = cursor.execute( - f"select count(source_name) from {self.TABLE_NAME} where status_target=={PIPELINE_STATUS.finished.value}" - ).fetchall()[0][0] - processing = cursor.execute( - f"select count(source_name) from {self.TABLE_NAME} where status_target=={PIPELINE_STATUS.processing.value}" - ).fetchall()[0][0] - error = cursor.execute( - f"select count(source_name) from {self.TABLE_NAME} where status_target=={PIPELINE_STATUS.error.value}" - ).fetchall()[0][0] - console.print("\nFields / science targets", style="bold") - console.print(f"= {not_started} targets not yet downloaded") - console.print(f"= {downloaded} targets downloaded", style="yellow") - console.print(f"= {processing} targets processing", style="cyan") - console.print(f"= {finished} targets finished", style="green") - console.print(f"= {error} targets failed", style="red") - - def update_db_statuses(self, running_fields): - if not running_fields: - return - print("== UPDATING DB STATUSES ==") - console = Console(highlight=False) - futures = running_fields.keys() - to_delete = set() - for future in futures: - field = running_fields[future] - console.print(f"[bold]Field: {field['name']}[/bold]") - console.print(f"Target SAS ID: {field['sasid']}") - console.print(f"Pipeline: {PIPELINE_NAMES[field['pipeline']]}") - if future.cancelled(): - console.print("Status: [bold yellow]cancelled[/bold yellow]") - del running_fields[future] - elif future.done(): - if future.exception(timeout=10): - console.print("Status: [bold red]failed[/bold red]") - print( - f"Processing {field['identifier']} for {field['name']} failed." - ) - print("Error was: ", future.exception(timeout=10)) - with sqlite3.connect(self.DATABASE) as db: - cursor = db.cursor() - cursor.execute( - f"update {self.TABLE_NAME} set status_{field['identifier']}={PIPELINE_STATUS.error.value} where source_name=='{field['name']}'" - ) - else: - result = future.result() - print(f"Result was {result}") - with sqlite3.connect(self.DATABASE) as db: - cursor = db.cursor() - if result: - console.print("Status: [bold green]finished[/bold green]") - cursor.execute( - f"update {self.TABLE_NAME} set status_{field['identifier']}={PIPELINE_STATUS.finished.value} where source_name=='{field['name']}'" - ) - if field["identifier"] == "target": - cursor.execute( - f"update {self.TABLE_NAME} set status_delay={PIPELINE_STATUS.downloaded.value} where source_name=='{field['name']}'" - ) - else: - console.print("Status: [bold red]failed[/bold red]") - cursor.execute( - f"update {self.TABLE_NAME} set status_{field['identifier']}={PIPELINE_STATUS.error.value} where source_name=='{field['name']}'" - ) - to_delete.add(future) - else: - console.print("Status: [bold cyan]running[/bold cyan]\n") - for f in to_delete: - print(f"Deleting future for {field['name']}") - del running_fields[f] - print("== UPDATING DB STATUSES FINISHED") + db = FlocsDB(dbname=dbname, db_table=table_name) + if pipeline == "all": + p_list = list(get_args(PIPELINES)) + p_list.remove("all") + for p in p_list: + logger.info(f"Setting pipeline {p} to status {set_status}") + match set_status: + case "nothing": + db.set_status_nothing(field_name, p, sas_id_target) + case "downloaded": + db.set_status_downloaded(field_name, sas_id_target) + case "processing": + db.set_status_processing(field_name, p, sas_id_target) + case "failed": + db.set_status_failed(field_name, p, sas_id_target) + case "success": + db.set_status_finished(field_name, p, sas_id_target) + else: + logger.info(f"Setting pipeline {pipeline} to status {set_status}") + match set_status: + case "nothing": + db.set_status_nothing(field_name, pipeline, sas_id_target) + case "downloaded": + db.set_status_downloaded(field_name, sas_id_target) + case "processing": + db.set_status_processing(field_name, pipeline, sas_id_target) + case "failed": + db.set_status_failed(field_name, pipeline, sas_id_target) + case "success": + db.set_status_finished(field_name, pipeline, sas_id_target) - def get_not_started(self, identifier: str): - not_started = self.get_db_columns(identifier, PIPELINE_STATUS.downloaded) - return not_started - def get_failed(self, identifier: str): - restart = self.get_db_columns(identifier, PIPELINE_STATUS.error) - return restart - - def get_db_columns(self, identifier: str, status: PIPELINE_STATUS): - with sqlite3.connect(self.DATABASE) as db: - cursor = db.cursor() - if "calibrator" in identifier: - columns = "source_name,sas_id_calibrator1,sas_id_calibrator2,sas_id_calibrator_final,sas_id_target" - elif "target" in identifier: - columns = "source_name,sas_id_calibrator_final,sas_id_target" - elif "delay" in identifier: - columns = "source_name,sas_id_target" - else: - columns = "*" - restart = cursor.execute( - f"select {columns} from {self.TABLE_NAME} where status_{identifier}=={status.value}" - ).fetchall() - return restart - - def is_processing(self, name, running_fields): - return name in [v["name"] for f, v in running_fields.items()] - - def set_status_processing(self, name, identifier, target): - with sqlite3.connect(self.DATABASE) as db: - cursor = db.cursor() - cursor.execute( - f"update {self.TABLE_NAME} set status_{identifier}={PIPELINE_STATUS.processing.value} where source_name=='{name}' and sas_id_target=='{target}'" - ) - - def check_fields_linc_calibrator(self, running_fields, tpe): - restart1 = self.get_failed("calibrator1") - restart2 = self.get_failed("calibrator2") - if restart1: - for name, cal1, cal2, cal_final, target in restart1: - if ( - not self.is_processing(name, running_fields) - and self.is_accepting_jobs - ): - print( - f"Re-starting LINC calibrator for calibrator 1 of field {name}" - ) - future = tpe.submit( - self.launch_calibrator, name, cal1, restart=True - ) - running_fields[future] = { - "name": name, - "pipeline": PIPELINE.linc_calibrator, - "identifier": "calibrator1", - "sasid": target, - } - self.set_status_processing(name, "calibrator1", target) - if restart2: - for name, cal1, cal2, cal_final, target in restart2: - if ( - not self.is_processing(name, running_fields) - and self.is_accepting_jobs - ): - print( - f"Re-starting LINC calibrator for calibrator 2 of field {name}" - ) - future = tpe.submit( - self.launch_calibrator, name, cal2, restart=True - ) - running_fields[future] = { - "name": name, - "pipeline": PIPELINE.linc_calibrator, - "identifier": "calibrator2", - "sasid": target, - } - self.set_status_processing(name, "calibrator2", target) - - not_started1 = self.get_not_started("calibrator1") - not_started2 = self.get_not_started("calibrator2") - if not_started1: - for name, cal1, cal2, cal_final, target in not_started1: - if ( - not self.is_processing(name, running_fields) - and self.is_accepting_jobs - ): - print( - f"Re-starting LINC calibrator for calibrator 1 of field {name}" - ) - future = tpe.submit(self.launch_calibrator, name, cal1) - running_fields[future] = { - "name": name, - "pipeline": PIPELINE.linc_calibrator, - "identifier": "calibrator1", - "sasid": target, - } - self.set_status_processing(name, "calibrator1", target) - if not_started2: - for name, cal1, cal2, cal_final, target in not_started2: - if ( - not self.is_processing(name, running_fields) - and self.is_accepting_jobs - ): - print( - f"Re-starting LINC calibrator for calibrator 2 of field {name}" - ) - future = tpe.submit(self.launch_calibrator, name, cal2) - running_fields[future] = { - "name": name, - "pipeline": PIPELINE.linc_calibrator, - "identifier": "calibrator2", - "sasid": target, - } - self.set_status_processing(name, "calibrator2", target) - - def check_fields_linc_target(self, running_fields, tpe): - restart = self.get_failed("target") - if restart: - for name, cal_final, target in restart: - if ( - not self.is_processing(name, running_fields) - and self.is_accepting_jobs - ): - print(f"Re-starting LINC target for field {name}") - future = tpe.submit( - self.launch_target, - name, - target, - cal_final, - restart=True, - ) - running_fields[future] = { - "name": name, - "pipeline": PIPELINE.linc_target, - "sasid": target, - "identifier": "target", - } - self.set_status_processing(name, "target", target) - self.is_accepting_jobs = len(running_fields) < self.MAX_RUNNING - print(f"Launched {name}") - - not_started = self.get_not_started("target") - if not_started: - for name, cal_final, target in not_started: - if ( - not self.is_processing(name, running_fields) - and self.is_accepting_jobs - ): - print(f"Starting LINC target for field {name}") - future = tpe.submit(self.launch_target, name, target, cal_final) - running_fields[future] = { - "name": name, - "pipeline": PIPELINE.linc_target, - "sasid": target, - "identifier": "target", - } - self.set_status_processing(name, "target", target) - self.is_accepting_jobs = len(running_fields) < self.MAX_RUNNING - print(f"Launched {name}") - - def check_fields_vlbi_delay(self, running_fields, tpe): - restart = self.get_failed("delay") - if restart: - for name, target in restart: - if ( - not self.is_processing(name, running_fields) - and self.is_accepting_jobs - ): - print(f"Re-starting VLBI delay for field {name}") - future = tpe.submit( - self.launch_vlbi_delay, - name, - target, - restart=True, - ) - running_fields[future] = { - "name": name, - "pipeline": PIPELINE.vlbi_delay, - "sasid": target, - "identifier": "delay", - } - self.set_status_processing(name, "delay", target) - print(f"Launched {name}") - - not_started = self.get_not_started("delay") - if not_started: - for name, target in not_started: - if ( - not self.is_processing(name, running_fields) - and self.is_accepting_jobs - ): - print(f"Starting VLBI delay for field {name}") - future = tpe.submit(self.launch_vlbi_delay, name, target) - running_fields[future] = { - "name": name, - "pipeline": PIPELINE.vlbi_delay, - "sasid": target, - "identifier": "delay", - } - self.set_status_processing(name, "delay", target) - print(f"Launched {name}") - - def start_processing_loop(self, allow_up_to=PIPELINE.linc_calibrator): - print("Starting processing loop") - allow_up_to = PIPELINE.vlbi_delay - self.MAX_RUNNING = 3 - max_noqueue = 5 - noqueue = 0 - lock = threading.RLock() - with ProcessPoolExecutor(max_workers=self.MAX_RUNNING + 1) as tpe: - running_fields = {} - - while True: - if len(running_fields) < 1: - noqueue += 1 - else: - noqueue = 0 - self.summarise_status() - if noqueue >= max_noqueue: - print( - f"No new jobs added in queue for {max_noqueue * 60} s, quitting processing loop." - ) - break - self.is_accepting_jobs = len(running_fields) < self.MAX_RUNNING - if allow_up_to >= PIPELINE.linc_calibrator: - with lock: - self.check_fields_linc_calibrator(running_fields, tpe) - if allow_up_to >= PIPELINE.linc_target: - with lock: - self.check_fields_linc_target(running_fields, tpe) - if allow_up_to >= PIPELINE.vlbi_delay: - with lock: - self.check_fields_vlbi_delay(running_fields, tpe) - with lock: - self.update_db_statuses(running_fields) - time.sleep(60) +@app.command() +def deploy_airflow( + airflow_cores: Annotated[ + int, Parameter(help="Number of cores to give to Airflow.") + ] = 6, + custom_airflow_home: Annotated[ + Optional[str], + Parameter( + help="Directory where Airflow stores its internal data. If given this overrides $AIRFLOW_HOME. If not given, $AIRFLOW_HOME is used to define other related variables." + ), + ] = "", +): + fp = FlocsAirflowProcessor(custom_airflow_home, airflow_cores) + fp.generate_jwt_secret() + fp.check_environment() + fp.check_airflow_init() + fp.write_source_file() + fp.deploy_airflow_tmux() def main(): diff --git a/flocs_processing/pipeline_runners.py b/flocs_processing/pipeline_runners.py new file mode 100644 index 0000000..de1e2ce --- /dev/null +++ b/flocs_processing/pipeline_runners.py @@ -0,0 +1,1485 @@ +import configparser +import os +import pathlib +import re +import subprocess +import time + +from airflow.exceptions import AirflowFailException +from airflow.sdk import get_current_context +from losoto.h5parm import h5parm + +from flocs_processing.db_utils import PIPELINE_STATUS, FlocsDB + +# Need to think of a way to centralise this and not read multiple times here and in the DAG +if "FLOCS_AIRFLOW_CONFIG" not in os.environ: + if not pathlib.Path(os.path.expandvars("$HOME/.flocs_airflow.cfg")).is_file(): + raise RuntimeError( + "FLOCS_AIRFLOW_CONFIG environment variable not set and no $HOME/.flocs_airflow.cfg exists. Please create a valid configuration file." + ) + else: + CONFIG_FILE = os.path.expandvars("$HOME/.flocs_airflow.cfg") +else: + CONFIG_FILE = os.getenv("FLOCS_AIRFLOW_CONFIG") or "" + +parser = configparser.ConfigParser() +parser.optionxform = str # ty: ignore[invalid-assignment] +with open(CONFIG_FILE, "r") as config: + parser.read_string("[DEFAULT]\n" + config.read()) + +print("Config summary:") +for k, v in parser["DEFAULT"].items(): + print(f"{k}: {v}") + +SLURM_ACCOUNT = parser["DEFAULT"]["SLURM_ACCOUNT"] +SLURM_QUEUE = parser["DEFAULT"]["SLURM_QUEUE"] +DATA_DIR = parser["DEFAULT"]["DATA_DIR"] +OUTPUT_DIR = parser["DEFAULT"]["OUTPUT_DIR"] +PROCESSING_DIR = parser["DEFAULT"]["PROCESSING_DIR"] +NN_MODEL_CACHE = parser["DEFAULT"]["NN_MODEL_CACHE"] +DDF_CONFIG = parser["DEFAULT"]["DDF_CONFIG"] +FLUX_CALIBRATOR_TEMPLATE = parser["DEFAULT"]["FLUX_CALIBRATOR_TEMPLATE"] +NEEDS_MANUAL_APPROVAL_DELAY = parser.getboolean( + "DEFAULT", "NEEDS_MANUAL_APPROVAL_DELAY" +) + + +def get_most_recent_run(searchpath: str, sas_id: str, pipeline: str) -> pathlib.Path: + rundirs = pathlib.Path(searchpath) + rundirs_sorted = sorted(rundirs.iterdir()) + if pipeline: + rundirs_sorted_filtered = [ + d + for d in rundirs_sorted + if ((sas_id in d.parts[-1]) and (pipeline in d.parts[-1])) and d.is_dir() + ] + else: + rundirs_sorted_filtered = [d for d in rundirs_sorted if sas_id in d.parts[-1]] + try: + rundir_final = rundirs_sorted_filtered[-1].absolute() + return rundir_final + except IndexError: + print(f"No {pipeline} run for {sas_id} found") + raise RuntimeError(f"No {pipeline} run for {sas_id} found") + + +def run_linc_calibrator_cwltool(field, calibrator_field: int, db: FlocsDB): + print( + f"Processing flux density calibrator {field[f'sas_id_calibrator{calibrator_field}']} for observation {field['target_name']} {field['sas_id_target']}" + ) + ms_folder = f"L{field[f'sas_id_calibrator{calibrator_field}']}" + db.set_status_processing( + field["target_name"], f"calibrator{calibrator_field}", field["sas_id_target"] + ) + outdir = os.path.join(OUTPUT_DIR, field["target_name"]) + logsdir = os.path.join(OUTPUT_DIR, field["target_name"], "logs") + if not os.path.isdir(outdir): + os.makedirs(outdir, exist_ok=True) + if not os.path.isdir(logsdir): + os.makedirs(logsdir, exist_ok=True) + cmd = f"flocs-run linc calibrator --runner cwltool --scheduler slurm --slurm-cores 32 --slurm-account {SLURM_ACCOUNT} --slurm-queue {SLURM_QUEUE} --rundir {PROCESSING_DIR} --outdir {outdir} {os.path.join(DATA_DIR, field['target_name'], 'calibrator', ms_folder)}" + with ( + open( + os.path.join( + logsdir, + f"log_LINC_calibrator_{field['target_name']}_{field[f'sas_id_calibrator{calibrator_field}']}.txt", + ), + "w+", + ) as f_out, + open( + os.path.join( + logsdir, + f"log_LINC_calibrator_{field['target_name']}_{field[f'sas_id_calibrator{calibrator_field}']}_err.txt", + ), + "w+", + ) as f_err, + ): + proc = subprocess.run(cmd, shell=True, text=True, stdout=f_out, stderr=f_err) + jobid = None + if not proc.returncode: + f_out.seek(0) + for line in f_out.readlines(): + if "Submitted batch job" in line: + jobid = line.strip().split()[-1] + else: + raise RuntimeError("Failed to submit job.") + + if not jobid: + raise RuntimeError("Failed to retrieve job id") + else: + while True: + print(f"Polling LINC calibrator job {jobid}") + poll_cmd = f"sacct -X -j {jobid} --format=State --noheader" + status = subprocess.run( + poll_cmd, shell=True, text=True, capture_output=True + ).stdout.strip() + if (status == "RUNNING") or (status == "PENDING"): + time.sleep(60) + elif status == "COMPLETED": + db.set_status_finished( + field["target_name"], + f"calibrator{calibrator_field}", + field["sas_id_target"], + ) + break + elif ( + (status == "FAILED") + or ("TIMEOUT" in status) + or ("CANCELLED" in status) + ): + raise RuntimeError( + f"LINC calibrator for {field['target_name']} {field['sas_id_target']} failed." + ) + + +def run_linc_calibrator_toil(field, calibrator_field: int, db: FlocsDB): + print( + f"Processing flux density calibrator {field[f'sas_id_calibrator{calibrator_field}']} for observation {field['target_name']} {field['sas_id_target']}" + ) + ms_folder = f"L{field[f'sas_id_calibrator{calibrator_field}']}" + db.set_status_processing( + field["target_name"], f"calibrator{calibrator_field}", field["sas_id_target"] + ) + outdir = os.path.join(OUTPUT_DIR, field["target_name"]) + logsdir = os.path.join(OUTPUT_DIR, field["target_name"], "logs") + if not os.path.isdir(outdir): + os.makedirs(outdir, exist_ok=True) + if not os.path.isdir(logsdir): + os.makedirs(logsdir, exist_ok=True) + cmd = f"flocs-run linc calibrator --runner toil --scheduler slurm --slurm-account {SLURM_ACCOUNT} --slurm-queue {SLURM_QUEUE} --rundir {PROCESSING_DIR} --outdir {outdir} {os.path.join(DATA_DIR, field['target_name'], 'calibrator', ms_folder)}" + print(cmd) + with ( + open( + os.path.join( + logsdir, + f"log_LINC_calibrator_{field['target_name']}_{field[f'sas_id_calibrator{calibrator_field}']}.txt", + ), + "w+", + ) as f_out, + open( + os.path.join( + logsdir, + f"log_LINC_calibrator_{field['target_name']}_{field[f'sas_id_calibrator{calibrator_field}']}_err.txt", + ), + "w+", + ) as f_err, + ): + proc = subprocess.run(cmd, shell=True, text=True, stdout=f_out, stderr=f_err) + success = False + pattern = re.compile(r"Workflow.* stopped. Success: True") + if not proc.returncode: + f_err.seek(0) + if pattern.search(f_err.read()): + success = True + if success: + db.set_status_finished( + field["target_name"], + f"calibrator{calibrator_field}", + field["sas_id_target"], + ) + else: + raise RuntimeError + + +def run_linc_target_cwltool(field, db: FlocsDB): + print( + f"Processing target observation {field['target_name']} {field['sas_id_target']} with calibrator {field['sas_id_calibrator_final']}" + ) + ms_folder = f"L{field['sas_id_target']}" + outdir = os.path.join(OUTPUT_DIR, field["target_name"]) + logsdir = os.path.join(OUTPUT_DIR, field["target_name"], "logs") + if not os.path.isdir(outdir): + os.makedirs(outdir, exist_ok=True) + if not os.path.isdir(logsdir): + os.makedirs(logsdir, exist_ok=True) + calibrator_path = get_most_recent_run( + outdir, field["sas_id_calibrator_final"], "LINC_calibrator" + ) + calibrator_solutions = ( + calibrator_path / "results_LINC_calibrator" / "cal_solutions.h5" + ) + db.set_status_processing(field["target_name"], "target", field["sas_id_target"]) + cmd = f"flocs-run linc target --runner cwltool --scheduler slurm --slurm-cores 64 --slurm-account {SLURM_ACCOUNT} --slurm-queue {SLURM_QUEUE} --rundir {PROCESSING_DIR} --outdir {outdir} --cal-solutions {calibrator_solutions} {os.path.join(DATA_DIR, field['target_name'], 'target', ms_folder)}" + print(cmd) + with ( + open( + os.path.join( + logsdir, + f"log_LINC_target_{field['target_name']}_{field['sas_id_target']}.txt", + ), + "w+", + ) as f_out, + open( + os.path.join( + logsdir, + f"log_LINC_target_{field['target_name']}_{field['sas_id_target']}_err.txt", + ), + "w+", + ) as f_err, + ): + proc = subprocess.run(cmd, shell=True, text=True, stdout=f_out, stderr=f_err) + jobid = None + if not proc.returncode: + f_out.seek(0) + for line in f_out.readlines(): + if "Submitted batch job" in line: + jobid = line.strip().split()[-1] + else: + raise RuntimeError("Failed to submit job.") + + if not jobid: + raise RuntimeError("Failed to retrieve job id") + else: + while True: + print(f"Polling LINC target job {jobid}") + poll_cmd = f"sacct -X -j {jobid} --format=State --noheader" + status = subprocess.run( + poll_cmd, shell=True, text=True, capture_output=True + ).stdout.strip() + if (status == "RUNNING") or (status == "PENDING"): + time.sleep(60) + elif status == "COMPLETED": + db.set_status_finished( + field["target_name"], + "target", + field["sas_id_target"], + ) + break + elif ( + (status == "FAILED") + or ("TIMEOUT" in status) + or ("CANCELLED" in status) + ): + raise RuntimeError( + f"LINC target for {field['target_name']} {field['sas_id_target']} failed." + ) + + +def run_linc_target_toil(field, db: FlocsDB): + print( + f"Processing target observation {field['target_name']} {field['sas_id_target']} with calibrator {field['sas_id_calibrator_final']}" + ) + ms_folder = f"L{field['sas_id_target']}" + outdir = os.path.join(OUTPUT_DIR, field["target_name"]) + logsdir = os.path.join(OUTPUT_DIR, field["target_name"], "logs") + if not os.path.isdir(outdir): + os.makedirs(outdir, exist_ok=True) + if not os.path.isdir(logsdir): + os.makedirs(logsdir, exist_ok=True) + calibrator_path = get_most_recent_run( + outdir, field["sas_id_calibrator_final"], "LINC_calibrator" + ) + calibrator_solutions = ( + calibrator_path / "results_LINC_calibrator" / "cal_solutions.h5" + ) + db.set_status_processing(field["target_name"], "target", field["sas_id_target"]) + cmd = f"flocs-run linc target --runner toil --scheduler slurm --slurm-account {SLURM_ACCOUNT} --slurm-queue {SLURM_QUEUE} --rundir {PROCESSING_DIR} --outdir {outdir} --cal-solutions {calibrator_solutions} {os.path.join(DATA_DIR, field['target_name'], 'target', ms_folder)}" + print(cmd) + with ( + open( + os.path.join( + logsdir, + f"log_LINC_target_{field['target_name']}_{field['sas_id_target']}.txt", + ), + "w+", + ) as f_out, + open( + os.path.join( + logsdir, + f"log_LINC_target_{field['target_name']}_{field['sas_id_target']}_err.txt", + ), + "w+", + ) as f_err, + ): + proc = subprocess.run(cmd, shell=True, text=True, stdout=f_out, stderr=f_err) + success = False + pattern = re.compile(r"Workflow.* stopped. Success: True") + if not proc.returncode: + f_err.seek(0) + if pattern.search(f_err.read()): + success = True + if success: + db.set_status_finished( + field["target_name"], "target", field["sas_id_target"] + ) + else: + raise RuntimeError + + +def run_pilot_delay_cwltool(field, db: FlocsDB): + print( + f"Processing delay calibration for {field['target_name']} {field['sas_id_target']}" + ) + outdir = os.path.join(OUTPUT_DIR, field["target_name"]) + logsdir = os.path.join(OUTPUT_DIR, field["target_name"], "logs") + if not os.path.isdir(outdir): + os.makedirs(outdir, exist_ok=True) + if not os.path.isdir(logsdir): + os.makedirs(logsdir, exist_ok=True) + target_path = get_most_recent_run(outdir, field["sas_id_target"], "LINC_target") + target_ms_path = target_path / "results_LINC_target" / "results" + db.set_status_processing(field["target_name"], "vlbi_delay", field["sas_id_target"]) + + delay_cat = os.path.join(outdir, "delay_calibrators.csv") + image_cat = os.path.join(outdir, "image_catalogue.csv") + + if not os.path.isfile(delay_cat) or not os.path.isfile(image_cat): + ms = list(target_ms_path.glob("*.dp3concat"))[0] + cmd = f"lofar-vlbi-plot --force --output_dir {outdir} --MS {ms}" + with ( + open( + os.path.join( + logsdir, + f"log_plot_field_{field['target_name']}_{field['sas_id_target']}.txt", + ), + "w+", + ) as f_out, + open( + os.path.join( + logsdir, + f"log_plot_field_{field['target_name']}_{field['sas_id_target']}_err.txt", + ), + "w+", + ) as f_err, + ): + proc = subprocess.run( + cmd, shell=True, text=True, stdout=f_out, stderr=f_err + ) + delay_cat = os.path.join(outdir, "delay_calibrators.csv") + image_cat = os.path.join(outdir, "image_catalogue.csv") + + if not os.path.isfile(delay_cat): + raise RuntimeError("Delay calibrator catalogue is missing or invalid.") + if not os.path.isfile(image_cat): + raise RuntimeError("Image source catalogue is missing or invalid.") + + proc = subprocess.run( + "detect_bad_slurm_nodes.sh", + shell=True, + text=True, + stdout=subprocess.PIPE, + ) + bad_nodes = proc.stdout.strip() + if bad_nodes: + print(f"Excluding the following bad nodes from scheduling: {bad_nodes}") + os.environ["TOIL_SLURM_ARGS"] = f"--exclude={bad_nodes}" + + cmd = f"flocs-run vlbi delay-calibration --record-toil-stats --runner cwltool --scheduler slurm --slurm-cores 64 --slurm-account {SLURM_ACCOUNT} --slurm-queue {SLURM_QUEUE} --rundir {PROCESSING_DIR} --outdir {outdir} --ms-suffix dp3concat --delay-calibrator {delay_cat} --image-catalogue {image_cat} --apply-delay-solutions {target_ms_path}" + if not os.path.isdir(outdir): + os.makedirs(outdir, exist_ok=True) + print(cmd) + with ( + open( + os.path.join( + logsdir, + f"log_VLBI_delay-calibration_{field['target_name']}_{field['sas_id_target']}.txt", + ), + "w+", + ) as f_out, + open( + os.path.join( + logsdir, + f"log_VLBI_delay-calibration_{field['target_name']}_{field['sas_id_target']}_err.txt", + ), + "w+", + ) as f_err, + ): + proc = subprocess.run(cmd, shell=True, text=True, stdout=f_out, stderr=f_err) + jobid = None + if not proc.returncode: + f_out.seek(0) + for line in f_out.readlines(): + if "Submitted batch job" in line: + jobid = line.strip().split()[-1] + else: + raise RuntimeError("Failed to submit job.") + + if not jobid: + raise RuntimeError("Failed to retrieve job id") + else: + while True: + print(f"Polling LINC target job {jobid}") + poll_cmd = f"sacct -X -j {jobid} --format=State --noheader" + status = subprocess.run( + poll_cmd, shell=True, text=True, capture_output=True + ).stdout.strip() + if (status == "RUNNING") or (status == "PENDING"): + time.sleep(60) + elif status == "COMPLETED": + if NEEDS_MANUAL_APPROVAL_DELAY: + db.set_status_await_approval( + field["target_name"], "vlbi_delay", field["sas_id_target"] + ) + else: + db.set_status_finished( + field["target_name"], "vlbi_delay", field["sas_id_target"] + ) + break + elif ( + (status == "FAILED") + or ("TIMEOUT" in status) + or ("CANCELLED" in status) + ): + raise RuntimeError( + f"PILOT delay calibration for {field['target_name']} {field['sas_id_target']} failed." + ) + + +def run_pilot_delay_toil(field, db: FlocsDB): + print( + f"Processing delay calibration for {field['target_name']} {field['sas_id_target']}" + ) + outdir = os.path.join(OUTPUT_DIR, field["target_name"]) + logsdir = os.path.join(OUTPUT_DIR, field["target_name"], "logs") + if not os.path.isdir(outdir): + os.makedirs(outdir, exist_ok=True) + if not os.path.isdir(logsdir): + os.makedirs(logsdir, exist_ok=True) + target_path = get_most_recent_run(outdir, field["sas_id_target"], "LINC_target") + target_ms_path = target_path / "results_LINC_target" / "results" + db.set_status_processing(field["target_name"], "vlbi_delay", field["sas_id_target"]) + + delay_cat = os.path.join(outdir, "delay_calibrators.csv") + image_cat = os.path.join(outdir, "image_catalogue.csv") + + if not os.path.isfile(delay_cat) or not os.path.isfile(image_cat): + ms = list(target_ms_path.glob("*.dp3concat"))[0] + cmd = f"lofar-vlbi-plot --force --output_dir {outdir} --MS {ms}" + with ( + open( + os.path.join( + logsdir, + f"log_plot_field_{field['target_name']}_{field['sas_id_target']}.txt", + ), + "w+", + ) as f_out, + open( + os.path.join( + logsdir, + f"log_plot_field_{field['target_name']}_{field['sas_id_target']}_err.txt", + ), + "w+", + ) as f_err, + ): + proc = subprocess.run( + cmd, shell=True, text=True, stdout=f_out, stderr=f_err + ) + delay_cat = os.path.join(outdir, "delay_calibrators.csv") + image_cat = os.path.join(outdir, "image_catalogue.csv") + + if not os.path.isfile(delay_cat): + raise RuntimeError("Delay calibrator catalogue is missing or invalid.") + if not os.path.isfile(image_cat): + raise RuntimeError("Image source catalogue is missing or invalid.") + + proc = subprocess.run( + "detect_bad_slurm_nodes.sh", + shell=True, + text=True, + stdout=subprocess.PIPE, + ) + bad_nodes = proc.stdout.strip() + if bad_nodes: + print(f"Excluding the following bad nodes from scheduling: {bad_nodes}") + os.environ["TOIL_SLURM_ARGS"] = f"--exclude={bad_nodes}" + + context = get_current_context() + if context["ti"].try_number == 1 or ( + not os.path.isfile( + os.path.join( + logsdir, + f"log_VLBI_delay-calibration_{field['target_name']}_{field['sas_id_target']}.txt", + ) + ) + ): + cmd = f"flocs-run vlbi delay-calibration --record-toil-stats --runner toil --scheduler slurm --slurm-account {SLURM_ACCOUNT} --slurm-queue {SLURM_QUEUE} --rundir {PROCESSING_DIR} --outdir {outdir} --ms-suffix dp3concat --delay-calibrator {delay_cat} --image-catalogue {image_cat} --apply-delay-solutions {target_ms_path}" + else: + # Extract the previous working directory + flocs_workdir = "" + print( + f"Scanning log_VLBI_delay-calibration_{field['target_name']}_{field['sas_id_target']}.txt for workdir." + ) + with open( + os.path.join( + logsdir, + f"log_VLBI_delay-calibration_{field['target_name']}_{field['sas_id_target']}.txt", + ) + ) as f_out: + for line in f_out.readlines(): + print(line) + if "Running workflow with" in line: + flocs_workdir = line.split(" ")[-1].strip() + break + if not flocs_workdir: + raise RuntimeError( + "Could not retrieve PILOT workdir. Flocs probably crashed before launching." + ) + print(f"Resuming failed PILOT run in {flocs_workdir}") + cmd = f"flocs-run vlbi delay-calibration --record-toil-stats --runner toil --scheduler slurm --slurm-account {SLURM_ACCOUNT} --slurm-queue {SLURM_QUEUE} --rundir {flocs_workdir} --restart --outdir {outdir} --ms-suffix dp3concat --delay-calibrator {delay_cat} --image-catalogue {image_cat} {target_ms_path}" + if not os.path.isdir(outdir): + os.makedirs(outdir, exist_ok=True) + print(cmd) + with ( + open( + os.path.join( + logsdir, + f"log_VLBI_delay-calibration_{field['target_name']}_{field['sas_id_target']}.txt", + ), + "w+", + ) as f_out, + open( + os.path.join( + logsdir, + f"log_VLBI_delay-calibration_{field['target_name']}_{field['sas_id_target']}_err.txt", + ), + "w+", + ) as f_err, + ): + proc = subprocess.run(cmd, shell=True, text=True, stdout=f_out, stderr=f_err) + success = False + pattern = re.compile(r"Workflow.* stopped. Success: True") + if not proc.returncode: + f_err.seek(0) + if pattern.search(f_err.read()): + success = True + + if success: + if NEEDS_MANUAL_APPROVAL_DELAY: + db.set_status_await_approval( + field["target_name"], "vlbi_delay", field["sas_id_target"] + ) + else: + db.set_status_finished( + field["target_name"], "vlbi_delay", field["sas_id_target"] + ) + else: + raise RuntimeError + + +def run_pilot_ddcal_cwltool(field, db: FlocsDB): + print( + f"Processing ILT dd calibration for {field['target_name']} {field['sas_id_target']}" + ) + outdir = os.path.join(OUTPUT_DIR, field["target_name"]) + logsdir = os.path.join(OUTPUT_DIR, field["target_name"], "logs") + if not os.path.isdir(outdir): + os.makedirs(outdir, exist_ok=True) + if not os.path.isdir(logsdir): + os.makedirs(logsdir, exist_ok=True) + if "status_ddf" not in field: + print("Not a widefield imaging run, checking LINC + delay calibration.") + target_path = get_most_recent_run(outdir, field["sas_id_target"], "LINC_target") + target_ms_path = target_path / "results_LINC_target" / "results" + print(f"Using LINC target run: {target_path}") + + sols_path = get_most_recent_run(outdir, field["sas_id_target"], "VLBI_delay") + sols_path = sols_path / "results_VLBI_delay-calibration" + sols = list(sols_path.glob("merged*selfcalcycle???_linearfulljones*.h5"))[0] + print(f"Using PILOT delay calibration solutions: {sols}") + + source_cat = os.path.join(DATA_DIR, field["target_name"], "vlbi_target.csv") + if not os.path.isfile(source_cat): + raise AirflowFailException(f"{source_cat} not found.") + + db.set_status_processing( + field["target_name"], "vlbi_dd", field["sas_id_target"] + ) + cmd = f"flocs-run vlbi dd-calibration --runner cwltool --scheduler slurm --slurm-cores 32 --slurm-time 24:00:00 --slurm-account {SLURM_ACCOUNT} --slurm-queue {SLURM_QUEUE} --rundir {PROCESSING_DIR} --outdir {outdir} --delay-solset {sols} --phasediff-score 10.0 --source-catalogue {source_cat} --model-cache {NN_MODEL_CACHE} --ms-suffix .dp3concat {target_ms_path}" + else: + print("Widefield imaging run, checking subtraction output.") + target_path = get_most_recent_run( + outdir, field["sas_id_target"], "VLBI_process-ddf" + ) + target_ms_path = target_path / "results_VLBI_process-ddf" + print(f"Using subtracted data at: {target_path}") + + source_cat = os.path.join(DATA_DIR, field["target_name"], "image_catalogue.csv") + if not os.path.isfile(source_cat): + raise AirflowFailException(f"{source_cat} not found.") + + db.set_status_processing( + field["target_name"], "vlbi_dd", field["sas_id_target"] + ) + + cmd = f"flocs-run vlbi dd-calibration --runner cwltool --scheduler slurm --slurm-cores 64 --slurm-time 24:00:00 --slurm-account {SLURM_ACCOUNT} --slurm-queue {SLURM_QUEUE} --rundir {PROCESSING_DIR} --outdir {outdir} --source-catalogue {source_cat} --model-cache {NN_MODEL_CACHE} --ms-suffix .dp3concat.sub.ms {target_ms_path}" + if not os.path.isdir(outdir): + os.makedirs(outdir, exist_ok=True) + print(cmd) + with ( + open( + os.path.join( + logsdir, + f"log_VLBI_dd-calibration_{field['target_name']}_{field['sas_id_target']}.txt", + ), + "w+", + ) as f_out, + open( + os.path.join( + logsdir, + f"log_VLBI_dd-calibration_{field['target_name']}_{field['sas_id_target']}_err.txt", + ), + "w+", + ) as f_err, + ): + proc = subprocess.run(cmd, shell=True, text=True, stdout=f_out, stderr=f_err) + jobid = None + if not proc.returncode: + f_out.seek(0) + for line in f_out.readlines(): + if "Submitted batch job" in line: + jobid = line.strip().split()[-1] + else: + raise RuntimeError("Failed to submit job.") + + if not jobid: + raise RuntimeError("Failed to retrieve job id") + else: + while True: + print(f"Polling LINC target job {jobid}") + poll_cmd = f"sacct -X -j {jobid} --format=State --noheader" + status = subprocess.run( + poll_cmd, shell=True, text=True, capture_output=True + ).stdout.strip() + if (status == "RUNNING") or (status == "PENDING"): + time.sleep(60) + elif status == "COMPLETED": + db.set_status_finished( + field["target_name"], "vlbi_dd", field["sas_id_target"] + ) + break + elif ( + (status == "FAILED") + or ("TIMEOUT" in status) + or ("CANCELLED" in status) + ): + raise RuntimeError( + f"PILOT direction-dependent calibration for {field['target_name']} {field['sas_id_target']} failed." + ) + + +def run_pilot_ddcal_toil(field, db: FlocsDB): + print( + f"Processing ILT dd calibration for {field['target_name']} {field['sas_id_target']}" + ) + outdir = os.path.join(OUTPUT_DIR, field["target_name"]) + logsdir = os.path.join(OUTPUT_DIR, field["target_name"], "logs") + if not os.path.isdir(outdir): + os.makedirs(outdir, exist_ok=True) + if not os.path.isdir(logsdir): + os.makedirs(logsdir, exist_ok=True) + if "status_ddf" not in field: + print("Not a widefield imaging run, checking LINC + delay calibration.") + target_path = get_most_recent_run(outdir, field["sas_id_target"], "LINC_target") + target_ms_path = target_path / "results_LINC_target" / "results" + print(f"Using LINC target run: {target_path}") + + sols_path = get_most_recent_run(outdir, field["sas_id_target"], "VLBI_delay") + sols_path = sols_path / "results_VLBI_delay-calibration" + sols = list(sols_path.glob("merged*selfcalcycle???_linearfulljones*.h5"))[0] + print(f"Using PILOT delay calibration solutions: {sols}") + + source_cat = os.path.join(DATA_DIR, field["target_name"], "vlbi_target.csv") + if not os.path.isfile(source_cat): + raise AirflowFailException(f"{source_cat} not found.") + + db.set_status_processing( + field["target_name"], "vlbi_dd", field["sas_id_target"] + ) + cmd = f"flocs-run vlbi dd-calibration --runner toil --scheduler slurm --slurm-time 24:00:00 --slurm-account {SLURM_ACCOUNT} --slurm-queue {SLURM_QUEUE} --rundir {PROCESSING_DIR} --outdir {outdir} --delay-solset {sols} --phasediff-score 10.0 --source-catalogue {source_cat} --model-cache {NN_MODEL_CACHE} --ms-suffix .dp3concat {target_ms_path}" + else: + print("Widefield imaging run, checking subtraction output.") + target_path = get_most_recent_run( + outdir, field["sas_id_target"], "VLBI_process-ddf" + ) + target_ms_path = target_path / "results_VLBI_process-ddf" + print(f"Using subtracted data at: {target_path}") + + source_cat = os.path.join(DATA_DIR, field["target_name"], "image_catalogue.csv") + if not os.path.isfile(source_cat): + raise AirflowFailException(f"{source_cat} not found.") + + db.set_status_processing( + field["target_name"], "vlbi_dd", field["sas_id_target"] + ) + + context = get_current_context() + if context["ti"].try_number == 1 or ( + not os.path.isfile( + os.path.join( + logsdir, + f"log_VLBI_dd-calibration_{field['target_name']}_{field['sas_id_target']}.txt", + ) + ) + ): + cmd = f"flocs-run vlbi dd-calibration --record-toil-stats --runner toil --scheduler slurm --slurm-time 24:00:00 --slurm-account {SLURM_ACCOUNT} --slurm-queue {SLURM_QUEUE} --rundir {PROCESSING_DIR} --outdir {outdir} --source-catalogue {source_cat} --model-cache {NN_MODEL_CACHE} --ms-suffix .dp3concat.sub.ms {target_ms_path}" + else: + if field["status_vlbi_dd"] == PIPELINE_STATUS.downloaded: + # This way we can force a clean restart in the database. + cmd = f"flocs-run vlbi dd-calibration --record-toil-stats --runner toil --scheduler slurm --slurm-time 24:00:00 --slurm-account {SLURM_ACCOUNT} --slurm-queue {SLURM_QUEUE} --rundir {PROCESSING_DIR} --outdir {outdir} --source-catalogue {source_cat} --model-cache {NN_MODEL_CACHE} --ms-suffix .dp3concat.sub.ms {target_ms_path}" + else: + # Extract the previous working directory + flocs_workdir = "" + print( + f"Scanning log_VLBI_dd-calibration_{field['target_name']}_{field['sas_id_target']}.txt for workdir." + ) + with open( + os.path.join( + logsdir, + f"log_VLBI_dd-calibration_{field['target_name']}_{field['sas_id_target']}.txt", + ) + ) as f_out: + for line in f_out.readlines(): + print(line) + if "Running workflow with" in line: + flocs_workdir = line.split(" ")[-1].strip() + break + if not flocs_workdir: + raise RuntimeError( + "Could not retrieve PILOT workdir. Flocs probably crashed before launching." + ) + print(f"Resuming failed PILOT run in {flocs_workdir}") + cmd = f"flocs-run vlbi dd-calibration --record-toil-stats --runner toil --scheduler slurm --slurm-time 24:00:00 --slurm-account {SLURM_ACCOUNT} --slurm-queue {SLURM_QUEUE} --rundir {flocs_workdir} --restart --outdir {outdir} --source-catalogue {source_cat} --model-cache {NN_MODEL_CACHE} --ms-suffix .dp3concat.sub.ms {target_ms_path}" + if not os.path.isdir(outdir): + os.makedirs(outdir, exist_ok=True) + print(cmd) + with ( + open( + os.path.join( + logsdir, + f"log_VLBI_dd-calibration_{field['target_name']}_{field['sas_id_target']}.txt", + ), + "w+", + ) as f_out, + open( + os.path.join( + logsdir, + f"log_VLBI_dd-calibration_{field['target_name']}_{field['sas_id_target']}_err.txt", + ), + "w+", + ) as f_err, + ): + proc = subprocess.run(cmd, shell=True, text=True, stdout=f_out, stderr=f_err) + success = False + pattern = re.compile(r"Workflow.* stopped. Success: True") + if not proc.returncode: + f_err.seek(0) + if pattern.search(f_err.read()): + success = True + if success: + db.set_status_finished( + field["target_name"], "vlbi_dd", field["sas_id_target"] + ) + db.set_field_finished(field["target_name"], field["sas_id_target"]) + else: + db.set_status_failed( + field["target_name"], "vlbi_dd", field["sas_id_target"] + ) + raise RuntimeError + + +def run_pilot_intermediate_image_toil(field, db: FlocsDB): + print( + f"Processing ILT intermediate resolution imaging for {field['target_name']} {field['sas_id_target']}" + ) + outdir = os.path.join(OUTPUT_DIR, field["target_name"]) + logsdir = os.path.join(OUTPUT_DIR, field["target_name"], "logs") + if not os.path.isdir(outdir): + os.makedirs(outdir, exist_ok=True) + if not os.path.isdir(logsdir): + os.makedirs(logsdir, exist_ok=True) + target_path = get_most_recent_run( + outdir, field["sas_id_target"], "VLBI_process-ddf" + ) + target_ms_path = target_path / "results_VLBI_process-ddf" + print(f"Using subtracted data at: {target_path}") + + dd_sols_path = get_most_recent_run( + outdir, field["sas_id_target"], "VLBI_dd-calibration" + ) + dd_sols = dd_sols_path / "results_VLBI_dd-calibration" / "merged.h5" + print(f"Using dd solutions: {dd_sols}") + + if not os.path.isfile(dd_sols): + raise AirflowFailException(f"{dd_sols} not found.") + + db.set_status_processing( + field["target_name"], "vlbi_intermediate_img", field["sas_id_target"] + ) + + context = get_current_context() + if context["ti"].try_number == 1 or ( + not os.path.isfile( + os.path.join( + logsdir, + f"log_VLBI_image_intermediate_resolution_{field['target_name']}_{field['sas_id_target']}.txt", + ) + ) + ): + cmd = f"flocs-run vlbi image-intermediate-resolution --record-toil-stats --runner toil --scheduler slurm --slurm-time 72:00:00 --slurm-account {SLURM_ACCOUNT} --slurm-queue {SLURM_QUEUE} --rundir {PROCESSING_DIR} --outdir {outdir} --dd-solutions {dd_sols} --ms-suffix .dp3concat.sub.ms {target_ms_path}" + else: + if field["status_vlbi_intermediate_img"] == PIPELINE_STATUS.downloaded: + # This way we can force a clean restart in the database. + cmd = f"flocs-run vlbi image-intermediate-resolution --record-toil-stats --runner toil --scheduler slurm --slurm-time 72:00:00 --slurm-account {SLURM_ACCOUNT} --slurm-queue {SLURM_QUEUE} --rundir {PROCESSING_DIR} --outdir {outdir} --dd-solutions {dd_sols} --ms-suffix .dp3concat.sub.ms {target_ms_path}" + else: + # Extract the previous working directory + flocs_workdir = "" + print( + f"Scanning log_VLBI_image_intermediate_resolution_{field['target_name']}_{field['sas_id_target']}.txt for workdir." + ) + with ( + open( + os.path.join( + logsdir, + f"log_VLBI_image_intermediate_resolution_{field['target_name']}_{field['sas_id_target']}.txt", + ) + ) as f_out, + open( + os.path.join( + logsdir, + f"log_VLBI_image_intermediate_resolution_{field['target_name']}_{field['sas_id_target']}_err.txt", + ) + ) as f_err, + ): + for line in f_out.readlines(): + print(line) + if "Running workflow with" in line: + flocs_workdir = line.split(" ")[-1].strip() + break + if not flocs_workdir: + raise RuntimeError( + "Could not retrieve PILOT workdir. Flocs probably crashed before launching." + ) + print(f"Resuming failed PILOT run in {flocs_workdir}") + cmd = f"flocs-run vlbi image-intermediate-resolution --record-toil-stats --runner toil --scheduler slurm --slurm-time 72:00:00 --slurm-account {SLURM_ACCOUNT} --slurm-queue {SLURM_QUEUE} --rundir {flocs_workdir} --restart --outdir {outdir} --dd-solutions {dd_sols} --ms-suffix .dp3concat.sub.ms {target_ms_path}" + if not os.path.isdir(outdir): + os.makedirs(outdir, exist_ok=True) + print(cmd) + with ( + open( + os.path.join( + logsdir, + f"log_VLBI_image_intermediate_resolution_{field['target_name']}_{field['sas_id_target']}.txt", + ), + "w+", + ) as f_out, + open( + os.path.join( + logsdir, + f"log_VLBI_image_intermediate_resolution_{field['target_name']}_{field['sas_id_target']}_err.txt", + ), + "w+", + ) as f_err, + ): + proc = subprocess.run(cmd, shell=True, text=True, stdout=f_out, stderr=f_err) + success = False + pattern = re.compile(r"Workflow.* stopped. Success: True") + if not proc.returncode: + f_err.seek(0) + if pattern.search(f_err.read()): + success = True + if success: + db.set_status_finished( + field["target_name"], + "vlbi_intermediate_img", + field["sas_id_target"], + ) + db.set_field_finished(field["target_name"], field["sas_id_target"]) + else: + db.set_status_failed( + field["target_name"], + "vlbi_intermediate_img", + field["sas_id_target"], + ) + raise RuntimeError + + +def run_pilot_facet_subtract_toil(field, db: FlocsDB): + print( + f"Processing ILT facet subtraction for {field['target_name']} {field['sas_id_target']}" + ) + outdir = os.path.join(OUTPUT_DIR, field["target_name"]) + logsdir = os.path.join(OUTPUT_DIR, field["target_name"], "logs") + if not os.path.isdir(outdir): + os.makedirs(outdir, exist_ok=True) + if not os.path.isdir(logsdir): + os.makedirs(logsdir, exist_ok=True) + target_path = get_most_recent_run( + outdir, field["sas_id_target"], "VLBI_process-ddf" + ) + target_ms_path = target_path / "results_VLBI_process-ddf" + print(f"Using subtracted data at: {target_path}") + + dd_sols_path = get_most_recent_run( + outdir, field["sas_id_target"], "VLBI_dd-calibration" + ) + dd_sols = dd_sols_path / "results_VLBI_dd-calibration" / "merged.h5" + print(f"Using dd solutions at: {dd_sols}") + + if not os.path.isfile(dd_sols): + raise AirflowFailException(f"{dd_sols} not found.") + + model_images_path = get_most_recent_run( + outdir, field["sas_id_target"], "VLBI_intermediate_resolution_imaging" + ) + model_images_path = ( + model_images_path / "results_VLBI_intermediate_resolution_imaging" + ) + model_images = list(model_images_path.glob("*-????-model-fpb.fits")) + print(f"Using model images at: {model_images_path}") + + if not model_images: + raise AirflowFailException( + "No suitable intermediate resolution model images found." + ) + + db.set_status_processing( + field["target_name"], "vlbi_facet_subtract", field["sas_id_target"] + ) + + context = get_current_context() + if context["ti"].try_number == 1 or ( + not os.path.isfile( + os.path.join( + logsdir, + f"log_VLBI_facet-subtract_{field['target_name']}_{field['sas_id_target']}.txt", + ) + ) + ): + cmd = f"flocs-run vlbi facet-subtract --record-toil-stats --runner toil --scheduler slurm --slurm-time 72:00:00 --slurm-account {SLURM_ACCOUNT} --slurm-queue {SLURM_QUEUE} --rundir {PROCESSING_DIR} --outdir {outdir} --dd-solutions {dd_sols} --model-image-directory {model_images_path} --ms-suffix .dp3concat.sub.ms {target_ms_path}" + else: + if field["status_vlbi_facet_subtract"] == PIPELINE_STATUS.downloaded: + # This way we can force a clean restart in the database. + cmd = f"flocs-run vlbi facet-subtract --record-toil-stats --runner toil --scheduler slurm --slurm-time 72:00:00 --slurm-account {SLURM_ACCOUNT} --slurm-queue {SLURM_QUEUE} --rundir {PROCESSING_DIR} --outdir {outdir} --dd-solutions {dd_sols} --model-image-directory {model_images_path} --ms-suffix .dp3concat.sub.ms {target_ms_path}" + else: + # Extract the previous working directory + flocs_workdir = "" + print( + f"Scanning log_VLBI_facet-subtract_{field['target_name']}_{field['sas_id_target']}.txt for workdir." + ) + with open( + os.path.join( + logsdir, + f"log_VLBI_facet-subtract_{field['target_name']}_{field['sas_id_target']}.txt", + ) + ) as f_out: + for line in f_out.readlines(): + print(line) + if "Running workflow with" in line: + flocs_workdir = line.split(" ")[-1].strip() + break + if not flocs_workdir: + raise RuntimeError( + "Could not retrieve PILOT workdir. Flocs probably crashed before launching." + ) + print(f"Resuming failed PILOT run in {flocs_workdir}") + cmd = f"flocs-run vlbi facet-subtract --record-toil-stats --runner toil --scheduler slurm --slurm-time 72:00:00 --slurm-account {SLURM_ACCOUNT} --slurm-queue {SLURM_QUEUE} --rundir {flocs_workdir} --restart --outdir {outdir} --dd-solutions {dd_sols} --model-image-directory {model_images_path} --ms-suffix .dp3concat.sub.ms {target_ms_path}" + if not os.path.isdir(outdir): + os.makedirs(outdir, exist_ok=True) + print(cmd) + with ( + open( + os.path.join( + logsdir, + f"log_VLBI_facet-subtract_{field['target_name']}_{field['sas_id_target']}.txt", + ), + "w+", + ) as f_out, + open( + os.path.join( + logsdir, + f"log_VLBI_facet-subtract_{field['target_name']}_{field['sas_id_target']}_err.txt", + ), + "w+", + ) as f_err, + ): + proc = subprocess.run(cmd, shell=True, text=True, stdout=f_out, stderr=f_err) + success = False + pattern = re.compile(r"Workflow.* stopped. Success: True") + if not proc.returncode: + f_err.seek(0) + if pattern.search(f_err.read()): + success = True + if success: + db.set_status_finished( + field["target_name"], + "vlbi_facet_subtract", + field["sas_id_target"], + ) + db.set_field_finished(field["target_name"], field["sas_id_target"]) + else: + db.set_status_failed( + field["target_name"], + "vlbi_facet_subtract", + field["sas_id_target"], + ) + raise RuntimeError + + +def run_pilot_facet_imaging_toil(field, db: FlocsDB): + print( + f"Processing ILT facet imaging for {field['target_name']} {field['sas_id_target']}" + ) + outdir = os.path.join(OUTPUT_DIR, field["target_name"]) + logsdir = os.path.join(OUTPUT_DIR, field["target_name"], "logs") + if not os.path.isdir(outdir): + os.makedirs(outdir, exist_ok=True) + if not os.path.isdir(logsdir): + os.makedirs(logsdir, exist_ok=True) + target_path = get_most_recent_run( + outdir, field["sas_id_target"], "VLBI_facet_subtract" + ) + target_ms_path = target_path / "results_VLBI_facet_subtract" + print(f"Using subtracted data at: {target_path}") + + facet_mses = list(target_ms_path.glob("facet*.ms")) + + if not facet_mses: + raise AirflowFailException( + "No suitable intermediate resolution model images found." + ) + + db.set_status_processing( + field["target_name"], "vlbi_facet_imaging", field["sas_id_target"] + ) + + context = get_current_context() + if context["ti"].try_number == 1 or ( + not os.path.isfile( + os.path.join( + logsdir, + f"log_VLBI_facet-imaging_{field['target_name']}_{field['sas_id_target']}.txt", + ) + ) + ): + cmd = f"flocs-run vlbi facet-imaging --record-toil-stats --runner toil --scheduler slurm --slurm-time 72:00:00 --slurm-account {SLURM_ACCOUNT} --slurm-queue {SLURM_QUEUE} --rundir {PROCESSING_DIR} --outdir {outdir} --resolution 0.3asec --pixel-scale 0.1 --ms-suffix .ms {target_ms_path}" + else: + if field["status_vlbi_facet_imaging"] == PIPELINE_STATUS.downloaded: + # This way we can force a clean restart in the database. + cmd = f"flocs-run vlbi facet-imaging --record-toil-stats --runner toil --scheduler slurm --slurm-time 72:00:00 --slurm-account {SLURM_ACCOUNT} --slurm-queue {SLURM_QUEUE} --rundir {PROCESSING_DIR} --outdir {outdir} --resolution 0.3asec --pixel-scale 0.1 --ms-suffix .ms {target_ms_path}" + else: + # Extract the previous working directory + flocs_workdir = "" + print( + f"Scanning log_VLBI_facet_imaging_{field['target_name']}_{field['sas_id_target']}.txt for workdir." + ) + with open( + os.path.join( + logsdir, + f"log_VLBI_facet_imaging_{field['target_name']}_{field['sas_id_target']}.txt", + ) + ) as f_out: + for line in f_out.readlines(): + print(line) + if "Running workflow with" in line: + flocs_workdir = line.split(" ")[-1].strip() + break + if not flocs_workdir: + raise RuntimeError( + "Could not retrieve PILOT workdir. Flocs probably crashed before launching." + ) + print(f"Resuming failed PILOT run in {flocs_workdir}") + cmd = f"flocs-run vlbi facet-imaging --record-toil-stats --runner toil --scheduler slurm --slurm-time 72:00:00 --slurm-account {SLURM_ACCOUNT} --slurm-queue {SLURM_QUEUE} --rundir {flocs_workdir} --restart --outdir {outdir} --resolution 0.3asec --pixel-scale 0.1 --ms-suffix .ms {target_ms_path}" + if not os.path.isdir(outdir): + os.makedirs(outdir, exist_ok=True) + print(cmd) + with ( + open( + os.path.join( + logsdir, + f"log_VLBI_facet_imaging_{field['target_name']}_{field['sas_id_target']}.txt", + ), + "w+", + ) as f_out, + open( + os.path.join( + logsdir, + f"log_VLBI_facet_imaging_{field['target_name']}_{field['sas_id_target']}_err.txt", + ), + "w+", + ) as f_err, + ): + proc = subprocess.run(cmd, shell=True, text=True, stdout=f_out, stderr=f_err) + success = False + pattern = re.compile(r"Workflow.* stopped. Success: True") + if not proc.returncode: + f_err.seek(0) + if pattern.search(f_err.read()): + success = True + if success: + db.set_status_finished( + field["target_name"], + "vlbi_facet_imaging", + field["sas_id_target"], + ) + db.set_field_finished(field["target_name"], field["sas_id_target"]) + else: + db.set_status_failed( + field["target_name"], + "vlbi_facet_imaging", + field["sas_id_target"], + ) + raise RuntimeError + + +def run_prepare_ddf_subtract(field): + print( + f"Preparing input for DDF subtract of {field['target_name']} {field['sas_id_target']}" + ) + outdir = os.path.join(OUTPUT_DIR, field["target_name"]) + logsdir = os.path.join(OUTPUT_DIR, field["target_name"], "logs") + if not os.path.isdir(outdir): + os.makedirs(outdir, exist_ok=True) + if not os.path.isdir(logsdir): + os.makedirs(logsdir, exist_ok=True) + target_path = get_most_recent_run(outdir, field["sas_id_target"], "VLBI_delay") + target_ms_path = target_path / "results_VLBI_delay-calibration" + mses_unaveraged = list(target_ms_path.glob("*.dp3concat")) + delay_sols = "" + if not mses_unaveraged: + print( + "No MSes found in delay-calibration output, will apply delay solutions to LINC." + ) + sols_path = get_most_recent_run(outdir, field["sas_id_target"], "VLBI_delay") + sols_path = sols_path / "results_VLBI_delay-calibration" + delay_sols = list(sols_path.glob("merged*selfcalcycle???_linearfulljones*.h5"))[ + 0 + ] + print(f"Using PILOT delay calibration solutions: {delay_sols}") + + linc_path = get_most_recent_run(outdir, field["sas_id_target"], "LINC_target") + print(f"Using LINC MSes at {linc_path}") + linc_ms_path = linc_path / "results_LINC_target" / "results" + mses_unaveraged = list(linc_ms_path.glob("*.dp3concat")) + mses_unaveraged_pilot = list(target_ms_path.glob("*.dp3concat")) + if not mses_unaveraged: + raise RuntimeError( + f"No unaveraged input MSes found at {linc_ms_path}/*.dp3concat" + ) + if mses_unaveraged_pilot and (len(mses_unaveraged_pilot) == len(mses_unaveraged)): + print("Appropriate input exists for ddf-pipeline.") + return mses_unaveraged_pilot + + jobids = [] + delay_corrected_mses = [] + for ms in mses_unaveraged: + out_ms = target_ms_path / f"{ms.stem}.dp3concat" + if out_ms.exists(): + print(f"Skipping {out_ms.name}, already exists.") + delay_corrected_mses.append(str(out_ms)) + continue + if delay_sols: + h5 = h5parm(str(delay_sols)) + ss = h5.getSolset("sol000") + # We only expect there to be one direction: the delay calibrator. + dirname = list(ss.getSou())[0] + sourcedir = ss.getSou()[dirname] + delaydir = f"[{sourcedir[0]},{sourcedir[1]}]" + dp3_cmd = f"apptainer exec $CWL_SINGULARITY_CACHE/astronrd_linc_latest.sif DP3 numthreads=2 msin={ms} msout={out_ms} msout.uvwcompression=False msout.antennacompression=False msout.scalarflags=False msout.storagemanager=Dysco steps=[applybeamdelay,applycal,applybeamtarget] applybeamdelay.type=applybeam applybeamdelay.beammode=full applybeamdelay.updateweights=True applybeamdelay.direction={delaydir} applycal.parmdb={delay_sols} applycal.correction=fulljones applycal.soltab=[amplitude000,phase000] applybeamtarget.type=applybeam applybeamtarget.beammode=full applybeamtarget.updateweights=True" + print(dp3_cmd) + else: + dp3_cmd = f"apptainer exec $CWL_SINGULARITY_CACHE/astronrd_linc_latest.sif DP3 numthreads=2 msin={ms} msout={out_ms} msout.uvwcompression=False msout.antennacompression=False msout.scalarflags=False msout.storagemanager=Dysco steps=[]" + submit_cmd = f'sbatch -A {SLURM_ACCOUNT} -p {SLURM_QUEUE} --time=08:00:00 -c 2 --job-name=dp3_avg_{ms.stem} --wrap="{dp3_cmd}"' + print(f"Submitting: {submit_cmd}") + proc = subprocess.run(submit_cmd, shell=True, text=True, capture_output=True) + if proc.returncode: + print(proc.stdout) + print(proc.stderr) + raise RuntimeError(f"Failed to submit SLURM job for {ms}") + jobid = proc.stdout.strip().split()[-1] + print(f"Submitted job {jobid} for {ms.name}") + jobids.append((jobid, out_ms)) + delay_corrected_mses.append(str(out_ms)) + + while jobids: + print(f"Polling {len(jobids)} SLURM jobs...") + remaining = [] + for jobid, out_ms in jobids: + poll_cmd = f"sacct -X -j {jobid} --format=State --noheader" + status = subprocess.run( + poll_cmd, shell=True, text=True, capture_output=True + ).stdout.strip() + if status == "COMPLETED": + print(f"Job {jobid} completed ({out_ms.name})") + elif ( + (status == "FAILED") or ("TIMEOUT" in status) or ("CANCELLED" in status) + ): + raise RuntimeError( + f"DP3 averaging job {jobid} failed for {out_ms.name}" + ) + elif status in ("PENDING", "RUNNING"): + remaining.append((jobid, out_ms)) + else: + remaining.append((jobid, out_ms)) + jobids = remaining + time.sleep(30) + + mses_delay_corrected = list(target_ms_path.glob("*.dp3concat")) + return mses_delay_corrected + + +def run_prepare_ddf(field): + print(f"Preparing DDF input for {field['target_name']} {field['sas_id_target']}") + outdir = os.path.join(OUTPUT_DIR, field["target_name"]) + logsdir = os.path.join(OUTPUT_DIR, field["target_name"], "logs") + if not os.path.isdir(outdir): + os.makedirs(outdir, exist_ok=True) + if not os.path.isdir(logsdir): + os.makedirs(logsdir, exist_ok=True) + target_path = get_most_recent_run(outdir, field["sas_id_target"], "VLBI_delay") + target_ms_path = target_path / "results_VLBI_delay-calibration" + mses_unaveraged = list(target_ms_path.glob("*.dp3concat")) + delay_sols = "" + if not mses_unaveraged: + print( + "No MSes found in delay-calibration output, will apply delay solutions to LINC." + ) + sols_path = get_most_recent_run(outdir, field["sas_id_target"], "VLBI_delay") + sols_path = sols_path / "results_VLBI_delay-calibration" + delay_sols = list(sols_path.glob("merged*selfcalcycle???_linearfulljones*.h5"))[ + 0 + ] + print(f"Using PILOT delay calibration solutions: {delay_sols}") + + linc_path = get_most_recent_run(outdir, field["sas_id_target"], "LINC_target") + print(f"Using LINC MSes at {linc_path}") + linc_ms_path = linc_path / "results_LINC_target" / "results" + mses_unaveraged = list(linc_ms_path.glob("*.dp3concat")) + else: + print("Found unaveraged data in delay calibration output.") + mses_averaged = list(target_ms_path.glob("*_pre-cal.ms")) + if not mses_unaveraged: + raise RuntimeError( + f"No unaveraged input MSes found at {linc_ms_path}/*.dp3concat" + ) + if mses_unaveraged and (len(mses_averaged) == len(mses_unaveraged)): + print("Appropriate input exists for ddf-pipeline.") + return mses_averaged + + jobids = [] + averaged_mses = [] + for ms in mses_unaveraged: + out_ms = target_ms_path / f"{ms.stem}_pre-cal.ms" + if out_ms.exists(): + print(f"Skipping {out_ms.name}, already exists.") + averaged_mses.append(str(out_ms)) + continue + if delay_sols: + h5 = h5parm(str(delay_sols)) + ss = h5.getSolset("sol000") + # We only expect there to be one direction: the delay calibrator. + dirname = list(ss.getSou())[0] + sourcedir = ss.getSou()[dirname] + delaydir = f"[{sourcedir[0]},{sourcedir[1]}]" + dp3_cmd = f"apptainer exec $CWL_SINGULARITY_CACHE/astronrd_linc_latest.sif DP3 numthreads=2 msin={ms} msout={out_ms} msout.uvwcompression=False msout.antennacompression=False msout.scalarflags=False msout.storagemanager=Dysco steps=[average,applybeamdelay,applycal,applybeamtarget,filter] average.timeresolution=8 average.freqresolution=97.64kHz applybeamdelay.type=applybeam applybeamdelay.beammode=full applybeamdelay.updateweights=True applybeamdelay.direction={delaydir} applycal.parmdb={delay_sols} applycal.correction=fulljones applycal.soltab=[amplitude000,phase000] applybeamtarget.type=applybeam applybeamtarget.beammode=full applybeamtarget.updateweights=True filter.remove=True filter.baseline='[CR]S*&&'" + print(dp3_cmd) + else: + dp3_cmd = f"apptainer exec $CWL_SINGULARITY_CACHE/astronrd_linc_latest.sif DP3 numthreads=2 msin={ms} msout={out_ms} msout.uvwcompression=False msout.antennacompression=False msout.scalarflags=False msout.storagemanager=Dysco steps=[filter,average] average.timeresolution=8 average.freqresolution=97.64kHz filter.remove=True filter.baseline='[CR]S*&&'" + submit_cmd = f'sbatch -A {SLURM_ACCOUNT} -p {SLURM_QUEUE} --time=02:00:00 -c 2 --job-name=dp3_avg_{ms.stem} --wrap="{dp3_cmd}"' + print(f"Submitting: {submit_cmd}") + proc = subprocess.run(submit_cmd, shell=True, text=True, capture_output=True) + if proc.returncode: + print(proc.stdout) + print(proc.stderr) + raise RuntimeError(f"Failed to submit SLURM job for {ms}") + jobid = proc.stdout.strip().split()[-1] + print(f"Submitted job {jobid} for {ms.name}") + jobids.append((jobid, out_ms)) + averaged_mses.append(str(out_ms)) + + while jobids: + print(f"Polling {len(jobids)} SLURM jobs...") + remaining = [] + for jobid, out_ms in jobids: + poll_cmd = f"sacct -X -j {jobid} --format=State --noheader" + status = subprocess.run( + poll_cmd, shell=True, text=True, capture_output=True + ).stdout.strip() + if status == "COMPLETED": + print(f"Job {jobid} completed ({out_ms.name})") + elif ( + (status == "FAILED") or ("TIMEOUT" in status) or ("CANCELLED" in status) + ): + raise RuntimeError( + f"DP3 averaging job {jobid} failed for {out_ms.name}" + ) + elif status in ("PENDING", "RUNNING"): + remaining.append((jobid, out_ms)) + else: + remaining.append((jobid, out_ms)) + jobids = remaining + time.sleep(30) + + mses_averaged = list(target_ms_path.glob("*_pre-cal.ms")) + return mses_averaged + + +def run_pilot_process_ddf_toil(field, db: FlocsDB): + print(f"Running ddf subtract for {field['target_name']} {field['sas_id_target']}") + outdir = os.path.join(OUTPUT_DIR, field["target_name"]) + logsdir = os.path.join(OUTPUT_DIR, field["target_name"], "logs") + if not os.path.isdir(outdir): + os.makedirs(outdir, exist_ok=True) + if not os.path.isdir(logsdir): + os.makedirs(logsdir, exist_ok=True) + target_path = get_most_recent_run( + outdir, field["sas_id_target"], "VLBI_delay-calibration" + ) + target_ms_path = target_path / "results_VLBI_delay-calibration" + print(f"Using data at: {target_path}/*.dp3concat") + + ddf_path = get_most_recent_run(outdir, field["sas_id_target"], "DDF-pipeline") + ddf_sols_path = ddf_path / "SOLSDIR" + print(f"Using DDF run at: {ddf_path}") + + context = get_current_context() + if context["ti"].try_number == 1 or ( + not os.path.isfile( + os.path.join( + logsdir, + f"log_VLBI_process-ddf_{field['target_name']}_{field['sas_id_target']}.txt", + ) + ) + ): + db.set_status_processing( + field["target_name"], "vlbi_ddf_subtract", field["sas_id_target"] + ) + cmd = f"flocs-run vlbi process-ddf --runner toil --record-toil-stats --scheduler slurm --slurm-time 24:00:00 --slurm-account {SLURM_ACCOUNT} --slurm-queue {SLURM_QUEUE} --rundir {PROCESSING_DIR} --outdir {outdir} --ms-suffix .dp3concat --ddf-rundir {ddf_path} --solsdir {ddf_sols_path} --do-subtraction {target_ms_path}" + else: + # Extract the previous working directory + flocs_workdir = "" + print( + f"Scanning log_VLBI_process-ddf_{field['target_name']}_{field['sas_id_target']}.txt for workdir." + ) + with ( + open( + os.path.join( + logsdir, + f"log_VLBI_process-ddf_{field['target_name']}_{field['sas_id_target']}.txt", + ) + ) as f_out, + open( + os.path.join( + logsdir, + f"log_VLBI_process-ddf_{field['target_name']}_{field['sas_id_target']}_err.txt", + ) + ) as f_err, + ): + for line in f_out.readlines(): + print(line) + if "Running workflow with" in line: + flocs_workdir = line.split(" ")[-1].strip() + break + if not flocs_workdir: + raise RuntimeError( + "Could not retrieve PILOT workdir. Flocs probably crashed before launching." + ) + print(f"Resuming failed PILOT run in {flocs_workdir}") + cmd = f"flocs-run vlbi process-ddf --runner toil --record-toil-stats --scheduler slurm --slurm-time 24:00:00 --slurm-account {SLURM_ACCOUNT} --slurm-queue {SLURM_QUEUE} --rundir {flocs_workdir} --restart --outdir {outdir} --ms-suffix .dp3concat --ddf-rundir {ddf_path} --solsdir {ddf_sols_path} --do-subtraction {target_ms_path}" + if not os.path.isdir(outdir): + os.makedirs(outdir, exist_ok=True) + print(cmd) + with ( + open( + os.path.join( + logsdir, + f"log_VLBI_process-ddf_{field['target_name']}_{field['sas_id_target']}.txt", + ), + "w+", + ) as f_out, + open( + os.path.join( + logsdir, + f"log_VLBI_process-ddf_{field['target_name']}_{field['sas_id_target']}_err.txt", + ), + "w+", + ) as f_err, + ): + proc = subprocess.run(cmd, shell=True, text=True, stdout=f_out, stderr=f_err) + success = False + pattern = re.compile(r"Workflow.* stopped. Success: True") + if not proc.returncode: + f_err.seek(0) + if pattern.search(f_err.read()): + success = True + if success: + db.set_status_finished( + field["target_name"], "vlbi_ddf_subtract", field["sas_id_target"] + ) + else: + db.set_status_failed( + field["target_name"], "vlbi_ddf_subtract", field["sas_id_target"] + ) + raise RuntimeError + + +def launch_ddf_pipeline(field, db: FlocsDB): + print(f"Starting ddf-pipeline for {field['target_name']} {field['sas_id_target']}") + outdir = os.path.join(OUTPUT_DIR, field["target_name"]) + logsdir = os.path.join(OUTPUT_DIR, field["target_name"], "logs") + if not os.path.isdir(outdir): + os.makedirs(outdir, exist_ok=True) + if not os.path.isdir(logsdir): + os.makedirs(logsdir, exist_ok=True) + target_path = get_most_recent_run( + outdir, field["sas_id_target"], "VLBI_delay-calibration" + ) + target_ms_path = target_path / "results_VLBI_delay-calibration" + if not list(target_ms_path.glob("*pre-cal.ms")): + target_path = get_most_recent_run(outdir, field["sas_id_target"], "LINC_target") + target_ms_path = target_path / "results_LINC_target" / "results" + + cmd = f"flocs-run ddf-pipeline --scheduler slurm --slurm-time 72:00:00 --slurm-cores 32 --slurm-account {SLURM_ACCOUNT} --slurm-queue {SLURM_QUEUE} --rundir {PROCESSING_DIR} --outdir {outdir} --config-file {DDF_CONFIG} {target_ms_path}" + print(cmd) + db.set_status_processing(field["target_name"], "ddf", field["sas_id_target"]) + with ( + open( + os.path.join( + logsdir, + f"log_DDF-pipeline_{field['target_name']}_{field['sas_id_target']}.txt", + ), + "w+", + ) as f_out, + open( + os.path.join( + logsdir, + f"log_DDF-pipeline_{field['target_name']}_{field['sas_id_target']}_err.txt", + ), + "w+", + ) as f_err, + ): + proc = subprocess.run(cmd, shell=True, text=True, stdout=f_out, stderr=f_err) + jobid = None + if not proc.returncode: + f_out.seek(0) + for line in f_out.readlines(): + if "Submitted batch job" in line: + jobid = line.strip().split()[-1] + else: + raise RuntimeError("Failed to submit job.") + + if not jobid: + raise RuntimeError("Failed to retrieve job id") + else: + while True: + print(f"Polling DDF-pipeine job {jobid}") + poll_cmd = f"sacct -X -j {jobid} --format=State --noheader" + status = subprocess.run( + poll_cmd, shell=True, text=True, capture_output=True + ).stdout.strip() + if (status == "RUNNING") or (status == "PENDING"): + time.sleep(60) + elif status == "COMPLETED": + db.set_status_finished( + field["target_name"], + "ddf", + field["sas_id_target"], + ) + break + elif ( + (status == "FAILED") + or ("TIMEOUT" in status) + or ("CANCELLED" in status) + ): + raise RuntimeError( + f"DDF-pipeline for {field['target_name']} {field['sas_id_target']} failed." + ) diff --git a/flocs_processing/processors.py b/flocs_processing/processors.py new file mode 100644 index 0000000..f86e867 --- /dev/null +++ b/flocs_processing/processors.py @@ -0,0 +1,190 @@ +import os +import shlex +import subprocess +import sys +from pathlib import Path +from typing import Optional + +import libtmux +import structlog + +logger = structlog.getLogger() + + +class FlocsAirflowProcessor: + def __init__(self, airflow_home: Optional[str], airflow_cores: int): + os.environ["AIRFLOW__CORE__LOAD_EXAMPLES"] = "False" + os.environ["AIRFLOW__CORE__PARALLELISM"] = str(airflow_cores) + self.REQUIRED_AIRFLOW_VARS = [ + "AIRFLOW_HOME", + "AIRFLOW__CORE__DAGS_FOLDER", + "AIRFLOW__CORE__LOAD_EXAMPLES", + "AIRFLOW__CORE__PARALLELISM", + "AIRFLOW__LOGGING__DAG_PROCESSOR_CHILD_PROCESS_LOG_DIRECTORY", + "AIRFLOW__CORE__PLUGINS_FOLDER", + "AIRFLOW__DATABASE__SQL_ALCHEMY_CONN", + "AIRFLOW__LOGGING__BASE_LOG_FOLDER", + ] + if airflow_home: + os.environ["AIRFLOW_HOME"] = airflow_home + os.environ["AIRFLOW__CORE__DAGS_FOLDER"] = os.path.join( + airflow_home, "dags" + ) + os.environ[ + "AIRFLOW__LOGGING__DAG_PROCESSOR_CHILD_PROCESS_LOG_DIRECTORY" + ] = os.path.join(airflow_home, "logs/dag_processor") + os.environ["AIRFLOW__CORE__PLUGINS_FOLDER"] = os.path.join( + airflow_home, "plugins" + ) + os.environ["AIRFLOW__DATABASE__SQL_ALCHEMY_CONN"] = ( + "sqlite:///$AIRFLOW_HOME/airflow.db" + ) + os.environ["AIRFLOW__LOGGING__BASE_LOG_FOLDER"] = os.path.join( + airflow_home, "logs" + ) + + def check_airflow_init(self): + """Check if the Airflow instance is initialised already. If not, runs `airflow db migrate` to intitialise it.""" + af_home = os.environ["AIRFLOW_HOME"] + if not os.path.isdir(af_home): + logger.info(f"Airflow home {af_home} does not exist, creating it") + os.makedirs(af_home) + if not os.listdir(af_home): + logger.info( + f"Airflow home {af_home} is empty. Running `airflow db migrate` to initialise." + ) + old_dir = os.getcwd() + os.chdir(af_home) + _ = subprocess.check_output( + "airflow db migrate", + text=True, + shell=True, + ) + os.chdir(old_dir) + else: + logger.info( + f"Airflow home {af_home} is populated. Leaving it alone and assuming it is operational." + ) + + def check_environment(self): + environment_ok = True + if "FLOCS_AIRFLOW_CONFIG" not in os.environ: + logger.warning( + "FLOCS_AIRFLOW_CONFIG environment variable not set. Please point this to a valid configuration file." + ) + environment_ok = False + else: + CONFIG_FILE = os.getenv("FLOCS_AIRFLOW_CONFIG") or "" + if not os.path.isfile(CONFIG_FILE): + logger.warning(f"{CONFIG_FILE} is not a valid file") + environment_ok = False + else: + logger.info(f"Found config file {os.getenv('FLOCS_AIRFLOW_CONFIG')}.") + + if "AIRFLOW__API_AUTH__JWT_SECRET" not in os.environ: + logger.critical("AIRFLOW__API_AUTH__JWT_SECRET not defined.") + environment_ok = False + + for kw in self.REQUIRED_AIRFLOW_VARS: + if kw not in os.environ: + logger.warning(f"Required variable {kw} is not set.") + environment_ok = False + + if environment_ok: + logger.info("Flocs automated processing environment appears ok") + else: + logger.critical( + "Flocs automated processing environment is not properly set up. Please see warnings." + ) + sys.exit(1) + + def check_airflow_services(self): + status = { + "api-server": False, + "dag-processor": False, + "scheduler": False, + "triggerer": False, + } + + tmux_server = libtmux.Server() + airflow_sessions = map( + lambda x: x.name, + filter(lambda x: "airflow" in x.name, tmux_server.sessions), + ) + logger.info("Checking Airflow status:") + for service in status.keys(): + if f"airflow-{service}" not in airflow_sessions: + logger.info(f"-- {service}: not running") + status[service] = False + else: + logger.info(f"-- {service}: running") + status[service] = True + return status + + def deploy_airflow_tmux(self): + af_status = self.check_airflow_services() + + if any(not up for up in af_status.values()): + source_file = os.path.abspath("source_flocs_airflow.sh") + logger.info("Not all Airflow services are online.") + tmux_server = libtmux.Server() + for service, running in af_status.items(): + if not running: + logger.info(f"-- deploying {service}") + tmux_server.new_session( + session_name=f"airflow-{service}", + attach=False, + window_command=( + f"bash -c 'source {source_file} && airflow {service}'; bash -i" + ), + ) + + def generate_jwt_secret(self): + if os.path.isfile(os.path.expandvars("$HOME/.config/airflow/jwt_secret")): + logger.info("JWT secret found, not (re)generating.") + else: + logger.info("No JWT secret found, generating one.") + _ = subprocess.check_output( + "mkdir -p $HOME/.config/airflow", + text=True, + shell=True, + ) + + _ = subprocess.check_output( + "chmod 700 $HOME/.config/airflow", + text=True, + shell=True, + ) + + _ = subprocess.check_output( + "openssl rand -hex 32 > $HOME/.config/airflow/jwt_secret", + text=True, + shell=True, + ) + + _ = subprocess.check_output( + "chmod 600 $HOME/.config/airflow/jwt_secret", + text=True, + shell=True, + ) + logger.info("JWT secret generated") + secret = Path.home().joinpath(".config/airflow/jwt_secret").read_text().strip() + os.environ["AIRFLOW__API_AUTH__JWT_SECRET"] = secret + + def write_source_file(self): + if os.path.exists("source_flocs_airflow.sh"): + logger.info("Found source_flocs_airflow.sh, not overwriting.") + else: + logger.info("Writing environment settings to source_flocs_airflow.sh") + with open("source_flocs_airflow.sh", "w") as f: + for kw in self.REQUIRED_AIRFLOW_VARS: + f.write(f"export {kw}={os.environ[kw]}\n") + + secret = ( + Path.home() + .joinpath(".config/airflow/jwt_secret") + .read_text() + .strip() + ) + f.write(f"export AIRFLOW__API_AUTH__JWT_SECRET={shlex.quote(secret)}\n") + os.chmod("source_flocs_airflow.sh", 0o600) diff --git a/pyproject.toml b/pyproject.toml index cc255e7..b100231 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -6,7 +6,7 @@ build-backend = "setuptools.build_meta" name = "flocs-processing" version = "0.0.0" description = "" -requires-python = ">3.9" +requires-python = ">=3.10" dependencies = ["cyclopts", "structlog", "packaging"] [project.scripts]