feat: Add JSTprove backend integration #40
Conversation
- Add JSTprove backend (dsperse/src/backends/JSTprove.py) with CLI-based compilation, witness generation, proving, and verification - Update compiler to support per-layer backend specifications and fallback logic (jstprove -> ezkl -> onnx) - Add JSTprove installation to install.sh (Open MPI + JSTprove CLI via uv) - Add mpi4py to requirements.txt for JSTprove support - Update compiler_utils to handle JSTprove compilation success checks (no payload subdirectory) - Add JSTPROVE_COMMAND constant to constants.py - Add JSTPROVE_BACKEND.md documentation - Fix: Enable fallback for unspecified layers in per-layer backend specs - Fix: Compile all layers when no --layers flag specified (default fallback behavior)
WalkthroughAdds a JSTprove backend and CLI/install integration: new JSTprove Python wrapper, compiler support for backend selection (default, single, or per-layer), per-slice fallback to EZKL/ONNX, updated metadata handling, installer steps for Open MPI/JSTprove, and updated requirements. Changes
Sequence Diagram(s)sequenceDiagram
autonumber
participant User
participant CLI as "dsperse CLI"
participant Compiler
participant JST as "JSTprove CLI"
participant EZKL as "EZKL backend"
participant FS as "Filesystem/Artifacts"
Note right of CLI: User invokes compile with --backend (optional)
User->>CLI: run compile (--backend / per-layer spec)
CLI->>Compiler: Compiler(backend=spec) / compile(...)
Compiler->>Compiler: parse per-layer backend spec
Compiler->>FS: read model slices / inputs
alt layer assigned to JSTprove
Compiler->>JST: compile_circuit (ONNX -> circuit)
JST-->>FS: write circuit, settings
Compiler->>JST: generate_witness (inputs -> witness)
JST-->>FS: write witness
Compiler->>JST: prove (witness,circuit -> proof)
JST-->>FS: write proof
else layer assigned to EZKL
Compiler->>EZKL: compile/prove pipeline
EZKL-->>FS: write artifacts
else fallback to ONNX
Compiler->>FS: keep ONNX artifacts (no ZK)
end
Compiler->>Compiler: update_slice_metadata(backend_name, files)
Compiler->>CLI: return compilation summary (backend_stats, paths)
CLI->>User: print summary / next steps
Estimated code review effort🎯 4 (Complex) | ⏱️ ~60 minutes
Poem
Pre-merge checks and finishing touches❌ Failed checks (1 warning)
✅ Passed checks (2 passed)
✨ Finishing touches
🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 6
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
dsperse/src/compile/compiler.py (1)
238-247: _compile_model silently ignores backend choice and always uses EZKL.The method calls
self.backend.compilation_pipeline()which unconditionally returns EZKL via thebackendproperty, regardless of the user-specifiedbackend_name. If a user specifiesbackend="jstprove"for whole-model compilation, it will silently use EZKL instead.Additionally, JSTprove backend lacks a
compilation_pipelinemethod entirely (it hascircuitization_pipelineinstead), so even sliced compilation with JSTprove will fail at runtime. The backend selection logic in_compile_sliced_modelattempts to usetry_backend.compilation_pipeline()for all backends, but this method doesn't exist on the JSTprove class.
🧹 Nitpick comments (5)
dsperse/src/compile/utils/compiler_utils.py (1)
105-125: Consider renaming function to reflect dual-backend support.The function
is_ezkl_compilation_successfulnow handles both EZKL and JSTprove backends, making the name misleading. Consider renaming tois_compilation_successfulfor clarity.@staticmethod - def is_ezkl_compilation_successful(compilation_data: Dict[str, Any]) -> bool: + def is_compilation_successful(compilation_data: Dict[str, Any]) -> bool: """ Determine if compilation was successful based on produced file paths. EZKL files are in payload subdirectories, JSTprove files are in backend directories. Supports both EZKL and JSTprove backends. """Note: This would require updating all call sites to use the new name.
install.sh (1)
277-279: Redundant PATH export.The PATH is already exported at line 259. This duplicate can be removed.
info "JSTprove installed successfully via uv" - # Add uv tool bin to PATH - export PATH="$HOME/.local/bin:$PATH" - info "Added $HOME/.local/bin to PATH for JSTprove"dsperse/src/backends/JSTprove.py (2)
433-435: Uselogger.exceptioninstead oflogger.errorin exception handlers.Within exception handlers,
logger.exceptionautomatically includes traceback information.except (KeyError, TypeError) as e: - logger.error(f"Could not process witness data: {e}") + logger.exception(f"Could not process witness data: {e}") return None
19-24: Consider annotating class attributes withClassVar.For better type checking, mutable class attributes should be annotated with
ClassVarto indicate they belong to the class, not instances.+from typing import ClassVar class JSTprove: """JSTprove backend for zero-knowledge proof generation using the JSTprove CLI.""" # Class constants - COMMAND = JSTPROVE_COMMAND - DEFAULT_FLAGS = ["--no-banner"] + COMMAND: ClassVar[str] = JSTPROVE_COMMAND + DEFAULT_FLAGS: ClassVar[List[str]] = ["--no-banner"]dsperse/src/compile/compiler.py (1)
82-100: Consider narrowing exception handling.Catching broad
Exceptioncould mask unexpected errors during development. However, given the variety of reasons a backend might fail to initialize (missing CLI, configuration issues, etc.), this defensive approach is acceptable for production robustness.If you want to improve debuggability, consider catching more specific exceptions or at least logging the exception type.
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (8)
docs/JSTPROVE_BACKEND.md(1 hunks)dsperse/src/backends/JSTprove.py(1 hunks)dsperse/src/cli/compile.py(3 hunks)dsperse/src/compile/compiler.py(7 hunks)dsperse/src/compile/utils/compiler_utils.py(1 hunks)dsperse/src/constants.py(1 hunks)install.sh(4 hunks)requirements.txt(1 hunks)
🧰 Additional context used
🧬 Code graph analysis (3)
dsperse/src/backends/JSTprove.py (1)
dsperse/src/backends/ezkl.py (1)
compilation_pipeline(578-671)
dsperse/src/cli/compile.py (1)
dsperse/src/compile/compiler.py (3)
backend(123-125)backend_name(128-129)Compiler(24-355)
dsperse/src/compile/compiler.py (4)
dsperse/src/backends/JSTprove.py (2)
JSTprove(19-458)get_version(438-458)dsperse/src/compile/utils/compiler_utils.py (4)
CompilerUtils(12-289)parse_layers(57-74)is_ezkl_compilation_successful(105-125)get_relative_paths(128-141)dsperse/src/backends/ezkl.py (2)
EZKL(107-693)get_version(137-158)dsperse/src/utils/utils.py (2)
Utils(12-359)save_metadata_file(18-42)
🪛 Ruff (0.14.7)
dsperse/src/backends/JSTprove.py
24-24: Mutable class attributes should be annotated with typing.ClassVar
(RUF012)
42-42: subprocess call: check for execution of untrusted input
(S603)
49-49: Avoid specifying long messages outside the exception class
(TRY003)
51-51: Within an except clause, raise exceptions with raise ... from err or raise ... from None to distinguish them from errors in exception handling
(B904)
51-51: Avoid specifying long messages outside the exception class
(TRY003)
75-75: Consider iterable unpacking instead of concatenation
Replace with iterable unpacking
(RUF005)
78-78: subprocess call: check for execution of untrusted input
(S603)
85-85: Consider moving this statement to an else block
(TRY300)
90-90: Use logging.exception instead of logging.error
Replace with exception
(TRY400)
102-102: Unused method argument: vk_path
(ARG002)
103-103: Unused method argument: settings_path
(ARG002)
126-126: Avoid specifying long messages outside the exception class
(TRY003)
139-139: Avoid specifying long messages outside the exception class
(TRY003)
141-141: Avoid specifying long messages outside the exception class
(TRY003)
156-156: Use logging.exception instead of logging.error
Replace with exception
(TRY400)
164-164: Consider moving this statement to an else block
(TRY300)
167-167: Use logging.exception instead of logging.error
Replace with exception
(TRY400)
175-175: Unused method argument: pk_path
(ARG002)
176-176: Unused method argument: check_mode
(ARG002)
177-177: Unused method argument: settings_path
(ARG002)
200-200: Avoid specifying long messages outside the exception class
(TRY003)
202-202: Avoid specifying long messages outside the exception class
(TRY003)
215-215: Use logging.exception instead of logging.error
Replace with exception
(TRY400)
227-227: Unused method argument: settings_path
(ARG002)
228-228: Unused method argument: vk_path
(ARG002)
256-256: Avoid specifying long messages outside the exception class
(TRY003)
266-266: Consider moving this statement to an else block
(TRY300)
268-268: Use logging.exception instead of logging.error
Replace with exception
(TRY400)
275-275: Unused method argument: settings_path
(ARG002)
294-294: Avoid specifying long messages outside the exception class
(TRY003)
304-304: Consider moving this statement to an else block
(TRY300)
305-305: Do not catch blind exception: Exception
(BLE001)
307-307: Use logging.exception instead of logging.error
Replace with exception
(TRY400)
315-315: Unused method argument: segment_details
(ARG002)
338-338: Avoid specifying long messages outside the exception class
(TRY003)
390-390: Use explicit conversion flag
Replace with conversion flag
(RUF010)
434-434: Use logging.exception instead of logging.error
Replace with exception
(TRY400)
446-446: subprocess call: check for execution of untrusted input
(S603)
456-456: Do not catch blind exception: Exception
(BLE001)
dsperse/src/cli/compile.py
146-146: f-string without any placeholders
Remove extraneous f prefix
(F541)
dsperse/src/compile/compiler.py
87-87: Do not catch blind exception: Exception
(BLE001)
97-97: Do not catch blind exception: Exception
(BLE001)
205-205: f-string without any placeholders
Remove extraneous f prefix
(F541)
307-307: Loop control variable idx not used within loop body
(B007)
🔇 Additional comments (17)
dsperse/src/constants.py (1)
16-17: LGTM!Clean addition of the JSTprove command constant following the existing pattern for EZKL configuration.
docs/JSTPROVE_BACKEND.md (1)
1-105: Well-structured documentation for JSTprove backend integration.The documentation thoroughly covers installation, usage modes, and backend comparison. Clear examples for per-layer backend assignment.
install.sh (1)
224-242: Linux installation uses sudo which may fail in non-interactive CI environments.The
apt-getandyumcommands usesudo, which could fail in containerized CI environments without sudo access. Consider documenting this limitation or adding a check for sudo availability.dsperse/src/cli/compile.py (1)
132-145: Clear backend-aware user messaging.Good handling of the different backend specification scenarios with appropriate user feedback.
dsperse/src/backends/JSTprove.py (6)
53-91: Well-implemented command execution helper.The
_run_commandmethod properly handles subprocess execution with good error capture. The subprocess security warnings (S603) are expected for a CLI wrapper that must execute external commands.
97-168: Good witness generation with auto-compilation support.The method correctly handles both pre-compiled circuits and ONNX models that need compilation. Unused parameters for backward compatibility are properly documented.
170-269: Consistent implementation for prove and verify methods.Both methods follow a consistent pattern with proper path normalization, file existence validation, and error handling.
299-308: Broad exception catch is acceptable for subprocess wrapper.The
except Exceptionat line 305 catches all exceptions from the subprocess call. While generally discouraged, it's pragmatic here since subprocess failures can raise various exception types. The error is properly logged and returned.
310-397: Well-designed pipeline with EZKL compatibility layer.The
circuitization_pipelinemethod provides EZKL-compatible output structure while handling JSTprove's simpler compilation model. The dummy settings file and alias tocompilation_pipelineensure smooth integration with existing infrastructure.
437-458: Robust version detection with proper timeout handling.Good use of timeout to prevent hanging, and graceful fallback to
Nonewhen version cannot be determined.requirements.txt (1)
12-12: mpi4py does not strictly require Open MPI pre-installation for typical pip installs.Prebuilt mpi4py wheels are available and work on most platforms without a separate Open MPI installation. However, if pip must compile from source, you will need an MPI implementation with development headers and
mpiccavailable. Consider documenting which installation method (wheels vs. source) is expected for this project.Likely an incorrect or invalid review comment.
dsperse/src/compile/compiler.py (6)
32-64: LGTM!The initialization logic cleanly handles the three backend specification modes (fallback, per-layer, and simple backend name). The lazy loading pattern for backends avoids initialization errors when a backend isn't needed.
66-80: LGTM!Good reuse of
CompilerUtils.parse_layersfor layer index parsing. The method correctly handles semicolon-separated backend specifications and normalizes backend names to lowercase.
121-133: LGTM!Properties correctly maintain backward compatibility. The
backendproperty returning EZKL ensures existing code continues to work.
186-226: LGTM on fallback logic!The fallback chain implementation is well-structured with proper exception handling and informative logging. The conditional re-raise at lines 224-225 correctly propagates errors when fallback is disabled.
282-294: LGTM!Good defensive coding with
hasattrchecks for version retrieval. The compilation metadata structure correctly records the backend used and its version undercompilation.<backend>.
316-355: LGTM!The
compilemethod properly orchestrates the compilation workflow with clear separation between sliced and whole-model paths. The updated logging at line 336 correctly describes the fallback behavior when no layers are specified.
| except FileNotFoundError: | ||
| raise RuntimeError("JSTprove CLI not found. Please install JSTprove: uv tool install jstprove") |
There was a problem hiding this comment.
Preserve exception chain with raise ... from e.
When catching FileNotFoundError and raising a new RuntimeError, preserve the exception chain for better debugging.
except FileNotFoundError:
- raise RuntimeError("JSTprove CLI not found. Please install JSTprove: uv tool install jstprove")
+ raise RuntimeError("JSTprove CLI not found. Please install JSTprove: uv tool install jstprove") from NoneUsing from None explicitly suppresses the chained exception if desired, or use from e to preserve it.
📝 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.
| except FileNotFoundError: | |
| raise RuntimeError("JSTprove CLI not found. Please install JSTprove: uv tool install jstprove") | |
| except FileNotFoundError: | |
| raise RuntimeError("JSTprove CLI not found. Please install JSTprove: uv tool install jstprove") from None |
🧰 Tools
🪛 Ruff (0.14.7)
51-51: Within an except clause, raise exceptions with raise ... from err or raise ... from None to distinguish them from errors in exception handling
(B904)
51-51: Avoid specifying long messages outside the exception class
(TRY003)
🤖 Prompt for AI Agents
In dsperse/src/backends/JSTprove.py around lines 50-51, the except block catches
FileNotFoundError and raises a new RuntimeError without preserving the original
exception chain; modify the except to capture the original exception (e) and
re-raise the RuntimeError using "raise RuntimeError(... ) from e" so the
underlying FileNotFoundError is retained for debugging (alternatively use "from
None" if you intentionally want to suppress chaining).
| compile_parser.add_argument('--backend', '-b', default=None, | ||
| help='Backend specification. Can be: "jstprove", "ezkl", or per-layer like "0,2:jstprove;3-4:ezkl". ' | ||
| 'Default: try both jstprove and ezkl, fallback to onnx. Note: requires --layers to compile.') |
There was a problem hiding this comment.
Misleading help text: --layers is not required.
The help text states "Note: requires --layers to compile" but the code at lines 135-137 shows that when no layers are specified, all layers are compiled with default fallback. Remove this misleading note.
compile_parser.add_argument('--backend', '-b', default=None,
help='Backend specification. Can be: "jstprove", "ezkl", or per-layer like "0,2:jstprove;3-4:ezkl". '
- 'Default: try both jstprove and ezkl, fallback to onnx. Note: requires --layers to compile.')
+ 'Default: try both jstprove and ezkl, fallback to onnx.')📝 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.
| compile_parser.add_argument('--backend', '-b', default=None, | |
| help='Backend specification. Can be: "jstprove", "ezkl", or per-layer like "0,2:jstprove;3-4:ezkl". ' | |
| 'Default: try both jstprove and ezkl, fallback to onnx. Note: requires --layers to compile.') | |
| compile_parser.add_argument('--backend', '-b', default=None, | |
| help='Backend specification. Can be: "jstprove", "ezkl", or per-layer like "0,2:jstprove;3-4:ezkl". ' | |
| 'Default: try both jstprove and ezkl, fallback to onnx.') |
🤖 Prompt for AI Agents
In dsperse/src/cli/compile.py around lines 118-120, the argument help string for
--backend incorrectly states "Note: requires --layers to compile." Remove that
clause and update the help text to reflect actual behavior: when --layers is
omitted all layers are compiled and the default behavior is to try jstprove and
ezkl then fall back to onnx. Replace the help string accordingly so it no longer
claims --layers is required.
| print(f"{Fore.CYAN}Compiling specified layers with {backend_name}...{Style.RESET_ALL}") | ||
| else: | ||
| print(f"{Fore.CYAN}Compiling specified layers (trying jstprove & ezkl, fallback to onnx)...{Style.RESET_ALL}") | ||
| logger.info(f"Starting slices compilation") |
There was a problem hiding this comment.
Remove extraneous f-prefix from string without placeholders.
As flagged by static analysis (Ruff F541), this f-string has no placeholders.
- logger.info(f"Starting slices compilation")
+ logger.info("Starting slices compilation")📝 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.
| logger.info(f"Starting slices compilation") | |
| logger.info("Starting slices compilation") |
🧰 Tools
🪛 Ruff (0.14.7)
146-146: f-string without any placeholders
Remove extraneous f prefix
(F541)
🤖 Prompt for AI Agents
In dsperse/src/cli/compile.py around line 146, the logger call uses an f-string
with no placeholders (logger.info(f"Starting slices compilation")); remove the
unnecessary leading f so it becomes a normal string literal
(logger.info("Starting slices compilation")) to satisfy Ruff F541 and eliminate
the extraneous f-prefix.
| calibration_input = os.path.join( | ||
| os.path.dirname(slice_path), | ||
| backend_dir, | ||
| f"calibration.json" | ||
| ) if os.path.exists(os.path.join(os.path.dirname(slice_path), backend_dir, "calibration.json")) else None |
There was a problem hiding this comment.
Remove extraneous f-string prefix.
The f-string on line 205 has no placeholders. This is a static analysis finding.
calibration_input = os.path.join(
os.path.dirname(slice_path),
backend_dir,
- f"calibration.json"
+ "calibration.json"
) if os.path.exists(os.path.join(os.path.dirname(slice_path), backend_dir, "calibration.json")) else 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.
| calibration_input = os.path.join( | |
| os.path.dirname(slice_path), | |
| backend_dir, | |
| f"calibration.json" | |
| ) if os.path.exists(os.path.join(os.path.dirname(slice_path), backend_dir, "calibration.json")) else None | |
| calibration_input = os.path.join( | |
| os.path.dirname(slice_path), | |
| backend_dir, | |
| "calibration.json" | |
| ) if os.path.exists(os.path.join(os.path.dirname(slice_path), backend_dir, "calibration.json")) else None |
🧰 Tools
🪛 Ruff (0.14.7)
205-205: f-string without any placeholders
Remove extraneous f prefix
(F541)
🤖 Prompt for AI Agents
In dsperse/src/compile/compiler.py around lines 202 to 206, the code builds
calibration_input using an f-string with no placeholders; remove the extraneous
f-string prefix so the filename is a normal string (e.g., "calibration.json")
while preserving the os.path.join and the conditional existence check, ensuring
behavior is unchanged.
| # Log summary | ||
| backend_summary = {} | ||
| for idx, backend in backend_stats.items(): | ||
| backend_summary[backend] = backend_summary.get(backend, 0) + 1 | ||
| summary_str = ", ".join(f"{k}: {v}" for k, v in backend_summary.items()) | ||
| if skipped_count > 0: | ||
| logger.info(f"Compilation completed. ZK compiled: {compiled_count} slices ({summary_str}). Skipped: {skipped_count} slices (will use pure ONNX at runtime)") | ||
| else: | ||
| logger.info(f"Compilation completed. ZK compiled: {compiled_count} slices. Backends used: {summary_str}") |
There was a problem hiding this comment.
Replace unused loop variable with underscore.
Static analysis correctly identified that idx is not used in the loop body at line 307.
# Log summary
backend_summary = {}
- for idx, backend in backend_stats.items():
+ for _, backend in backend_stats.items():
backend_summary[backend] = backend_summary.get(backend, 0) + 1📝 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.
| # Log summary | |
| backend_summary = {} | |
| for idx, backend in backend_stats.items(): | |
| backend_summary[backend] = backend_summary.get(backend, 0) + 1 | |
| summary_str = ", ".join(f"{k}: {v}" for k, v in backend_summary.items()) | |
| if skipped_count > 0: | |
| logger.info(f"Compilation completed. ZK compiled: {compiled_count} slices ({summary_str}). Skipped: {skipped_count} slices (will use pure ONNX at runtime)") | |
| else: | |
| logger.info(f"Compilation completed. ZK compiled: {compiled_count} slices. Backends used: {summary_str}") | |
| # Log summary | |
| backend_summary = {} | |
| for _, backend in backend_stats.items(): | |
| backend_summary[backend] = backend_summary.get(backend, 0) + 1 | |
| summary_str = ", ".join(f"{k}: {v}" for k, v in backend_summary.items()) | |
| if skipped_count > 0: | |
| logger.info(f"Compilation completed. ZK compiled: {compiled_count} slices ({summary_str}). Skipped: {skipped_count} slices (will use pure ONNX at runtime)") | |
| else: | |
| logger.info(f"Compilation completed. ZK compiled: {compiled_count} slices. Backends used: {summary_str}") |
🧰 Tools
🪛 Ruff (0.14.7)
307-307: Loop control variable idx not used within loop body
(B007)
🤖 Prompt for AI Agents
In dsperse/src/compile/compiler.py around lines 305 to 313, the loop declares an
unused variable `idx` when iterating backend_stats.items(); replace `for idx,
backend in backend_stats.items():` with either `for _, backend in
backend_stats.items():` or, clearer, iterate values via `for backend in
backend_stats.values():` so the unused variable is removed and the loop intent
is explicit.
| @@ -1,9 +1,11 @@ | |||
| #!/usr/bin/env bash | |||
| # install.sh - Installer for Dsperse CLI and EZKL (with lookup tables) | |||
| # install.sh - Installer for Dsperse CLI, JSTproveand EZKL (with lookup tables) | |||
There was a problem hiding this comment.
Fix typo: missing space between "JSTprove" and "and".
-# install.sh - Installer for Dsperse CLI, JSTproveand EZKL (with lookup tables)
+# install.sh - Installer for Dsperse CLI, JSTprove and EZKL (with lookup tables)📝 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.
| # install.sh - Installer for Dsperse CLI, JSTproveand EZKL (with lookup tables) | |
| # install.sh - Installer for Dsperse CLI, JSTprove and EZKL (with lookup tables) |
🤖 Prompt for AI Agents
In install.sh around line 2, the comment header contains a typo with
"JSTproveand" missing a space; update the string to "JSTprove and" so it reads
"# install.sh - Installer for Dsperse CLI, JSTprove and EZKL (with lookup
tables)".
- Bug 1: update_slice_metadata() now accepts backend_name parameter and writes under correct key (jstprove/ezkl/onnx) instead of always 'ezkl' - Bug 2: _compile_model() now uses proper backend selection with fallback logic instead of always using EZKL
Resolved conflicts: - compiler.py: Combined idx parameter with backend_name tracking - compiler_utils.py: Merged idx-based slice lookup with backend-specific metadata - requirements.txt: Kept simple format with mpi4py for JSTprove
There was a problem hiding this comment.
Actionable comments posted: 2
♻️ Duplicate comments (2)
dsperse/src/compile/compiler.py (2)
368-376: Replace unused loop variable with underscore.The loop variable
idxon line 370 is not used in the loop body, as previously noted. Consider using underscore or iterating over values directly.# Log summary backend_summary = {} - for idx, backend in backend_stats.items(): + for _, backend in backend_stats.items(): backend_summary[backend] = backend_summary.get(backend, 0) + 1Alternatively, iterate values directly for clarity:
# Log summary backend_summary = {} - for idx, backend in backend_stats.items(): + for backend in backend_stats.values(): backend_summary[backend] = backend_summary.get(backend, 0) + 1
207-211: Remove extraneous f-string prefix.The f-string on line 210 has no placeholders, as previously noted.
calibration_input = os.path.join( os.path.dirname(slice_path), backend_dir, - f"calibration.json" + "calibration.json" ) if os.path.exists(os.path.join(os.path.dirname(slice_path), backend_dir, "calibration.json")) else None
🧹 Nitpick comments (1)
requirements.txt (1)
12-15: Clear documentation for system and tool dependencies.The comments correctly document that jstprove is installed via
uv tool install(not pip) and that Open MPI must be pre-installed. This is helpful context. Thempi4py>=3.1.0constraint uses a loose lower bound; consider pinning a tested version range (e.g.,mpi4py>=3.1.0,<4.0.0) if mpi4py introduces breaking changes frequently, though this is optional.
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (3)
dsperse/src/compile/compiler.py(5 hunks)dsperse/src/compile/utils/compiler_utils.py(6 hunks)requirements.txt(1 hunks)
🧰 Additional context used
🧬 Code graph analysis (1)
dsperse/src/compile/utils/compiler_utils.py (3)
dsperse/src/compile/compiler.py (1)
backend_name(128-129)dsperse/src/backends/JSTprove.py (2)
JSTprove(19-458)get_version(438-458)dsperse/src/backends/ezkl.py (2)
get_version(137-158)EZKL(107-693)
🪛 Ruff (0.14.7)
dsperse/src/compile/compiler.py
87-87: Do not catch blind exception: Exception
(BLE001)
97-97: Do not catch blind exception: Exception
(BLE001)
198-198: Undefined name layer_idx
(F821)
210-210: f-string without any placeholders
Remove extraneous f prefix
(F541)
214-214: Undefined name layer_idx
(F821)
223-223: Undefined name layer_idx
(F821)
226-226: Undefined name layer_idx
(F821)
228-228: Undefined name layer_idx
(F821)
248-248: Avoid specifying long messages outside the exception class
(TRY003)
288-288: Avoid specifying long messages outside the exception class
(TRY003)
310-310: Avoid specifying long messages outside the exception class
(TRY003)
370-370: Loop control variable idx not used within loop body
(B007)
🔇 Additional comments (4)
dsperse/src/compile/utils/compiler_utils.py (1)
190-258: Multi-backend metadata update implementation looks good.The function correctly:
- Retrieves backend-specific versions
- Nests compilation info under the backend name for proper organization
- Handles missing slices gracefully with root-level fallback
- Maintains backward compatibility with the default "ezkl" backend
The conditional import of JSTprove (line 210) is acceptable, though importing at the module level would be slightly more efficient if this function is called frequently.
dsperse/src/compile/compiler.py (3)
32-134: Flexible backend selection and initialization looks good.The implementation correctly:
- Supports multiple backend specification formats (None, simple name, per-layer specs)
- Uses lazy initialization to avoid errors when backends aren't needed
- Maintains backward compatibility with properties
- Handles initialization failures gracefully with warnings
The broad
Exceptioncatching in lines 87 and 97 is acceptable here since the failures are logged and handled by returningNone.
243-310: Backend fallback for single model compilation looks good.The implementation correctly mirrors the fallback logic from
_compile_sliceand properly tries backends in the configured order. The exception messages on lines 288 and 310 are descriptive and appropriate for debugging compilation failures.
379-419: Main compile method integration looks good.The updated documentation clearly explains the default fallback behavior when no layers are specified, and the logic correctly delegates to the appropriate compilation methods based on model type.
| torch==2.6.0 | ||
| onnx==1.18.0 | ||
| pytest==8.3.5 | ||
| ezkl==22.2.1 | ||
| networkx~=3.4.2 | ||
| matplotlib~=3.10.1 | ||
| numpy~=2.2.3 | ||
| tqdm~=4.67.1 | ||
| onnxruntime==1.21.0 | ||
| colorama~=0.4.6 |
There was a problem hiding this comment.
🧩 Analysis chain
🌐 Web query:
numpy 2.2.3 onnxruntime 1.21.0 compatibility
💡 Result:
Short answer: Probably not — ONNX Runtime 1.21.0 wheels are built against the older NumPy (1.x) ABI, so they can be incompatible with NumPy 2.x (including 2.2.3). If you try to use a prebuilt onnxruntime 1.21.0 wheel with NumPy 2.2.3 you may hit the “A module that was compiled using NumPy 1.x cannot be run in NumPy 2.0…” error. [1][2]
Recommendations:
- Use numpy<2 (e.g., pip install "numpy<2") with onnxruntime==1.21.0. [2]
- Or rebuild ONNX Runtime from source against NumPy 2.x (or use an ONNX Runtime release/wheel that explicitly documents NumPy‑2 support). [1][3]
Sources:
- ONNX Runtime GitHub issue about NumPy‑2 compatibility (discussion & error text). [microsoft/onnxruntime issue]. [2]
- NumPy 2.x release notes (NumPy 2.2.3). [numpy release notes]. [1]
- ONNX Runtime installation/compatibility docs and release issues referencing build-time NumPy ABI concerns. [onnxruntime install docs / release issues]. [3]
If you want, I can check your environment (pip list / python -c "import onnxruntime; import numpy; print(numpy.version)") and give a precise command to fix it.
🌐 Web query:
torch 2.6.0 numpy 2.2.3 compatibility issues
💡 Result:
Short answer: PyTorch 2.6.0 binary builds that were compiled against NumPy 1.x are not ABI-compatible with NumPy 2.x (including 2.2.3), so importing torch can raise the “A module that was compiled using NumPy 1.x cannot be run in NumPy 2.x …” error. Workarounds: install a NumPy 1.x wheel (pip install "numpy<2") or build/reinstall a PyTorch wheel compiled against NumPy ≥2.0. See sources below.
- PyTorch 2.6 release notes (PyTorch 2.6). [1]
- NumPy 2.2.3 release notes (NumPy 2.x series). [2]
- Ongoing PyTorch issues/threads showing the ABI mismatch and recommended workarounds (downgrade numpy or rebuild PyTorch). [3]
Suggested quick fix:
- pip uninstall numpy && pip install "numpy<2" (or pin to a known 1.26.x).
🌐 Web query:
onnxruntime 1.21.0 numpy 2.0 support
💡 Result:
Short answer: ONNX Runtime 1.21.0 is not compatible with NumPy 2.0 (you will see the "A module that was compiled using NumPy 1.x cannot be run in NumPy 2.0…" error). Use numpy<2 or rebuild ONNX Runtime with NumPy 2.0 (and updated pybind11) to get compatibility. [1][2][3]
Sources:
- ONNX Runtime GitHub issue tracking NumPy 2.0 support (reports the incompatible-module error and discussion). [1]
- Multiple downstream projects reporting the runtime error when ONNX Runtime built against NumPy 1.x is used with NumPy 2.x. [2]
- ONNX Runtime build CMake shows a warning that binaries built against older NumPy will not be compatible with NumPy 2.0 (and indicates rebuilding is required). [3]
References:
[1] microsoft/onnxruntime#21063
[2] example reports: dusty-nv/jetson-containers#561 and chroma-core/chroma#2359
[3] onnxruntime packaging/CMake snippet showing NumPy 2.0 compatibility warning (Debian/Fedora packaging view).
Resolve numpy 2.2.3 incompatibility with torch and onnxruntime before merging.
The combination of torch==2.6.0, onnxruntime==1.21.0, and numpy~=2.2.3 will fail at runtime with the error: "A module that was compiled using NumPy 1.x cannot be run in NumPy 2.0…" Both PyTorch and ONNX Runtime 1.21.0 wheels are compiled against NumPy 1.x ABI and are not forward-compatible with NumPy 2.x. Either pin numpy<2 (e.g., numpy<2 or numpy~=1.26.0) or upgrade to versions with explicit NumPy 2.x support (ONNX Runtime 1.18.0+ requires verification or later releases that document NumPy 2.x compatibility).
🤖 Prompt for AI Agents
requirements.txt lines 1-10: the combination of torch==2.6.0 and
onnxruntime==1.21.0 is incompatible with numpy~=2.2.3 (wheels were built against
NumPy 1.x); to fix, pin numpy to a 1.x release (e.g., change to numpy<2 or
numpy~=1.26.0) OR upgrade torch/onnxruntime to versions that explicitly support
NumPy 2.x, then run a fresh install and smoke-test imports to confirm the
runtime error is resolved.
…atements - Fixed NameError: logger statements in _compile_slice were referencing undefined variable layer_idx - Changed all logger statements to use idx parameter (lines 198, 214, 223, 226, 228)
There was a problem hiding this comment.
Actionable comments posted: 0
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
dsperse/src/compile/compiler.py (1)
394-399: Log message at line 399 is misleading for non-default backend configurations.The message "Will compile all layers with default fallback (jstprove -> ezkl -> onnx)" is printed whenever
layersis not specified, but it doesn't account for the Compiler's backend configuration. When instantiated with a specific backend likeCompiler(backend="ezkl"), theuse_fallbackflag is set toFalse, yet this log message still claims fallback will be used. The log should be conditional on the actual backend configuration and fallback state, or reflect the configured behavior accurately (e.g., "Will compile all layers with ezkl backend only").
♻️ Duplicate comments (2)
dsperse/src/compile/compiler.py (2)
207-211: Remove extraneous f-string prefix.This is a duplicate of a previous review comment. The f-string has no placeholders.
Apply this diff:
calibration_input = os.path.join( os.path.dirname(slice_path), backend_dir, - f"calibration.json" + "calibration.json" ) if os.path.exists(os.path.join(os.path.dirname(slice_path), backend_dir, "calibration.json")) else None
368-376: Replace unused loop variable with underscore.This is a duplicate of a previous review comment. The
idxvariable in the loop is not used.Apply this diff:
# Log summary backend_summary = {} - for idx, backend in backend_stats.items(): + for _, backend in backend_stats.items(): backend_summary[backend] = backend_summary.get(backend, 0) + 1Or better yet, iterate over values directly since the key is unused:
# Log summary backend_summary = {} - for idx, backend in backend_stats.items(): + for backend in backend_stats.values(): backend_summary[backend] = backend_summary.get(backend, 0) + 1
🧹 Nitpick comments (1)
dsperse/src/compile/compiler.py (1)
82-90: Consider catching more specific exceptions.The broad
Exceptioncatch provides graceful degradation but may mask unexpected errors during initialization.If JSTprove has specific initialization exceptions (e.g.,
ImportError,RuntimeError), catch those explicitly:def _get_jstprove(self): """Lazy initialization of JSTprove backend""" if self._jstprove is None: try: self._jstprove = JSTprove() - except Exception as e: + except (ImportError, RuntimeError) as e: logger.warning(f"Failed to initialize JSTprove: {e}") return None return self._jstprove
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (1)
dsperse/src/compile/compiler.py(5 hunks)
🧰 Additional context used
🪛 Ruff (0.14.7)
dsperse/src/compile/compiler.py
87-87: Do not catch blind exception: Exception
(BLE001)
97-97: Do not catch blind exception: Exception
(BLE001)
210-210: f-string without any placeholders
Remove extraneous f prefix
(F541)
248-248: Avoid specifying long messages outside the exception class
(TRY003)
288-288: Avoid specifying long messages outside the exception class
(TRY003)
310-310: Avoid specifying long messages outside the exception class
(TRY003)
370-370: Loop control variable idx not used within loop body
(B007)
🔇 Additional comments (7)
dsperse/src/compile/compiler.py (7)
32-64: LGTM! Backend specification logic is well-structured.The three-mode backend specification (None for fallback, simple name, or per-layer spec) is clearly implemented with appropriate fallback flags.
66-80: LGTM! Per-layer backend parsing reuses existing utilities.The implementation correctly parses the layer specification format and delegates layer index parsing to CompilerUtils.
102-119: LGTM! Layer-specific backend resolution is correct.The logic properly prioritizes per-layer assignments over the default backend, and returns None appropriately for fallback scenarios.
158-231: LGTM! Fallback chain logic is well-implemented.The multi-backend fallback mechanism correctly builds a prioritized list of backends and attempts each until one succeeds, with proper error handling and logging.
243-310: LGTM! Model compilation with fallback is correctly implemented.The fallback chain for whole-model compilation parallels the slice compilation logic, ensuring consistent behavior across both compilation paths.
329-366: LGTM! Per-slice backend tracking is comprehensive.The implementation correctly tracks which backend was used for each slice, retrieves version information, and nests compilation metadata appropriately under the backend name.
121-133: Verify ifbackendproperty misleads callers expecting dynamic backend detection.The
backendproperty always returns EZKL regardless of which backend is active, andbackend_namedefaults to "ezkl". While the comment acknowledges backward compatibility, code that inspects these properties to determine the active backend will always see EZKL even when JSTprove is in use. Verify whether external code relies on these properties for backend detection or if this static return value causes confusion in practice.
Summary by CodeRabbit
New Features
Documentation
Chores
✏️ Tip: You can customize this high-level summary in your review settings.