merge dev: resolve conflicts, keep dev changes

Author: masoud@anyscale.com

.github/workflows/ci.yml

@@ -55,6 +55,9 @@ jobs:
           
           echo "Running chunk execution tests..."
           cargo test --test chunk_execution_test -- --nocapture
+          
+          echo "Running chunk proof generation test (release mode for speed)..."
+          cargo test --test chunk_proof_test --release -- --nocapture
       
       - name: Run clippy (warnings only)
         working-directory: zkml

.github/workflows/slow-tests.yml

@@ -0,0 +1,67 @@
+name: Slow Tests
+
+# Triggers:
+# 1. Manual trigger via GitHub UI ("Run workflow" button)
+# 2. When 'run-slow-tests' label is added to a PR
+# 3. Nightly at 2am UTC (optional, commented out)
+
+on:
+  workflow_dispatch:  # Manual trigger
+  pull_request:
+    types: [labeled]
+
+# schedule:
+#   - cron: '0 2 * * *'  # Uncomment for nightly runs
+
+env:
+  CARGO_TERM_COLOR: always
+  RUST_BACKTRACE: 1
+
+jobs:
+  slow-tests:
+    name: Slow Tests (KZG Proofs)
+    runs-on: ubuntu-latest
+    # Only run if manually triggered OR if the 'run-slow-tests' label was added
+    if: |
+      github.event_name == 'workflow_dispatch' ||
+      (github.event_name == 'pull_request' && github.event.label.name == 'run-slow-tests')
+    
+    steps:
+      - name: Checkout code
+        uses: actions/checkout@v4
+      
+      - name: Install Rust toolchain
+        uses: dtolnay/rust-toolchain@nightly
+        with:
+          components: rustfmt, clippy
+      
+      - name: Cache cargo registry
+        uses: actions/cache@v4
+        with:
+          path: |
+            ~/.cargo/registry
+            ~/.cargo/git
+            zkml/target
+          key: ${{ runner.os }}-cargo-slow-${{ hashFiles('zkml/Cargo.lock') }}
+          restore-keys: |
+            ${{ runner.os }}-cargo-slow-
+            ${{ runner.os }}-cargo-
+      
+      - name: Build zkml library (release)
+        working-directory: zkml
+        run: cargo build --lib --release
+      
+      - name: Run chunk proof generation test
+        working-directory: zkml
+        run: |
+          echo "Running KZG proof generation test (this takes ~1-2 minutes)..."
+          cargo test --test chunk_proof_test --release -- --nocapture
+      
+      - name: Summary
+        run: |
+          echo "## Slow Tests Complete ✅" >> $GITHUB_STEP_SUMMARY
+          echo "- KZG chunk proof generation: PASSED" >> $GITHUB_STEP_SUMMARY
+
+  # Note: AWS/GPU tests are excluded from CI entirely.
+  # Run them manually: cd tests/aws && python gpu_test.py
+

README.md

@@ -342,9 +342,16 @@ Distributed Proving Simulation: PASS
 #### "PyTorch CUDA not available"
 - Install PyTorch with CUDA: `pip install torch --index-url https://download.pytorch.org/whl/cu118`
 
+## CI
+
+Lightweight CI runs on every PR to `main` and `dev`:
+- Builds zkml library (nightly Rust)
+- Runs `zkml/testing/` tests (~3-4 min total)
+- AWS/GPU tests excluded to save costs
+
 ## TODO: Next Steps
 
-1. **Make Merkle root public**: Add root to public values so next chunk can verify it
+1. ~~**Make Merkle root public**: Add root to public values so next chunk can verify it~~ ✅ Done
 2. **Complete proof generation**: Connect chunk execution to actual proof generation
 3. **Ray-Rust integration**: Connect Python Ray workers to Rust proof generation
 4. **GPU acceleration**: Current implementation is CPU-based. GPU acceleration for proof generation requires additional work (Halo2 GPU support or custom GPU kernels)

