framework: cache decoded SecretKey instead of decode/encode bytes for every sign (leanEthereum#777)

Author: unnawut

packages/testing/src/consensus_testing/keys.py

@@ -180,7 +180,7 @@ class XmssKeyManager:
 
     - Raw JSON (lightweight hex strings, ~2.7 KB per validator)
     - Deserialized public keys only (avoids the heavy secret key objects)
-    - Advanced secret key state as compact SSZ bytes
+    - Advanced secret key state as live Python objects
     """
 
     __slots__ = (
@@ -267,10 +267,8 @@ def __init__(
         # Populated lazily on first directory scan.
         self._available_indices: set[ValidatorIndex] | None = None
 
-        # Advanced secret key state cached as raw SSZ bytes.
-        # Raw bytes (~2.7 KB each) instead of deserialized objects (~370 MB each)
-        # to avoid holding massive Pydantic model trees in memory.
-        self._secret_state: dict[tuple[ValidatorIndex, KeyRole], bytes] = {}
+        # Advanced secret-key state held as live Python objects.
+        self._secret_state: dict[tuple[ValidatorIndex, KeyRole], SecretKey] = {}
 
     def _scan_indices(self) -> set[ValidatorIndex]:
         """
@@ -413,9 +411,9 @@ def _sign_with_secret(
 
         Memory strategy:
 
-        1. Deserialize the secret key from cached bytes or disk
-        2. Advance and sign (only one full key object in memory)
-        3. Re-serialize to compact bytes (~2.7 KB) for caching
+        1. On cache miss, deserialize the key once from disk
+        2. Advance and sign
+        3. Keep the advanced object for the next sign
 
         Args:
             validator_id: Which validator's key to use.
@@ -428,9 +426,11 @@ def _sign_with_secret(
         """
         cache_key = (validator_id, role)
 
-        # Deserialize the secret key from either the byte cache or disk.
+        # Reuse the cached object directly when present, else decode from disk.
+        # Holding the object avoids the bytes-to-object round-trip on every sign.
+        # That round-trip dominated prod-scheme runtime under the compact-bytes cache.
         if cache_key in self._secret_state:
-            sk = SecretKey.decode_bytes(self._secret_state[cache_key])
+            sk = self._secret_state[cache_key]
         else:
             sk = self._get_secret_key(validator_id, role)
 
@@ -452,9 +452,8 @@ def _sign_with_secret(
         # Produce the signature for the target slot.
         signature = self.scheme.sign(sk, slot, message)
 
-        # Re-serialize the advanced key state to compact bytes for caching.
-        # This drops the full Python object tree from memory immediately.
-        self._secret_state[cache_key] = sk.encode_bytes()
+        # Park the advanced object back in the cache for the next sign.
+        self._secret_state[cache_key] = sk
 
         return signature