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,

src/documents/tests/test_classifier.py

@@ -1,5 +1,8 @@
+import pickle
 import re
 import warnings
+from datetime import UTC
+from datetime import datetime
 from pathlib import Path
 from unittest import mock
 
@@ -10,6 +13,7 @@ from django.db import connection
 from django.test import TestCase
 from django.test import override_settings
 from django.test.utils import CaptureQueriesContext
+from pytest_django.fixtures import Settings
 from pytest_mock import MockerFixture
 
 from documents.classifier import ClassifierModelCorruptError
@@ -914,6 +918,85 @@ class TestClassifier(DirectoriesMixin, TestCase):
             load_classifier(raise_exception=True)
 
 
+class TestClassifierSave:
+    @pytest.fixture
+    def model_file(self, tmp_path: Path, settings: Settings) -> Path:
+        settings.MODEL_FILE = tmp_path / "classifier.pickle"
+        return settings.MODEL_FILE
+
+    @pytest.fixture
+    def classifier(self) -> DocumentClassifier:
+        classifier = DocumentClassifier()
+        classifier.last_doc_change_time = datetime(2026, 1, 2, 3, 4, 5, tzinfo=UTC)
+        classifier.last_auto_type_hash = b"\x01" * 32
+        return classifier
+
+    def test_save_writes_signed_pickle(
+        self,
+        model_file: Path,
+        classifier: DocumentClassifier,
+    ) -> None:
+        """
+        GIVEN:
+            - A classifier with training state
+        WHEN:
+            - The classifier is saved
+        THEN:
+            - The file is the HMAC of the pickled state followed by that pickle
+            - The pickle uses the highest protocol
+            - The saved state loads back into a new classifier
+            - No temporary file is left behind
+        """
+        classifier.save()
+
+        raw = model_file.read_bytes()
+        signature = raw[: DocumentClassifier.HMAC_SIZE]
+        data = raw[DocumentClassifier.HMAC_SIZE :]
+        assert signature == DocumentClassifier._compute_hmac(data)
+        # A pickle opens with the PROTO opcode followed by the protocol number
+        assert data[:2] == bytes([pickle.PROTO[0], pickle.HIGHEST_PROTOCOL])
+        assert pickle.loads(data)[:3] == (
+            DocumentClassifier.FORMAT_VERSION,
+            classifier.last_doc_change_time,
+            classifier.last_auto_type_hash,
+        )
+
+        loaded = DocumentClassifier()
+        loaded.load()
+        assert loaded.last_doc_change_time == classifier.last_doc_change_time
+        assert loaded.last_auto_type_hash == classifier.last_auto_type_hash
+
+        assert not model_file.with_name(f"{model_file.name}.part").exists()
+
+    def test_save_failure_removes_partial_file(
+        self,
+        model_file: Path,
+        classifier: DocumentClassifier,
+        mocker: MockerFixture,
+    ) -> None:
+        """
+        GIVEN:
+            - An existing classifier model file
+        WHEN:
+            - Saving a new classifier fails part way through writing
+        THEN:
+            - The error is raised
+            - The partially written temporary file is removed
+            - The existing model file is left untouched
+        """
+        model_file.write_bytes(b"existing model")
+        mocker.patch(
+            "documents.classifier.pickle.dump",
+            side_effect=RuntimeError("disk full"),
+        )
+
+        with pytest.raises(RuntimeError, match="disk full"):
+            classifier.save()
+
+        assert not model_file.with_name(f"{model_file.name}.part").exists()
+        assert model_file.read_bytes() == b"existing model"
+
+
 class _StubProbaClassifier:
     """
     A fake scikit-learn classifier exposing just enough of the API for