Performance: Streams the classifier pickle file during save to file (#14121)

src/documents/classifier.py

@@ -13,6 +13,9 @@ if TYPE_CHECKING:
     from collections.abc import Callable
     from collections.abc import Iterator
     from datetime import datetime
+    from types import TracebackType
+    from typing import BinaryIO
+    from typing import Self
 
     from numpy import ndarray
 
@@ -68,6 +71,49 @@ RE_DIGIT = re.compile(r"\d")
 RE_WORD = re.compile(r"\b[\w]+\b")  # words that may contain digits
 
 
+class _SignedFileWriter:
+    """
+    Atomically writes a file made of an HMAC signature followed by the data,
+    signing the data as it streams to disk rather than holding it in memory.
+
+    The signature is only known once everything is written, so its space is
+    reserved at the start of the file and filled in on exit. The target is only
+    replaced once the file is complete; on error the partial file is removed.
+    """
+
+    def __init__(self, target: Path, mac: hmac.HMAC) -> None:
+        self._target = target
+        self._temp = target.with_name(f"{target.name}.part")
+        self._mac = mac
+        self._file: BinaryIO
+
+    def __enter__(self) -> Self:
+        self._file = self._temp.open("wb")
+        self._file.write(bytes(self._mac.digest_size))
+        return self
+
+    def write(self, data: bytes | memoryview) -> int:
+        self._mac.update(data)
+        return self._file.write(data)
+
+    def __exit__(
+        self,
+        exc_type: type[BaseException] | None,
+        exc_value: BaseException | None,
+        traceback: TracebackType | None,
+    ) -> None:
+        try:
+            with self._file:
+                if exc_type is None:
+                    self._file.seek(0)
+                    self._file.write(self._mac.digest())
+            if exc_type is None:
+                self._temp.rename(self._target)
+        finally:
+            # A no-op after a successful rename, otherwise removes the partial file
+            self._temp.unlink(missing_ok=True)
+
+
 class IncompatibleClassifierVersionError(Exception):
     def __init__(self, message: str, *args: object) -> None:
         self.message: str = message
@@ -159,13 +205,15 @@ class DocumentClassifier:
             pickle.dumps(self.data_vectorizer),
         ).hexdigest()
 
+    @staticmethod
+    def _new_hmac() -> hmac.HMAC:
+        return hmac.new(settings.SECRET_KEY.encode(), digestmod=sha256)
+
     @staticmethod
     def _compute_hmac(data: bytes | memoryview) -> bytes:
-        return hmac.new(
-            settings.SECRET_KEY.encode(),
-            data,
-            sha256,
-        ).digest()
+        mac = DocumentClassifier._new_hmac()
+        mac.update(data)
+        return mac.digest()
 
     def load(self) -> None:
         from sklearn.exceptions import InconsistentVersionWarning
@@ -226,29 +274,24 @@ class DocumentClassifier:
                 raise IncompatibleClassifierVersionError("sklearn version update")
 
     def save(self) -> None:
-        target_file: Path = settings.MODEL_FILE
-        target_file_temp: Path = target_file.with_suffix(".pickle.part")
-
-        data = pickle.dumps(
-            (
-                self.FORMAT_VERSION,
-                self.last_doc_change_time,
-                self.last_auto_type_hash,
-                self.data_vectorizer,
-                self.tags_binarizer,
-                self.tags_classifier,
-                self.correspondent_classifier,
-                self.document_type_classifier,
-                self.storage_path_classifier,
-            ),
-        )
-
-        signature = self._compute_hmac(data)
-
-        with target_file_temp.open("wb") as f:
-            f.write(signature + data)
-
-        target_file_temp.rename(target_file)
+        # Stream to disk instead of building the payload in memory. Protocol 5+
+        # pickles numpy arrays without copying them (the default is 4 before 3.14).
+        with _SignedFileWriter(settings.MODEL_FILE, self._new_hmac()) as f:
+            pickle.dump(
+                (
+                    self.FORMAT_VERSION,
+                    self.last_doc_change_time,
+                    self.last_auto_type_hash,
+                    self.data_vectorizer,
+                    self.tags_binarizer,
+                    self.tags_classifier,
+                    self.correspondent_classifier,
+                    self.document_type_classifier,
+                    self.storage_path_classifier,
+                ),
+                f,
+                protocol=pickle.HIGHEST_PROTOCOL,
+            )
 
     def train(
         self,