complete todo 1: make merkle root public, add lightweight CI (#5)

a set of commits:

* updated readme with current and future work

* feat: restructure as separate crate with distributed proving

- Create distributed-zkml as separate Rust crate
- Move distributed/ray modules from zkml to distributed-zkml/src
- Add Merkle tree integration tests
- Add AWS GPU test infrastructure
- Update documentation and testing guides

* Reorganize tests: move simple_distributed to tests, add tests README, add AWS test infrastructure

* Make AWS resource names configurable via env vars

* added applictions

* Consolidate QUICK_START.md and remove NEXT_STEPS.md

* Remove QUICK_START.md (consolidated into README.md)

* readme cleanups

* complete todo 1: make merkle root public, add lightweight CI for PRs

* add mnist example files for CI tests

* fix CI: use nightly rust for int_roundings feature

Author: Masoud

.github/workflows/ci.yml

@@ -0,0 +1,66 @@
+name: CI
+
+on:
+  pull_request:
+    branches: [main, dev]
+  push:
+    branches: [main, dev]
+
+env:
+  CARGO_TERM_COLOR: always
+  RUST_BACKTRACE: 1
+
+jobs:
+  build-and-test:
+    name: Build & Test
+    runs-on: ubuntu-latest
+    
+    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-${{ hashFiles('zkml/Cargo.lock') }}
+          restore-keys: |
+            ${{ runner.os }}-cargo-
+      
+      - name: Check formatting
+        working-directory: zkml
+        run: cargo fmt --check || echo "::warning::Code formatting issues found (non-blocking)"
+        continue-on-error: true
+      
+      - name: Build zkml library
+        working-directory: zkml
+        run: cargo build --lib
+      
+      - name: Run zkml/testing tests
+        working-directory: zkml
+        run: |
+          echo "Running Merkle root public values test..."
+          cargo test --test test_merkle_root_public -- --nocapture
+          
+          echo "Running Merkle tree tests..."
+          cargo test --test merkle_tree_test -- --nocapture
+          
+          echo "Running chunk execution tests..."
+          cargo test --test chunk_execution_test -- --nocapture
+      
+      - name: Run clippy (warnings only)
+        working-directory: zkml
+        run: cargo clippy --lib -- -W clippy::all || echo "::warning::Clippy warnings found (non-blocking)"
+        continue-on-error: true
+
+  # Note: AWS tests are intentionally excluded from CI to save costs.
+  # Run them manually with: cd tests/aws && python gpu_test.py
+

README.md

@@ -4,9 +4,7 @@ Extension of [zkml](https://github.com/uiuc-kang-lab/zkml) for distributed provi
 
 ## Overview
 
-This repository extends zkml (see [ZKML: An Optimizing System for ML Inference in Zero-Knowledge Proofs](https://ddkang.github.io/papers/2024/zkml-eurosys.pdf)) with distributed proving capabilities. The zkml repository is included as a git submodule in the `zkml/` directory and modified to support Merkle tree commitments for intermediate layer outputs.
-
-zkml provides: An optimizing compiler from TensorFlow to halo2 ZK-SNARK circuits for single-machine proof generation.
+This repository extends zkml (see [ZKML: An Optimizing System for ML Inference in Zero-Knowledge Proofs](https://ddkang.github.io/papers/2024/zkml-eurosys.pdf)) with distributed proving capabilities. The zkml repository is included as a git submodule in the `zkml/` directory and modified to support Merkle tree commitments for intermediate layer outputs required in a distributed setting. zkml provides an optimizing compiler from TensorFlow to halo2 ZK-SNARK circuits for single-machine proof generation. High-stakes AI applications in biology or robotics are more practical with trustless verification using [ZKPs (Zero Knowledge Proofs](https://toc.csail.mit.edu/node/218), [SNARKs (Succient Non-interactive Arguments of Knowledge), and zk-SNARKs](https://cs251.stanford.edu/lectures/lecture15.pdf).
 
 distributed-zkml adds:
 - Layer-wise partitioning: Split ML models into chunks for parallel proving across multiple GPUs
@@ -127,26 +125,12 @@ python3 tests/simple_distributed.py \
 pytest tests/
 ```
 
-#### Run Specific Test File
+#### Run specific GPU and AWS tests
 ```bash
 pytest tests/aws/gpu_test.py
-```
-
-#### Run Specific Test
-```bash
 pytest tests/aws/gpu_test.py::test_aws_credentials
 ```
 
-#### Run with Verbose Output
-```bash
-pytest tests/ -v
-```
-
-#### Run Directly (without pytest)
-```bash
-python3 tests/aws/gpu_test.py
-```
-
 ### Rust Tests (Cargo)
 
 #### Run All Tests in zkml
@@ -184,7 +168,7 @@ cd zkml
 cargo check --lib
 ```
 
-Note: Broken example files have been moved to `zkml/examples/broken/` to prevent compilation errors. Use `--test` flags when running tests.
+Broken example files are moved to `zkml/examples/broken/` to prevent compilation errors. Use `--test` flags when running tests.
 
 ### Test Files
 
@@ -278,13 +262,6 @@ pip install ray torch
 nvidia-smi  # Verify GPU is available
 ```
 
-### Running Tests
-
-```bash
-# From distributed-zkml root
-python3 tests/aws/gpu_test.py
-```
-
 ### Test Suite
 
 The test suite includes:
@@ -365,7 +342,7 @@ Distributed Proving Simulation: PASS
 #### "PyTorch CUDA not available"
 - Install PyTorch with CUDA: `pip install torch --index-url https://download.pytorch.org/whl/cu118`
 
-## Next Steps
+## TODO: Next Steps
 
 1. **Make Merkle root public**: Add root to public values so next chunk can verify it
 2. **Complete proof generation**: Connect chunk execution to actual proof generation

tests/README.md tests/README_tests.mdtests/README.md renamed to tests/README_tests.md

@@ -35,10 +35,8 @@ export AWS_SESSION_TOKEN=your_token
 - `setup_aws_resources.sh` - Creates/gets key pair and security group
 - `manage_aws_instance.sh` - Launches instance, runs tests, shuts down
 - `find_ami.sh` - Lists available Deep Learning AMIs
-- `README.md` - AWS testing documentation
-- `QUICK_START.md` - Detailed step-by-step guide
+- `README.md` - AWS testing documentation (includes quick start guide)
 - `aws_setup_guide.md` - Comprehensive AWS setup guide
-- `NEXT_STEPS.md` - Post-setup instructions
 
 See `aws/README.md` for detailed documentation.
 
@@ -59,3 +57,34 @@ pytest tests/aws/gpu_test.py
 python3 tests/aws/gpu_test.py
 ```
 
+## Testing Merkle Root in Public Values
+
+Tests to verify that Merkle root is correctly added to public values when chunk execution is enabled.
+
+### Run Test Binary
+```bash
+cd zkml
+
+# Test with chunk execution (layers 0-2) and Merkle root
+cargo run --bin test_merkle_public -- \
+    examples/mnist/config.msgpack \
+    examples/mnist/inp.msgpack \
+    0 2
+
+# Test with full model execution (no chunk)
+cargo run --bin test_merkle_public -- \
+    examples/mnist/config.msgpack \
+    examples/mnist/inp.msgpack
+```
+
+### Run Test Suite
+```bash
+cd zkml
+cargo test --test test_merkle_root_public -- --nocapture
+```
+
+**What it verifies:**
+- Merkle root is computed and added to public values when chunk execution with Merkle is enabled
+- Circuit verification passes with MockProver
+- Backward compatibility: full model execution still works
+