feat: Add custom_image input to debug Spice Cloud workflow (#238)

phillipleblanc · Jeadie · commit f3960291f840 · 2026-04-02T11:23:13.000+10:00
* feat: Add custom_image input to debug Spice Cloud workflow

Add a custom_image workflow input to run_spicebench_debug_spice_cloud
that allows specifying a custom runtime container image (e.g.
ghcr.io/spiceai/spiceai-dev:spicebench-sf10) instead of the default
nightly image.

When set, the image reference is parsed into registry, image name, and
tag components and passed through to spidapter as SPIDAPTER_IMAGE_REGISTRY,
SPIDAPTER_IMAGE_NAME, and SPIDAPTER_IMAGE_TAG env vars. The channel is
automatically switched to internal.

Also adds executor_memory_limit input and fixes NUM_QUERY_CLIENTS to
match the main workflow (2 instead of 8).

* fix: Run row count validation first in checkpoint validation

Move table row count validation to run as the first phase (Phase 0)
before the probe query. Row count queries are cheap SELECT COUNT(*)
and immediately surface data loss or duplication without waiting for
expensive analytical queries to converge.

* chore: Add per-batch operation row count logging for data reconciliation

Log insert/update/delete row counts for each batch written through
write_segments_for_batch. This covers both the initialization phase
and the main ETL run pipeline.

Example log output:
  INFO etl: Writing segments for batch table=customer batch_id=5
    segments=3 insert_rows=8192 update_rows=512 delete_rows=128

This allows post-hoc reconciliation: summing insert_rows - delete_rows
per table should match the expected row count at each checkpoint.
If there is a mismatch, the per-batch logs pinpoint which batch_id
had unexpected operation counts.
diff --git a/.github/workflows/run_spicebench_debug_spice_cloud.yml b/.github/workflows/run_spicebench_debug_spice_cloud.yml
@@ -46,6 +46,11 @@ on:
         required: false
         default: 'latest'
         type: string
+      num_query_clients:
+        description: 'Number of concurrent query clients'
+        required: false
+        default: '2'
+        type: string
       enable_module_debug_logging:
         description: 'Enable debug logs'
         required: false
@@ -61,11 +66,26 @@ on:
         required: false
         default: false
         type: boolean
+      executor_replicas:
+        description: 'Number of executor replicas'
+        required: false
+        default: '4'
+        type: string
       app_memory_limit:
         description: 'Memory limit for the scheduler (app) pod (e.g. 16Gi, 20Gi)'
         required: false
         default: '16Gi'
         type: string
+      executor_memory_limit:
+        description: 'Memory limit for the executor pod (e.g. 16Gi, 62Gi)'
+        required: false
+        default: '16Gi'
+        type: string
+      custom_image:
+        description: 'Custom runtime container image (e.g. ghcr.io/spiceai/spiceai-dev:spicebench-sf10). Overrides the default channel image. Requires the internal update channel on the target SCP environment.'
+        required: false
+        default: ''
+        type: string
 jobs:
   run-spicebench:
     name: Run spicebench
@@ -134,7 +154,7 @@ jobs:
           SCENARIO: ${{ github.event.inputs.scenario || 'tpch' }}
           SYSTEM_UNDER_TEST: ${{ github.event.inputs.system_under_test || 'spice_cloud' }}
           SYSTEM_ADAPTER: ${{ github.event.inputs.system_under_test || 'spice_cloud' }}
-          NUM_QUERY_CLIENTS: '8'
+          NUM_QUERY_CLIENTS: ${{ github.event.inputs.num_query_clients || '2' }}
           ETL_BUCKET: 'spicebench'
           ETL_PREFIX: ${{ github.event.inputs.etl_type == 'changes' && 'data-gen-mutable' || 'data-gen' }}
           SCALE_FACTOR: ${{ github.event.inputs.scale_factor || '1' }}
@@ -154,7 +174,11 @@ jobs:
           SPIDAPTER_ICEBERG_CATALOG_FROM: iceberg:https://glue.us-west-1.amazonaws.com/iceberg/v1/catalogs/211125479522/namespaces
           DISABLE_TEARDOWN: ${{ github.event.inputs.disable_teardown || 'false' }}
           ENABLE_PVC: ${{ github.event.inputs.enable_pvc || 'false' }}
+          EXECUTOR_REPLICAS: ${{ github.event.inputs.executor_replicas || '4' }}
           SPIDAPTER_APP_MEMORY_LIMIT: ${{ github.event.inputs.app_memory_limit || '16Gi' }}
+          SPIDAPTER_EXECUTOR_MEMORY_LIMIT: ${{ github.event.inputs.executor_memory_limit || '16Gi' }}
+          SPIDAPTER_ORGANIZATION_TAG: 'spicehq'
+          CUSTOM_IMAGE: ${{ github.event.inputs.custom_image || '' }}
         run: |
           set -euo pipefail
           if [ "${ENABLE_MODULE_DEBUG_LOGGING}" = "true" ]; then
@@ -198,16 +222,47 @@ jobs:
           export SPICEBENCH_ADBC_FLUSH_STREAM_BEFORE_UPSERT=true
           export SPICEBENCH_ADBC_DELETE_BATCH_SIZE=50000
           ADAPTER_CMD="docker"
-          ADAPTER_DOCKER_OPTS="run -i -e SPIDAPTER_EXECUTOR_REPLICAS=4 -e SPICEAI_API_KEY -e SPICE_CLOUD_API_URL -e AWS_ACCESS_KEY_ID=${S3_AWS_ACCESS_KEY_ID} -e AWS_SECRET_ACCESS_KEY=${S3_AWS_SECRET_ACCESS_KEY} -e SPIDAPTER_ICEBERG_REGION -e SPIDAPTER_ICEBERG_CATALOG_FROM -e SCHEDULER_STATE_LOCATION"
+          ADAPTER_DOCKER_OPTS="run -i -e SPIDAPTER_EXECUTOR_REPLICAS=${EXECUTOR_REPLICAS} -e SPICEAI_API_KEY -e SPICE_CLOUD_API_URL -e AWS_ACCESS_KEY_ID=${S3_AWS_ACCESS_KEY_ID} -e AWS_SECRET_ACCESS_KEY=${S3_AWS_SECRET_ACCESS_KEY} -e SPIDAPTER_ICEBERG_REGION -e SPIDAPTER_ICEBERG_CATALOG_FROM -e SCHEDULER_STATE_LOCATION"
 
           ADAPTER_DOCKER_OPTS="${ADAPTER_DOCKER_OPTS} -e SPIDAPTER_APP_MEMORY_LIMIT=${SPIDAPTER_APP_MEMORY_LIMIT}"
+          ADAPTER_DOCKER_OPTS="${ADAPTER_DOCKER_OPTS} -e SPIDAPTER_EXECUTOR_MEMORY_LIMIT=${SPIDAPTER_EXECUTOR_MEMORY_LIMIT}"
+          ADAPTER_DOCKER_OPTS="${ADAPTER_DOCKER_OPTS} -e SPIDAPTER_ORGANIZATION_TAG"
+
+          # Parse custom image into registry, name, and tag for spidapter.
+          # Expected format: registry/image:tag (e.g. ghcr.io/spiceai/spiceai-dev:spicebench-sf10)
+          # Parses as: registry=ghcr.io, image=spiceai/spiceai-dev, tag=spicebench-sf10
+          if [ -n "${CUSTOM_IMAGE}" ]; then
+            IMAGE_WITH_TAG="${CUSTOM_IMAGE}"
+            if [[ "${IMAGE_WITH_TAG}" == *":"* ]]; then
+              SPIDAPTER_IMAGE_TAG="${IMAGE_WITH_TAG##*:}"
+              IMAGE_WITHOUT_TAG="${IMAGE_WITH_TAG%:*}"
+            else
+              SPIDAPTER_IMAGE_TAG=""
+              IMAGE_WITHOUT_TAG="${IMAGE_WITH_TAG}"
+            fi
+            # Registry is the first path component (hostname), image is the rest
+            SPIDAPTER_IMAGE_REGISTRY="${IMAGE_WITHOUT_TAG%%/*}"
+            SPIDAPTER_IMAGE_NAME="${IMAGE_WITHOUT_TAG#*/}"
+
+            echo "Custom image: registry=${SPIDAPTER_IMAGE_REGISTRY}, name=${SPIDAPTER_IMAGE_NAME}, tag=${SPIDAPTER_IMAGE_TAG}"
+            ADAPTER_DOCKER_OPTS="${ADAPTER_DOCKER_OPTS} -e SPIDAPTER_IMAGE_REGISTRY=${SPIDAPTER_IMAGE_REGISTRY} -e SPIDAPTER_IMAGE_NAME=${SPIDAPTER_IMAGE_NAME}"
+            if [ -n "${SPIDAPTER_IMAGE_TAG}" ]; then
+              ADAPTER_DOCKER_OPTS="${ADAPTER_DOCKER_OPTS} -e SPIDAPTER_IMAGE_TAG=${SPIDAPTER_IMAGE_TAG}"
+            fi
+          fi
 
           if [ "${ENABLE_PVC}" = "true" ]; then
             echo "PVC enabled: app=3GB, executor=2GB"
             ADAPTER_DOCKER_OPTS="${ADAPTER_DOCKER_OPTS} -e SPIDAPTER_APP_STORAGE_SIZE_GB=3 -e SPIDAPTER_EXECUTOR_STORAGE_SIZE_GB=2 -e SPIDAPTER_CAYENNE_DATA_DIR=/data/data -e SPIDAPTER_CAYENNE_METADATA_DIR=/data/metadata"
           fi
 
-          ADAPTER_ARGS="${ADAPTER_DOCKER_OPTS} ghcr.io/spiceai/spidapter:${{ github.event.inputs.spidapter_version || 'latest' }} stdio --verbose --channel nightly"
+          # Use internal channel when a custom image is specified, otherwise nightly.
+          SPIDAPTER_CHANNEL="nightly"
+          if [ -n "${CUSTOM_IMAGE}" ]; then
+            SPIDAPTER_CHANNEL="internal"
+          fi
+
+          ADAPTER_ARGS="${ADAPTER_DOCKER_OPTS} ghcr.io/spiceai/spidapter:${{ github.event.inputs.spidapter_version || 'latest' }} stdio --verbose --channel ${SPIDAPTER_CHANNEL}"
           ADAPTER_ENVS=""
 
           NO_TEARDOWN_ARG=""
diff --git a/crates/data-generation/src/dataset/mod.rs b/crates/data-generation/src/dataset/mod.rs
@@ -31,6 +31,31 @@ use crate::dataset::simple_sequence::SimpleSequenceDataset;
 use crate::dataset::tpch::TpchDataset;
 use crate::storage::DataStorage;
 
+/// ETL mode derived from mutation ratios.
+#[derive(Debug, Clone, Copy, PartialEq, Eq)]
+pub enum EtlType {
+    /// Append-only data (no updates/deletes).
+    Events,
+    /// Change-data mode (updates and/or deletes present).
+    Changes,
+}
+
+impl EtlType {
+    #[must_use]
+    pub fn as_str(self) -> &'static str {
+        match self {
+            Self::Events => "events",
+            Self::Changes => "changes",
+        }
+    }
+}
+
+impl std::fmt::Display for EtlType {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        f.write_str(self.as_str())
+    }
+}
+
 /// Metadata about a table in a dataset.
 #[derive(Debug, Clone)]
 pub struct DatasetTable {
@@ -156,6 +181,15 @@ impl MutationConfig {
             delete_ratio,
         }
     }
