From 45e69798aaa6980f25cf11a9877d03e6dccef967 Mon Sep 17 00:00:00 2001 From: Umut Date: Mon, 26 Jun 2023 11:20:11 +0200 Subject: [PATCH] feat(frontend-python): support ctrl+c during compilation and key generation --- .../lib/Bindings/Python/CompilerAPIModule.cpp | 3 +++ .../concrete/fhe/compilation/keys.py | 13 ++++++++----- .../concrete/fhe/compilation/server.py | 4 ++-- 3 files changed, 13 insertions(+), 7 deletions(-) diff --git a/compilers/concrete-compiler/compiler/lib/Bindings/Python/CompilerAPIModule.cpp b/compilers/concrete-compiler/compiler/lib/Bindings/Python/CompilerAPIModule.cpp index 3d6769444..1d7d64355 100644 --- a/compilers/concrete-compiler/compiler/lib/Bindings/Python/CompilerAPIModule.cpp +++ b/compilers/concrete-compiler/compiler/lib/Bindings/Python/CompilerAPIModule.cpp @@ -118,6 +118,7 @@ void mlir::concretelang::python::populateCompilerAPISubmodule( .def("compile", [](JITSupport_Py &support, std::string mlir_program, CompilationOptions options) { + pybind11::gil_scoped_release release; return jit_compile(support, mlir_program.c_str(), options); }) .def("load_client_parameters", @@ -168,6 +169,7 @@ void mlir::concretelang::python::populateCompilerAPISubmodule( .def("compile", [](LibrarySupport_Py &support, std::string mlir_program, mlir::concretelang::CompilationOptions options) { + pybind11::gil_scoped_release release; return library_compile(support, mlir_program.c_str(), options); }) .def("load_client_parameters", @@ -211,6 +213,7 @@ void mlir::concretelang::python::populateCompilerAPISubmodule( [](clientlib::ClientParameters clientParameters, clientlib::KeySetCache *cache, uint64_t seedMsb, uint64_t seedLsb) { + pybind11::gil_scoped_release release; auto optCache = cache == nullptr ? std::nullopt : std::optional(*cache); diff --git a/frontends/concrete-python/concrete/fhe/compilation/keys.py b/frontends/concrete-python/concrete/fhe/compilation/keys.py index ee2270301..243be1668 100644 --- a/frontends/concrete-python/concrete/fhe/compilation/keys.py +++ b/frontends/concrete-python/concrete/fhe/compilation/keys.py @@ -11,6 +11,7 @@ from typing import Optional, Union from concrete.compiler import ClientSupport, EvaluationKeys, KeySet, KeySetCache from .specs import ClientSpecs +from .utils import interruptable_native_call # pylint: enable=import-error,no-name-in-module @@ -64,11 +65,13 @@ class Keys: seed_msb = (seed >> 64) & ((2**64) - 1) if self._keyset is None or force: - self._keyset = ClientSupport.key_set( - self.client_specs.client_parameters, - self._keyset_cache, - seed_msb, - seed_lsb, + self._keyset = interruptable_native_call( + lambda: ClientSupport.key_set( + self.client_specs.client_parameters, + self._keyset_cache, + seed_msb, + seed_lsb, + ) ) def save(self, location: Union[str, Path]): diff --git a/frontends/concrete-python/concrete/fhe/compilation/server.py b/frontends/concrete-python/concrete/fhe/compilation/server.py index 54d74feb8..5e851b655 100644 --- a/frontends/concrete-python/concrete/fhe/compilation/server.py +++ b/frontends/concrete-python/concrete/fhe/compilation/server.py @@ -147,7 +147,7 @@ class Server: output_dir = None support = JITSupport.new() - compilation_result = support.compile(mlir, options) + compilation_result = interruptable_native_call(lambda: support.compile(mlir, options)) server_lambda = support.load_server_lambda(compilation_result) else: @@ -159,7 +159,7 @@ class Server: support = LibrarySupport.new( str(output_dir_path), generateCppHeader=False, generateStaticLib=False ) - compilation_result = support.compile(mlir, options) + compilation_result = interruptable_native_call(lambda: support.compile(mlir, options)) server_lambda = support.load_server_lambda(compilation_result) client_parameters = support.load_client_parameters(compilation_result)