From f919c981e10d30d1165e176bdfa43f54bf4cf62a Mon Sep 17 00:00:00 2001 From: Trenton H <797416+stumpylog@users.noreply.github.com> Date: Tue, 15 Sep 2026 13:35:44 -0700 Subject: [PATCH] Peformance: Improves the memory efficency of classifier training (#14124) --- src/documents/classifier.py | 71 +++++++++++++++++--------- src/documents/tests/test_classifier.py | 65 +++++++++++++++++++++++ 2 files changed, 111 insertions(+), 25 deletions(-) diff --git a/src/documents/classifier.py b/src/documents/classifier.py index bac810cb8..1e078c9c9 100644 --- a/src/documents/classifier.py +++ b/src/documents/classifier.py @@ -70,6 +70,9 @@ read_cache = caches["read-cache"] RE_DIGIT = re.compile(r"\d") RE_WORD = re.compile(r"\b[\w]+\b") # words that may contain digits +# Documents whose content is fetched per query while training +_CONTENT_CHUNK_SIZE = 1000 + class _SignedFileWriter: """ @@ -300,10 +303,30 @@ class DocumentClassifier: notify = status_callback if status_callback is not None else lambda _: None # Get non-inbox documents - docs_queryset = ( - Document.objects.exclude( - tags__is_inbox_tag=True, - ) + docs_queryset = Document.objects.exclude( + tags__is_inbox_tag=True, + ).order_by("pk") + + # No documents exit to train against + doc_count = docs_queryset.count() + if doc_count == 0: + raise ValueError("No training data available.") + + labels_tags = [] + labels_correspondent = [] + labels_document_type = [] + labels_storage_path = [] + # Content is fetched separately later, for exactly these documents in this + # order, so it never all has to be in memory at once + doc_pks: list[int] = [] + latest_doc_change: datetime | None = None + + # Step 1: Extract and preprocess training data from the database. + logger.debug("Gathering data from database...") + notify(f"Gathering data from {doc_count} document(s)...") + hasher = sha256() + for doc in ( + docs_queryset.defer("content") .select_related("document_type", "correspondent", "storage_path") .prefetch_related( Prefetch( @@ -316,23 +339,12 @@ class DocumentClassifier: to_attr="auto_tags", ), ) - .order_by("pk") - ) + .iterator(chunk_size=2000) + ): + doc_pks.append(doc.pk) + if latest_doc_change is None or doc.modified > latest_doc_change: + latest_doc_change = doc.modified - # No documents exit to train against - if docs_queryset.count() == 0: - raise ValueError("No training data available.") - - labels_tags = [] - labels_correspondent = [] - labels_document_type = [] - labels_storage_path = [] - - # Step 1: Extract and preprocess training data from the database. - logger.debug("Gathering data from database...") - notify(f"Gathering data from {docs_queryset.count()} document(s)...") - hasher = sha256() - for doc in docs_queryset: y = -1 dt = doc.document_type if dt and dt.matching_algorithm == MatchingModel.MATCH_AUTO: @@ -366,7 +378,6 @@ class DocumentClassifier: # Check if retraining is actually required. # A document has been updated since the classifier was trained # New auto tags, types, correspondent, storage paths exist - latest_doc_change = docs_queryset.latest("modified").modified if ( self.last_doc_change_time is not None and self.last_doc_change_time >= latest_doc_change @@ -393,7 +404,7 @@ class DocumentClassifier: num_storage_paths: int = len(set(labels_storage_path) | {-1}) - 1 logger.debug( - f"{docs_queryset.count()} documents, {num_tags} tag(s), {num_correspondents} correspondent(s), " + f"{len(doc_pks)} documents, {num_tags} tag(s), {num_correspondents} correspondent(s), " f"{num_document_types} document type(s). {num_storage_paths} storage path(s)", ) @@ -415,10 +426,20 @@ class DocumentClassifier: def content_generator() -> Iterator[str]: """ - Generates the content for documents, but once at a time + Generates the content for documents, in the same order as the labels, + fetching it a chunk at a time """ - for doc in docs_queryset: - yield self.preprocess_content(doc.content, shared_cache=False) + for start in range(0, len(doc_pks), _CONTENT_CHUNK_SIZE): + chunk = doc_pks[start : start + _CONTENT_CHUNK_SIZE] + docs = Document.objects.only("content").order_by().in_bulk(chunk) + for pk in chunk: + # A document deleted since its labels were gathered still + # needs a row, so labels and content stay aligned + doc = docs.get(pk) + yield self.preprocess_content( + doc.content if doc is not None else "", + shared_cache=False, + ) self.data_vectorizer = CountVectorizer( analyzer="word", diff --git a/src/documents/tests/test_classifier.py b/src/documents/tests/test_classifier.py index 99d7196f3..673c18fe5 100644 --- a/src/documents/tests/test_classifier.py +++ b/src/documents/tests/test_classifier.py @@ -1172,3 +1172,68 @@ class TestClassifierTrainTagLabels: assert list(classifier.tags_binarizer.classes_) == sorted( tag.pk for tag in auto_tags ) + + +@pytest.mark.django_db +class TestClassifierTrainContent: + def test_train_content_follows_label_order_across_chunks( + self, + mocker: MockerFixture, + ) -> None: + """ + GIVEN: + - More documents than fit in one content chunk + WHEN: + - The classifier is trained + THEN: + - Every document's content is preprocessed once, in document order + """ + mocker.patch("documents.classifier._CONTENT_CHUNK_SIZE", 2) + docs = DocumentFactory.create_batch(5) + preprocess = mocker.patch.object( + DocumentClassifier, + "preprocess_content", + side_effect=dummy_preprocess, + ) + + DocumentClassifier().train() + + assert [call.args[0] for call in preprocess.call_args_list] == [ + doc.content for doc in sorted(docs, key=lambda doc: doc.pk) + ] + + def test_train_document_deleted_while_training( + self, + mocker: MockerFixture, + ) -> None: + """ + GIVEN: + - Two documents + WHEN: + - The second document is deleted after its labels were gathered, but + before its content is fetched + THEN: + - Training completes + - The deleted document is trained with empty content, keeping labels + and content aligned + """ + mocker.patch("documents.classifier._CONTENT_CHUNK_SIZE", 1) + first, second = DocumentFactory.create_batch(2) + + def delete_second_then_preprocess(content: str, **kwargs) -> str: + if content == first.content: + second.delete() + return dummy_preprocess(content) + + preprocess = mocker.patch.object( + DocumentClassifier, + "preprocess_content", + side_effect=delete_second_then_preprocess, + ) + + assert DocumentClassifier().train() + + assert [call.args[0] for call in preprocess.call_args_list] == [ + first.content, + "", + ]