Skip to content

feat: Add JSTprove backend integration #40

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
shirin-shahabi wants to merge 5 commits into
base: dan/dslice
Choose a base branch
from
Open

feat: Add JSTprove backend integration #40

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
5 commits
Select commit Hold shift + click to select a range
f9400f9
feat: Add JSTprove backend integration
shirin-shahabi Dec 2, 2025
78fd8d4
fix: Correct backend tracking in metadata and model compilation
shirin-shahabi Dec 2, 2025
a9a3733
Merge origin/dan/dslice into JSTprove_integration
shirin-shahabi Dec 2, 2025
7d91b00
fix: Replace undefined layer_idx with idx in _compile_slice logger st…
shirin-shahabi Dec 2, 2025
191d2e6
feat: Add JSTprove optimizations and runner integration
shirin-shahabi Dec 13, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
105 changes: 105 additions & 0 deletions docs/JSTPROVE_BACKEND.md
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,105 @@
# JSTprove Backend Integration

## Overview

This document describes the integration of JSTprove as an additional ZK proof backend alongside EZKL in the Dsperse compilation pipeline.

## Features

### 1. JSTprove Backend Support
- New backend class `JSTprove` in `dsperse/src/backends/JSTprove.py`
- Uses JSTprove CLI (`jst` command) for circuit compilation, witness generation, proof generation, and verification
- Compatible with existing EZKL interface for seamless integration

### 2. Flexible Backend Selection
The compiler now supports three modes:

**Default (Fallback Mode):**
```bash
dsperse compile --path model/slices
```
- Tries JSTprove first
- Falls back to EZKL if JSTprove fails
- Falls back to ONNX (skip ZK compilation) if both fail

**Single Backend:**
```bash
dsperse compile --path model/slices --backend jstprove
dsperse compile --path model/slices --backend ezkl
```

**Per-Layer Backend Assignment:**
```bash
dsperse compile --path model/slices --backend "0,2:jstprove;3-4:ezkl"
```
- Layer 0 and 2: Use JSTprove
- Layer 3 and 4: Use EZKL
- Unspecified layers use default backend

## Installation

1. Install Open MPI (required for JSTprove):
```bash
brew install open-mpi # macOS
# or apt-get install openmpi-bin libopenmpi-dev # Linux
```

2. Install JSTprove:
```bash
uv tool install jstprove
# or: pip install jstprove
```

3. Verify installation:
```bash
jst --help
```

The `install.sh` script has been updated to automatically install these dependencies.

## File Changes

### New Files
- `dsperse/src/backends/JSTprove.py` - JSTprove backend implementation

### Modified Files
- `dsperse/src/cli/compile.py` - Added `--backend` argument
- `dsperse/src/compile/compiler.py` - Backend selection and fallback logic
- `dsperse/src/compile/utils/compiler_utils.py` - Support for JSTprove compilation success check
- `dsperse/src/constants.py` - Added JSTprove command constant
- `install.sh` - Added Open MPI and JSTprove installation
- `requirements.txt` - Added mpi4py dependency

## Usage Examples

**Compile all layers with default fallback:**
```bash
dsperse compile --path model/slices
```

**Compile specific layers with mixed backends:**
```bash
dsperse compile --path model/slices --layers "0-4" --backend "0,2:jstprove;3-4:ezkl"
```

**Compile with single backend:**
```bash
dsperse compile --path model/slices --backend jstprove
```

## Backend Comparison

| Feature | JSTprove | EZKL |
|---------|----------|------|
| Circuit Format | `.txt` | `.compiled` |
| Keys | Not required | `vk.key`, `pk.key` |
| Settings | Dummy JSON | Full settings.json |
| CLI Command | `jst` | `ezkl` |

## Notes

- JSTprove uses CLI-only interface (no Python package import)
- Fallback logic ensures compilation continues even if preferred backend fails
- Metadata tracks which backend was used for each slice
- All changes maintain backward compatibility with existing EZKL workflows

78 changes: 55 additions & 23 deletions dsperse/src/analyzers/runner_analyzer.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -95,16 +95,29 @@ def _process_slices_model(slices_dir: Path, slices_list: list[dict]) -> dict:
dependencies = item.get("dependencies") or {}
parameters = item.get("parameters", 0)

