JSTprove 1.0.0__py3-none-macosx_11_0_arm64.whl

python/frontend/commands/args.py

@@ -0,0 +1,100 @@
+"""Argument specifications for CLI commands."""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    import argparse
+
+from dataclasses import dataclass, field
+from typing import Any
+
+
+@dataclass(frozen=True)
+class ArgSpec:
+    """Specification for a command-line argument."""
+
+    name: str
+    flag: str
+    help_text: str
+    short: str = ""
+    arg_type: type | None = None
+    extra_kwargs: dict[str, Any] = field(default_factory=dict)
+
+    @property
+    def positional(self) -> str:
+        """Return the positional argument name."""
+        return f"pos_{self.name}"
+
+    def add_to_parser(
+        self,
+        parser: argparse.ArgumentParser,
+        help_override: str | None = None,
+    ) -> None:
+        """Add both positional and flag arguments to the parser."""
+        help_text = help_override or self.help_text
+        kwargs = {"help": help_text, **self.extra_kwargs}
+        if self.arg_type is not None:
+            kwargs["type"] = self.arg_type
+
+        if self.short:
+            parser.add_argument(
+                self.positional,
+                nargs="?",
+                metavar=self.name,
+                **kwargs,
+            )
+            parser.add_argument(
+                self.short,
+                self.flag,
+                **kwargs,
+            )
+        else:
+            parser.add_argument(
+                self.flag,
+                **kwargs,
+            )
+
+
+MODEL_PATH = ArgSpec(
+    name="model_path",
+    flag="--model-path",
+    short="-m",
+    help_text="Path to the original ONNX model.",
+)
+
+CIRCUIT_PATH = ArgSpec(
+    name="circuit_path",
+    flag="--circuit-path",
+    short="-c",
+    help_text="Path to the compiled circuit.",
+)
+
+INPUT_PATH = ArgSpec(
+    name="input_path",
+    flag="--input-path",
+    short="-i",
+    help_text="Path to input JSON.",
+)
+
+OUTPUT_PATH = ArgSpec(
+    name="output_path",
+    flag="--output-path",
+    short="-o",
+    help_text="Path to write model outputs JSON.",
+)
+
+WITNESS_PATH = ArgSpec(
+    name="witness_path",
+    flag="--witness-path",
+    short="-w",
+    help_text="Path to write witness.",
+)
+
+PROOF_PATH = ArgSpec(
+    name="proof_path",
+    flag="--proof-path",
+    short="-p",
+    help_text="Path to write proof.",
+)

python/frontend/commands/base.py

