msc.py 2.45 KB
Newer Older
BO ZHANG's avatar
tweaks    
BO ZHANG committed
1
2
3
"""
Aim
---
BO ZHANG's avatar
tweaks    
BO ZHANG committed
4
Process an MSC dataset, given a set of parameters.
BO ZHANG's avatar
tweaks    
BO ZHANG committed
5
6
7

Example
-------
BO ZHANG's avatar
BO ZHANG committed
8
9
python -m csst_dag.cli.msc -h

BO ZHANG's avatar
tweaks    
BO ZHANG committed
10
11
python -m csst_dag.cli.msc \
    --batch-id=csci-test-20250507 \
BO ZHANG's avatar
BO ZHANG committed
12
13
    --dataset=csst-msc-c9-25sqdeg-v3 \
    --obs-group=W1 \
BO ZHANG's avatar
tweaks    
BO ZHANG committed
14
15
    --priority=1 \
    --initial-prc-status=-1024 \
BO ZHANG's avatar
BO ZHANG committed
16
    --final-prc-status=-2 \
BO ZHANG's avatar
BO ZHANG committed
17
    --demo
BO ZHANG's avatar
tweaks    
BO ZHANG committed
18
19
"""

BO ZHANG's avatar
BO ZHANG committed
20
from csst_dag.dag import CsstDAGs
BO ZHANG's avatar
BO ZHANG committed
21
22
23
24
25
26
27
import argparse

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

BO ZHANG's avatar
BO ZHANG committed
28
29
30
31
# task related parameters
parser.add_argument("--batch-id", type=str, help="Batch ID", default="default_batch")
parser.add_argument("--priority", type=str, help="Task priority", default=1)
# data related parameters
BO ZHANG's avatar
BO ZHANG committed
32
33
parser.add_argument("--dataset", type=str, help="Dataset name")
# parser.add_argument("--instrument", type=str, help="Instrument name", default="MSC")
BO ZHANG's avatar
BO ZHANG committed
34
parser.add_argument("--obs-group", type=str, help="Observation group", default="none")
BO ZHANG's avatar
tweaks    
BO ZHANG committed
35
parser.add_argument("--obs-type", type=str, help="Observation type", default="")
BO ZHANG's avatar
BO ZHANG committed
36
# status related parameters
BO ZHANG's avatar
tweaks    
BO ZHANG committed
37
38
39
40
41
42
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
)
BO ZHANG's avatar
BO ZHANG committed
43
parser.add_argument(
BO ZHANG's avatar
BO ZHANG committed
44
    "--demo", action="store_true", help="Enable demo mode", default=False
BO ZHANG's avatar
BO ZHANG committed
45
)
BO ZHANG's avatar
BO ZHANG committed
46

BO ZHANG's avatar
BO ZHANG committed
47
args = parser.parse_args()
BO ZHANG's avatar
BO ZHANG committed
48
print("CLI parameters: ", args)
BO ZHANG's avatar
BO ZHANG committed
49

BO ZHANG's avatar
BO ZHANG committed
50
51
52
53
# define DAP LIST in this group
DAG_MAP = {
    "WIDE": ["csst-msc-l1-mbi", "csst-msc-l1-sls", "csst-msc-l1-ast"],
    "DEEP": ["csst-msc-l1-mbi", "csst-msc-l1-sls", "csst-msc-l1-ast"],
BO ZHANG's avatar
BO ZHANG committed
54
55
56
57
    "BIAS": ["csst-msc-l1-qc0"],
    "DARK": ["csst-msc-l1-qc0"],
    "FLAT": ["csst-msc-l1-qc0"],
}
BO ZHANG's avatar
BO ZHANG committed
58
59

# if obs_type is set
BO ZHANG's avatar
BO ZHANG committed
60
if args.obs_type:
BO ZHANG's avatar
BO ZHANG committed
61
62
    assert args.obs_type in DAG_MAP.keys(), f"Unknown obs_type: {args.obs_type}"
    DAG_MAP = {args.obs_type: DAG_MAP[args.obs_type]}
BO ZHANG's avatar
BO ZHANG committed
63

BO ZHANG's avatar
BO ZHANG committed
64
for obs_type, dags in DAG_MAP.items():
BO ZHANG's avatar
tweaks    
BO ZHANG committed
65
    print(f"* Processing `{obs_type}`")
BO ZHANG's avatar
BO ZHANG committed
66
    for this_dag in dags:
BO ZHANG's avatar
tweaks    
BO ZHANG committed
67
        print(f"  - Scheduling `{this_dag}` -> ", end="")
BO ZHANG's avatar
tweaks    
BO ZHANG committed
68
        dag = CsstDAGs().get(this_dag)
BO ZHANG's avatar
BO ZHANG committed
69
        msgs = dag.schedule(
BO ZHANG's avatar
BO ZHANG committed
70
71
            batch_id=args.batch_id,
            priority=args.priority,
BO ZHANG's avatar
BO ZHANG committed
72
            dataset=args.dataset,
BO ZHANG's avatar
BO ZHANG committed
73
            obs_type=obs_type,
BO ZHANG's avatar
BO ZHANG committed
74
            obs_group=args.obs_group,
BO ZHANG's avatar
tweaks    
BO ZHANG committed
75
76
            initial_prc_status=args.initial_prc_status,
            final_prc_status=args.final_prc_status,
BO ZHANG's avatar
BO ZHANG committed
77
            demo=args.demo,
BO ZHANG's avatar
BO ZHANG committed
78
        )
BO ZHANG's avatar
BO ZHANG committed
79
        print(f"{len(msgs)} tasks.")