added query_knobs param to time_query

Author: patrickhyw

env/pg_conn.py

@@ -110,11 +110,16 @@ def force_statement_timeout(self, timeout: float) -> None:
                 retry = True
 
     def time_query(
-        self, query: str, add_explain: bool = False, timeout: float = 0
+        self,
+        query: str,
+        query_knobs: list[str] = [],
+        add_explain: bool = False,
+        timeout: float = 0,
     ) -> tuple[float, bool, Optional[dict[str, Any]]]:
         """
-        Run a query with a timeout (in seconds). If you want to attach per-query knobs, attach them to the query string
-        itself. Following Postgres's convention, timeout=0 indicates "disable timeout"
+        Run a query with a timeout (in seconds). Following Postgres's convention, timeout=0 indicates "disable timeout".
+
+        Use query_knobs to pass query knobs. An example input is query_knobs=["SET (enable_sort on)", "IndexOnlyScan(it)"].
 
         It returns the runtime, whether the query timed out, and the explain data if add_explain is True. Note that if
         the query timed out, it won't have any explain data and thus explain_data will be None.
@@ -133,6 +138,9 @@ def time_query(
         explain_data = None
 
         try:
+            if query_knobs:
+                query = f"/*+ {' '.join(query_knobs)} */ {query}"
+
             if add_explain:
                 assert (
                     "explain" not in query.lower()

scripts/pat_test.sh

@@ -5,9 +5,7 @@ set -euxo pipefail
 . ./experiments/load_per_machine_envvars.sh
 
 # space for testing. uncomment this to run individual commands from the script (copy pasting is harder because there are envvars)
-python3 task.py tune protox agent hpo job --workload-name-suffix demo --num-samples 2 --max-concurrent 2 --workload-timeout 15 --query-timeout 2 --tune-duration-during-hpo 0.03  --intended-dbdata-hardware $INTENDED_DBDATA_HARDWARE --dbdata-parent-dpath $DBDATA_PARENT_DPATH --build-space-good-for-boot
 python3 task.py tune protox agent tune job --workload-name-suffix demo
-python3 task.py tune protox agent replay job --workload-name-suffix demo
 exit 0
 
 # benchmark

tune/protox/env/util/execute.py

@@ -24,6 +24,7 @@
 def _acquire_metrics_around_query(
     pg_conn: PostgresConn,
     query: str,
+    query_knobs: list[str],
     query_timeout: float = 0.0,
     observation_space: Optional[StateSpace] = None,
 ) -> tuple[float, bool, Optional[dict[str, Any]], Any]:
@@ -32,7 +33,7 @@ def _acquire_metrics_around_query(
         initial_metrics = observation_space.construct_online(pg_conn.conn())
 
     qid_runtime, did_time_out, explain_data = pg_conn.time_query(
-        query, add_explain=True, timeout=query_timeout
+        query, query_knobs=query_knobs, add_explain=True, timeout=query_timeout
     )
 
     if observation_space and observation_space.require_metrics():
@@ -61,31 +62,24 @@ def execute_variations(
     best_qr = BestQueryRun(None, None, True, None, None)
 
     for qr in runs:
-        # Attach the specific per-query knobs.
-        pqk_query = (
-            "/*+ "
-            + " ".join(
-                [
-                    knob.resolve_per_query_knob(
-                        value,
-                        all_knobs=sysknobs if sysknobs else KnobSpaceContainer({}),
-                    )
-                    for knob, value in qr.qknobs.items()
-                ]
+        # Get the per-query knobs for this query in the form list[str].
+        query_knobs = [
+            knob.resolve_per_query_knob(
+                value,
+                all_knobs=sysknobs if sysknobs else KnobSpaceContainer({}),
             )
-            + " */"
-            + query
-        )
+            for knob, value in qr.qknobs.items()
+        ]
 
         # Log out the knobs that we are using.
-        pqkk = [(knob.name(), val) for knob, val in qr.qknobs.items()]
         logging.getLogger(DBGYM_LOGGER_NAME).debug(
-            f"{qr.prefix_qid} executing with {pqkk}"
+            f"{qr.prefix_qid} executing with {query_knobs}"
         )
 
         runtime, did_time_out, explain_data, metric = _acquire_metrics_around_query(
             pg_conn=pg_conn,
-            query=pqk_query,
+            query=query,
+            query_knobs=query_knobs,
             query_timeout=timeout_limit,
             observation_space=observation_space,
         )