Add SUT metrics for Databricks SQL Warehouse: disk I/O, num_active_sessions, compute nodes count (#180)

* Databricks: resolve UTF8 opaque types to fix schema mismatch

* Add Databricks `num_active_sessions` metric

* Enable scrape_sut_metrics

* Add num_compute_nodes

* Update

* Update

* disk_read_bytes and disk_write_bytes metrics

* Update

---------

Co-authored-by: William <98815791+peasee@users.noreply.github.com>

Author: sgrebnov

.github/workflows/run_spicebench.yml

@@ -53,6 +53,11 @@ on:
         required: false
         default: false
         type: boolean
+      scrape_sut_metrics:
+        description: 'Enable periodic SUT metrics collection via the system adapter (maps to spicebench --scrape-sut-metrics)'
+        required: false
+        default: true
+        type: boolean
       scheduler_state_location:
         description: 'S3 URI for shared scheduler state (e.g. s3://bucket/scheduler-state/). If empty, scheduler state is not configured.'
         required: false
@@ -248,6 +253,7 @@ jobs:
           SCHEDULER_STATE_LOCATION: ${{ github.event.inputs.scheduler_state_location || '' }}
           VALIDATE_CHECKPOINT_RESULTS: ${{ github.event.inputs.validate_checkpoint_results || 'false' }}
           ENABLE_MODULE_DEBUG_LOGGING: ${{ github.event.inputs.enable_module_debug_logging || 'false' }}
+          SCRAPE_SUT_METRICS: ${{ github.event.inputs.scrape_sut_metrics || 'true' }}
           SPICEAI_BENCHMARK_METRICS_KEY: ${{ secrets.SPICEAI_BENCHMARK_METRICS_KEY }}
           MINIO_ENDPOINT: ${{ secrets.MINIO_ENDPOINT }}
           AWS_ACCESS_KEY_ID: ${{ secrets.MINIO_ACCESS_KEY_ID }}
@@ -300,6 +306,11 @@ jobs:
             SCHEDULER_ARGS="--scheduler-state-location ${SCHEDULER_STATE_LOCATION}"
           fi
 
+          SUT_METRICS_ARGS=""
+          if [ "${SCRAPE_SUT_METRICS}" = "true" ]; then
+            SUT_METRICS_ARGS="--scrape-sut-metrics"
+          fi
+
           if [ "${SYSTEM_UNDER_TEST_PREFIX}" = "databricks" ]; then
             ADAPTER_CMD="${HOME}/.spice/bin/databricks-system-adapter"
             ADAPTER_ARGS="stdio"
@@ -335,11 +346,11 @@ jobs:
             --concurrency "${NUM_QUERY_CLIENTS}"  \
             --scenario "${SCENARIO}" \
             --executor-instance-type "${EXECUTOR_INSTANCE_TYPE}" \
-            --scrape-sut-metrics \
             ${ETL_ARGS} \
             ${ETL_SINK_ARGS} \
             ${VALIDATION_ARGS} \
             ${SCHEDULER_ARGS} \
+            ${SUT_METRICS_ARGS} \
             --system-adapter-stdio-cmd "${ADAPTER_CMD}" \
             --system-adapter-stdio-args "${ADAPTER_ARGS}" \
             ${ADAPTER_ENVS} \

crates/system-adapter-protocol/src/client.rs

@@ -212,8 +212,15 @@ impl Client {
     }
 
     /// Collect current metrics from the system under test
-    pub async fn metrics(&mut self, run_id: uuid::Uuid) -> Result<crate::MetricsResponse> {
-        let request = crate::MetricsRequest { run_id };
+    pub async fn metrics(
+        &mut self,
+        run_id: uuid::Uuid,
+        final_scrape: bool,
+    ) -> Result<crate::MetricsResponse> {
+        let request = crate::MetricsRequest {
+            run_id,
+            final_scrape,
+        };
         let rpc_request = JsonRpcRequest::new(1, crate::methods::METRICS, request);
         let response = self.call_typed(rpc_request).await?;
         response

crates/system-adapter-protocol/src/lib.rs

@@ -217,6 +217,11 @@ pub struct TeardownResponse {
 pub struct MetricsRequest {
     /// Unique identifier for the benchmark run
     pub run_id: Uuid,
+    /// When true, this is the final metrics collection after the benchmark run
+    /// has finished. Adapters may perform heavier queries (e.g. Query History)
+    /// that would be too slow for periodic scraping.
+    #[serde(default)]
+    pub final_scrape: bool,
 }
 
 /// Resource utilization snapshot from the system under test
@@ -240,6 +245,9 @@ pub struct ResourceMetrics {
     /// Cumulative disk write operations
     #[serde(skip_serializing_if = "Option::is_none")]
     pub disk_write_iops: Option<u64>,
+    /// Number of active compute nodes / clusters backing the SUT
+    #[serde(skip_serializing_if = "Option::is_none")]
+    pub num_compute_nodes: Option<u64>,
 }
 
 /// Ingestion progress snapshot from the system under test

crates/system-adapter-protocol/src/server.rs

@@ -86,9 +86,15 @@ pub trait Handler: Send + Sync {
     ///
     /// Called periodically by spicebench when `--scrape-sut-metrics` is enabled.
     /// Returns a snapshot of resource utilization and ingestion progress.
+    /// When `final_scrape` is true, the benchmark run has finished and the
+    /// adapter may perform heavier queries (e.g. Query History aggregation).
     /// Default implementation returns empty metrics.
-    async fn metrics(&mut self, run_id: Uuid) -> std::result::Result<MetricsResponse, String> {
-        let _ = run_id;
+    async fn metrics(
+        &mut self,
+        run_id: Uuid,
+        final_scrape: bool,
+    ) -> std::result::Result<MetricsResponse, String> {
+        let _ = (run_id, final_scrape);
         Ok(MetricsResponse::default())
     }
 
@@ -263,7 +269,7 @@ impl<H: Handler> Server<H> {
             Ok(r) => r,
             Err(e) => return e,
         };
-        Self::handler_response(self.handler.metrics(req.run_id).await, id)
+        Self::handler_response(self.handler.metrics(req.run_id, req.final_scrape).await, id)
     }
 
     async fn handle_rpc_methods(&mut self, id: serde_json::Value) -> serde_json::Value {
@@ -302,7 +308,11 @@ mod tests {
             Ok(TeardownResponse { ok: true })
         }
 
-        async fn metrics(&mut self, _run_id: Uuid) -> std::result::Result<MetricsResponse, String> {
+        async fn metrics(
+            &mut self,
+            _run_id: Uuid,
+            _final_scrape: bool,
+        ) -> std::result::Result<MetricsResponse, String> {
             Ok(MetricsResponse::default())
         }
     }

src/commands/load/mod.rs

@@ -153,6 +153,9 @@ fn record_sut_metrics(
     if let Some(v) = response.ingestion.active_connections {
         crate::metrics::ACTIVE_CONNECTIONS.record(v, attributes);
     }
+    if let Some(v) = response.resource.num_compute_nodes {
+        crate::metrics::NUM_COMPUTE_NODES.record(v, attributes);
+    }
 }
 
 /// Spawn a task that periodically scrapes SUT metrics from the system adapter.
@@ -180,7 +183,7 @@ fn spawn_sut_metrics_scraper(
         loop {
             tokio::select! {
                 _ = ticker.tick() => {
-                    let metrics_result = adapter.lock().await.metrics(run_id).await;
+                    let metrics_result = adapter.lock().await.metrics(run_id, false).await;
                     match metrics_result {
                         Ok(resp) => {
                             log_sut_metrics_snapshot(&resp);
@@ -205,7 +208,7 @@ fn spawn_sut_metrics_scraper(
                 }
                 () = token.cancelled() => {
                     // Final scrape before exiting
-                    if let Ok(resp) = adapter.lock().await.metrics(run_id).await {
+                    if let Ok(resp) = adapter.lock().await.metrics(run_id, true).await {
                         log_sut_metrics_snapshot(&resp);
                         record_sut_metrics(
                             &resp,

src/metrics.rs

@@ -141,6 +141,14 @@ pub static ACTIVE_CONNECTIONS: LazyLock<Gauge<u64>> = LazyLock::new(|| {
         .build()
 });
 
+pub static NUM_COMPUTE_NODES: LazyLock<Gauge<u64>> = LazyLock::new(|| {
+    meter()
+        .u64_gauge("num_compute_nodes")
+        .with_description("Number of active compute nodes / clusters backing the SUT.")
+        .with_unit("nodes")
+        .build()
+});
+
 // --- Efficiency ---
 
 pub static EFFICIENCY_QUERIES_PER_CORE: LazyLock<Gauge<f64>> = LazyLock::new(|| {