Use locmemcache for ci

src/documents/caching.py

@@ -3,13 +3,16 @@ from __future__ import annotations
 import hashlib
 import logging
 import pickle
+import time
 import uuid
 from binascii import hexlify
 from collections import OrderedDict
 from dataclasses import dataclass
+from hashlib import sha256
 from typing import TYPE_CHECKING
 from typing import Any
 from typing import Final
+from uuid import uuid4
 
 from django.conf import settings
 from django.core.cache import cache
@@ -21,6 +24,7 @@ from paperless.signed_pickle import signed_pickle_dumps
 from paperless.signed_pickle import signed_pickle_loads
 
 if TYPE_CHECKING:
+    from django.contrib.auth.models import User
     from django.core.cache.backends.base import BaseCache
 
     from documents.classifier import DocumentClassifier
@@ -59,6 +63,9 @@ CLASSIFIER_MODIFIED_KEY: Final[str] = "classifier_modified"
 #          validated separately, so candidate-anchored 1001 results are stale
 LLM_CACHE_CLASSIFIER_VERSION: Final[int] = 1002
 
+# How often a request waiting on llm generation re-checks the cache
+LLM_SUGGESTION_POLL_INTERVAL: Final[float] = 0.5
+
 CACHE_1_MINUTE: Final[int] = 60
 CACHE_5_MINUTES: Final[int] = 5 * CACHE_1_MINUTE
 CACHE_50_MINUTES: Final[int] = 50 * CACHE_1_MINUTE
