diff --git a/compilers/concrete-compiler/compiler/lib/Bindings/Python/CompilerAPIModule.cpp b/compilers/concrete-compiler/compiler/lib/Bindings/Python/CompilerAPIModule.cpp index 3fe1c05b1..d69942796 100644 --- a/compilers/concrete-compiler/compiler/lib/Bindings/Python/CompilerAPIModule.cpp +++ b/compilers/concrete-compiler/compiler/lib/Bindings/Python/CompilerAPIModule.cpp @@ -1675,6 +1675,21 @@ void mlir::concretelang::python::populateCompilerAPISubmodule( return pybind11::bytes(keySetSerialize(keySet)); }, "Serialize a Keyset to bytes.") + .def( + "serialize_to_file", + [](Keyset &keySet, const std::string path) { + std::ofstream ofs; + ofs.open(path); + if (!ofs.good()) { + throw std::runtime_error("Failed to open keyset file " + path); + } + auto keysetProto = keySet.toProto(); + auto maybeBuffer = keysetProto.writeBinaryToOstream(ofs); + if (maybeBuffer.has_failure()) { + throw std::runtime_error("Failed to serialize keys."); + } + }, + "Serialize a Keyset to bytes.") .def( "serialize_lwe_secret_key_as_glwe", [](Keyset &keyset, size_t keyIndex, size_t glwe_dimension, diff --git a/frontends/concrete-python/concrete/fhe/compilation/keys.py b/frontends/concrete-python/concrete/fhe/compilation/keys.py index 76356ca7c..9ef2a698c 100644 --- a/frontends/concrete-python/concrete/fhe/compilation/keys.py +++ b/frontends/concrete-python/concrete/fhe/compilation/keys.py @@ -115,7 +115,7 @@ def save(self, location: Union[str, Path]): message = f"Unable to save keys to {location} because it already exists" raise ValueError(message) - location.write_bytes(self.serialize()) + self.serialize_to_file(location) def load(self, location: Union[str, Path]): """ @@ -171,6 +171,8 @@ def serialize(self) -> bytes: Serialize keys into bytes. Serialized keys are not encrypted, so be careful how you store/transfer them! + `serialize_to_file` is supposed to be more performant as it avoid copying the buffer + between the Compiler and the Frontend. Returns: bytes: @@ -184,6 +186,23 @@ def serialize(self) -> bytes: serialized_keyset = self._keyset.serialize() return serialized_keyset + def serialize_to_file(self, path: Path): + """ + Serialize keys into a file. + + Serialized keys are not encrypted, so be careful how you store/transfer them! + This is supposed to be more performant than `serialize` as it avoid copying the buffer + between the Compiler and the Frontend. + + Args: + path (Path): where to save serialized keys + """ + if self._keyset is None: + message = "Keys cannot be serialized before they are generated" + raise RuntimeError(message) + + self._keyset.serialize_to_file(str(path)) + @staticmethod def deserialize(serialized_keys: Union[Path, bytes]) -> "Keys": """ diff --git a/frontends/concrete-python/tests/compilation/test_keys.py b/frontends/concrete-python/tests/compilation/test_keys.py index 94d56ef85..2a2486807 100644 --- a/frontends/concrete-python/tests/compilation/test_keys.py +++ b/frontends/concrete-python/tests/compilation/test_keys.py @@ -175,6 +175,13 @@ def f(x): expected_message = "Keys cannot be serialized before they are generated" helpers.check_str(expected_message, str(excinfo.value)) + with pytest.raises(RuntimeError) as excinfo: + # path doesn't matter as it will fail + circuit.keys.serialize_to_file(Path("_keys_file")) + + expected_message = "Keys cannot be serialized before they are generated" + helpers.check_str(expected_message, str(excinfo.value)) + def test_keys_generate_manual_seed(helpers): """