fixed some issues from review

Author: milindsrivastava1997

asap-tools/experiments/experiment_only_ingest_path.py

@@ -289,16 +289,7 @@ def main(cfg: DictConfig):
             f"http://localhost:{prometheus_service.get_query_endpoint_port()}"
         )
 
-        print("Waiting for Prometheus to become ready...")
-        prometheus_ready_timeout = 60
-        prometheus_ready_start = time.time()
-        while not prometheus_service.is_healthy():
-            if time.time() - prometheus_ready_start > prometheus_ready_timeout:
-                raise RuntimeError(
-                    f"Prometheus did not become ready within {prometheus_ready_timeout}s"
-                )
-            time.sleep(2)
-        print("Prometheus is ready.")
+        prometheus_service.wait_until_ready()
 
         label_discovery_wait = prometheus_scrape_interval * 2
         print(

asap-tools/experiments/experiment_run_e2e.py

@@ -342,17 +342,7 @@ def main(cfg: DictConfig):
             # Poll until Prometheus is actually accepting connections before sleeping
             # for scrape data. Prometheus takes a few seconds to bind the port after
             # its process starts, so a fixed sleep alone can race.
-            print("Waiting for Prometheus to become ready...")
-            prometheus_ready_timeout = 60
-            prometheus_ready_start = time.time()
-            while not prometheus_service.is_healthy():
-                if time.time() - prometheus_ready_start > prometheus_ready_timeout:
-                    raise RuntimeError(
-                        "Prometheus did not become ready within "
-                        f"{prometheus_ready_timeout}s"
-                    )
-                time.sleep(2)
-            print("Prometheus is ready.")
+            prometheus_service.wait_until_ready()
 
             # Wait for two scrape intervals so Prometheus has series to return.
             label_discovery_wait = prometheus_scrape_interval * 2
@@ -484,9 +474,7 @@ def main(cfg: DictConfig):
                         pipeline_id=arroyosketch_pipeline_id,
                         experiment_output_dir=experiment_output_dir,
                     )
-            elif args.streaming_engine == "precompute":
-                pass  # precompute engine handles sketch computation internally; no external pipeline needed
-            else:
+            elif args.streaming_engine not in ("precompute",):
                 raise ValueError(
                     "Invalid streaming engine: {}. Supported engines are 'flink', 'arroyo', and 'precompute'".format(
                         args.streaming_engine

asap-tools/experiments/experiment_run_grafana_demo.py

@@ -350,16 +350,7 @@ def main(cfg: DictConfig):
             f"http://localhost:{prometheus_service.get_query_endpoint_port()}"
         )
 
-        print("Waiting for Prometheus to become ready...")
-        prometheus_ready_timeout = 60
-        prometheus_ready_start = time.time()
-        while not prometheus_service.is_healthy():
-            if time.time() - prometheus_ready_start > prometheus_ready_timeout:
-                raise RuntimeError(
-                    f"Prometheus did not become ready within {prometheus_ready_timeout}s"
-                )
-            time.sleep(2)
-        print("Prometheus is ready.")
+        prometheus_service.wait_until_ready()
 
         label_discovery_wait = prometheus_scrape_interval * 2
         print(
@@ -448,9 +439,7 @@ def main(cfg: DictConfig):
                     pipeline_id=arroyosketch_pipeline_id,
                     experiment_output_dir=experiment_output_dir,
                 )
-        elif args.streaming_engine == "precompute":
-            pass  # precompute engine handles sketch computation internally; no external pipeline needed
-        else:
+        elif args.streaming_engine not in ("precompute",):
             raise ValueError(
                 "Invalid streaming engine: {}. Supported engines are 'flink', 'arroyo', and 'precompute'".format(
                     args.streaming_engine

asap-tools/experiments/experiment_utils/services/base.py

@@ -57,6 +57,18 @@ def is_healthy(self) -> bool:
         """
         return True
 
+    def wait_until_ready(self, timeout: int = 60) -> None:
+        """Block until is_healthy() returns True, or raise RuntimeError on timeout."""
+        print(f"Waiting for {self.__class__.__name__} to become ready...")
+        start = time.time()
+        while not self.is_healthy():
+            if time.time() - start > timeout:
+                raise RuntimeError(
+                    f"{self.__class__.__name__} did not become ready within {timeout}s"
+                )
+            time.sleep(2)
+        print(f"{self.__class__.__name__} is ready.")
+
     def restart(self, **kwargs) -> None:
         """
         Restart the service. Default implementation stops then starts.

asap-tools/experiments/experiment_utils/services/query_engine.py

@@ -8,6 +8,7 @@
 import yaml
 
 import constants
+import utils
 from .base import BaseService
 from experiment_utils.providers.base import InfrastructureProvider
 
@@ -126,8 +127,9 @@ def _build_engine_config(
                 "forward_unsupported_queries": forward_unsupported_queries,
             }
 
-        # Ingest config depends on the streaming engine
-        if streaming_engine == "arroyo":
+        # Ingest config depends on the streaming engine.
+        # Both flink and arroyo produce to the same Kafka topic.
+        if streaming_engine in ("arroyo", "flink"):
             ingest: dict = {
                 "type": "kafka",
                 "broker": kafka_broker,
@@ -143,7 +145,7 @@ def _build_engine_config(
         else:
             raise ValueError(
                 f"streaming_engine='{streaming_engine}' is not supported by the Rust query engine. "
-                "Use 'arroyo' or 'precompute'."
+                "Use 'flink', 'arroyo', or 'precompute'."
             )
 
         return {
@@ -179,8 +181,6 @@ def _write_engine_config_to_remote(
             local_path: Local path to write the YAML file to
             remote_path: Absolute path on the remote node where the file should land
         """
-        import utils  # top-level module in the experiments/ tree
-
         os.makedirs(os.path.dirname(local_path), exist_ok=True)
         config_yaml = yaml.dump(
             config_dict, default_flow_style=False, allow_unicode=True