diff --git a/vector_search/conftest.py b/vector_search/conftest.py index 24ae8df991..98e62d25cc 100644 --- a/vector_search/conftest.py +++ b/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() diff --git a/vector_search/constants.py b/vector_search/constants.py index d5aaca7097..2302fae0bf 100644 --- a/vector_search/constants.py +++ b/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 = { diff --git a/vector_search/tasks.py b/vector_search/tasks.py index 8711cd2734..015958f02e 100644 --- a/vector_search/tasks.py +++ b/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) diff --git a/vector_search/tasks_test.py b/vector_search/tasks_test.py index 06c1565e49..5a5a1c7597 100644 --- a/vector_search/tasks_test.py +++ b/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 diff --git a/vector_search/utils.py b/vector_search/utils.py index 7f147ea93a..afe0e79824 100644 --- a/vector_search/utils.py +++ b/vector_search/utils.py @@ -184,6 +184,12 @@ def tune_collection(client, collection_name): ) +@cache +def ensure_qdrant_collections() -> None: + """Ensure Qdrant collections exist, at most once per worker process.""" + create_qdrant_collections(force_recreate=False) + + def create_qdrant_collections(force_recreate): """ Create or recreate QDrant collections @@ -798,7 +804,7 @@ def embed_learning_resources(ids, resource_type, overwrite): # noqa: PLR0915, C return client = qdrant_client() - create_qdrant_collections(force_recreate=False) + ensure_qdrant_collections() if resource_type != CONTENT_FILE_TYPE: serialized_resources = list(serialize_bulk_learning_resources(ids)) points = [ diff --git a/vector_search/utils_test.py b/vector_search/utils_test.py index 0eae55645e..1e8339957d 100644 --- a/vector_search/utils_test.py +++ b/vector_search/utils_test.py @@ -11,6 +11,7 @@ from qdrant_client import models from qdrant_client.models import PointStruct +import vector_search.utils as vs_utils from learning_resources.constants import GROUP_CONTENT_FILE_CONTENT_VIEWERS from learning_resources.factories import ( ContentFileFactory, @@ -2066,3 +2067,21 @@ def test_best_run_ids_for_resources_no_published_run(): ) assert best_run_ids_for_resources([course.readable_id]) == [] + + +def test_ensure_qdrant_collections_runs_once(mocker): + create = mocker.patch("vector_search.utils.create_qdrant_collections") + vs_utils.ensure_qdrant_collections() + vs_utils.ensure_qdrant_collections() + create.assert_called_once_with(force_recreate=False) + + +def test_embed_learning_resources_uses_collection_guard(mocker): + """embed_learning_resources delegates collection-ensuring to the guard + (not a direct create_qdrant_collections call). + """ + ensure = mocker.patch("vector_search.utils.ensure_qdrant_collections") + mocker.patch("vector_search.utils.qdrant_client") + mocker.patch("vector_search.utils.serialize_bulk_content_files", return_value=[]) + vs_utils.embed_learning_resources([1], CONTENT_FILE_TYPE, overwrite=True) + ensure.assert_called_once()