From 6d86ceabb5989c78fb2534f89ae2ef3489ff4474 Mon Sep 17 00:00:00 2001 From: Furkan Date: Fri, 8 Nov 2024 20:12:37 +0300 Subject: [PATCH] feat: use literal class names for configs and deprecate `DAQ_JOB_TYPE_TO_CLASS` --- configs/examples/alert_slack.toml | 2 +- configs/examples/handle_alerts.toml | 2 +- configs/examples/handle_stats.toml | 2 +- configs/examples/healthcheck.toml | 2 +- configs/examples/n1081b.toml | 2 +- configs/examples/remote.toml | 2 +- configs/examples/serve_http.toml | 2 +- configs/examples/store_csv.toml | 2 +- configs/examples/store_root.toml | 2 +- configs/examples/test.toml | 2 +- src/daq/daq_job.py | 21 +++++++++++++++++---- src/daq/store/models.py | 22 ++++++++++++++++++++-- src/utils/subclasses.py | 4 ++++ 13 files changed, 51 insertions(+), 16 deletions(-) create mode 100644 src/utils/subclasses.py diff --git a/configs/examples/alert_slack.toml b/configs/examples/alert_slack.toml index cd98d45..0d5a1bb 100644 --- a/configs/examples/alert_slack.toml +++ b/configs/examples/alert_slack.toml @@ -1,2 +1,2 @@ -daq_job_type = "alert_slack" +daq_job_type = "DAQJobAlertSlack" slack_webhook_url = "https://hooks.slack.com/services/xxxxxxxxxxxxxxxxxxxxxxxxx" diff --git a/configs/examples/handle_alerts.toml b/configs/examples/handle_alerts.toml index e3cf271..a696549 100644 --- a/configs/examples/handle_alerts.toml +++ b/configs/examples/handle_alerts.toml @@ -1,4 +1,4 @@ -daq_job_type = "handle_alerts" +daq_job_type = "DAQJobHandleAlerts" [store_config.csv] file_path = "alerts.csv" diff --git a/configs/examples/handle_stats.toml b/configs/examples/handle_stats.toml index b34c56d..dbe8bf9 100644 --- a/configs/examples/handle_stats.toml +++ b/configs/examples/handle_stats.toml @@ -1,4 +1,4 @@ -daq_job_type = "handle_stats" +daq_job_type = "DAQJobHandleStats" [store_config.csv] file_path = "stats.csv" diff --git a/configs/examples/healthcheck.toml b/configs/examples/healthcheck.toml index 0f599e0..ef590e9 100644 --- a/configs/examples/healthcheck.toml +++ b/configs/examples/healthcheck.toml @@ -1,4 +1,4 @@ -daq_job_type = "healthcheck" +daq_job_type = "DAQJobHealthcheck" # Alert if restarted [[healthcheck_stats]] diff --git a/configs/examples/n1081b.toml b/configs/examples/n1081b.toml index 221d48e..0c1e343 100644 --- a/configs/examples/n1081b.toml +++ b/configs/examples/n1081b.toml @@ -1,4 +1,4 @@ -daq_job_type = "n1081b" +daq_job_type = "DAQJobN1081B" host = "1.2.3.4" port = "8080" password = "password" diff --git a/configs/examples/remote.toml b/configs/examples/remote.toml index 7e13f32..c248f31 100644 --- a/configs/examples/remote.toml +++ b/configs/examples/remote.toml @@ -1,3 +1,3 @@ -daq_job_type = "remote" +daq_job_type = "DAQJobRemote" zmq_local_url = "tcp://localhost:10191" zmq_remote_urls = ["tcp://1.2.3.4:10191"] diff --git a/configs/examples/serve_http.toml b/configs/examples/serve_http.toml index ac83869..1b2a063 100644 --- a/configs/examples/serve_http.toml +++ b/configs/examples/serve_http.toml @@ -1,4 +1,4 @@ -daq_job_type = "serve_http" +daq_job_type = "DAQJobServeHTTP" serve_path = "out/" host = "0.0.0.0" port = 8000 diff --git a/configs/examples/store_csv.toml b/configs/examples/store_csv.toml index c624f71..4f36512 100644 --- a/configs/examples/store_csv.toml +++ b/configs/examples/store_csv.toml @@ -1 +1 @@ -daq_job_type = "store_csv" +daq_job_type = "DAQJobStoreCSV" diff --git a/configs/examples/store_root.toml b/configs/examples/store_root.toml index 52a61ee..4b9b1cb 100644 --- a/configs/examples/store_root.toml +++ b/configs/examples/store_root.toml @@ -1 +1 @@ -daq_job_type = "store_root" +daq_job_type = "DAQJobStoreROOT" diff --git a/configs/examples/test.toml b/configs/examples/test.toml index ca77a16..74453ca 100644 --- a/configs/examples/test.toml +++ b/configs/examples/test.toml @@ -1,4 +1,4 @@ -daq_job_type = "test" +daq_job_type = "DAQJobTest" rand_min = 1 rand_max = 100 diff --git a/src/daq/daq_job.py b/src/daq/daq_job.py index 0f2d5b8..dc39135 100644 --- a/src/daq/daq_job.py +++ b/src/daq/daq_job.py @@ -8,16 +8,29 @@ from daq.base import DAQJob, DAQJobThread from daq.models import DAQJobConfig from daq.types import DAQ_JOB_TYPE_TO_CLASS +from utils.subclasses import all_subclasses + +ALL_DAQ_JOBS = all_subclasses(DAQJob) def build_daq_job(toml_config: bytes) -> DAQJob: generic_daq_job_config = msgspec.toml.decode(toml_config, type=DAQJobConfig) - - if generic_daq_job_config.daq_job_type not in DAQ_JOB_TYPE_TO_CLASS: + daq_job_class = None + + if generic_daq_job_config.daq_job_type in DAQ_JOB_TYPE_TO_CLASS: + daq_job_class = DAQ_JOB_TYPE_TO_CLASS[generic_daq_job_config.daq_job_type] + logging.warning( + f"DAQ job type '{generic_daq_job_config.daq_job_type}' is deprecated, please use the '{daq_job_class.__name__}' instead" + ) + else: + for daq_job in ALL_DAQ_JOBS: + if daq_job.__name__ == generic_daq_job_config.daq_job_type: + daq_job_class = daq_job + + if daq_job_class is None: raise Exception(f"Invalid DAQ job type: {generic_daq_job_config.daq_job_type}") - # Get DAQ and DAQ config clasess based on daq_job_type - daq_job_class = DAQ_JOB_TYPE_TO_CLASS[generic_daq_job_config.daq_job_type] + # Get DAQ config clase based on daq_job_type daq_job_config_class: DAQJobConfig = daq_job_class.config_type # Load the config in diff --git a/src/daq/store/models.py b/src/daq/store/models.py index 7c7c3be..6b1cb81 100644 --- a/src/daq/store/models.py +++ b/src/daq/store/models.py @@ -36,12 +36,30 @@ class StorableDAQJobConfig(DAQJobConfig): store_config: DAQJobStoreConfig -class DAQJobStoreConfigCSV(Struct): +class DAQJobStoreTarget(Struct): + instances: Optional[list["DAQJobStoreTargetInstance"]] = None + jobs: Optional[list["DAQJobStoreTargetJob"]] = None + + +class DAQJobStoreTargetInstance(Struct): + supervisor_id: Optional[str] = None + is_self: Optional[bool] = None + + +class DAQJobStoreTargetJob(Struct): + job_name: str + + +class DAQJobStoreConfigBase(Struct, kw_only=True): + target: Optional[DAQJobStoreTarget] = None + + +class DAQJobStoreConfigCSV(DAQJobStoreConfigBase): file_path: str add_date: bool overwrite: Optional[bool] = None -class DAQJobStoreConfigROOT(Struct): +class DAQJobStoreConfigROOT(DAQJobStoreConfigBase): file_path: str add_date: bool diff --git a/src/utils/subclasses.py b/src/utils/subclasses.py new file mode 100644 index 0000000..de9906c --- /dev/null +++ b/src/utils/subclasses.py @@ -0,0 +1,4 @@ +def all_subclasses(cls): + return set(cls.__subclasses__()).union( + [s for c in cls.__subclasses__() for s in all_subclasses(c)] + )