Adding specific error messaging and enforced maximum duration to WASM calls (#62)

Preventing infinite looping of WASM calls and providing additional details to users when WASM errors occur. 

---------

Co-authored-by: Ciarán Ryan-Anderson <70174051+qciaran@users.noreply.github.com>

Author: nealerickson-qtm

.github/workflows/python-app.yml

@@ -32,7 +32,7 @@ jobs:
     - name: Pre-commit checks
       run: |
         python -m pip install pre-commit
-        pre-commit run --all-files
+        pre-commit run --all-files --show-diff-on-failure
     - name: Test with pytest
       run: |
         pip install pytest-cov

.pre-commit-config.yaml

@@ -1,6 +1,6 @@
 repos:
   - repo: https://github.com/pre-commit/pre-commit-hooks
-    rev: v4.5.0
+    rev: v4.6.0
     hooks:
       - id: trailing-whitespace
         exclude: ^docs/reference/_autosummary/
@@ -15,19 +15,23 @@ repos:
       - id: debug-statements
 
   - repo: https://github.com/crate-ci/typos
-    rev: v1.19.0
+    rev: v1.21.0
     hooks:
       - id: typos
         args: []
+        exclude: |
+            (?x)^(
+                python/pecos/simulators/cuquantum_old/.*|
+            )$
 
   - repo: https://github.com/astral-sh/ruff-pre-commit
-    rev: v0.3.3
+    rev: v0.4.3
     hooks:
       - id: ruff
         args: [--fix, --exit-non-zero-on-fix]
 
   - repo: https://github.com/psf/black
-    rev: 24.3.0
+    rev: 24.4.2
     hooks:
       - id: black
         exclude: |

.typos.toml

@@ -3,3 +3,5 @@ ba = "ba"
 datas = "datas"
 ket = "ket"
 wqs = "wqs"
+thr = "thr"
+IY = "IY"

docs/api_guide/quantum_circuits.rst

@@ -167,7 +167,7 @@ A ``tick`` keyword can be used to specify which tick the gate is discarded from.
 Retrieving Information
 ----------------------
 
-Next, how to retrieve information from a ``QuantumCircuit`` will be dicuss, for example, through attributes or for
+Next, how to retrieve information from a ``QuantumCircuit`` will be discussed, for example, through attributes or for
 loops.
 
 Number of Ticks

pyproject.toml

@@ -15,7 +15,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "quantum-pecos"
-version = "0.5.0.dev11"
+version = "0.6.0.dev1"
 authors = [
     {name = "The PECOS Developers"},
 ]

python/pecos/engines/cvm/wasm.py

@@ -95,7 +95,7 @@ def eval_cfunc(runner, params, output):
         ccop_type = runner.circuit.metadata["ccop_type"]
 
         if ccop is None:
-            msg = "Wasm not supplied but requested!"
+            msg = f"Wasm ({ccop_type}) function not found: {func} with args: {args}"
             raise MissingCCOPError(msg) from AttributeError
 
         msg = f"Classical coprocessor object not assigned or missing exec method. Wasm-type = {ccop_type}"

python/pecos/engines/cvm/wasm_vms/pywasm.py

@@ -42,4 +42,7 @@ def exec(self, func, args, debug=False):
             args = [int(b) for _, b in args]
             return self.p.exec(func, args)
 
+        def teardown(self):
+            pass  # Only needed for wasmtime
+
     return PywasmReader(p)

python/pecos/engines/cvm/wasm_vms/pywasm3.py

@@ -28,4 +28,7 @@ def exec(self, func, args, debug=False):
             args = [int(b) for _, b in args]
             return self.rt.find_function(func)(*args)
 
+        def teardown(self):
+            pass  # Only needed for wasmtime
+
     return Reader(rt)

python/pecos/engines/cvm/wasm_vms/wasmer.py

@@ -72,4 +72,7 @@ def exec(self, func_name, args, debug=False):
                 args = [int(b) for _, b in args]
                 return method(*args)
 
+        def teardown(self):
+            pass  # Only needed for wasmtime
+
     return WasmerInstance(path, compiler)

python/pecos/engines/cvm/wasm_vms/wasmtime.py

@@ -41,4 +41,7 @@ def exec(self, func_name, args, debug=False):
                 args = [int(b) for _, b in args]
                 return self.wasmtime.exec(func_name, args)
 
+        def teardown(self):
+            self.wasmtime.teardown()
+
     return WASM(path)

python/pecos/engines/hybrid_engine_old.py

@@ -124,6 +124,9 @@ def run(
         if output_export:
             output = output_export
 
+        if self.ccop:
+            self.ccop.teardown()  # Tear down WASM execution context
+
         return output, error_circuits
 
     def run_circuit(self, state, output, output_export, circuit, error_gen, removed_locations=None):

python/pecos/errors.py

@@ -17,5 +17,13 @@ class NotSupportedGateError(PECOSError):
     """Indicates a gate not supported by a simulator."""
 
 
-class MissingCCOPError(PECOSError):
+class WasmError(PECOSError):
+    """Base WASM-related exception type"""
+
+
+class MissingCCOPError(WasmError):
     """Indicates missing a classical function library."""
+
+
+class WasmRuntimeError(WasmError):
+    """Indicates a runtime WASM error."""

python/pecos/foreign_objects/wasm_execution_timer_thread.py

@@ -0,0 +1,16 @@
+from threading import Event, Thread
+
+# These values multiplied should equal the intended maximum execution time
+WASM_EXECUTION_TICK_LENGTH_S: float = 0.25
+WASM_EXECUTION_MAX_TICKS: int = 4
+
+
+class WasmExecutionTimerThread(Thread):
+    def __init__(self, stop_event: Event, func) -> None:
+        Thread.__init__(self, daemon=True)
+        self._stop_event = stop_event
+        self._func = func
+
+    def run(self):
+        while not self._stop_event.wait(WASM_EXECUTION_TICK_LENGTH_S):
+            self._func()

python/pecos/foreign_objects/wasmer.py

@@ -17,6 +17,7 @@
 from wasmer import FunctionType, Instance, Module, Store, engine
 from wasmer_compiler_cranelift import Compiler as Cranelift
 
+from pecos.errors import MissingCCOPError, WasmRuntimeError
 from pecos.foreign_objects.foreign_object_abc import ForeignObject
 
 if TYPE_CHECKING:
@@ -88,12 +89,21 @@ def get_funcs(self) -> list[str]:
         return self.func_names
 
     def exec(self, func_name: str, args: Sequence) -> tuple:
-        func = getattr(self.instance.exports, func_name)
+        try:
+            func = getattr(self.instance.exports, func_name)
+        except AttributeError as e:
+            message = f"Func {func_name} not found in WASM"
+            raise MissingCCOPError(message) from e
+
         params = func.type.params
         if len(args) != len(params):
             msg = f"Wasmer function `{func_name}` takes {len(params)} args and {len(args)} were given!"
-            raise TypeError(msg)
-        return func(*args)
+            raise WasmRuntimeError(msg)
+
+        try:
+            return func(*args)
+        except Exception as ex:
+            raise WasmRuntimeError(ex.args[0]) from ex
 
     def to_dict(self) -> dict:
         return {"fobj_class": WasmerObj, "wasm_bytes": self.wasm_bytes}

python/pecos/foreign_objects/wasmtime.py

@@ -12,11 +12,18 @@
 from __future__ import annotations
 
 from pathlib import Path
+from threading import Event
 from typing import TYPE_CHECKING
 
-from wasmtime import FuncType, Instance, Module, Store
+from wasmtime import Config, Engine, FuncType, Instance, Module, Store, Trap, TrapCode
 
+from pecos.errors import MissingCCOPError, WasmRuntimeError
 from pecos.foreign_objects.foreign_object_abc import ForeignObject
+from pecos.foreign_objects.wasm_execution_timer_thread import (
+    WASM_EXECUTION_MAX_TICKS,
+    WASM_EXECUTION_TICK_LENGTH_S,
+    WasmExecutionTimerThread,
+)
 
 if TYPE_CHECKING:
     from collections.abc import Sequence
@@ -65,8 +72,14 @@ def new_instance(self) -> None:
         self.instance = Instance(self.store, self.module, [])
 
     def spin_up_wasm(self) -> None:
-        self.store = Store()
+        config = Config()
+        config.epoch_interruption = True
+        engine = Engine(config)
+        self.store = Store(engine)
         self.module = Module(self.store.engine, self.wasm_bytes)
+        self.stop_flag = Event()
+        self.inc_thread_handle = WasmExecutionTimerThread(self.stop_flag, self._increment_engine)
+        self.inc_thread_handle.start()
         self.new_instance()
 
     def get_funcs(self) -> list[str]:
@@ -80,9 +93,42 @@ def get_funcs(self) -> list[str]:
 
         return self.func_names
 
+    def _increment_engine(self):
+        self.store.engine.increment_epoch()
+
     def exec(self, func_name: str, args: Sequence) -> tuple:
-        func = self.instance.exports(self.store)[func_name]
-        return func(self.store, *args)
+        try:
+            func = self.instance.exports(self.store)[func_name]
+        except KeyError as e:
+            message = f"No method found with name {func_name} in WASM"
+            raise MissingCCOPError(message) from e
+
+        try:
+            self.store.engine.increment_epoch()
+            self.store.set_epoch_deadline(WASM_EXECUTION_MAX_TICKS)
+            output = func(self.store, *args)
+            return output  # noqa: TRY300
+        except Trap as t:
+            if t.trap_code is TrapCode.INTERRUPT:
+                message = (
+                    f"WASM error: WASM failed during run-time. Execution time of "
+                    f"function '{func_name}' exceeded maximum "
+                    f"{WASM_EXECUTION_MAX_TICKS * WASM_EXECUTION_TICK_LENGTH_S}s"
+                )
+            else:
+                message = (
+                    f"WASM error: WASM failed during run-time. Execution of "
+                    f"function '{func_name}' resulted in {t.trap_code}\n"
+                    f"{t.message}"
+                )
+            raise WasmRuntimeError(message) from t
+        except Exception as e:
+            message = f"Error during execution of function '{func_name}' with args {args}"
+            raise WasmRuntimeError(message) from e
+
+    def teardown(self) -> None:
+        self.stop_flag.set()
+        self.inc_thread_handle.join()
 
     def to_dict(self) -> dict:
         return {"fobj_class": WasmtimeObj, "wasm_bytes": self.wasm_bytes}

ruff.toml

@@ -112,6 +112,9 @@ ignore = [
     # TODO: Remove to improve error handling...
     "TRY002", # Custom exceptions
     "TRY003", # Avoid specifying long messages outside the exception class
+
+    # UP031 should be enabled after fixing all suggestions in a separate PR
+    "UP031",
 ]
 
 [lint.per-file-ignores]

tests/integration/state_sim_tests/test_cointoss.py

@@ -1,4 +1,4 @@
-# Copyright 2023 The PECOS Developers
+# Copyright 2024 The PECOS Developers
 #
 # Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except in compliance with
 # the License.You may obtain a copy of the License at

tests/integration/state_sim_tests/test_statevec.py

@@ -1,4 +1,4 @@
-# Copyright 2023 The PECOS Developers
+# Copyright 2024 The PECOS Developers
 #
 # Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except in compliance with
 # the License.You may obtain a copy of the License at