zkml/Cargo.toml

@@ -92,4 +92,8 @@ path = "testing/chunk_execution_test.rs"
 
 [[test]]
 name = "test_merkle_root_public"
-path = "testing/test_merkle_root_public.rs"
+path = "testing/test_merkle_root_public.rs"
+
+[[test]]
+name = "chunk_proof_test"
+path = "testing/chunk_proof_test.rs"

zkml/src/utils/proving_kzg.rs

@@ -168,6 +168,128 @@ pub fn time_circuit_kzg(circuit: ModelCircuit<Fr>) {
   println!("Verifying time: {:?}", verify_duration - proof_duration);
 }
 
+/// Result of proving a chunk
+pub struct ChunkProofResult {
+  /// The serialized proof bytes
+  pub proof: Vec<u8>,
+  /// Public values including Merkle root (if use_merkle=true)
+  pub public_vals: Vec<Fr>,
+  /// The Merkle root (last public value if use_merkle=true)
+  pub merkle_root: Option<Fr>,
+  /// Proving time in milliseconds
+  pub proving_time_ms: u128,
+  /// Verification time in milliseconds
+  pub verify_time_ms: u128,
+}
+
+/// Generate a KZG proof for a chunk of the model
+/// 
+/// # Arguments
+/// * `config_path` - Path to model config msgpack file
+/// * `input_path` - Path to input msgpack file
+/// * `chunk_start` - Starting layer index (inclusive)
+/// * `chunk_end` - Ending layer index (exclusive)
+/// * `use_merkle` - Whether to compute and include Merkle root in public values
+/// * `params_dir` - Directory for KZG params (will be created if needed)
+/// 
+/// # Returns
+/// ChunkProofResult containing proof, public values, and timing info
+pub fn prove_chunk_kzg(
+  config_path: &str,
+  input_path: &str,
+  chunk_start: usize,
+  chunk_end: usize,
+  use_merkle: bool,
+  params_dir: &str,
+) -> ChunkProofResult {
+  use crate::utils::loader::load_model_msgpack;
+  
+  let start = Instant::now();
+  
+  // Load and configure circuit for chunk execution
+  let config = load_model_msgpack(config_path, input_path);
+  let mut circuit = ModelCircuit::<Fr>::generate_from_file(config_path, input_path);
+  
+  // If using Merkle, ensure Poseidon hasher is configured
+  if use_merkle && circuit.commit_after.is_empty() && circuit.commit_before.is_empty() {
+    // Set commit_after to enable Poseidon hasher
+    if let Some(layer) = config.layers.get(chunk_end.saturating_sub(1)) {
+      if !layer.out_idxes.is_empty() {
+        circuit.commit_after = vec![layer.out_idxes.clone()];
+      }
+    }
+  }
+  
+  // Configure for chunk execution
+  circuit.set_chunk_config(chunk_start, chunk_end, use_merkle);
+  
+  let degree = circuit.k as u32;
+  let params = get_kzg_params(params_dir, degree);
+  
+  // Generate keys
+  let vk = keygen_vk(&params, &circuit).unwrap();
+  let pk = keygen_pk(&params, vk, &circuit).unwrap();
+  
+  // First run to get public values
+  let _mock = MockProver::run(degree, &circuit, vec![vec![]]).unwrap();
+  let public_vals = get_public_values();
+  
+  // Extract Merkle root (last public value if use_merkle)
+  let merkle_root = if use_merkle && !public_vals.is_empty() {
+    Some(public_vals[public_vals.len() - 1])
+  } else {
+    None
+  };
+  
+  // Generate proof
+  let rng = rand::thread_rng();
+  let prove_start = Instant::now();
+  let mut transcript = Blake2bWrite::<_, G1Affine, Challenge255<_>>::init(vec![]);
+  create_proof::<
+    KZGCommitmentScheme<Bn256>,
+    ProverSHPLONK<'_, Bn256>,
+    Challenge255<G1Affine>,
+    _,
+    Blake2bWrite<Vec<u8>, G1Affine, Challenge255<G1Affine>>,
+    ModelCircuit<Fr>,
+  >(
+    &params,
+    &pk,
+    &[circuit],
+    &[&[&public_vals]],
+    rng,
+    &mut transcript,
+  )
+  .unwrap();
+  let proof = transcript.finalize();
+  let proving_time_ms = prove_start.elapsed().as_millis();
+  
+  // Verify the proof
+  let verify_start = Instant::now();
+  let strategy = SingleStrategy::new(&params);
+  let transcript_read = Blake2bRead::<_, _, Challenge255<_>>::init(&proof[..]);
+  verify_kzg(&params, pk.get_vk(), strategy, &public_vals, transcript_read);
+  let verify_time_ms = verify_start.elapsed().as_millis();
+  
+  println!(
+    "Chunk [{}, {}): proof generated in {}ms, verified in {}ms, {} public vals{}",
+    chunk_start,
+    chunk_end,
+    proving_time_ms,
+    verify_time_ms,
+    public_vals.len(),
+    if use_merkle { ", includes Merkle root" } else { "" }
+  );
+  
+  ChunkProofResult {
+    proof,
+    public_vals,
+    merkle_root,
+    proving_time_ms,
+    verify_time_ms,
+  }
+}
+
 // Standalone verification
 pub fn verify_circuit_kzg(
   circuit: ModelCircuit<Fr>,

zkml/testing/chunk_proof_test.rs

@@ -0,0 +1,58 @@
+//! Tests for chunk proof generation
+//!
+//! These tests verify that:
+//! 1. Real KZG proofs can be generated for model chunks
+//! 2. Proofs verify correctly
+//!
+//! Note: These are slow tests (~50s each) as they generate real cryptographic proofs.
+//! Run with: cargo test --test chunk_proof_test --release -- --nocapture
+
+#[cfg(test)]
+mod tests {
+    use zkml::utils::proving_kzg::prove_chunk_kzg;
+    use std::fs;
+
+    /// Test: Generate and verify a real KZG proof for a chunk
+    /// 
+    /// This is the main test - it proves that we can generate real ZK proofs
+    /// for a subset of model layers.
+    #[test]
+    fn test_chunk_proof_generation() {
+        let config_file = "examples/mnist/model.msgpack";
+        let input_file = "examples/mnist/inp.msgpack";
+        
+        if !std::path::Path::new(config_file).exists() {
+            eprintln!("Skipping test: example files not found");
+            return;
+        }
+        
+        // Use unique params directory to avoid race conditions
+        let params_dir = "./params_kzg_chunk_test";
+        fs::create_dir_all(params_dir).ok();
+        
+        // Generate proof for first 2 layers WITHOUT Merkle
+        // This tests the core proof generation functionality
+        let result = prove_chunk_kzg(
+            config_file,
+            input_file,
+            0,  // chunk_start
+            2,  // chunk_end
+            false,  // use_merkle = false (simpler case)
+            params_dir,
+        );
+        
+        // Verify we got a valid proof
+        assert!(!result.proof.is_empty(), "Proof should not be empty");
+        assert!(!result.public_vals.is_empty(), "Public values should not be empty");
+        assert!(result.merkle_root.is_none(), "Merkle root should be None when use_merkle=false");
+        assert!(result.proving_time_ms > 0, "Should have non-zero proving time");
+        assert!(result.verify_time_ms > 0, "Should have non-zero verify time");
+        
+        println!("✓ Chunk proof generated and verified successfully");
+        println!("  Proof size: {} bytes", result.proof.len());
+        println!("  Public values: {}", result.public_vals.len());
+        println!("  Proving time: {}ms", result.proving_time_ms);
+        println!("  Verify time: {}ms", result.verify_time_ms);
+    }
+}
+