# EZKL compilation info
comp = ((item.get("compilation") or {}).get("ezkl") or {})
files = (comp.get("files") or {})
compiled_flag = bool(comp.get("compiled", False))

# Accept both keys: 'compiled_circuit' and legacy 'compiled'
compiled_rel = files.get("compiled_circuit") or files.get("compiled")
settings_rel = files.get("settings")
pk_rel = files.get("pk_key")
vk_rel = files.get("vk_key")
# Check for compilation info (JSTprove first, then EZKL)
compilation = item.get("compilation") or {}
jst_comp = compilation.get("jstprove") or {}
ezkl_comp = compilation.get("ezkl") or {}

# Prefer JSTprove if available, otherwise EZKL
if jst_comp.get("compiled"):
backend = "jstprove"
files = jst_comp.get("files") or {}
compiled_flag = True
compiled_rel = files.get("circuit")
settings_rel = None
pk_rel = None
vk_rel = None
else:
backend = "ezkl"
files = ezkl_comp.get("files") or {}
compiled_flag = bool(ezkl_comp.get("compiled", False))
# Accept both keys: 'compiled_circuit' and legacy 'compiled'
compiled_rel = files.get("compiled_circuit") or files.get("compiled")
settings_rel = files.get("settings")
pk_rel = files.get("pk_key")
vk_rel = files.get("vk_key")

def _norm(rel: Optional[str]) -> Optional[str]:
if not rel:
Expand All @@ -124,6 +137,7 @@ def _norm(rel: Optional[str]) -> Optional[str]:
"output_shape": output_shape,
"ezkl_compatible": True,
"ezkl": bool(compiled_flag),
"backend": backend,
"circuit_size": 0, # unknown without touching filesystem; keep 0
"dependencies": dependencies,
"parameters": parameters,
Expand Down Expand Up @@ -171,28 +185,44 @@ def _process_slices_per_slice(slices_dir: Path, slices_data_list: list[dict]) ->
dependencies = slice.get("dependencies") or {}
parameters = slice.get("parameters", 0)

# EZKL compilation info
comp = ((slice.get("compilation") or {}).get("ezkl") or {})
files = (comp.get("files") or {})
compiled_flag = bool(comp.get("compiled", False))

if files:
circuit_path = os.path.join(parent_dir, files.get("compiled_circuit") or files.get("compiled"))
settings_path = os.path.join(parent_dir, files.get("settings"))
pk_path = os.path.join(parent_dir, files.get("pk_key"))
vk_path = os.path.join(parent_dir, files.get("vk_key"))
else:
circuit_path = None
# Check for compilation info (JSTprove first, then EZKL)
compilation = slice.get("compilation") or {}
jst_comp = compilation.get("jstprove") or {}
ezkl_comp = compilation.get("ezkl") or {}

if jst_comp.get("compiled"):
backend = "jstprove"
files = jst_comp.get("files") or {}
compiled_flag = True
if files:
circuit_path = os.path.join(parent_dir, files.get("circuit"))
else:
circuit_path = None
Comment on lines +197 to +200
Copy link

@coderabbitai coderabbitai bot Dec 13, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

⚠️ Potential issue | 🟡 Minor

Potential TypeError if files.get("circuit") returns None.

When files is truthy but files.get("circuit") returns None, os.path.join(parent_dir, None) will raise a TypeError.

             if jst_comp.get("compiled"):
                 backend = "jstprove"
                 files = jst_comp.get("files") or {}
                 compiled_flag = True
-                if files:
-                    circuit_path = os.path.join(parent_dir, files.get("circuit"))
+                circuit_rel = files.get("circuit") if files else None
+                if circuit_rel:
+                    circuit_path = os.path.join(parent_dir, circuit_rel)
                 else:
                     circuit_path = None
📝 Committable suggestion

‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.