+
+    #[must_use]
+    pub fn etl_type(&self) -> EtlType {
+        if self.update_ratio == 0.0 && self.delete_ratio == 0.0 {
+            EtlType::Events
+        } else {
+            EtlType::Changes
+        }
+    }
 }
 
 #[async_trait]
diff --git a/crates/data-generation/src/version.rs b/crates/data-generation/src/version.rs
@@ -27,7 +27,7 @@ use serde::{Deserialize, Serialize};
 use serde_json::Value;
 
 use crate::config::DatasetConfig;
-use crate::dataset::MutationConfig;
+use crate::dataset::{EtlType, MutationConfig};
 
 /// Converts an Arrow [`SchemaRef`] to a JSON-compatible representation
 /// using Arrow's built-in IPC JSON serialization (the "Schema" portion
@@ -105,16 +105,9 @@ impl VersionMetadata {
     }
 
     /// Returns the ETL type based on the mutation configuration.
-    ///
-    /// - `"events"` — append-only data (no updates or deletes).
-    /// - `"changes"` — data with mutations (updates and/or deletes).
     #[must_use]
-    pub fn etl_type(&self) -> &'static str {
-        if self.mutations.update_ratio == 0.0 && self.mutations.delete_ratio == 0.0 {
-            "events"
-        } else {
-            "changes"
-        }
+    pub fn etl_type(&self) -> EtlType {
+        self.mutation_config().etl_type()
     }
 }
 
