github/tfhe-rs
1
1
Fork 0
mirror of https://github.com/zama-ai/tfhe-rs.git synced 2026-01-09 14:47:56 -05:00
Code Issues Packages Projects Releases Wiki Activity

chore: pre-generate keyswitching keys for shortint tests

Browse Source 
- we run in a cross process race condition which fucks up the key file
- no rust crate seems to help and linux locks are just a fucking mess
- also avoid truncating file when we are going to write to it, get a lock
first
This commit is contained in:
Arthur Meyre
2025-03-07 11:29:42 +01:00
parent 396f30ff5d
commit b0d7bb9f95
2 changed files with 31 additions and 16 deletions
Download Patch File Download Diff File Expand all files Collapse all files

25
tfhe/examples/utilities/generates_test_keys.rs
View File

@@ -27,6 +27,16 @@ use tfhe::shortint::parameters::{
};
use tfhe::shortint::MultiBitPBSParameters;
const KSK_PARAMS: [(
ClassicPBSParameters,
ClassicPBSParameters,
ShortintKeySwitchingParameters,
); 1] = [(
V1_0_PARAM_MESSAGE_1_CARRY_1_KS_PBS_GAUSSIAN_2M128,
V1_0_PARAM_MESSAGE_2_CARRY_2_KS_PBS_GAUSSIAN_2M128,
V1_0_PARAM_KEYSWITCH_1_1_KS_PBS_TO_2_2_KS_PBS_GAUSSIAN_2M128,
)];
fn client_server_keys() {
let matches = Command::new("test key gen")
.arg(
@@ -45,6 +55,9 @@ fn client_server_keys() {
)
.get_matches();
// Always generate those as they may be used in the different cases
generate_ksk_keys(&KSK_PARAMS);
// If set using the command line flag "--ladner-fischer" this algorithm will be used in
// additions
let multi_bit_only: bool = matches.get_flag("multi_bit_only");
@@ -81,18 +94,6 @@ fn client_server_keys() {
generate_pbs_multi_bit_keys(&MULTI_BIT_PARAMS);
}
const KSK_PARAMS: [(
ClassicPBSParameters,
ClassicPBSParameters,
ShortintKeySwitchingParameters,
); 1] = [(
V1_0_PARAM_MESSAGE_1_CARRY_1_KS_PBS_GAUSSIAN_2M128,
V1_0_PARAM_MESSAGE_2_CARRY_2_KS_PBS_GAUSSIAN_2M128,
V1_0_PARAM_KEYSWITCH_1_1_KS_PBS_TO_2_2_KS_PBS_GAUSSIAN_2M128,
)];
generate_ksk_keys(&KSK_PARAMS);
#[cfg(feature = "experimental")]
{
const WOPBS_PARAMS: [(ClassicPBSParameters, WopbsParameters); 1] = [(

22
tfhe/src/keycache/mod.rs
View File

@@ -7,7 +7,7 @@ pub mod utils {
use fs2::FileExt;
use serde::de::DeserializeOwned;
use serde::Serialize;
use std::fs::File;
use std::fs::{File, OpenOptions};
use std::io::{BufReader, BufWriter};
use std::ops::Deref;
use std::path::PathBuf;
@@ -110,7 +110,10 @@ pub mod utils {
if path_buf.exists() {
let file = File::open(&path_buf).unwrap();
// Lock for reading
// TODO Manage file locking for inter process stuff, unfortunately linux locks are a
// mess and nothing seems to work
//
// Lock for reading this only works for our process not inter process
fs2::FileExt::lock_shared(&file).unwrap();
let file_reader = BufReader::new(file);
bincode::deserialize_from::<_, (P, K)>(file_reader)
@@ -128,9 +131,20 @@ pub mod utils {
path_buf.push(param.name());
path_buf.set_extension("bin");
let file = File::create(&path_buf).unwrap();
// Lock for writing
let file = OpenOptions::new()
.create(true)
.write(true)
.truncate(false)
.open(&path_buf)
.unwrap();
// TODO Manage file locking for inter process stuff, unfortunately linux locks are a
// mess and nothing seems to work
//
// Lock for writing this only works for our process not inter process
file.lock_exclusive().unwrap();
// Truncate manually
file.set_len(0).unwrap();
let file_writer = BufWriter::new(file);
bincode::serialize_into(file_writer, &(param, key)).unwrap();