@@ -0,0 +1,199 @@
+from __future__ import annotations
+
+import argparse
+import functools
+import importlib
+from abc import ABC, abstractmethod
+from pathlib import Path
+from typing import TYPE_CHECKING, Any, ClassVar
+
+if TYPE_CHECKING:
+    from collections.abc import Callable
+
+    from python.frontend.commands.args import ArgSpec
+
+from python.frontend.commands.constants import (
+    DEFAULT_CIRCUIT_CLASS,
+    DEFAULT_CIRCUIT_MODULE,
+)
+
+
+class HiddenPositionalHelpFormatter(argparse.HelpFormatter):
+    def _format_usage(
+        self,
+        usage: str | None,
+        actions: list,
+        groups: list,
+        prefix: str | None,
+    ) -> str:
+        filtered_actions = [
+            action
+            for action in actions
+            if not (
+                isinstance(action, argparse._StoreAction)  # noqa: SLF001
+                and action.dest.startswith("pos_")
+            )
+        ]
+        return super()._format_usage(usage, filtered_actions, groups, prefix)
+
+    def _format_action(self, action: argparse.Action) -> str:
+        if isinstance(
+            action,
+            argparse._StoreAction,  # noqa: SLF001
+        ) and action.dest.startswith(
+            "pos_",
+        ):
+            return ""
+        return super()._format_action(action)
+
+
+class BaseCommand(ABC):
+    """Base class for CLI commands."""
+
+    name: ClassVar[str]
+    aliases: ClassVar[list[str]] = []
+    help: ClassVar[str]
+
+    @classmethod
+    @abstractmethod
+    def configure_parser(
+        cls: type[BaseCommand],
+        parser: argparse.ArgumentParser,
+    ) -> None:
+        """Configure the argument parser for this command."""
+
+    @classmethod
+    @abstractmethod
+    def run(cls: type[BaseCommand], args: argparse.Namespace) -> None:
+        """Execute the command."""
+
+    @staticmethod
+    def validate_required(*required: ArgSpec) -> Callable:
+        def decorator(func: Callable) -> Callable:
+            @functools.wraps(func)
+            def wrapper(cls: type[BaseCommand], args: argparse.Namespace) -> None:
+                for arg_spec in required:
+                    flag_val = getattr(args, arg_spec.name, None)
+                    pos_val = getattr(args, arg_spec.positional, None)
+                    merged = flag_val if flag_val is not None else pos_val
+                    if not merged:
+                        msg = f"Missing required argument: {arg_spec.name}"
+                        raise ValueError(msg)
+                    setattr(args, arg_spec.name, merged)
+                return func(cls, args)
+
+            return wrapper
+
+        return decorator
+
+    @staticmethod
+    def validate_paths(*paths: ArgSpec) -> Callable:
+        def decorator(func: Callable) -> Callable:
+            @functools.wraps(func)
+            def wrapper(cls: type[BaseCommand], args: argparse.Namespace) -> None:
+                for arg_spec in paths:
+                    cls._ensure_file_exists(getattr(args, arg_spec.name))
+                return func(cls, args)
+
+            return wrapper
+
+        return decorator
+
+    @staticmethod
+    def validate_parent_paths(*paths: ArgSpec) -> Callable:
+        def decorator(func: Callable) -> Callable:
+            @functools.wraps(func)
+            def wrapper(cls: type[BaseCommand], args: argparse.Namespace) -> None:
+                for arg_spec in paths:
+                    cls._ensure_parent_dir(getattr(args, arg_spec.name))
+                return func(cls, args)
+
+            return wrapper
+
+        return decorator
+
+    @staticmethod
+    def validate_optional_paths(*paths: ArgSpec) -> Callable:
+        def decorator(func: Callable) -> Callable:
+            @functools.wraps(func)
+            def wrapper(cls: type[BaseCommand], args: argparse.Namespace) -> None:
+                for arg_spec in paths:
+                    flag_val = getattr(args, arg_spec.name, None)
+                    pos_val = getattr(args, arg_spec.positional, None)
+                    merged = flag_val if flag_val is not None else pos_val
+                    if merged is not None:
+                        cls._ensure_file_exists(merged)
+                    setattr(args, arg_spec.name, merged)
+                return func(cls, args)
+
+            return wrapper
+
+        return decorator
+
+    @staticmethod
+    def _ensure_file_exists(path: str) -> None:
+        p = Path(path)
+        if not p.is_file():
+            msg = f"Required file not found: {path}"
+            raise FileNotFoundError(msg)
+        if not p.stat().st_mode & 0o444:
+            msg = f"Cannot read file: {path}"
+            raise PermissionError(msg)
+
+    @staticmethod
+    def _ensure_parent_dir(path: str) -> None:
+        Path(path).parent.mkdir(parents=True, exist_ok=True)
+
+    @staticmethod
+    def append_arg(cmd: list[str], flag: str, val: object | None) -> None:
+        if val is None:
+            return
+        if isinstance(val, str) and not val.strip():
+            return
+        cmd.extend([flag, str(val)])
+
+    @staticmethod
+    def append_args_from_specs(
+        cmd: list[str],
+        specs: tuple[tuple[ArgSpec, str], ...],
+    ) -> None:
+        for spec, value in specs:
+            cmd.extend([spec.flag, value])
+
+    @staticmethod
+    def append_args_from_namespace(
+        cmd: list[str],
+        args: argparse.Namespace,
+        specs: tuple[ArgSpec, ...],
+    ) -> None:
+        for spec in specs:
+            value = getattr(args, spec.name, None)
+            BaseCommand.append_arg(cmd, spec.flag, value)
+
+    @staticmethod
+    def _build_circuit(model_name_hint: str | None = None) -> Any:  # noqa: ANN401
+        mod = importlib.import_module(DEFAULT_CIRCUIT_MODULE)
+        try:
+            cls = getattr(mod, DEFAULT_CIRCUIT_CLASS)
+        except AttributeError as e:
+            msg = (
+                f"Default circuit class '{DEFAULT_CIRCUIT_CLASS}' "
+                f"not found in '{DEFAULT_CIRCUIT_MODULE}'"
+            )
+            raise RuntimeError(msg) from e
+
+        name = model_name_hint or "cli"
+
+        for attempt in (
+            lambda: cls(model_name=name),
+            lambda: cls(name=name),
+            lambda: cls(name),
+            lambda: cls(),
+        ):
+            try:
+                return attempt()
+            except TypeError:  # noqa: PERF203
+                continue
+
+        msg = f"Could not construct {cls.__name__} with/without name '{name}'"
+        raise RuntimeError(msg)

