9393from urllib .parse import urlparse
9494
9595import hydra
96+ import yaml
9697from omegaconf import DictConfig , OmegaConf
9798
9899import constants
106107from experiment_utils .services .misc import ControllerService
107108from experiment_utils .services .query_engine import QueryEngineRustService
108109
110+
111+ def _inline_sql_queries_in_experiment_config (local_experiment_root_dir : str ) -> None :
112+ """Enrich the saved experiment_params.yaml by inlining SQL from sql_file references.
113+
114+ Downstream analysis scripts expect query_groups[i]["queries"] to be a list
115+ of query strings. The clickhouse config stores sql_file paths instead, so
116+ we read each file and add the queries in-place before the scripts run.
117+ """
118+ config_path = os .path .join (
119+ local_experiment_root_dir , "experiment_config" , "experiment_params.yaml"
120+ )
121+ if not os .path .exists (config_path ):
122+ return
123+ with open (config_path ) as f :
124+ data = yaml .safe_load (f )
125+
126+ def _expand_groups (groups ):
127+ if not groups :
128+ return
129+ for group in groups :
130+ sql_file = group .get ("sql_file" )
131+ if sql_file and "queries" not in group :
132+ with open (sql_file ) as fq :
133+ content = fq .read ()
134+ group ["queries" ] = [s .strip () for s in content .split (";" ) if s .strip ()]
135+
136+ _expand_groups (data .get ("query_groups" ))
137+ _expand_groups (data .get ("sketchdb_query_groups" ))
138+
139+ with open (config_path , "w" ) as f :
140+ yaml .dump (data , f , allow_unicode = True )
141+
142+
109143# Register resolvers used by config.yaml interpolation.
110144OmegaConf .register_new_resolver (
111145 "local_experiment_dir" , lambda : constants .LOCAL_EXPERIMENT_DIR
@@ -234,6 +268,7 @@ def main(cfg: DictConfig) -> None:
234268 )
235269
236270 sync .copy_experiment_config (cfg .experiment_params , local_experiment_root_dir )
271+ _inline_sql_queries_in_experiment_config (local_experiment_root_dir )
237272
238273 # --- dataset config ---
239274 ep = cfg .experiment_params
@@ -339,6 +374,11 @@ def main(cfg: DictConfig) -> None:
339374
340375 if experiment_mode == constants .SKETCHDB_EXPERIMENT_NAME :
341376 # --- sketchdb mode: precompute engine + JSON ingest + ClickHouse fallback ---
377+ # Kill any leftover query_engine_rust from a previous run (mirrors
378+ # query_engine_service.stop() at the top of the e2e mode loop).
379+ QueryEngineRustService (
380+ provider = provider , use_container = False , node_offset = node_offset
381+ ).stop ()
342382 # Mirrors experiment_run_e2e.py: planner runs first and generates
343383 # streaming_config.yaml + inference_config.yaml into controller_output_dir,
344384 # then the query engine starts reading from that same directory.
@@ -438,7 +478,9 @@ def main(cfg: DictConfig) -> None:
438478 provider = provider ,
439479 node_offset = node_offset ,
440480 config_file = controller_client_config ,
441- output_dir = experiment_output_dir ,
481+ output_dir = os .path .join (
482+ experiment_output_dir , "prometheus_client_output"
483+ ),
442484 use_container = use_container ,
443485 parallel = parallel ,
444486 )
@@ -465,7 +507,9 @@ def main(cfg: DictConfig) -> None:
465507 provider = provider ,
466508 node_offset = node_offset ,
467509 config_file = controller_client_config ,
468- output_dir = experiment_output_dir ,
510+ output_dir = os .path .join (
511+ experiment_output_dir , "prometheus_client_output"
512+ ),
469513 use_container = use_container ,
470514 parallel = parallel ,
471515 )
0 commit comments