@@ -262,6 +269,68 @@ def get_llm_suggestion_cache(
     return None
 
 
+def retrieve_llm_suggestions(
+    document: Document,
+    user: User | None,
+    output_language: str | None,
+    *,
+    backend: str,
+    lock_timeout: int,
+) -> dict:
+    """Return cached LLM suggestions, generating them once across workers."""
+    # Lazy import to avoid pulling in the whole AI stuff
+    from paperless_ai.ai_classifier import get_ai_document_classification
+    from paperless_ai.exceptions import LLMTimeoutError
+
+    lock_key = (
+        f"{get_suggestion_cache_key(document.pk)}_llm_lock_"
+        f"{sha256(backend.encode()).hexdigest()}"
+    )
+    waited = False
+
+    while True:
+        cached = get_llm_suggestion_cache(document.pk, backend=backend)
+        if cached is not None:
+            refresh_suggestions_cache(document.pk)
+            return cached.suggestions
+
+        lock_token = uuid4().hex
+        if cache.add(lock_key, lock_token, lock_timeout):
+            if waited:
+                # The generation we were waiting on has ended without caching
+                # anything so it either failed or outlived its lock. Give up
+                # rather than re-running it
+                cache.delete(lock_key)
+                raise LLMTimeoutError
+
+            try:
+                # The cache may have been populated while acquiring the lock.
+                cached = get_llm_suggestion_cache(document.pk, backend=backend)
+                if cached is not None:
+                    refresh_suggestions_cache(document.pk)
+                    return cached.suggestions
+
+                suggestions = get_ai_document_classification(
+                    document,
+                    user,
+                    output_language,
+                )
+                set_llm_suggestions_cache(
+                    document.pk,
+                    suggestions,
+                    backend=backend,
+                )
+                return suggestions
+            finally:
+                # Don't remove lock if this one expired while generation was still running
+                if cache.get(lock_key) == lock_token:
+                    cache.delete(lock_key)
+
+        waited = True
+        # Another worker is generating suggestions, poll to avoid another LLM request
+        time.sleep(LLM_SUGGESTION_POLL_INTERVAL)
+
+
 def set_llm_suggestions_cache(
     document_id: int,
     suggestions: dict,

src/documents/tests/test_api_documents.py

@@ -2608,7 +2608,7 @@ class TestDocumentApi(DirectoriesMixin, ConsumeTaskMixin, APITestCase):
         response = self.client.get("/api/documents/34676/suggestions/")
         self.assertEqual(response.status_code, status.HTTP_404_NOT_FOUND)
 
-    @mock.patch("documents.views.get_ai_document_classification")
+    @mock.patch("paperless_ai.ai_classifier.get_ai_document_classification")
     @override_settings(AI_ENABLED=True)
     def test_suggestions_still_uses_classifier_when_ai_enabled(
         self,

src/documents/tests/test_caching.py

@@ -1,7 +1,17 @@
+from concurrent.futures import ThreadPoolExecutor
+from threading import Event
+from threading import Lock
+from uuid import uuid4
+
+import pytest
+from django.core.cache.backends.locmem import LocMemCache
+
 from documents.caching import StoredLRUCache
+from documents.caching import retrieve_llm_suggestions
 from paperless.signed_pickle import HMAC_SIZE
 from paperless.signed_pickle import signed_pickle_dumps
 from paperless.signed_pickle import signed_pickle_loads
+from paperless_ai.exceptions import LLMTimeoutError
 
 
 def test_lru_cache_entries() -> None:
@@ -56,3 +66,131 @@ def test_stored_lru_cache_rejects_tampered_data(mocker) -> None:
     cache.load()
 
     assert cache.get("x") is None
+
+
+def test_llm_suggestions_are_generated_once_for_concurrent_requests(mocker) -> None:
+    mocker.patch(
+        "documents.caching.cache",
+        LocMemCache(uuid4().hex, {}),
+    )
+    generation_started = Event()
+    finish_generation = Event()
+    waiter_started = Event()
+    release_waiter = Event()
+    call_lock = Lock()
+    calls = 0
+    suggestions = {"title": "Generated once"}
+    document = mocker.Mock(pk=42)
+    user = mocker.Mock()
+
+    def generate(*args) -> dict:
+        nonlocal calls
+        with call_lock:
+            calls += 1
+        generation_started.set()
+        assert finish_generation.wait(timeout=2)
+        return suggestions
+
+    def wait_for_generation(_interval: float) -> None:
+        waiter_started.set()
+        assert release_waiter.wait(timeout=2)
+
+    mock_get_classification = mocker.patch(
+        "paperless_ai.ai_classifier.get_ai_document_classification",
+        side_effect=generate,
+    )
+    mocker.patch("documents.caching.time.sleep", side_effect=wait_for_generation)
+
+    with ThreadPoolExecutor(max_workers=2) as executor:
+        first = executor.submit(
+            retrieve_llm_suggestions,
+            document,
+            user,
+            None,
+            backend="ollama:model",
+            lock_timeout=10,
+        )
+        assert generation_started.wait(timeout=2)
+        second = executor.submit(
+            retrieve_llm_suggestions,
+            document,
+            user,
+            None,
+            backend="ollama:model",
+            lock_timeout=10,
+        )
+        assert waiter_started.wait(timeout=2)
+        finish_generation.set()
+
+        assert first.result(timeout=2) == suggestions
+        release_waiter.set()
+        assert second.result(timeout=2) == suggestions
+
+    assert calls == 1
+    mock_get_classification.assert_called_once_with(document, user, None)
+
+
+def test_llm_suggestions_waiter_does_not_rerun_a_failed_generation(mocker) -> None:
+    """
+    A request queued behind a generation that fails should give up, not take
+    its turn at re-running a query that just failed.
+    """
+    mocker.patch(
+        "documents.caching.cache",
+        LocMemCache(uuid4().hex, {}),
+    )
+    generation_started = Event()
+    fail_generation = Event()
+    waiter_started = Event()
+    release_waiter = Event()
+    call_lock = Lock()
+    calls = 0
+    document = mocker.Mock(pk=43)
+    user = mocker.Mock()
+
+    def generate(*args) -> dict:
+        nonlocal calls
+        with call_lock:
+            calls += 1
+        generation_started.set()
+        assert fail_generation.wait(timeout=2)
+        raise ValueError("Unknown model")
+
+    def wait_for_generation(_interval: float) -> None:
+        waiter_started.set()
+        assert release_waiter.wait(timeout=2)
+
+    mocker.patch(
+        "paperless_ai.ai_classifier.get_ai_document_classification",
+        side_effect=generate,
+    )
+    mocker.patch("documents.caching.time.sleep", side_effect=wait_for_generation)
+
+    with ThreadPoolExecutor(max_workers=2) as executor:
+        first = executor.submit(
+            retrieve_llm_suggestions,
+            document,
+            user,
+            None,
+            backend="ollama:model",
+            lock_timeout=10,
+        )
+        assert generation_started.wait(timeout=2)
+        second = executor.submit(
+            retrieve_llm_suggestions,
+            document,
+            user,
+            None,
+            backend="ollama:model",
+            lock_timeout=10,
+        )
+        assert waiter_started.wait(timeout=2)
+        fail_generation.set()
+
+        with pytest.raises(ValueError, match="Unknown model"):
+            first.result(timeout=2)
+        release_waiter.set()
+        with pytest.raises(LLMTimeoutError):
+            second.result(timeout=2)
+
+    assert calls == 1

src/documents/tests/test_views.py

@@ -446,7 +446,7 @@ class TestAISuggestions(DirectoriesMixin, TestCase):
         self.assertEqual(response.json()["tags"], [])
         self.assertEqual(response.json()["suggested_tags"], [])
 
-    @patch("documents.views.get_ai_document_classification")
+    @patch("paperless_ai.ai_classifier.get_ai_document_classification")
     @override_settings(
         AI_ENABLED=True,
         LLM_BACKEND="mock_backend",
@@ -496,7 +496,7 @@ class TestAISuggestions(DirectoriesMixin, TestCase):
             None,
         )
 
-    @patch("documents.views.get_ai_document_classification")
+    @patch("paperless_ai.ai_classifier.get_ai_document_classification")
     @override_settings(
         AI_ENABLED=True,
         LLM_BACKEND="mock_backend",
@@ -534,7 +534,7 @@ class TestAISuggestions(DirectoriesMixin, TestCase):
             "KI Title",
         )
 
-    @patch("documents.views.get_ai_document_classification")
+    @patch("paperless_ai.ai_classifier.get_ai_document_classification")
     @override_settings(
         AI_ENABLED=True,
         LLM_BACKEND="mock_backend",
@@ -573,7 +573,7 @@ class TestAISuggestions(DirectoriesMixin, TestCase):
             "Titre IA",
         )
 
-    @patch("documents.views.get_ai_document_classification")
+    @patch("paperless_ai.ai_classifier.get_ai_document_classification")
     @override_settings(
         AI_ENABLED=True,
         LLM_BACKEND="mock_backend",
@@ -609,7 +609,7 @@ class TestAISuggestions(DirectoriesMixin, TestCase):
             ),
         )
 
-    @patch("documents.views.get_ai_document_classification")
+    @patch("paperless_ai.ai_classifier.get_ai_document_classification")
     @override_settings(
         AI_ENABLED=True,
         LLM_BACKEND="mock_backend",
@@ -681,7 +681,7 @@ class TestAISuggestions(DirectoriesMixin, TestCase):
             ),
         )
 
-    @patch("documents.views.get_ai_document_classification")
+    @patch("paperless_ai.ai_classifier.get_ai_document_classification")
     @override_settings(
         AI_ENABLED=True,
         LLM_BACKEND="openai-like",
@@ -710,7 +710,7 @@ class TestAISuggestions(DirectoriesMixin, TestCase):
             get_llm_suggestion_cache(self.document.pk, backend="openai-like"),
         )
 
-    @patch("documents.views.get_ai_document_classification")
+    @patch("paperless_ai.ai_classifier.get_ai_document_classification")
     @override_settings(
         AI_ENABLED=True,
         LLM_BACKEND="openai-like",
@@ -737,7 +737,7 @@ class TestAISuggestions(DirectoriesMixin, TestCase):
             get_llm_suggestion_cache(self.document.pk, backend="openai-like"),
         )
 
-    @patch("documents.views.get_ai_document_classification")
+    @patch("paperless_ai.ai_classifier.get_ai_document_classification")
     @override_settings(
         AI_ENABLED=True,
         LLM_BACKEND="mock_backend",
@@ -775,7 +775,7 @@ class TestAISuggestions(DirectoriesMixin, TestCase):
         self.assertEqual(response.json()["tags"], [self.tag1.pk])
         self.assertEqual(response.json()["suggested_tags"], ["Follow-up"])
 
-    @patch("documents.views.get_ai_document_classification")
+    @patch("paperless_ai.ai_classifier.get_ai_document_classification")
     @override_settings(
         AI_ENABLED=True,
         LLM_BACKEND="mock_backend",
@@ -814,7 +814,7 @@ class TestAISuggestions(DirectoriesMixin, TestCase):
         self.assertEqual(response.json()["tags"], [self.tag1.pk])
         self.assertEqual(response.json()["suggested_tags"], [])
 
-    @patch("documents.views.get_ai_document_classification")
+    @patch("paperless_ai.ai_classifier.get_ai_document_classification")
     @override_settings(
         AI_ENABLED=True,
         LLM_BACKEND="mock_backend",

src/documents/views.py

@@ -116,7 +116,7 @@ from documents.caching import get_suggestion_cache
 from documents.caching import refresh_llm_suggestions_cache
 from documents.caching import refresh_metadata_cache
 from documents.caching import refresh_suggestions_cache
-from documents.caching import set_llm_suggestions_cache
+from documents.caching import retrieve_llm_suggestions
 from documents.caching import set_metadata_cache
 from documents.caching import set_suggestions_cache
 from documents.classifier import load_classifier
@@ -249,7 +249,6 @@ from paperless.parsers.remote import RemoteEngineConfig
 from paperless.serialisers import GroupSerializer
 from paperless.serialisers import UserSerializer
 from paperless.views import StandardPagination
-from paperless_ai.ai_classifier import get_ai_document_classification
 from paperless_ai.ai_classifier import get_llm_output_language
 from paperless_ai.chat import stream_chat_with_documents
 from paperless_ai.exceptions import LLMTimeoutError
@@ -1575,10 +1574,13 @@ class DocumentViewSet(
             llm_suggestions = cached_llm_suggestions.suggestions
         else:
             try:
-                llm_suggestions = get_ai_document_classification(
-                    doc,
-                    request.user,
-                    output_language,
+                llm_suggestions = retrieve_llm_suggestions(
+                    document=doc,
+                    user=request.user,
+                    output_language=output_language,
+                    backend=llm_cache_backend,
+                    # Classification, localization + 30s
+                    lock_timeout=(2 * ai_config.llm_request_timeout) + 30,
                 )
             except ValueError as exc:
                 logger.exception(
@@ -1603,11 +1605,6 @@ class DocumentViewSet(
                     {"ai": [_("AI backend request timed out.")]},
                     status=status.HTTP_503_SERVICE_UNAVAILABLE,
                 )
-            set_llm_suggestions_cache(
-                doc.pk,
-                llm_suggestions,
-                backend=llm_cache_backend,
-            )
 
         tags_choice: TaxonomyChoiceDict = llm_suggestions["tags"]
         correspondents_choice: TaxonomyChoiceDict = llm_suggestions["correspondents"]

src/paperless_mail/tasks.py

@@ -1,9 +1,7 @@
 import logging
 
-from celery import Task
 from celery import shared_task
 
-from documents.models import PaperlessTask
 from paperless_mail.mail import MailAccountHandler
 from paperless_mail.mail import MailError
 from paperless_mail.models import MailAccount
@@ -12,27 +10,8 @@ from paperless_mail.models import MailRule
 logger = logging.getLogger("paperless.mail.tasks")
 
 
-@shared_task(bind=True)
-def process_mail_accounts(self: Task, account_ids: list[int] | None = None) -> str:
-    # A scheduled check can still be running (or queued) when the next one
-    # fires, e.g. a large attachment batch that takes longer to process than
-    # the check interval. ProcessedMail dedup only records a message once its
-    # handling has finished, so an overlapping run can still pick up the same
-    # not-yet-recorded message. Skip outright rather than race it.
-    other_mail_fetch_running = (
-        PaperlessTask.objects.filter(
-            task_type=PaperlessTask.TaskType.MAIL_FETCH,
-            status__in=[PaperlessTask.Status.PENDING, PaperlessTask.Status.STARTED],
-        )
-        .exclude(task_id=self.request.id)
-        .exists()
-    )
-    if other_mail_fetch_running:
-        logger.info(
-            "Mail account processing is already running; skipping this run.",
-        )
-        return "Skipped: mail account processing already in progress."
-
+@shared_task
+def process_mail_accounts(account_ids: list[int] | None = None) -> str:
     total_new_documents = 0
     accounts = (
         MailAccount.objects.filter(pk__in=account_ids)

src/paperless_mail/tests/test_process_mail_accounts_overlap.py

@@ -1,89 +0,0 @@
-from unittest import mock
-
-import pytest
-
-from documents.models import PaperlessTask
-from paperless_mail import tasks
-from paperless_mail.tests.factories import MailAccountFactory
-from paperless_mail.tests.factories import MailRuleFactory
-
-
-@pytest.mark.django_db
-class TestProcessMailAccountsOverlap:
-    def test_skips_when_another_mail_fetch_task_is_running(self) -> None:
-        account = MailAccountFactory.create()
-        MailRuleFactory.create(account=account, enabled=True)
-
-        PaperlessTask.objects.create(
-            task_id="other-running-task",
-            task_type=PaperlessTask.TaskType.MAIL_FETCH,
-            trigger_source=PaperlessTask.TriggerSource.SCHEDULED,
-            status=PaperlessTask.Status.STARTED,
-        )
-
-        with mock.patch.object(
-            tasks.MailAccountHandler,
-            "handle_mail_account",
-        ) as mocked_handle:
-            result = tasks.process_mail_accounts()
-
-        mocked_handle.assert_not_called()
-        assert result == "Skipped: mail account processing already in progress."
-
-    def test_runs_when_no_other_mail_fetch_task_is_running(self) -> None:
-        account = MailAccountFactory.create()
-        MailRuleFactory.create(account=account, enabled=True)
-
-        with mock.patch.object(
-            tasks.MailAccountHandler,
-            "handle_mail_account",
-            return_value=0,
-        ) as mocked_handle:
-            result = tasks.process_mail_accounts()
-
-        mocked_handle.assert_called_once()
-        assert result == "No new documents were added."
-
-    def test_ignores_completed_mail_fetch_tasks(self) -> None:
-        account = MailAccountFactory.create()
-        MailRuleFactory.create(account=account, enabled=True)
-
-        PaperlessTask.objects.create(
-            task_id="finished-task",
-            task_type=PaperlessTask.TaskType.MAIL_FETCH,
-            trigger_source=PaperlessTask.TriggerSource.SCHEDULED,
-            status=PaperlessTask.Status.SUCCESS,
-        )
-
-        with mock.patch.object(
-            tasks.MailAccountHandler,
-            "handle_mail_account",
-            return_value=0,
-        ) as mocked_handle:
-            result = tasks.process_mail_accounts()
-
-        mocked_handle.assert_called_once()
-        assert result == "No new documents were added."
-
-    def test_does_not_skip_due_to_its_own_task_row(self) -> None:
-        account = MailAccountFactory.create()
-        MailRuleFactory.create(account=account, enabled=True)
-
-        PaperlessTask.objects.create(
-            task_id="self-task-id",
-            task_type=PaperlessTask.TaskType.MAIL_FETCH,
-            trigger_source=PaperlessTask.TriggerSource.SCHEDULED,
-            status=PaperlessTask.Status.STARTED,
-        )
-
-        with mock.patch.object(
-            tasks.MailAccountHandler,
-            "handle_mail_account",
-            return_value=0,
-        ) as mocked_handle:
-            result = tasks.process_mail_accounts.apply(
-                task_id="self-task-id",
-            ).result
-
-        mocked_handle.assert_called_once()
-        assert result == "No new documents were added."