feat(vector_search): embeddings deadline resilience (retry + finalize tail)

vector_search/conftest.py

@@ -55,3 +55,13 @@ def _use_test_qdrant_settings(settings, mocker):
         "vector_search.utils.qdrant_client",
         return_value=mock_qdrant,
     )
+
+
+@pytest.fixture(autouse=True)
+def _reset_qdrant_collections_guard():
+    """Per-process embed guard must not leak across tests."""
+    from vector_search.utils import ensure_qdrant_collections
+
+    ensure_qdrant_collections.cache_clear()
+    yield
+    ensure_qdrant_collections.cache_clear()

vector_search/constants.py

@@ -162,7 +162,7 @@
 QDRANT_OPTIMIZER_FLUSH_INTERVAL_XLARGE = 30
 
 # Indexing threshold ratio
-QDRANT_OPTIMIZER_INDEXING_THRESHOLD_RATIO = 0.4
+QDRANT_OPTIMIZER_INDEXING_THRESHOLD_RATIO = 0.8
 
 
 VECTOR_SEARCH_SCORE_BOOST = {

vector_search/tasks.py

@@ -5,7 +5,9 @@
 import grpc
 import sentry_sdk
 from celery.exceptions import Ignore
+from celery.utils.time import get_exponential_backoff_interval
 from django.conf import settings
+from django.core.cache import caches
 from django.db.models import Q
 
 from learning_resources.content_summarizer import ContentSummarizer
@@ -54,6 +56,18 @@
 
 log = logging.getLogger(__name__)
 
+EMBED_FAILURE_TTL = 60 * 60 * 24  # 24h defensive cleanup for the per-run counter
+
+
+def _record_embedding_failure(failure_key: str) -> None:
+    """Bump the per-invocation embedding-failure counter in the shared redis cache."""
+    cache = caches["redis"]
+    key = f"embed_errors:{failure_key}"
+    try:
+        cache.incr(key)
+    except ValueError:  # key absent
+        cache.set(key, 1, EMBED_FAILURE_TTL)
+
 
 @app.task
 def tune_qdrant_collections():
@@ -73,6 +87,25 @@ def _replace_with_chain(task, task_signatures):
     return task.replace(celery.chain(*task_signatures))
 
 
+def _replace_with_finalized_chain(
+    task: celery.Task, content_file_ids: list[int], *, overwrite: bool
+) -> None:
+    """
+    Chain of content-file embedding chunks + a finalize tail that fails the parent
+    if any chunk failed. Returns None when there is nothing to embed.
+    """
+    failure_key = task.request.id
+    sigs = [
+        generate_embeddings.si(
+            ids, CONTENT_FILE_TYPE, overwrite=overwrite, failure_key=failure_key
+        )
+        for ids in chunks(content_file_ids, chunk_size=settings.QDRANT_CHUNK_SIZE)
+    ]
+    if not sigs:
+        return None
+    return task.replace(celery.chain(*sigs, finalize_embeddings.si(failure_key)))
+
+
 def _queue_program_content_file_embedding_tasks(index_tasks, program_ids, overwrite):
     """Queue content file embedding tasks for programs using a single bulk query."""
     if not program_ids:
@@ -97,33 +130,56 @@ def _queue_program_content_file_embedding_tasks(index_tasks, program_ids, overwr
     )
 
 
+def _retry_countdown(retries: int) -> int:
+    """Full-jitter exponential backoff (mirrors retry_backoff=True), capped at 10m."""
+    return get_exponential_backoff_interval(
+        factor=1, retries=retries, maximum=600, full_jitter=True
+    )
+
+
 @app.task(
+    bind=True,
     acks_late=True,
     reject_on_worker_lost=True,
-    autoretry_for=(RetryError,),
-    retry_backoff=True,
+    max_retries=3,
     rate_limit="200/m",
 )
-def generate_embeddings(ids, resource_type, overwrite):
+def generate_embeddings(
+    self,
+    ids: list[int],
+    resource_type: str,
+    overwrite: bool,  # noqa: FBT001
+    failure_key: str | None = None,
+) -> None:
     """
-    Generate learning resource embeddings and index in Qdrant
-
-    Args:
-        ids(list of int): List of resource id's
-        resource_type (string): resource_type value for the learning resource objects
+    Generate learning resource embeddings and index in Qdrant.
 
+    Retries transient Qdrant/search errors with jittered backoff. On exhaustion or a
+    non-transient error: if failure_key is set, log + record the failure and return so
+    the chain continues (finalize_embeddings fails the parent); otherwise propagate.
     """
     try:
         with wrap_retry_exception(*SEARCH_CONN_EXCEPTIONS):
             embed_learning_resources(ids, resource_type, overwrite)
-    except (RetryError, Ignore):
+    except Ignore:
         raise
-    except SystemExit as err:
-        raise RetryError(SystemExit.__name__) from err
-    except grpc.RpcError as err:
-        if err.code() == grpc.StatusCode.DEADLINE_EXCEEDED:
-            raise RetryError(str(err)) from err
+    except SystemExit as err:  # worker shutdown: transient; propagate if exhausted
+        if self.request.retries < self.max_retries:
+            raise self.retry(exc=err, countdown=_retry_countdown(self.request.retries))  # noqa: B904
         raise
+    except Exception as err:
+        is_deadline = (
+            isinstance(err, grpc.RpcError)
+            and err.code() == grpc.StatusCode.DEADLINE_EXCEEDED
+        )
+        if (isinstance(err, RetryError) or is_deadline) and (
+            self.request.retries < self.max_retries
+        ):
+            raise self.retry(exc=err, countdown=_retry_countdown(self.request.retries))  # noqa: B904
+        if failure_key is None:
+            raise  # generic callers: propagate terminal failure (current behavior)
+        log.exception("generate_embeddings failed for %s", resource_type)
+        _record_embedding_failure(failure_key)
 
 
 @app.task(
@@ -155,6 +211,19 @@ def remove_embeddings(ids, resource_type):
         raise
 
 
+@app.task
+def finalize_embeddings(failure_key: str) -> None:
+    """Chain tail: fail the parent task if any chunk recorded a failure."""
+    cache = caches["redis"]
+    key = f"embed_errors:{failure_key}"
+    failures = cache.get(key, 0)
+    cache.delete(key)
+    if failures:
+        msg = f"{failures} embedding chunk(s) failed for {failure_key}"
+        log.error(msg)
+        raise RuntimeError(msg)
+
+
 @app.task(bind=True)
 def start_embed_resources(self, indexes, skip_content_files, overwrite):  # noqa: C901
     """
@@ -383,14 +452,11 @@ def embed_new_content_files(self):
         .exclude(learning_resource__published=False, learning_resource__test_mode=False)
     )
 
-    tasks = [
-        generate_embeddings.si(ids, CONTENT_FILE_TYPE, overwrite=False)
-        for ids in chunks(
-            new_content_files.values_list("id", flat=True),
-            chunk_size=settings.QDRANT_CHUNK_SIZE,
-        )
-    ]
-    return _replace_with_chain(self, tasks)
+    return _replace_with_finalized_chain(
+        self,
+        list(new_content_files.values_list("id", flat=True)),
+        overwrite=False,
+    )
 
 
 @app.task(bind=True)
@@ -402,11 +468,7 @@ def embed_run_content_files(self, run_id):
         ContentFile.objects.filter(run__id=run_id).values_list("id", flat=True)
     )
 
-    tasks = [
-        generate_embeddings.si(ids, CONTENT_FILE_TYPE, overwrite=True)
-        for ids in chunks(content_file_ids, chunk_size=settings.QDRANT_CHUNK_SIZE)
-    ]
-    return _replace_with_chain(self, tasks)
+    return _replace_with_finalized_chain(self, content_file_ids, overwrite=True)
 
 
 @app.task(bind=True)

vector_search/tasks_test.py

@@ -3,7 +3,9 @@
 
 import grpc
 import pytest
+from celery.exceptions import Retry
 from django.conf import settings
+from django.core.cache.backends.locmem import LocMemCache
 
 from learning_resources.etl.constants import (
     RESOURCE_FILE_ETL_SOURCES,
@@ -28,11 +30,13 @@
 from learning_resources_search.exceptions import RetryError
 from main.utils import now_in_utc
 from vector_search.tasks import (
+    _record_embedding_failure,
     embed_learning_resources_by_id,
     embed_new_content_files,
     embed_new_learning_resources,
     embed_run_content_files,
     embeddings_healthcheck,
+    finalize_embeddings,
     generate_embeddings,
     remove_embeddings,
     remove_run_content_files,
@@ -51,6 +55,15 @@ def _rpc_error(code):
     return err
 
 
+@pytest.fixture
+def embed_cache(mocker):
+    """Real (LocMem) backing store for the redis-alias counter in tasks under test."""
+    cache = LocMemCache("embed-test", {})
+    cache.clear()
+    mocker.patch("vector_search.tasks.caches", {"redis": cache})
+    return cache
+
+
 @pytest.mark.parametrize("index", list(LEARNING_RESOURCE_TYPES))
 def test_start_embed_resources(mocker, mocked_celery, index):
     """
@@ -222,16 +235,29 @@ def test_embed_new_content_files(mocker, mocked_celery):
     generate_embeddings_mock = mocker.patch(
         "vector_search.tasks.generate_embeddings", autospec=True
     )
+    finalize_embeddings_mock = mocker.patch(
+        "vector_search.tasks.finalize_embeddings", autospec=True
+    )
 
     with pytest.raises(mocked_celery.replace_exception_class):
         embed_new_content_files.delay()
 
     embedded_ids = generate_embeddings_mock.si.mock_calls[0].args[0]
     assert sorted(new_content_file_ids) == sorted(embedded_ids)
-    assert mocked_celery.chain.call_args.args == tuple(
+    assert all(
+        mock_call.kwargs.get("overwrite") is False and "failure_key" in mock_call.kwargs
+        for mock_call in generate_embeddings_mock.si.mock_calls
+    )
+    assert (
+        finalize_embeddings_mock.si.call_args.args[0]
+        == generate_embeddings_mock.si.mock_calls[0].kwargs["failure_key"]
+    )
+    chain_args = mocked_celery.chain.call_args.args
+    assert chain_args[:-1] == tuple(
         generate_embeddings_mock.si.return_value
         for _ in generate_embeddings_mock.si.mock_calls
     )
+    assert chain_args[-1] == finalize_embeddings_mock.si.return_value
 
 
 def test_remove_run_content_files(mocker, mocked_celery, settings):
@@ -638,6 +664,9 @@ def test_embed_run_content_files(mocker, mocked_celery, settings):
     generate_embeddings_mock = mocker.patch(
         "vector_search.tasks.generate_embeddings", autospec=True
     )
+    finalize_embeddings_mock = mocker.patch(
+        "vector_search.tasks.finalize_embeddings", autospec=True
+    )
 
     with pytest.raises(mocked_celery.replace_exception_class):
         embed_run_content_files.delay(run.id)
@@ -650,13 +679,21 @@ def test_embed_run_content_files(mocker, mocked_celery, settings):
     assert sorted(embedded_ids) == sorted(content_file_ids)
     assert all(
         mock_call.args[1:] == (CONTENT_FILE_TYPE,)
-        and mock_call.kwargs == {"overwrite": True}
+        and mock_call.kwargs["overwrite"] is True
+        and "failure_key" in mock_call.kwargs
         for mock_call in generate_embeddings_mock.si.mock_calls
     )
-    assert mocked_celery.chain.call_args.args == tuple(
+    # chain = all chunk sigs, then the finalize tail
+    chain_args = mocked_celery.chain.call_args.args
+    assert chain_args[:-1] == tuple(
         generate_embeddings_mock.si.return_value
         for _ in generate_embeddings_mock.si.mock_calls
     )
+    assert chain_args[-1] == finalize_embeddings_mock.si.return_value
+    assert (
+        finalize_embeddings_mock.si.call_args.args[0]
+        == generate_embeddings_mock.si.mock_calls[0].kwargs["failure_key"]
+    )
     assert mocked_celery.replace.call_count == 1
 
 
@@ -676,6 +713,15 @@ def test_embed_run_content_files_no_content_files(mocker, mocked_celery):
     mocked_celery.replace.assert_not_called()
 
 
+def test_embed_run_content_files_no_files_returns_none(mocker, mocked_celery):
+    """No content files → no chain, no replace, returns None."""
+    run = LearningResourceRunFactory.create()  # no content files
+    mocker.patch("vector_search.tasks.generate_embeddings", autospec=True)
+    mocker.patch("vector_search.tasks.finalize_embeddings", autospec=True)
+    assert embed_run_content_files(run.id) is None
+    mocked_celery.chain.assert_not_called()
+
+
 def test_embeddings_healthcheck_no_missing_embeddings(mocker):
     """
     Test embeddings_healthcheck when there are no missing embeddings
@@ -793,14 +839,48 @@ def test_embeddings_healthcheck_missing_summaries(mocker):
     )
 
 
-def test_generate_embeddings_raises_retryerror_on_grpc_deadline(mocker):
-    """A DEADLINE_EXCEEDED gRPC error becomes a RetryError (autoretry_for picks it up)."""
+def test_generate_embeddings_retries_on_deadline(mocker):
+    """A deadline with retry budget left calls self.retry (jittered backoff)."""
     mocker.patch(
         "vector_search.tasks.embed_learning_resources",
         side_effect=_rpc_error(grpc.StatusCode.DEADLINE_EXCEEDED),
     )
-    with pytest.raises(RetryError):
-        generate_embeddings([1], COURSE_TYPE, overwrite=False)
+    retry = mocker.patch.object(generate_embeddings, "retry", side_effect=Retry())
+    with pytest.raises(Retry):
+        generate_embeddings([1], CONTENT_FILE_TYPE, overwrite=True, failure_key="k")
+    retry.assert_called_once()
+    assert retry.call_args.kwargs["countdown"] >= 0
+
+
+def test_generate_embeddings_records_on_exhaustion(mocker):
+    """Exhausted deadline + failure_key: record + return, do not raise (chain continues)."""
+    mocker.patch(
+        "vector_search.tasks.embed_learning_resources",
+        side_effect=_rpc_error(grpc.StatusCode.DEADLINE_EXCEEDED),
+    )
+    record = mocker.patch("vector_search.tasks._record_embedding_failure")
+    generate_embeddings.push_request(retries=3)
+    try:
+        assert (
+            generate_embeddings([1], CONTENT_FILE_TYPE, overwrite=True, failure_key="k")
+            is None
+        )
+    finally:
+        generate_embeddings.pop_request()
+    record.assert_called_once_with("k")
+
+
+def test_generate_embeddings_records_non_transient_with_key(mocker):
+    """Non-transient error + failure_key: record + return, do not raise."""
+    mocker.patch(
+        "vector_search.tasks.embed_learning_resources", side_effect=ValueError("boom")
+    )
+    record = mocker.patch("vector_search.tasks._record_embedding_failure")
+    assert (
+        generate_embeddings([1], CONTENT_FILE_TYPE, overwrite=True, failure_key="k")
+        is None
+    )
+    record.assert_called_once_with("k")
 
 
 def test_generate_embeddings_reraises_other_grpc_errors(mocker):
@@ -851,3 +931,21 @@ def test_remove_embeddings_does_not_swallow_errors(mocker):
     )
     with pytest.raises(ValueError, match="boom"):
         remove_embeddings([1], COURSE_TYPE)
+
+
+def test_record_embedding_failure_increments(embed_cache):
+    _record_embedding_failure("run-1")
+    _record_embedding_failure("run-1")
+    assert embed_cache.get("embed_errors:run-1") == 2
+
+
+def test_finalize_embeddings_raises_and_clears_on_failures(embed_cache):
+    embed_cache.set("embed_errors:run-1", 3)
+    with pytest.raises(RuntimeError, match="3 embedding chunk"):
+        finalize_embeddings("run-1")
+    assert embed_cache.get("embed_errors:run-1") is None
+
+
+def test_finalize_embeddings_succeeds_when_clean(embed_cache):
+    assert finalize_embeddings("run-1") is None
+    assert embed_cache.get("embed_errors:run-1") is None