Switching Context of OpenFHE In and Out of memory in a large scale application

I am working on a project with OpenFHE that requires a large amount of rotation keys so I am looking for ways to reduce my memory usage by switching context of my program in and out of memory. I am using the the following steps:

  1. Generate Context with application public and private keys and save in disk.
  2. Generate rotation keys for different sections of the program.
  3. Clear Context from memory to remove all keys and data not required
  4. Load Context into memory
  5. Load just the rotation keys needed for each section. After computing the section, I can clear these keys and context from memory as well. Load the context and keys for the next section and so forth.

For some weird reason, when I clear the context from memory and load it back from disk, I can not use it.

I get the following error

pke/cryptocontext.h:l.344:ValidateKey<std::shared_ptr<lbcrypto::PublicKeyImpl<lbcrypto::DCRTPolyImpl<bigintdyn::mubintvec<bigintdyn::ubint<long unsigned int> > > > > >(): Key was not generated with the same crypto context
Aborted (core dumped)

This suggest that my public key was not generated with the same context which is not true since I do not regenerate the context anywhere in the program.

My implementation looks something like this:


CryptoContext<DCRTPoly> context; // public variable and shared across multiple controllers. 
int main(int argc, char *argv[]) {
       -----------------
    FHEController fheController(context); // creating my fhe controller  
	fheController.generate_context(15, 16, 12); // generating fhe context with public and private keys and storing it. 
	-----------------
    fheController.generate_bootstrapping_and_rotation_keys("section1_keys.bin");
    fheController.clear_bootstrapping_and_rotation_keys();
    -----------------
    fheController.clear_context(num_slots);
    fheController.load_context(true);
    -----------------
    Ctext enData = fheController.encrypt_packedData(data);
}

From the code snippet above, if I remove the following lines, I won’t have any issues but they are needed for scalability and for the system to work.

 fheController.clear_context(num_slots);
 fheController.load_context(true);

I believe my issue should be coming from either how I generate, serialization and store the context or how I read the context back into my program. Below are the functions to do all of this;

void FHEController::generate_context(bool serialize, int numSlots, int ringDim, int mlevelBootstrap) {
    CCParams<CryptoContextCKKSRNS> parameters;

    num_slots = 1 << numSlots;
    int dcrtBits               = 47;
    int firstMod               = 52;

    auto secretKeyDist = SPARSE_TERNARY;
    parameters.SetSecretKeyDist(SPARSE_TERNARY);
    parameters.SetSecurityLevel(lbcrypto::HEStd_NotSet);
    parameters.SetNumLargeDigits(3);
    parameters.SetRingDim(1 << ringDim);
    parameters.SetBatchSize(num_slots);
    level_budget = {4, 4}; 
    ScalingTechnique rescaleTech = FLEXIBLEAUTO; 
    parameters.SetScalingModSize(dcrtBits);
    parameters.SetFirstModSize(firstMod);
    parameters.SetScalingTechnique(rescaleTech);

    uint32_t levelsAvailableAfterBootstrap = mlevelBootstrap;

    circuit_depth = levelsAvailableAfterBootstrap + FHECKKSRNS::GetBootstrapDepth(level_budget, secretKeyDist);

    parameters.SetMultiplicativeDepth(circuit_depth);
    context = GenCryptoContext(parameters);

    context->Enable(PKE);
    context->Enable(KEYSWITCH);
    context->Enable(LEVELEDSHE);
    context->Enable(ADVANCEDSHE);
    context->Enable(FHE);

    keyPair = context->KeyGen();
    context->EvalMultKeyGen(keyPair.secretKey);
    context->EvalSumKeyGen(keyPair.secretKey);

    std::vector<uint32_t> bsgsDim     = {0, 0};
    numSlots = context->GetRingDimension() / 2;
    context->EvalBootstrapSetup(level_budget, bsgsDim, numSlots);
    context->EvalBootstrapKeyGen(keyPair.secretKey, numSlots);

    serialize_context();
}