diff --git a/crates/etl/src/lib.rs b/crates/etl/src/lib.rs
@@ -666,13 +666,11 @@ fn strip_internal_columns(batch: &RecordBatch) -> anyhow::Result<RecordBatch> {
 fn batch_is_insert_only(batch: &RecordBatch, mutations: &MutationConfig) -> anyhow::Result<bool> {
     // Naive check: if the dataset was generated with non-zero update or delete
     // ratios, batches may contain `_op` values other than "c".
-    if mutations.update_ratio == 0.0 && mutations.delete_ratio == 0.0 {
+    if mutations.etl_type() != data_generation::dataset::EtlType::Changes {
         return Ok(true);
     }
 
-    let schema = batch.schema();
-
-    let op_idx = match schema.index_of("_op") {
+    let op_idx = match batch.schema().index_of("_op") {
         Ok(idx) => idx,
         Err(_) => return Ok(true),
     };
@@ -1197,7 +1195,12 @@ impl ETLPipeline {
             .map(|(name, table)| {
                 let schema =
                     schema_with_created_at(&schema_without_internal_columns(&table.schema));
-                let primary_key_columns = dataset.primary_key(&name);
+                let primary_key_columns =
+                    if mutations.etl_type() == data_generation::dataset::EtlType::Changes {
+                        dataset.primary_key(&name)
+                    } else {
+                        Vec::new()
+                    };
                 let config = ProtocolDatasetConfig {
                     schema,
                     primary_key_columns,
@@ -2009,6 +2012,9 @@ mod tests {
     use arrow::array::{RecordBatch, StringViewArray};
     use arrow::datatypes::{DataType, Field, Schema};
     use arrow::ipc::writer::StreamWriter;
+    use data_generation::config::DatasetConfig as GenerationDatasetConfig;
+    use data_generation::dataset::MutationConfig;
+    use data_generation::storage::file::FileStorage;
     use std::sync::Arc;
 
     /// Measures the IPC-serialized size of a [`RecordBatch`].
@@ -2090,4 +2096,63 @@ mod tests {
             );
         }
     }
+
+    #[test]
+    fn create_tables_request_datasets_omits_primary_keys_for_events() {
+        let config = GenerationDatasetConfig {
+            dataset_type: "simple_sequence".to_string(),
+            scale_factor: 0.01,
+            num_steps: 1,
+        };
+        let tempdir = tempfile::tempdir().expect("tempdir");
+        let storage = Arc::new(FileStorage::new(tempdir.path().to_path_buf()));
+        let mutations = MutationConfig::new(0.0, 0.0);
+
+        let datasets = ETLPipeline::create_tables_request_datasets(
+            DatasetSource::SimpleSequence,
+            &config,
+            storage,
+            &mutations,
+            None,
+        )
+        .expect("create tables datasets");
+
+        let table_cfg = datasets
+            .get("integer_sequence")
+            .expect("integer_sequence config");
+        assert!(
+            table_cfg.primary_key_columns.is_empty(),
+            "events ETL should not provide primary keys"
+        );
+    }
+
+    #[test]
+    fn create_tables_request_datasets_includes_primary_keys_for_changes() {
+        let config = GenerationDatasetConfig {
+            dataset_type: "simple_sequence".to_string(),
+            scale_factor: 0.01,
+            num_steps: 1,
+        };
+        let tempdir = tempfile::tempdir().expect("tempdir");
+        let storage = Arc::new(FileStorage::new(tempdir.path().to_path_buf()));
+        let mutations = MutationConfig::new(0.1, 0.0);
+
+        let datasets = ETLPipeline::create_tables_request_datasets(
+            DatasetSource::SimpleSequence,
+            &config,
+            storage,
+            &mutations,
+            None,
+        )
+        .expect("create tables datasets");
+
+        let table_cfg = datasets
+            .get("integer_sequence")
+            .expect("integer_sequence config");
+        assert_eq!(
+            table_cfg.primary_key_columns,
+            vec!["id".to_string()],
+            "changes ETL should provide primary keys"
+        );
+    }
 }
diff --git a/src/commands/load/mod.rs b/src/commands/load/mod.rs
@@ -653,8 +653,48 @@ async fn run_checkpoint_validation(
         return CheckpointValidationResult::TimedOut;
     };
 
+    // Phase 0: validate table row counts first as a fast correctness probe.
+    // Row count queries are cheap and immediately surface data loss/duplication
+    // without waiting for expensive analytical queries to converge.
+    println!("Checkpoint {checkpoint_idx}: validating table row counts before probing queries",);
+    {
+        let mut row_count_ticker = tokio::time::interval(probe_period);
+        let mut row_count_attempt = 0u64;
+        loop {
+            if tokio::time::Instant::now() >= deadline {
+                println!(
+                    "Checkpoint {checkpoint_idx}: row count validation timed out after {} attempts",
+                    row_count_attempt
+                );
+                return CheckpointValidationResult::TimedOut;
+            }
+
+            tokio::select! {
+                biased;
+                _ = signal::ctrl_c() => return CheckpointValidationResult::Interrupted,
+                _ = row_count_ticker.tick() => {}
+            };
+
+            row_count_attempt += 1;
+            if validate_checkpoint_table_row_counts(
+                executor,
+                expected_row_counts,
+                checkpoint_idx,
+                query_catalog_namespace,
+            )
+            .await
+            {
+                break;
+            }
+
+            println!(
+                "Checkpoint {checkpoint_idx}: row count validation attempt {row_count_attempt} failed, retrying",
+            );
+        }
+    }
+
     println!(
-        "Checkpoint {checkpoint_idx}: probing '{}' every {}s",
+        "Checkpoint {checkpoint_idx}: row counts passed, probing '{}' every {}s",
         probe_query.name,
         probe_period.as_secs()
     );
@@ -663,7 +703,7 @@ async fn run_checkpoint_validation(
     let mut probe_count = 0u64;
 
     loop {
-        // Phase 1: probe Q1 until it passes
+        // Phase 1: probe query until it passes
         let e2e_send_time = match probe_until_pass(
             executor,
             probe_query,
@@ -677,7 +717,7 @@ async fn run_checkpoint_validation(
         {
             ProbeOutcome::Passed(send_time) => {
                 println!(
-                    "Checkpoint {checkpoint_idx}: probe query '{}' passed, validating table row counts",
+                    "Checkpoint {checkpoint_idx}: probe query '{}' passed, validating full query set",
                     probe_query.name
                 );
                 send_time
@@ -691,17 +731,6 @@ async fn run_checkpoint_validation(
             return CheckpointValidationResult::TimedOut;
         }
 
-        if !validate_checkpoint_table_row_counts(
-            executor,
-            expected_row_counts,
-            checkpoint_idx,
-            query_catalog_namespace,
-        )
-        .await
-        {
-            continue;
-        }
-
         if validate_full_query_set(
             executor,
             queries,
