From dde40cfc54977eb2294e9a777117c5a6d99f37d7 Mon Sep 17 00:00:00 2001 From: Matt Bertrand Date: Thu, 25 Jun 2026 10:21:20 -0400 Subject: [PATCH 1/9] feat(vector_search): add embedding failure counter + finalize_embeddings tail --- vector_search/tasks.py | 26 ++++++++++++++++++++++++++ vector_search/tasks_test.py | 30 ++++++++++++++++++++++++++++++ 2 files changed, 56 insertions(+) diff --git a/vector_search/tasks.py b/vector_search/tasks.py index 8711cd2734..21186f485e 100644 --- a/vector_search/tasks.py +++ b/vector_search/tasks.py @@ -6,6 +6,7 @@ import sentry_sdk from celery.exceptions import Ignore 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 +55,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): + """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(): @@ -155,6 +168,19 @@ def remove_embeddings(ids, resource_type): raise +@app.task(bind=True) +def finalize_embeddings(self, failure_key): # noqa: ARG001 + """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 """ diff --git a/vector_search/tasks_test.py b/vector_search/tasks_test.py index 06c1565e49..903d6782be 100644 --- a/vector_search/tasks_test.py +++ b/vector_search/tasks_test.py @@ -4,6 +4,7 @@ import grpc import pytest 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 +29,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 +54,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): """ @@ -851,3 +863,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 From bd539688ea4ba747e5972c5e45b7e7f436642443 Mon Sep 17 00:00:00 2001 From: Matt Bertrand Date: Thu, 25 Jun 2026 10:29:55 -0400 Subject: [PATCH 2/9] feat(vector_search): manual jittered retry + gated skip-on-exhaustion for generate_embeddings --- vector_search/tasks.py | 46 ++++++++++++++++++++++++++----------- vector_search/tasks_test.py | 43 ++++++++++++++++++++++++++++++---- 2 files changed, 71 insertions(+), 18 deletions(-) diff --git a/vector_search/tasks.py b/vector_search/tasks.py index 21186f485e..b8d46c666c 100644 --- a/vector_search/tasks.py +++ b/vector_search/tasks.py @@ -5,6 +5,7 @@ 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 @@ -110,33 +111,50 @@ def _queue_program_content_file_embedding_tasks(index_tasks, program_ids, overwr ) +def _retry_countdown(retries): + """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, resource_type, overwrite, failure_key=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( diff --git a/vector_search/tasks_test.py b/vector_search/tasks_test.py index 903d6782be..a271aaefac 100644 --- a/vector_search/tasks_test.py +++ b/vector_search/tasks_test.py @@ -3,6 +3,7 @@ import grpc import pytest +from celery.exceptions import Retry from django.conf import settings from django.core.cache.backends.locmem import LocMemCache @@ -805,14 +806,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): From 6153bd3c86270e90fdfc38756faaaaa4699251c1 Mon Sep 17 00:00:00 2001 From: Matt Bertrand Date: Thu, 25 Jun 2026 10:39:22 -0400 Subject: [PATCH 3/9] feat(vector_search): finalize tail + failure_key for run/new content-file embedding Co-Authored-By: Claude Opus 4.8 (1M context) --- vector_search/tasks.py | 36 +++++++++++++++++++++++------------- vector_search/tasks_test.py | 31 ++++++++++++++++++++++++++++--- 2 files changed, 51 insertions(+), 16 deletions(-) diff --git a/vector_search/tasks.py b/vector_search/tasks.py index b8d46c666c..e1e843593d 100644 --- a/vector_search/tasks.py +++ b/vector_search/tasks.py @@ -87,6 +87,23 @@ def _replace_with_chain(task, task_signatures): return task.replace(celery.chain(*task_signatures)) +def _replace_with_finalized_chain(task, content_file_ids, overwrite): + """ + 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: @@ -427,14 +444,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) @@ -446,11 +460,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 a271aaefac..2ad02a5df1 100644 --- a/vector_search/tasks_test.py +++ b/vector_search/tasks_test.py @@ -235,16 +235,21 @@ 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( + 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): @@ -651,6 +656,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) @@ -663,13 +671,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 @@ -689,6 +705,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 From 7fcc43de2fe3e6e027538cb1422d4a97feee3d42 Mon Sep 17 00:00:00 2001 From: Matt Bertrand Date: Thu, 25 Jun 2026 10:43:42 -0400 Subject: [PATCH 4/9] test(vector_search): assert failure_key threading on new-content-file embedding path Co-Authored-By: Claude Opus 4.8 (1M context) --- vector_search/tasks_test.py | 8 ++++++++ 1 file changed, 8 insertions(+) diff --git a/vector_search/tasks_test.py b/vector_search/tasks_test.py index 2ad02a5df1..5a5a1c7597 100644 --- a/vector_search/tasks_test.py +++ b/vector_search/tasks_test.py @@ -244,6 +244,14 @@ def test_embed_new_content_files(mocker, mocked_celery): embedded_ids = generate_embeddings_mock.si.mock_calls[0].args[0] assert sorted(new_content_file_ids) == sorted(embedded_ids) + 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 From 7c76423c98a80d40bb8cd2d2ea19c1adfd2515ed Mon Sep 17 00:00:00 2001 From: Matt Bertrand Date: Thu, 25 Jun 2026 10:51:02 -0400 Subject: [PATCH 5/9] perf(vector_search): ensure Qdrant collections once per process Co-Authored-By: Claude Opus 4.8 (1M context) --- vector_search/utils.py | 14 +++++++++++++- vector_search/utils_test.py | 28 ++++++++++++++++++++++++++++ 2 files changed, 41 insertions(+), 1 deletion(-) diff --git a/vector_search/utils.py b/vector_search/utils.py index 7f147ea93a..2e8cbf3239 100644 --- a/vector_search/utils.py +++ b/vector_search/utils.py @@ -184,6 +184,18 @@ def tune_collection(client, collection_name): ) +_collections_ensured = False + + +def ensure_qdrant_collections(): + """Ensure Qdrant collections exist, at most once per worker process.""" + global _collections_ensured # noqa: PLW0603 + if _collections_ensured: + return + create_qdrant_collections(force_recreate=False) + _collections_ensured = True + + def create_qdrant_collections(force_recreate): """ Create or recreate QDrant collections @@ -798,7 +810,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..f53c5adb18 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,30 @@ def test_best_run_ids_for_resources_no_published_run(): ) assert best_run_ids_for_resources([course.readable_id]) == [] + + +@pytest.fixture +def reset_collections_guard(): + vs_utils._collections_ensured = False # noqa: SLF001 + yield + vs_utils._collections_ensured = False # noqa: SLF001 + + +def test_ensure_qdrant_collections_runs_once(mocker, reset_collections_guard): + 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, reset_collections_guard +): + """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() From 6233f5eaab8fba6e3f5fbec215a92af48c1e062d Mon Sep 17 00:00:00 2001 From: Matt Bertrand Date: Thu, 25 Jun 2026 12:10:54 -0400 Subject: [PATCH 6/9] test(vector_search): autouse-reset Qdrant collection guard between tests Prevents the process-global _collections_ensured flag from leaking across tests when an earlier test runs embed_learning_resources. Co-Authored-By: Claude Opus 4.8 (1M context) --- vector_search/conftest.py | 10 ++++++++++ 1 file changed, 10 insertions(+) diff --git a/vector_search/conftest.py b/vector_search/conftest.py index 24ae8df991..d227050760 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(): + """Module-global embed guard must not leak across tests.""" + import vector_search.utils as vs_utils + + vs_utils._collections_ensured = False # noqa: SLF001 + yield + vs_utils._collections_ensured = False # noqa: SLF001 From 0a2e24b8109537e0ba503cb35560285c806954a6 Mon Sep 17 00:00:00 2001 From: Matt Bertrand Date: Thu, 25 Jun 2026 12:13:37 -0400 Subject: [PATCH 7/9] perf(vector_search): raise Qdrant indexing_threshold ratio to 0.8 for bulk writes Co-Authored-By: Claude Opus 4.8 (1M context) --- vector_search/constants.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) 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 = { From ca00cea334a7ef2218deb887649a81831313f8a0 Mon Sep 17 00:00:00 2001 From: Matt Bertrand Date: Thu, 25 Jun 2026 13:41:36 -0400 Subject: [PATCH 8/9] Simplification --- vector_search/conftest.py | 8 ++++---- vector_search/tasks.py | 4 ++-- vector_search/utils.py | 8 +------- vector_search/utils_test.py | 13 ++----------- 4 files changed, 9 insertions(+), 24 deletions(-) diff --git a/vector_search/conftest.py b/vector_search/conftest.py index d227050760..98e62d25cc 100644 --- a/vector_search/conftest.py +++ b/vector_search/conftest.py @@ -59,9 +59,9 @@ def _use_test_qdrant_settings(settings, mocker): @pytest.fixture(autouse=True) def _reset_qdrant_collections_guard(): - """Module-global embed guard must not leak across tests.""" - import vector_search.utils as vs_utils + """Per-process embed guard must not leak across tests.""" + from vector_search.utils import ensure_qdrant_collections - vs_utils._collections_ensured = False # noqa: SLF001 + ensure_qdrant_collections.cache_clear() yield - vs_utils._collections_ensured = False # noqa: SLF001 + ensure_qdrant_collections.cache_clear() diff --git a/vector_search/tasks.py b/vector_search/tasks.py index e1e843593d..7d5df165b9 100644 --- a/vector_search/tasks.py +++ b/vector_search/tasks.py @@ -203,8 +203,8 @@ def remove_embeddings(ids, resource_type): raise -@app.task(bind=True) -def finalize_embeddings(self, failure_key): # noqa: ARG001 +@app.task +def finalize_embeddings(failure_key): """Chain tail: fail the parent task if any chunk recorded a failure.""" cache = caches["redis"] key = f"embed_errors:{failure_key}" diff --git a/vector_search/utils.py b/vector_search/utils.py index 2e8cbf3239..3d75b2cc5d 100644 --- a/vector_search/utils.py +++ b/vector_search/utils.py @@ -184,16 +184,10 @@ def tune_collection(client, collection_name): ) -_collections_ensured = False - - +@cache def ensure_qdrant_collections(): """Ensure Qdrant collections exist, at most once per worker process.""" - global _collections_ensured # noqa: PLW0603 - if _collections_ensured: - return create_qdrant_collections(force_recreate=False) - _collections_ensured = True def create_qdrant_collections(force_recreate): diff --git a/vector_search/utils_test.py b/vector_search/utils_test.py index f53c5adb18..1e8339957d 100644 --- a/vector_search/utils_test.py +++ b/vector_search/utils_test.py @@ -2069,23 +2069,14 @@ def test_best_run_ids_for_resources_no_published_run(): assert best_run_ids_for_resources([course.readable_id]) == [] -@pytest.fixture -def reset_collections_guard(): - vs_utils._collections_ensured = False # noqa: SLF001 - yield - vs_utils._collections_ensured = False # noqa: SLF001 - - -def test_ensure_qdrant_collections_runs_once(mocker, reset_collections_guard): +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, reset_collections_guard -): +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). """ From 63618b1aa7402566b2aa9805cc3ce02d3d40781b Mon Sep 17 00:00:00 2001 From: Matt Bertrand Date: Thu, 25 Jun 2026 13:57:57 -0400 Subject: [PATCH 9/9] type hints --- vector_search/tasks.py | 18 +++++++++++++----- vector_search/utils.py | 2 +- 2 files changed, 14 insertions(+), 6 deletions(-) diff --git a/vector_search/tasks.py b/vector_search/tasks.py index 7d5df165b9..015958f02e 100644 --- a/vector_search/tasks.py +++ b/vector_search/tasks.py @@ -59,7 +59,7 @@ EMBED_FAILURE_TTL = 60 * 60 * 24 # 24h defensive cleanup for the per-run counter -def _record_embedding_failure(failure_key): +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}" @@ -87,7 +87,9 @@ def _replace_with_chain(task, task_signatures): return task.replace(celery.chain(*task_signatures)) -def _replace_with_finalized_chain(task, content_file_ids, overwrite): +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. @@ -128,7 +130,7 @@ def _queue_program_content_file_embedding_tasks(index_tasks, program_ids, overwr ) -def _retry_countdown(retries): +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 @@ -142,7 +144,13 @@ def _retry_countdown(retries): max_retries=3, rate_limit="200/m", ) -def generate_embeddings(self, ids, resource_type, overwrite, failure_key=None): +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. @@ -204,7 +212,7 @@ def remove_embeddings(ids, resource_type): @app.task -def finalize_embeddings(failure_key): +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}" diff --git a/vector_search/utils.py b/vector_search/utils.py index 3d75b2cc5d..afe0e79824 100644 --- a/vector_search/utils.py +++ b/vector_search/utils.py @@ -185,7 +185,7 @@ def tune_collection(client, collection_name): @cache -def ensure_qdrant_collections(): +def ensure_qdrant_collections() -> None: """Ensure Qdrant collections exist, at most once per worker process.""" create_qdrant_collections(force_recreate=False)