void FHEController::serialize_context(){
 if (fs::exists(sskeys_folder) && fs::is_directory(sskeys_folder)) {
        for (const auto& entry : fs::directory_iterator(sskeys_folder)) {
            if (fs::is_regular_file(entry)) {
                cout << "Deleting file: " << entry.path().filename() << "\n";
                fs::remove(entry);
            }
        }
        cout << "All files deleted from " << sskeys_folder << endl;;
    } 
    else {
        if (!fs::create_directory(sskeys_folder)) {context->EvalMultKeyGen(keyPair.secretKey);
            std::cerr << "Failed to create directory: " << sskeys_folder << endl;
            return;
        }
        std::cout << sskeys_folder << " created to store keys " << endl; 
    }

    ofstream multKeyFile(sskeys_folder + "/mult-keys.txt", ios::out | ios::binary);
    if (multKeyFile.is_open()) {
        if (!context->SerializeEvalMultKey(multKeyFile, SerType::BINARY)) {
            cerr << "Error writing eval mult keys" << std::endl;
            exit(1);
        }
        cout << "Relinearization Keys have been serialized" << std::endl;
        multKeyFile.close();
    }
    else {
        cerr << "Error serializing EvalMult keys in \"" << sskeys_folder + "/mult-keys.txt" << "\"" << endl;
        exit(1);
    }

    if (!Serial::SerializeToFile(sskeys_folder + "/crypto-context.txt", context, SerType::BINARY)) {
        cerr << "Error writing serialization of the crypto context to crypto-context.txt" << endl;
    } else {
        cout << "Crypto Context have been serialized" << std::endl;
    }

    if (!Serial::SerializeToFile(sskeys_folder + "/public-key.txt", keyPair.publicKey, SerType::BINARY)) {
        cerr << "Error writing serialization of public key to public-key.txt" << endl;
    } else {
        cout << "Public Key has been serialized" << std::endl;
    }

    if (!Serial::SerializeToFile(sskeys_folder + "/secret-key.txt", keyPair.secretKey, SerType::BINARY)) {
        cerr << "Error writing serialization of public key to secret-key.txt" << endl;
    } else {
        cout << "Secret Key has been serialized" << std::endl;
    }
}

void FHEController::load_context(bool verbose) {
    context->ClearEvalMultKeys();
    context->ClearEvalAutomorphismKeys();
    CryptoContextFactory<lbcrypto::DCRTPoly>::ReleaseAllContexts();

    if (verbose) cout << "Reading serialized context..." << endl;

    if (!Serial::DeserializeFromFile(sskeys_folder + "/crypto-context.txt", context, SerType::BINARY)) {
        cerr << "I cannot read serialized data from: " << sskeys_folder + "/crypto-context.txt" << endl;
        exit(1);
    }

    PublicKey<DCRTPoly> clientPublicKey;
    if (!Serial::DeserializeFromFile(sskeys_folder + "/public-key.txt", clientPublicKey, SerType::BINARY)) {
        cerr << "I cannot read serialized data from public-key.txt" << endl;
        exit(1);
    }

    PrivateKey<DCRTPoly> serverSecretKey;
    if (!Serial::DeserializeFromFile(sskeys_folder + "/secret-key.txt", serverSecretKey, SerType::BINARY)) {
        cerr << "I cannot read serialized data from public-key.txt" << endl;
        exit(1);
    }

    keyPair.publicKey = clientPublicKey;
    keyPair.secretKey = serverSecretKey;

    std::ifstream multKeyIStream(sskeys_folder + "/mult-keys.txt", ios::in | ios::binary);
    if (!multKeyIStream.is_open()) {
        cerr << "Cannot read serialization from " << "mult-keys.txt" << endl;
        exit(1);
    }
    if (!context->DeserializeEvalMultKey(multKeyIStream, SerType::BINARY)) {
        cerr << "Could not deserialize eval mult key file" << endl;
        exit(1);
    }

    uint32_t approxBootstrapDepth = 4 + 4;
   	uint32_t levelsUsedBeforeBootstrap = 12;
    circuit_depth = levelsUsedBeforeBootstrap + 2 + FHECKKSRNS::GetBootstrapDepth(approxBootstrapDepth, level_budget, SPARSE_TERNARY);
    if (verbose) cout << "Circuit depth: " << circuit_depth << ", available multiplications: " << levelsUsedBeforeBootstrap - 2 << endl;
}

Ctext FHEController::encrypt_packedImage(vector<double> packedVector, int level, int num_slots) {
    Ptext plaintext = context->MakeCKKSPackedPlaintext(packedVector, 1, level, nullptr, num_slots);
    plaintext->SetLength(num_slots);
    auto encryptImage = context->Encrypt(plaintext, keyPair.publicKey);
    return encryptImage;
}

I have been stocked on this issue for days and I am out of ideas. Any feedback, pointers, ideas etc will be greatly appreciated.

Could you check the following example: openfhe-development/src/pke/examples/simple-real-numbers-serial.cpp at v1.2.4 · openfheorg/openfhe-development · GitHub? It includes serialization and deserialization, along with clearing the keys in between. You can compare this example with your code to see what is missing.

In particular, please make sure the following header files are included

// header files needed for serialization
#include "ciphertext-ser.h"
#include "cryptocontext-ser.h"
#include "key/key-ser.h"
#include "scheme/ckksrns/ckksrns-ser.h"