python/frontend/commands/bench/__init__.py

@@ -0,0 +1,54 @@
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, ClassVar
+
+if TYPE_CHECKING:
+    import argparse
+
+from python.frontend.commands.base import BaseCommand
+from python.frontend.commands.bench.list import ListCommand
+from python.frontend.commands.bench.model import ModelCommand
+from python.frontend.commands.bench.sweep import SweepCommand
+
+
+class BenchCommand(BaseCommand):
+    """Benchmark JSTprove models with various configurations."""
+
+    name: ClassVar[str] = "bench"
+    aliases: ClassVar[list[str]] = []
+    help: ClassVar[str] = "Benchmark JSTprove models with various configurations."
+
+    SUBCOMMANDS: ClassVar[list[type[BaseCommand]]] = [
+        ListCommand,
+        ModelCommand,
+        SweepCommand,
+    ]
+
+    @classmethod
+    def configure_parser(
+        cls: type[BenchCommand],
+        parser: argparse.ArgumentParser,
+    ) -> None:
+        subparsers = parser.add_subparsers(
+            dest="bench_subcommand",
+            required=True,
+            help="Benchmark subcommands",
+        )
+
+        for subcommand_cls in cls.SUBCOMMANDS:
+            subparser = subparsers.add_parser(
+                subcommand_cls.name,
+                help=subcommand_cls.help,
+                aliases=subcommand_cls.aliases,
+            )
+            subcommand_cls.configure_parser(subparser)
+
+    @classmethod
+    def run(cls: type[BenchCommand], args: argparse.Namespace) -> None:
+        for subcommand_cls in cls.SUBCOMMANDS:
+            if args.bench_subcommand in [subcommand_cls.name, *subcommand_cls.aliases]:
+                subcommand_cls.run(args)
+                return
+
+        msg = f"Unknown bench subcommand: {args.bench_subcommand}"
+        raise ValueError(msg)

python/frontend/commands/bench/list.py

@@ -0,0 +1,42 @@
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, ClassVar
+
+if TYPE_CHECKING:
+    import argparse
+
+from python.frontend.commands.args import ArgSpec
+from python.frontend.commands.base import BaseCommand
+
+LIST_MODELS = ArgSpec(
+    name="list_models",
+    flag="--list-models",
+    help_text="List all available circuit models.",
+    extra_kwargs={"action": "store_true", "default": False},
+)
+
+
+class ListCommand(BaseCommand):
+    """List all available circuit models for benchmarking."""
+
+    name: ClassVar[str] = "list"
+    aliases: ClassVar[list[str]] = []
+    help: ClassVar[str] = "List all available circuit models for benchmarking."
+
+    @classmethod
+    def configure_parser(
+        cls: type[ListCommand],
+        parser: argparse.ArgumentParser,
+    ) -> None:
+        LIST_MODELS.add_to_parser(parser)
+
+    @classmethod
+    def run(cls: type[ListCommand], args: argparse.Namespace) -> None:  # noqa: ARG003
+        from python.core.utils.model_registry import (  # noqa: PLC0415
+            list_available_models,
+        )
+
+        available_models = list_available_models()
+        print("\nAvailable Circuit Models:")  # noqa: T201
+        for model in available_models:
+            print(f"- {model}")  # noqa: T201

python/frontend/commands/bench/model.py