Suggested change
if files:
circuit_path = os.path.join(parent_dir, files.get("circuit"))
else:
circuit_path = None
circuit_rel = files.get("circuit") if files else None
if circuit_rel:
circuit_path = os.path.join(parent_dir, circuit_rel)
else:
circuit_path = None
🤖 Prompt for AI Agents 
In dsperse/src/analyzers/runner_analyzer.py around lines 197 to 200, the current
code calls os.path.join(parent_dir, files.get("circuit")) when files is truthy
but files.get("circuit") may be None which will raise a TypeError; change the
logic to first extract the circuit name (e.g. circuit_name =
files.get("circuit")) and only call os.path.join(parent_dir, circuit_name) if
circuit_name is not None/empty, otherwise set circuit_path = None (or handle
absent value accordingly).

settings_path = None
pk_path = None
vk_path = None
else:
backend = "ezkl"
files = ezkl_comp.get("files") or {}
compiled_flag = bool(ezkl_comp.get("compiled", False))
if files:
circuit_path = os.path.join(parent_dir, files.get("compiled_circuit") or files.get("compiled"))
settings_path = os.path.join(parent_dir, files.get("settings"))
pk_path = os.path.join(parent_dir, files.get("pk_key"))
vk_path = os.path.join(parent_dir, files.get("vk_key"))
else:
circuit_path = None
settings_path = None
pk_path = None
vk_path = None
Comment on lines +208 to +217
Copy link

@coderabbitai coderabbitai bot Dec 13, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

⚠️ Potential issue | 🟡 Minor

Same None handling issue for EZKL file paths.

Similar to JSTprove, files.get() can return None for any key, causing os.path.join to fail.

             else:
                 backend = "ezkl"
                 files = ezkl_comp.get("files") or {}
                 compiled_flag = bool(ezkl_comp.get("compiled", False))
                 if files:
-                    circuit_path = os.path.join(parent_dir, files.get("compiled_circuit") or files.get("compiled"))
-                    settings_path = os.path.join(parent_dir, files.get("settings"))
-                    pk_path = os.path.join(parent_dir, files.get("pk_key"))
-                    vk_path = os.path.join(parent_dir, files.get("vk_key"))
+                    circuit_rel = files.get("compiled_circuit") or files.get("compiled")
+                    circuit_path = os.path.join(parent_dir, circuit_rel) if circuit_rel else None
+                    settings_rel = files.get("settings")
+                    settings_path = os.path.join(parent_dir, settings_rel) if settings_rel else None
+                    pk_rel = files.get("pk_key")
+                    pk_path = os.path.join(parent_dir, pk_rel) if pk_rel else None
+                    vk_rel = files.get("vk_key")
+                    vk_path = os.path.join(parent_dir, vk_rel) if vk_rel else None
                 else:
                     circuit_path = None
                     settings_path = None
                     pk_path = None
                     vk_path = None
🤖 Prompt for AI Agents 
In dsperse/src/analyzers/runner_analyzer.py around lines 208 to 217, calling
os.path.join(parent_dir, files.get(...)) can raise if files.get(...) returns
None; change the assignments to first retrieve each value into a local variable
and only call os.path.join when that value is non-empty (e.g., if val: path =
os.path.join(parent_dir, val) else: path = None) for compiled_circuit/compiled,
settings, pk_key and vk_key so None values are preserved instead of passed into
os.path.join.


slices[slice_key] = {
"path": onnx_path,
"input_shape": input_shape,
"output_shape": output_shape,
"ezkl_compatible": True,
"ezkl": bool(compiled_flag),
"backend": backend,
"circuit_size": 0,
"dependencies": dependencies,
"parameters": parameters,
Expand Down Expand Up @@ -243,9 +273,11 @@ def _build_execution_chain(slices: dict):
meta = slices.get(slice_key, {})
circuit_path = meta.get('circuit_path')
onnx_path = meta.get('path')
backend = meta.get('backend', 'ezkl')
has_circuit = circuit_path is not None and circuit_path != ""
has_keys = (meta.get('pk_path') is not None) and (meta.get('vk_path') is not None)
use_circuit = bool(meta.get('ezkl')) and has_circuit and has_keys
# JSTprove doesn't require pk/vk keys; EZKL does
use_circuit = bool(meta.get('ezkl')) and has_circuit and (backend == 'jstprove' or has_keys)

next_slice = ordered_keys[i + 1] if i < len(ordered_keys) - 1 else None
execution_chain["nodes"][slice_key] = {
Expand Down
Loading