fix: use env variable

zama-ai · Jun 13, 2024 · ffbea40 · ffbea40
1 parent cabc34b
commit ffbea40
Show file tree

Hide file tree

Showing 6 changed files with 11 additions and 7 deletions.
diff --git a/conftest.py b/conftest.py
@@ -155,7 +155,6 @@ def default_configuration():
         fhe_simulation=False,
         fhe_execution=True,
         compress_input_ciphertexts=os.environ.get("USE_INPUT_COMPRESSION", "1") == "1",
-        compress_evaluation_keys=True,
     )
 
 
@@ -174,7 +173,6 @@ def simulation_configuration():
         fhe_simulation=True,
         fhe_execution=False,
         compress_input_ciphertexts=os.environ.get("USE_INPUT_COMPRESSION", "1") == "1",
-        compress_eval_keys=True,
     )
 
 

diff --git a/docker/release_resources/sanity_check.py b/docker/release_resources/sanity_check.py
@@ -76,7 +76,6 @@ def ml_check(args, keyring_dir_as_str):
         use_insecure_key_cache=is_fast,
         insecure_key_cache_location=keyring_dir_as_str,
         compress_input_ciphertexts=True,
-        compress_eval_keys=True,
     )
 
     # We first compile the model with some data, here the training set
@@ -123,7 +122,7 @@ def function_to_compile(x):
         use_insecure_key_cache=is_fast,
         insecure_key_cache_location=keyring_dir_as_str,
         compress_input_ciphertexts=True,
-        compress_eval_keys=True,
+        compress_evaluation_keys=True,
     )
 
     print("Compiling...")

diff --git a/src/concrete/ml/common/utils.py b/src/concrete/ml/common/utils.py
@@ -53,6 +53,11 @@
 # However, for internal testing purposes, we retain the capability to disable this feature
 os.environ["USE_INPUT_COMPRESSION"] = os.environ.get("USE_INPUT_COMPRESSION", "1")
 
+# Enable PBS evaluation key compression (~4x size reduction)
+# Note: This setting is fixed and cannot be altered by users
+# However, for internal testing purposes, we retain the capability to disable this feature
+os.environ["USE_KEY_COMPRESSION"] = os.environ.get("USE_KEY_COMPRESSION", "1")
+
 
 class FheMode(str, enum.Enum):
     """Enum representing the execution mode.

diff --git a/src/concrete/ml/quantization/quantized_module.py b/src/concrete/ml/quantization/quantized_module.py
@@ -876,6 +876,7 @@ def compile(
 
         # Enable input ciphertext compression
         enable_input_compression = os.environ.get("USE_INPUT_COMPRESSION", "1") == "1"
+        enable_key_compression = os.environ.get("USE_KEY_COMPRESSION", "1") == "1"
 
         self.fhe_circuit = compiler.compile(
             inputset,
@@ -889,7 +890,7 @@ def compile(
             fhe_simulation=False,
             fhe_execution=True,
             compress_input_ciphertexts=enable_input_compression,
-            compress_evaluation_keys=True,
+            compress_evaluation_keys=enable_key_compression,
         )
 
         self._is_compiled = True

diff --git a/src/concrete/ml/sklearn/base.py b/src/concrete/ml/sklearn/base.py
@@ -569,6 +569,8 @@ def compile(
 
         # Enable input ciphertext compression
         enable_input_compression = os.environ.get("USE_INPUT_COMPRESSION", "1") == "1"
+        # Enable evaluation key compression
+        enable_key_compression = os.environ.get("USE_KEY_COMPRESSION", "1") == "1"
 
         self.fhe_circuit_ = module_to_compile.compile(
             inputset,
@@ -582,7 +584,7 @@ def compile(
             fhe_simulation=False,
             fhe_execution=True,
             compress_input_ciphertexts=enable_input_compression,
-            compress_evaluation_keys=True,
+            compress_evaluation_keys=enable_key_compression,
         )
 
         self._is_compiled = True

diff --git a/tests/torch/test_hybrid_converter.py b/tests/torch/test_hybrid_converter.py
@@ -41,7 +41,6 @@ def run_hybrid_llm_test(
     configuration = Configuration(
         single_precision=False,
         compress_input_ciphertexts=True,
-        compress_evaluation_keys=True,
     )
 
     # Create a hybrid model