@@ -0,0 +1,172 @@
+from __future__ import annotations
+
+import os
+import sys
+import tempfile
+from pathlib import Path
+from typing import TYPE_CHECKING, ClassVar
+
+if TYPE_CHECKING:
+    import argparse
+
+from python.core.utils.constants import MODEL_SOURCE_CLASS, MODEL_SOURCE_ONNX
+from python.core.utils.helper_functions import (
+    ensure_parent_dir,
+    run_subprocess,
+    to_json,
+)
+from python.frontend.commands.args import ArgSpec
+from python.frontend.commands.base import BaseCommand
+
+SOURCE_CHOICES: tuple[str, ...] = (MODEL_SOURCE_CLASS, MODEL_SOURCE_ONNX)
+
+BENCH_MODEL_PATH = ArgSpec(
+    name="model_path",
+    flag="--model-path",
+    short="-m",
+    help_text="Direct path to ONNX model file to benchmark.",
+)
+
+MODEL = ArgSpec(
+    name="model",
+    flag="--model",
+    help_text=(
+        "Model name(s) from registry to benchmark. "
+        "Use multiple times to test more than one."
+    ),
+    extra_kwargs={"action": "append", "default": None},
+)
+
+SOURCE = ArgSpec(
+    name="source",
+    flag="--source",
+    help_text="Restrict registry models to a specific source: class or onnx.",
+    extra_kwargs={"choices": list(SOURCE_CHOICES), "default": None},
+)
+
+ITERATIONS = ArgSpec(
+    name="iterations",
+    flag="--iterations",
+    help_text="E2E loops per model (default 5)",
+    arg_type=int,
+)
+
+RESULTS = ArgSpec(
+    name="results",
+    flag="--results",
+    help_text="Path to JSONL results (e.g., benchmarking/model_name.jsonl)",
+)
+
+
+class ModelCommand(BaseCommand):
+    """Benchmark specific models from registry or file path."""
+
+    name: ClassVar[str] = "model"
+    aliases: ClassVar[list[str]] = []
+    help: ClassVar[str] = "Benchmark specific models from registry or file path."
+
+    DEFAULT_ITERATIONS: ClassVar[int] = 5
+    SCRIPT_BENCHMARK_RUNNER: ClassVar[str] = "python.scripts.benchmark_runner"
+
+    @classmethod
+    def configure_parser(
+        cls: type[ModelCommand],
+        parser: argparse.ArgumentParser,
+    ) -> None:
+        BENCH_MODEL_PATH.add_to_parser(parser)
+        MODEL.add_to_parser(parser)
+        SOURCE.add_to_parser(parser)
+        ITERATIONS.add_to_parser(parser)
+        RESULTS.add_to_parser(parser)
+
+    @classmethod
+    @BaseCommand.validate_optional_paths(BENCH_MODEL_PATH)
+    def run(cls: type[ModelCommand], args: argparse.Namespace) -> None:
+        if args.model_path:
+            name = Path(args.model_path).stem
+            cls._run_bench_single_model(args, args.model_path, name)
+            return
+
+        if args.model or args.source:
+            cls._run_bench_on_models(args)
+            return
+
+        msg = "Specify --model-path, --model, or --source"
+        raise ValueError(msg)
+
+    @classmethod
+    def _run_bench_on_models(cls: type[ModelCommand], args: argparse.Namespace) -> None:
+        from python.core.utils.model_registry import get_models_to_test  # noqa: PLC0415
+
+        models = get_models_to_test(args.model, args.source or SOURCE_CHOICES[1])
+        if not models:
+            msg = "No models selected for benchmarking."
+            raise ValueError(msg)
+
+        for model_entry in models:
+            instance = model_entry.loader()
+            cls._run_bench_single_model(
+                args,
+                instance.model_file_name,
+                model_entry.name,
+            )
+
+    @classmethod
+    def _generate_model_input(
+        cls: type[ModelCommand],
+        model_path: str,
+        input_file: Path,
+        name: str,
+    ) -> None:
+        instance = BaseCommand._build_circuit()  # noqa: SLF001
+        instance.model_file_name = model_path
+
+        try:
+            instance.load_model(model_path)
+        except Exception as e:
+            msg = f"Failed to load model {model_path}: {e}"
+            raise RuntimeError(msg) from e
+
+        try:
+            inputs = instance.get_inputs()
+            formatted_inputs = instance.format_inputs(inputs)
+            to_json(formatted_inputs, str(input_file))
+        except Exception as e:
+            msg = f"Failed to generate input for {name}: {e}"
+            raise RuntimeError(msg) from e
+
+    @classmethod
+    def _run_bench_single_model(
+        cls: type[ModelCommand],
+        args: argparse.Namespace,
+        model_path: str,
+        name: str,
+    ) -> None:
+        with tempfile.TemporaryDirectory() as tmpdir:
+            tmp_path = Path(tmpdir)
+            input_file = tmp_path / "input.json"
+
+            cls._generate_model_input(model_path, input_file, name)
+
+            iterations = str(args.iterations or cls.DEFAULT_ITERATIONS)
+            results = args.results or f"benchmarking/{name}.jsonl"
+            ensure_parent_dir(results)
+
+            cmd = [
+                sys.executable,
+                "-m",
+                cls.SCRIPT_BENCHMARK_RUNNER,
+                "--model",
+                model_path,
+                "--input",
+                str(input_file),
+                "--iterations",
+                iterations,
+                "--output",
+                results,
+                "--summarize",
+            ]
+            if os.environ.get("JSTPROVE_DEBUG") == "1":
+                print(f"[debug] bench {name} cmd:", " ".join(cmd))  # noqa: T201
+
+            run_subprocess(cmd)