cpic.py

"""
Aim
---
Process an CPIC dataset, given a set of parameters.

Example
-------
python -m csst_dag.cli.cpic -h

python -m csst_dag.cli.cpic \
    --dataset=csst-cpic-c11-hip71681-v1 \
    --obs-group=hip71681 \
    --batch-id=csci-test-20250507 \
    --priority=1 \
    --initial-prc-status=-1024 \
    --final-prc-status=-2 \
    --demo
"""

from csst_dag.dag import CsstDAG
import argparse

parser = argparse.ArgumentParser(
    description="Scheduler for CPIC L1 pipeline.",
    formatter_class=argparse.ArgumentDefaultsHelpFormatter,
)

parser.add_argument("--dataset", type=str, help="Dataset name")
# parser.add_argument("--instrument", type=str, help="Instrument name", default="CPIC")
parser.add_argument("--obs-group", type=str, help="Observation group", default="none")
parser.add_argument("--obs-type", type=str, help="Observation type", default="")
parser.add_argument("--batch-id", type=str, help="Batch ID", default="default_batch")
parser.add_argument("--priority", type=str, help="Task priority", default=1)
parser.add_argument(
    "--initial-prc-status", type=int, help="Initial processing status", default=-1024
)
parser.add_argument(
    "--final-prc-status", type=int, help="Final processing status", default=-2
)
parser.add_argument(
    "--demo", action="store_true", help="Enable demo mode", default=False
)

args = parser.parse_args()
print("CLI parameters: ", args)

DAG_LOOP_MAP = {
    "SCI": ["csst-cpic-l1"],
    "DSF": ["csst-cpic-l1"],
    "CALS": ["csst-cpic-l1"],
    "BIAS": ["csst-cpic-l1-qc0"],
    "DARK": ["csst-cpic-l1-qc0"],
    "FLAT": ["csst-cpic-l1-qc0"],
    "BKG": ["csst-cpic-l1-qc0"],
    "LASER": ["csst-cpic-l1-qc0"],
}
if args.obs_type:
    assert args.obs_type in DAG_LOOP_MAP.keys(), f"Unknown obs_type: {args.obs_type}"
    DAG_LOOP_MAP = {args.obs_type: DAG_LOOP_MAP[args.obs_type]}

for obs_type, dag_ids in DAG_LOOP_MAP.items():
    print(f"* Processing {obs_type}")
    for dag_id in dag_ids:
        print(f"   - Scheduling `{dag_id}` -> ", end="")
        dag = CsstDAG.get_dag(dag_id=dag_id)
        msgs = dag.schedule(
            dataset=args.dataset,
            obs_type=obs_type,
            obs_group=args.obs_group,
            batch_id=args.batch_id,
            initial_prc_status=args.initial_prc_status,
            final_prc_status=args.final_prc_status,
            demo=True,
            priority=args.priority,
        )
        print(f"{len(msgs)} tasks.")