fix: Add num steps input for datagen (#108)

Author: peasee

.github/workflows/data_generation_run.yml

@@ -34,6 +34,14 @@ on:
         required: false
         default: false
         type: boolean
+      num_steps:
+        required: false
+        default: '25'
+        type: string
+      checkpoint_interval_steps:
+        required: false
+        default: '100'
+        type: string
       runner_type:
         required: false
         default: 'ubuntu-latest'
@@ -88,6 +96,16 @@ on:
         required: false
         default: false
         type: boolean
+      num_steps:
+        description: 'Number of data generation steps (partitions for TPC-H dbgen)'
+        required: false
+        default: '5'
+        type: string
+      checkpoint_interval_steps:
+        description: 'Every N steps to take a checkpoint'
+        required: false
+        default: '3'
+        type: string
 
 jobs:
   run-data-generation:
@@ -130,6 +148,7 @@ jobs:
           MAX_CONCURRENCY: ${{ inputs.max_concurrency || github.event.inputs.max_concurrency || '8' }}
           REGION: ${{ inputs.region || github.event.inputs.region || 'us-east-1' }}
           SKIP_INITIAL: ${{ inputs.skip_initial || github.event.inputs.skip_initial || false }}
+          NUM_STEPS: ${{ inputs.num_steps || github.event.inputs.num_steps || '25' }}
           AWS_ACCESS_KEY_ID: ${{ secrets.AWS_ACCESS_KEY_ID }}
           AWS_SECRET_ACCESS_KEY: ${{ secrets.AWS_SECRET_ACCESS_KEY }}
           RUST_LOG: info
@@ -142,6 +161,7 @@ jobs:
           ARGS="${ARGS} --prefix ${PREFIX}"
           ARGS="${ARGS} --max-concurrency ${MAX_CONCURRENCY}"
           ARGS="${ARGS} --region ${REGION}"
+          ARGS="${ARGS} --num-steps ${NUM_STEPS}"
 
           if [ "${SKIP_INITIAL}" = "true" ]; then
             ARGS="${ARGS} --skip-initial"
@@ -173,6 +193,7 @@ jobs:
           BUCKET: ${{ inputs.bucket || github.event.inputs.bucket }}
           PREFIX: ${{ inputs.prefix || github.event.inputs.prefix || 'data-gen' }}
           REGION: ${{ inputs.region || github.event.inputs.region || 'us-east-1' }}
+          CHECKPOINT_INTERVAL_STEPS: ${{ inputs.checkpoint_interval_steps || github.event.inputs.checkpoint_interval_steps || '100' }}
           AWS_ACCESS_KEY_ID: ${{ secrets.AWS_ACCESS_KEY_ID }}
           AWS_SECRET_ACCESS_KEY: ${{ secrets.AWS_SECRET_ACCESS_KEY }}
           RUST_LOG: info
@@ -188,5 +209,7 @@ jobs:
             ARGS="${ARGS} --region ${REGION}"
           fi
 
+          ARGS="${ARGS} --checkpoint-interval-steps ${CHECKPOINT_INTERVAL_STEPS}"
+
           echo "Running: checkpointer ${ARGS}"
           ~/.spice/bin/checkpointer ${ARGS}