From d7f1c925933458cfd6c42577a04131b1166fd326 Mon Sep 17 00:00:00 2001 From: Matt Bertrand Date: Sat, 18 Jul 2026 09:07:01 -0400 Subject: [PATCH 1/6] Embedding pipeline mechanics: broker safety, chunking, dispatch Phase 1 (mechanics) of the content-file embedding load reduction: - Raise Redis broker visibility_timeout to 6h so long-running summarization+embed tasks are not redelivered mid-flight during a backlog (the primary redelivery-storm driver). - Split QDRANT_CHUNK_SIZE (10->100 for resources) from a new QDRANT_CONTENT_FILE_CHUNK_SIZE (25) so content-file tasks stay small. - ignore_result on fire-and-forget embedding tasks; drop the finalize_embeddings tail and embed_errors counter plumbing (failure signal is log.exception + the weekly presence-based healthcheck). - Convert whole-catalog backfill from self.replace(chain) to bounded apply_async batches; per-run dispatch becomes a group of chunk sigs (no O(N^2) chain payloads). - embed_run_content_files filters to published files. - generate_embeddings rate_limit from CELERY_EMBEDDINGS_RATE_LIMIT. - try_with_retry_as_task dispatches via apply_async with publish-time retry_policy instead of re-.delay() on any exception. Addresses https://github.com/mitodl/hq/issues/12453 Related: mitodl/hq#12015 (infra concurrency/grace-period counterpart), mitodl/hq#12008 (Qdrant performance settings), mitodl/hq#12172 (embedding pipeline failures). Co-Authored-By: Claude Opus 4.8 --- learning_resources_search/plugins.py | 22 ++- main/settings.py | 14 +- main/settings_celery.py | 10 ++ vector_search/tasks.py | 109 +++++++------- vector_search/tasks_test.py | 213 ++++++++++----------------- vector_search/utils.py | 104 ++++++++++--- vector_search/utils_test.py | 31 ++++ 7 files changed, 282 insertions(+), 221 deletions(-) diff --git a/learning_resources_search/plugins.py b/learning_resources_search/plugins.py index 0f3f5abfa6..59af690d00 100644 --- a/learning_resources_search/plugins.py +++ b/learning_resources_search/plugins.py @@ -2,7 +2,7 @@ import logging -from celery import chain +from celery import Task, chain from django.apps import apps from django.conf import settings as django_settings @@ -23,12 +23,22 @@ def try_with_retry_as_task(function, *args): """ - Try running the task, if it errors, run it as a celery task. + Dispatch a task/signature to the broker with publish-time retry. + + Accepts a bare task object (plus positional args) or an already-built + signature (e.g. a chain). Publish-time retry absorbs transient broker + blips without a second manual publish that could double-dispatch. """ - try: - function(*args) - except Exception: # noqa: BLE001 - function.delay(*args) + signature = function.si(*args) if isinstance(function, Task) else function + signature.apply_async( + retry=True, + retry_policy={ + "max_retries": 3, + "interval_start": 0.2, + "interval_step": 0.5, + "interval_max": 2, + }, + ) class SearchIndexPlugin: diff --git a/main/settings.py b/main/settings.py index 378b0891ea..22cc15b2d3 100644 --- a/main/settings.py +++ b/main/settings.py @@ -778,9 +778,21 @@ def get_all_config_keys(): default="vector_search.encoders.sparse_hash.SparseHashEncoder", ) +# Number of resource ids per generate_embeddings task. Resource-metadata +# embedding is cheap per item, so a larger chunk means far fewer tasks. QDRANT_CHUNK_SIZE = get_int( name="QDRANT_CHUNK_SIZE", - default=10, + default=100, +) + +# Number of content-file ids per generate_embeddings task. Kept smaller than +# QDRANT_CHUNK_SIZE because content-file tasks do inline LLM summarization + +# embedding per file, so a large chunk risks long runtimes (visibility_timeout) +# and high worker memory. Raise toward QDRANT_CHUNK_SIZE once summarization is +# decoupled from embedding. +QDRANT_CONTENT_FILE_CHUNK_SIZE = get_int( + name="QDRANT_CONTENT_FILE_CHUNK_SIZE", + default=25, ) QDRANT_ENCODER = get_string( diff --git a/main/settings_celery.py b/main/settings_celery.py index 62f11fce73..026408608c 100644 --- a/main/settings_celery.py +++ b/main/settings_celery.py @@ -25,6 +25,14 @@ # (the default here) that saturates memory. Keep results only long enough for # chord callbacks to consume them. CELERY_RESULT_EXPIRES = get_int("CELERY_RESULT_EXPIRES", 60 * 60) +# visibility_timeout must exceed the longest possible task runtime (inline +# summarization + embedding) plus max retry backoff (600s), or in-flight tasks +# get redelivered while still running, amplifying load during a backlog. +CELERY_BROKER_TRANSPORT_OPTIONS = { + "visibility_timeout": get_int( + name="CELERY_BROKER_VISIBILITY_TIMEOUT", default=6 * 60 * 60 + ), +} CELERY_BEAT_SCHEDULER = RedBeatScheduler redbeat_redis_url = CELERY_BROKER_URL CELERY_TASK_ALWAYS_EAGER = get_bool("CELERY_TASK_ALWAYS_EAGER", False) # noqa: FBT003 @@ -221,3 +229,5 @@ CELERY_VECTOR_SEARCH_RATE_LIMIT = get_string( "CELERY_VECTOR_SEARCH_RATE_LIMIT", CELERY_RATE_LIMIT ) +# Per-worker execution rate for generate_embeddings; tunable without a deploy. +CELERY_EMBEDDINGS_RATE_LIMIT = get_string("CELERY_EMBEDDINGS_RATE_LIMIT", "200/m") diff --git a/vector_search/tasks.py b/vector_search/tasks.py index 1d933c37c1..916148856e 100644 --- a/vector_search/tasks.py +++ b/vector_search/tasks.py @@ -7,7 +7,6 @@ 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.models import ( @@ -58,18 +57,6 @@ 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(): @@ -89,23 +76,41 @@ def _replace_with_chain(task, task_signatures): return task.replace(celery.chain(*task_signatures)) -def _replace_with_finalized_chain( +def _replace_with_content_file_group( 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. + Replace the parent with a group of content-file embedding chunks. + + A group (not a chain) so chunks are not serialized behind one another and the + remaining work is not carried in every message; returns None when there is + nothing to embed. Failures surface via generate_embeddings' own logging and + the periodic embeddings_healthcheck, not a chain tail. """ - failure_key = task.request.id sigs = [ - generate_embeddings.si( - ids, CONTENT_FILE_TYPE, overwrite=overwrite, failure_key=failure_key + generate_embeddings.si(ids, CONTENT_FILE_TYPE, overwrite=overwrite) + for ids in chunks( + content_file_ids, chunk_size=settings.QDRANT_CONTENT_FILE_CHUNK_SIZE ) - 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))) + return task.replace(celery.group(sigs)) + + +def _dispatch_signatures(task_signatures) -> int: + """ + Fire-and-forget dispatch of already-built chunk signatures (backfill paths). + + Publishes one message per chunk rather than a single catalog-wide group/chain, + so a full-catalog backfill does not spike the broker with one huge message. + Nothing waits on the results (embedding tasks are ignore_result). + """ + count = 0 + for sig in task_signatures: + sig.apply_async() + count += 1 + return count def _queue_program_content_file_embedding_tasks(index_tasks, program_ids, overwrite): @@ -147,21 +152,21 @@ def _retry_countdown(retries: int) -> int: acks_late=True, reject_on_worker_lost=True, max_retries=3, - rate_limit="200/m", + rate_limit=settings.CELERY_EMBEDDINGS_RATE_LIMIT, + ignore_result=True, ) 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. 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. + non-transient error, logs and propagates the failure (surfaced via Sentry and the + periodic embeddings_healthcheck). """ try: with wrap_retry_exception(*SEARCH_CONN_EXCEPTIONS): @@ -171,6 +176,7 @@ def generate_embeddings( 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 + log.exception("generate_embeddings exhausted retries for %s", resource_type) raise except Exception as err: is_transient_grpc = isinstance(err, grpc.RpcError) and err.code() in ( @@ -181,10 +187,8 @@ def generate_embeddings( 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) + raise @app.task( @@ -193,6 +197,7 @@ def generate_embeddings( autoretry_for=(RetryError,), retry_backoff=True, rate_limit=settings.CELERY_VECTOR_SEARCH_RATE_LIMIT, + ignore_result=True, ) def remove_embeddings(ids, resource_type): """ @@ -217,20 +222,7 @@ def remove_embeddings(ids, resource_type): @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 +def start_embed_resources(indexes, skip_content_files, overwrite): # noqa: C901 """ Celery task to embed all learning resources for given indexes @@ -329,13 +321,13 @@ def start_embed_resources(self, indexes, skip_content_files, overwrite): # noqa log.exception(error) return error - # Use self.replace so that code waiting on this task will also wait on the embedding - # and finish tasks - return _replace_with_chain(self, index_tasks) + # Fire-and-forget per-chunk dispatch rather than one catalog-wide chain, which + # would spike the broker with a single huge message. Nothing waits on results. + return _dispatch_signatures(index_tasks) -@app.task(bind=True) -def embed_learning_resources_by_id(self, ids, skip_content_files, overwrite): +@app.task +def embed_learning_resources_by_id(ids, skip_content_files, overwrite): """ Celery task to embed specific resources @@ -394,17 +386,16 @@ def embed_learning_resources_by_id(self, ids, skip_content_files, overwrite): ] except: # noqa: E722 - error = "start_embed_resources threw an error" + error = "embed_learning_resources_by_id threw an error" log.exception(error) return error - # Use self.replace so that code waiting on this task will also wait on the embedding - # and finish tasks - - return _replace_with_chain(self, index_tasks) + # Fire-and-forget per-chunk dispatch rather than one catalog-wide chain, which + # would spike the broker with a single huge message. Nothing waits on results. + return _dispatch_signatures(index_tasks) -@app.task(bind=True) +@app.task(bind=True, ignore_result=True) def embed_new_learning_resources(self): """ Embed new resources from QDRANT_EMBEDDINGS_TASK_LOOKBACK_WINDOW minutes ago @@ -440,7 +431,7 @@ def embed_new_learning_resources(self): return self.replace(embed_tasks) -@app.task(bind=True) +@app.task(bind=True, ignore_result=True) def embed_new_content_files(self): """ Embed new content files from QDRANT_EMBEDDINGS_TASK_LOOKBACK_WINDOW minutes ago @@ -457,14 +448,14 @@ def embed_new_content_files(self): .exclude(learning_resource__published=False, learning_resource__test_mode=False) ) - return _replace_with_finalized_chain( + return _replace_with_content_file_group( self, list(new_content_files.values_list("id", flat=True)), overwrite=False, ) -@app.task(bind=True, max_retries=3) +@app.task(bind=True, max_retries=3, ignore_result=True) def embed_run_content_files(self, run_id): """ Embed the run's published content files whose Qdrant points are missing or @@ -568,10 +559,10 @@ def is_stale(pid, checksum, meta): remove_qdrant_records(leftover_ids, CONTENT_FILE_TYPE) if not ids: return None - return _replace_with_finalized_chain(self, ids, overwrite=True) + return _replace_with_content_file_group(self, ids, overwrite=True) -@app.task(bind=True) +@app.task(bind=True, ignore_result=True) def remove_run_content_files(self, run_id): """ Remove content files associated with a run from Qdrant @@ -586,7 +577,7 @@ def remove_run_content_files(self, run_id): return _replace_with_chain(self, tasks) -@app.task(bind=True) +@app.task(bind=True, ignore_result=True) def remove_unpublished_run_content_files(self, run_id): """ Remove unpublished content files associated with a run from Qdrant diff --git a/vector_search/tasks_test.py b/vector_search/tasks_test.py index e821d5bc16..5b8e1754b9 100644 --- a/vector_search/tasks_test.py +++ b/vector_search/tasks_test.py @@ -5,7 +5,6 @@ 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, @@ -32,14 +31,12 @@ from main.utils import now_in_utc from vector_search.constants import CONTENT_FILE_PREPASS_PAYLOAD_FIELDS from vector_search.tasks import ( - _record_embedding_failure, _retry_countdown, 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, @@ -58,15 +55,6 @@ 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): """ @@ -104,16 +92,16 @@ def test_start_embed_resources(mocker, mocked_celery, index): "vector_search.tasks.generate_embeddings", autospec=True ) - with pytest.raises(mocked_celery.replace_exception_class): - start_embed_resources.delay([index], skip_content_files=True, overwrite=True) + # Backfill now dispatches each chunk signature directly (no self.replace). + start_embed_resources.delay([index], skip_content_files=True, overwrite=True) generate_embeddings_mock.si.assert_called_once_with( resource_ids, index, True, # noqa: FBT003 ) - assert mocked_celery.replace.call_count == 1 - assert mocked_celery.replace.call_args[0][1] == mocked_celery.chain.return_value + generate_embeddings_mock.si.return_value.apply_async.assert_called_once_with() + assert mocked_celery.replace.call_count == 0 @pytest.mark.parametrize( @@ -163,10 +151,8 @@ def test_start_embed_resources_excludes_blocklisted_courses(mocker, mocked_celer "vector_search.tasks.generate_embeddings", autospec=True ) - with pytest.raises(mocked_celery.replace_exception_class): - start_embed_resources.delay( - [COURSE_TYPE], skip_content_files=False, overwrite=True - ) + # Backfill now dispatches each chunk signature directly (no self.replace). + start_embed_resources.delay([COURSE_TYPE], skip_content_files=False, overwrite=True) embedded_resource_ids = { resource_id @@ -240,7 +226,7 @@ def test_embed_new_learning_resources(mocker, mocked_celery): list(mocked_celery.group.call_args[0][0]) assert generate_embeddings_mock.si.call_count == 1 - embedded_ids = generate_embeddings_mock.si.mock_calls[0].args[0] + embedded_ids = generate_embeddings_mock.si.call_args_list[0].args[0] assert sorted(new_resource_ids) == sorted(embedded_ids) @@ -282,29 +268,27 @@ 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] + embedded_ids = [ + content_file_id + for mock_call in generate_embeddings_mock.si.call_args_list + for content_file_id in mock_call.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"] + mock_call.kwargs.get("overwrite") is False + for mock_call in generate_embeddings_mock.si.call_args_list ) - chain_args = mocked_celery.chain.call_args.args - assert chain_args[:-1] == tuple( + # content files now dispatch as a group of chunk signatures (no finalize tail) + group_args = list(mocked_celery.group.call_args.args[0]) + assert group_args == [ generate_embeddings_mock.si.return_value - for _ in generate_embeddings_mock.si.mock_calls - ) - assert chain_args[-1] == finalize_embeddings_mock.si.return_value + for _ in generate_embeddings_mock.si.call_args_list + ] + assert mocked_celery.replace.call_args[0][1] == mocked_celery.group.return_value def test_remove_run_content_files(mocker, mocked_celery, settings): @@ -327,13 +311,13 @@ def test_remove_run_content_files(mocker, mocked_celery, settings): removed_ids = [ content_file_id - for mock_call in remove_embeddings_mock.si.mock_calls + for mock_call in remove_embeddings_mock.si.call_args_list for content_file_id in mock_call.args[0] ] assert sorted(removed_ids) == sorted(content_file_ids) assert all( mock_call.args[1] == CONTENT_FILE_TYPE - for mock_call in remove_embeddings_mock.si.mock_calls + for mock_call in remove_embeddings_mock.si.call_args_list ) assert mocked_celery.chain.call_count == 1 assert mocked_celery.replace.call_count == 1 @@ -407,14 +391,13 @@ def test_embed_learning_resources_by_id(mocker, mocked_celery): ) content_ids.append(cf.id) - with pytest.raises(mocked_celery.replace_exception_class): - embed_learning_resources_by_id.delay( - resource_ids, skip_content_files=False, overwrite=True - ) - for mock_call in generate_embeddings_mock.si.mock_calls[1:]: + embed_learning_resources_by_id.delay( + resource_ids, skip_content_files=False, overwrite=True + ) + for mock_call in generate_embeddings_mock.si.call_args_list[1:]: assert mock_call.args[0][0] in content_ids assert mock_call.args[1] == "content_file" - embedded_resource_ids = generate_embeddings_mock.si.mock_calls[0].args[0] + embedded_resource_ids = generate_embeddings_mock.si.call_args_list[0].args[0] assert sorted(resource_ids) == sorted(embedded_resource_ids) @@ -455,10 +438,7 @@ def test_embedded_content_from_all_runs(mocker, mocked_celery): "vector_search.tasks.generate_embeddings", autospec=True ) - with pytest.raises(mocked_celery.replace_exception_class): - start_embed_resources.delay( - ["course"], skip_content_files=False, overwrite=True - ) + start_embed_resources.delay(["course"], skip_content_files=False, overwrite=True) assert all_contentfiles <= _embedded_content_file_ids(generate_embeddings_mock) @@ -492,10 +472,9 @@ def test_embed_by_id_all_runs_excludes_unpublished(mocker, mocked_celery): "vector_search.tasks.generate_embeddings", autospec=True ) - with pytest.raises(mocked_celery.replace_exception_class): - embed_learning_resources_by_id.delay( - [course.learning_resource.id], skip_content_files=False, overwrite=True - ) + embed_learning_resources_by_id.delay( + [course.learning_resource.id], skip_content_files=False, overwrite=True + ) embedded = _embedded_content_file_ids(generate_embeddings_mock) assert published_ids <= embedded @@ -527,11 +506,8 @@ def test_embedded_content_file_without_runs(mocker, mocked_celery): "vector_search.tasks.generate_embeddings", autospec=True ) - with pytest.raises(mocked_celery.replace_exception_class): - start_embed_resources.delay( - ["course"], skip_content_files=False, overwrite=True - ) - embedded_ids = generate_embeddings_mock.mock_calls[-1].args[0] + start_embed_resources.delay(["course"], skip_content_files=False, overwrite=True) + embedded_ids = generate_embeddings_mock.si.call_args_list[-1].args[0] for contentfile_id in contentfiles_with_no_run: assert contentfile_id in embedded_ids @@ -556,14 +532,13 @@ def test_start_embed_resources_program_content_files(mocker, mocked_celery): "vector_search.tasks.generate_embeddings", autospec=True ) - with pytest.raises(mocked_celery.replace_exception_class): - start_embed_resources.delay( - [PROGRAM_TYPE], skip_content_files=False, overwrite=True - ) + start_embed_resources.delay( + [PROGRAM_TYPE], skip_content_files=False, overwrite=True + ) content_file_calls = [ call - for call in generate_embeddings_mock.si.mock_calls + for call in generate_embeddings_mock.si.call_args_list if call.args[1] == "content_file" ] embedded_content_ids = [] @@ -592,14 +567,13 @@ def test_embed_learning_resources_by_id_program_content_files(mocker, mocked_cel "vector_search.tasks.generate_embeddings", autospec=True ) - with pytest.raises(mocked_celery.replace_exception_class): - embed_learning_resources_by_id.delay( - resource_ids, skip_content_files=False, overwrite=True - ) + embed_learning_resources_by_id.delay( + resource_ids, skip_content_files=False, overwrite=True + ) content_file_calls = [ call - for call in generate_embeddings_mock.si.mock_calls + for call in generate_embeddings_mock.si.call_args_list if call.args[1] == "content_file" ] embedded_content_ids = [] @@ -628,19 +602,18 @@ def test_program_embedding_includes_test_mode_unpublished_programs( "vector_search.tasks.generate_embeddings", autospec=True ) - with pytest.raises(mocked_celery.replace_exception_class): - start_embed_resources.delay( - [PROGRAM_TYPE], skip_content_files=False, overwrite=True - ) + start_embed_resources.delay( + [PROGRAM_TYPE], skip_content_files=False, overwrite=True + ) program_calls = [ call - for call in generate_embeddings_mock.si.mock_calls + for call in generate_embeddings_mock.si.call_args_list if call.args[1] == PROGRAM_TYPE ] content_calls = [ call - for call in generate_embeddings_mock.si.mock_calls + for call in generate_embeddings_mock.si.call_args_list if call.args[1] == CONTENT_FILE_TYPE ] assert any(resource_id in call.args[0] for call in program_calls) @@ -648,19 +621,18 @@ def test_program_embedding_includes_test_mode_unpublished_programs( generate_embeddings_mock.reset_mock() - with pytest.raises(mocked_celery.replace_exception_class): - embed_learning_resources_by_id.delay( - [resource_id], skip_content_files=False, overwrite=True - ) + embed_learning_resources_by_id.delay( + [resource_id], skip_content_files=False, overwrite=True + ) program_calls = [ call - for call in generate_embeddings_mock.si.mock_calls + for call in generate_embeddings_mock.si.call_args_list if call.args[1] == PROGRAM_TYPE ] content_calls = [ call - for call in generate_embeddings_mock.si.mock_calls + for call in generate_embeddings_mock.si.call_args_list if call.args[1] == CONTENT_FILE_TYPE ] assert any(resource_id in call.args[0] for call in program_calls) @@ -692,7 +664,7 @@ def test_embed_new_content_files_without_runs(mocker, mocked_celery): with pytest.raises(mocked_celery.replace_exception_class): embed_new_content_files.delay() - embedded_ids = generate_embeddings_mock.si.mock_calls[0].args[0] + embedded_ids = generate_embeddings_mock.si.call_args_list[0].args[0] for contentfile_id in content_files_without_run: assert contentfile_id in embedded_ids @@ -702,7 +674,7 @@ def test_embed_run_content_files(mocker, mocked_celery, settings): embed_run_content_files should replace itself with embedding tasks for all content files associated with the run. """ - settings.QDRANT_CHUNK_SIZE = 2 + settings.QDRANT_CONTENT_FILE_CHUNK_SIZE = 2 run = LearningResourceRunFactory.create() content_file_ids = [ content_file.id @@ -712,37 +684,29 @@ 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) embedded_ids = [ content_file_id - for mock_call in generate_embeddings_mock.si.mock_calls + for mock_call in generate_embeddings_mock.si.call_args_list for content_file_id in mock_call.args[0] ] assert sorted(embedded_ids) == sorted(content_file_ids) assert all( mock_call.args[1:] == (CONTENT_FILE_TYPE,) and mock_call.kwargs["overwrite"] is True - and "failure_key" in mock_call.kwargs - for mock_call in generate_embeddings_mock.si.mock_calls + for mock_call in generate_embeddings_mock.si.call_args_list ) - # chain = all chunk sigs, then the finalize tail - chain_args = mocked_celery.chain.call_args.args - assert chain_args[:-1] == tuple( + # content files dispatch as a group of chunk sigs (no finalize tail) + group_args = list(mocked_celery.group.call_args.args[0]) + assert group_args == [ 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"] - ) + for _ in generate_embeddings_mock.si.call_args_list + ] assert mocked_celery.replace.call_count == 1 + assert mocked_celery.replace.call_args[0][1] == mocked_celery.group.return_value def test_embed_run_content_files_no_content_files(mocker, mocked_celery): @@ -762,12 +726,11 @@ def test_embed_run_content_files_no_content_files(mocker, mocked_celery): def test_embed_run_content_files_no_files_returns_none(mocker, mocked_celery): - """No content files → no chain, no replace, returns None.""" + """No content files → no group, 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() + mocked_celery.group.assert_not_called() def test_embed_run_content_files_skips_unpublished(mocker, mocked_celery, settings): @@ -1213,7 +1176,7 @@ def test_generate_embeddings_retries_on_deadline(mocker): ) 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") + generate_embeddings([1], CONTENT_FILE_TYPE, overwrite=True) retry.assert_called_once() assert retry.call_args.kwargs["countdown"] >= 0 @@ -1226,7 +1189,7 @@ def test_generate_embeddings_retries_on_unavailable(mocker): ) 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") + generate_embeddings([1], CONTENT_FILE_TYPE, overwrite=True) retry.assert_called_once() @@ -1241,35 +1204,31 @@ def test_retry_countdown_is_minutes_scale(mocker): ) -def test_generate_embeddings_records_on_exhaustion(mocker): - """Exhausted deadline + failure_key: record + return, do not raise (chain continues).""" +def test_generate_embeddings_logs_and_raises_on_exhaustion(mocker): + """Exhausted deadline: log the failure and propagate (no swallowing).""" 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") + log_exc = mocker.patch("vector_search.tasks.log.exception") generate_embeddings.push_request(retries=3) try: - assert ( - generate_embeddings([1], CONTENT_FILE_TYPE, overwrite=True, failure_key="k") - is None - ) + with pytest.raises(grpc.RpcError): + generate_embeddings([1], CONTENT_FILE_TYPE, overwrite=True) finally: generate_embeddings.pop_request() - record.assert_called_once_with("k") + log_exc.assert_called_once() -def test_generate_embeddings_records_non_transient_with_key(mocker): - """Non-transient error + failure_key: record + return, do not raise.""" +def test_generate_embeddings_logs_and_raises_non_transient(mocker): + """Non-transient error: log the failure and propagate.""" 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") + log_exc = mocker.patch("vector_search.tasks.log.exception") + with pytest.raises(ValueError, match="boom"): + generate_embeddings([1], CONTENT_FILE_TYPE, overwrite=True) + log_exc.assert_called_once() def test_generate_embeddings_reraises_other_grpc_errors(mocker): @@ -1320,21 +1279,3 @@ 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 750aab6f27..3a1b26c273 100644 --- a/vector_search/utils.py +++ b/vector_search/utils.py @@ -502,6 +502,33 @@ def _content_file_embedding_context(document): return document.get("content", "") +# Sentinel distinguishing "caller did not supply a stored point" from an +# explicit None (point genuinely absent from Qdrant). +_UNSET = object() + +# Cap ids per Qdrant retrieve so large (backfill / healthcheck) batches stay +# under the server's request-size limits. +QDRANT_RETRIEVE_BATCH_SIZE = 256 + + +def _batch_retrieve_points(point_ids, collection_name): + """ + Retrieve points by id in sub-batches, returning ``{str(point_id): point}``. + + One (sub-batched) call replaces a per-document ``client.retrieve`` so a batch + of N documents costs ⌈N/256⌉ round trips instead of N. + """ + client = qdrant_client() + found = {} + unique_ids = list(dict.fromkeys(str(pid) for pid in point_ids)) + for id_batch in chunks(unique_ids, chunk_size=QDRANT_RETRIEVE_BATCH_SIZE): + for point in client.retrieve( + collection_name=collection_name, ids=list(id_batch) + ): + found[str(point.id)] = point + return found + + def _process_resource_embeddings(serialized_resources): docs = [] metadata = [] @@ -509,12 +536,19 @@ def _process_resource_embeddings(serialized_resources): encoder_dense = dense_encoder() encoder_sparse = sparse_encoder() + stored_points = _batch_retrieve_points( + [vector_point_id(vector_point_key(doc)) for doc in serialized_resources], + RESOURCES_COLLECTION_NAME, + ) + for doc in serialized_resources: - if not should_generate_resource_embeddings(doc): - update_learning_resource_payload(doc) + point_id = vector_point_id(vector_point_key(doc)) + stored_point = stored_points.get(point_id) + if not should_generate_resource_embeddings(doc, stored_point=stored_point): + update_learning_resource_payload(doc, stored_point=stored_point) continue metadata.append(doc) - ids.append(vector_point_id(vector_point_key(doc))) + ids.append(point_id) docs.append(_learning_resource_embedding_context(doc)) if len(docs) > 0: embeddings = encoder_dense.embed_documents(docs) @@ -528,10 +562,20 @@ def _process_resource_embeddings(serialized_resources): return None -def update_learning_resource_payload(serialized_document): +def update_learning_resource_payload(serialized_document, stored_point=_UNSET): """ Refresh a resource's Qdrant payload without re-embedding. + + When the caller supplies the already-retrieved stored point and its payload + matches, the write is skipped -- avoids a Qdrant write per resource on every + metadata sweep even when nothing changed. """ + if ( + stored_point is not _UNSET + and stored_point is not None + and stored_point.payload == serialized_document + ): + return point_id = vector_point_id(vector_point_key(serialized_document)) qdrant_client().overwrite_payload( collection_name=RESOURCES_COLLECTION_NAME, @@ -593,20 +637,24 @@ def _set_payload(points, document, param_map, collection_name): ) -def should_generate_resource_embeddings(serialized_document): +def should_generate_resource_embeddings(serialized_document, stored_point=_UNSET): """ - Determine if we should generate embeddings for a learning resource + Determine if we should generate embeddings for a learning resource. + + Pass stored_point (the already-retrieved point, or None if absent) to reuse a + batched retrieve instead of issuing one retrieve per document. """ - client = qdrant_client() - point_id = vector_point_id(vector_point_key(serialized_document)) - response = client.retrieve( - collection_name=RESOURCES_COLLECTION_NAME, - ids=[point_id], - ) - if len(response) > 0: - resource_payload = response[0].payload + if stored_point is _UNSET: + client = qdrant_client() + point_id = vector_point_id(vector_point_key(serialized_document)) + response = client.retrieve( + collection_name=RESOURCES_COLLECTION_NAME, + ids=[point_id], + ) + stored_point = response[0] if response else None + if stored_point is not None: stored_embedding_content = _learning_resource_embedding_context( - resource_payload + stored_point.payload ) current_embedding_content = _learning_resource_embedding_context( serialized_document @@ -662,12 +710,19 @@ def _stored_content_payloads( def should_generate_content_embeddings( - serialized_document: dict, point_id: str | None = None + serialized_document: dict, point_id: str | None = None, stored_point=_UNSET ) -> bool: """ - Determine if we should generate embeddings for a content file + Determine if we should generate embeddings for a content file. + + Pass stored_point (the already-retrieved chunk-0 point, or None if absent) to + reuse a batched retrieve instead of issuing one retrieve per document. """ - point = _retrieve_content_file_point(serialized_document, point_id=point_id) + point = ( + _retrieve_content_file_point(serialized_document, point_id=point_id) + if stored_point is _UNSET + else stored_point + ) if not point: return True qdrant_checksum = (point.payload or {}).get("checksum") @@ -688,6 +743,15 @@ def _embed_course_metadata_as_contentfile(serialized_resources): metadata = [] ids = [] docs = [] + # One batched retrieve of the course-info chunk-0 points, keyed by point id, + # instead of a per-resource retrieve inside the loop. + stored_points = _batch_retrieve_points( + [ + vector_point_id(vector_point_key(doc, document_type="course_information")) + for doc in serialized_resources + ], + CONTENT_FILES_COLLECTION_NAME, + ) for doc in serialized_resources: resource_vector_point_id = str(vector_point_id(vector_point_key(doc))) serializer = LearningResourceMetadataDisplaySerializer( @@ -705,7 +769,9 @@ def _embed_course_metadata_as_contentfile(serialized_resources): vector_point_key(doc, document_type="course_information") ) if not should_generate_content_embeddings( - serialized_document, document_point_id + serialized_document, + document_point_id, + stored_point=stored_points.get(str(document_point_id)), ): continue diff --git a/vector_search/utils_test.py b/vector_search/utils_test.py index 37ff207cc7..c5607689c7 100644 --- a/vector_search/utils_test.py +++ b/vector_search/utils_test.py @@ -59,6 +59,8 @@ ) from vector_search.encoders.utils import dense_encoder, sparse_encoder from vector_search.utils import ( + QDRANT_RETRIEVE_BATCH_SIZE, + _batch_retrieve_points, _chunk_documents, _chunk_markdown_documents, _embed_course_metadata_as_contentfile, @@ -1593,6 +1595,9 @@ def test_embed_course_metadata_as_contentfile_uploads_points_on_change(mocker): record that matches the checksum of metadata doc """ mock_point = mocker.Mock() + mock_point.id = vector_point_id( + vector_point_key(serialized_resource, document_type="course_information") + ) mock_point.payload = {"checksum": "checksum2"} mock_client.retrieve.return_value = [mock_point] @@ -2657,3 +2662,29 @@ def test_check_missing_content_file_ids_skips_unimportant_block_types(mocker): mock_present.assert_not_called() mock_client.count.assert_not_called() mock_log.assert_not_called() + + +def test_batch_retrieve_points_subbatches_at_cap(mocker): + """_batch_retrieve_points splits ids into retrieve calls capped at 256.""" + mock_client = mocker.patch("vector_search.utils.qdrant_client") + mock_client.return_value.retrieve.return_value = [] + ids = [str(i) for i in range(QDRANT_RETRIEVE_BATCH_SIZE * 2 + 5)] + + _batch_retrieve_points(ids, "some_collection") + + # ceil((2*256+5)/256) == 3 retrieve calls, none exceeding the cap + assert mock_client.return_value.retrieve.call_count == 3 + for call in mock_client.return_value.retrieve.call_args_list: + assert len(call.kwargs["ids"]) <= QDRANT_RETRIEVE_BATCH_SIZE + + +def test_update_learning_resource_payload_skips_write_when_stored_matches(mocker): + """Unchanged resource payload on a sweep → no overwrite_payload write.""" + doc = {"readable_id": "r1", "title": "t"} + stored = mocker.MagicMock() + stored.payload = dict(doc) + mock_client = mocker.patch("vector_search.utils.qdrant_client") + + update_learning_resource_payload(doc, stored_point=stored) + + mock_client.return_value.overwrite_payload.assert_not_called() From 8ae701ae06012f8fe4c1b89ddae571468a58dba6 Mon Sep 17 00:00:00 2001 From: Matt Bertrand Date: Sun, 19 Jul 2026 10:16:06 -0400 Subject: [PATCH 2/6] Add embed_run_content_files test for a run with only unpublished files Covers the published-only filter returning None (no group), per review feedback. Co-Authored-By: Claude Opus 4.8 --- vector_search/tasks_test.py | 9 +++++++++ 1 file changed, 9 insertions(+) diff --git a/vector_search/tasks_test.py b/vector_search/tasks_test.py index 5b8e1754b9..827ebeb5b6 100644 --- a/vector_search/tasks_test.py +++ b/vector_search/tasks_test.py @@ -961,6 +961,15 @@ def test_embed_run_content_files_all_unchanged_dispatches_nothing( mocked_celery.chain.assert_not_called() +def test_embed_run_content_files_only_unpublished_returns_none(mocker, mocked_celery): + """A run whose content files are all unpublished embeds nothing (published filter).""" + run = LearningResourceRunFactory.create() + ContentFileFactory.create_batch(2, run=run, published=False) + mocker.patch("vector_search.tasks.generate_embeddings", autospec=True) + assert embed_run_content_files(run.id) is None + mocked_celery.group.assert_not_called() + + def test_embeddings_healthcheck_no_missing_embeddings(mocker): """ Test embeddings_healthcheck when there are no missing embeddings From ea5db544297cb8b61ed3d2717e6f3289b92df0fa Mon Sep 17 00:00:00 2001 From: Matt Bertrand Date: Tue, 28 Jul 2026 09:04:50 -0400 Subject: [PATCH 3/6] Purge before embed; drop dead resource payload write-skip - Reorder the content_files_loaded chain to purge unpublished files' points before embedding (cherry-picked from the change-detection follow-up). celery uplifts a group replacement to a chord, so with embed first the purge task becomes the chord body; if the gap between two chunk completions exceeds result_expires (1h) the chord's Redis bookkeeping expires and the body is dropped with no error logged. Purging first makes the body a bare celery.accumulate, whose loss costs nothing. The two tasks touch disjoint sets (published=False vs published=True). - Drop the payload-equality write-skip in update_learning_resource_payload. It never fired: Qdrant returns resource_age_date as a string while the serialized doc holds a datetime, so the payloads never compared equal. Its test asserted against a MagicMock echoing its own input, so the no-op looked covered. The batched retrieve it rode on is unaffected and still saves N retrieves per sweep. Co-Authored-By: Claude Opus 5 (1M context) --- vector_search/utils.py | 14 ++------------ vector_search/utils_test.py | 12 ------------ 2 files changed, 2 insertions(+), 24 deletions(-) diff --git a/vector_search/utils.py b/vector_search/utils.py index 3a1b26c273..57f47e75e3 100644 --- a/vector_search/utils.py +++ b/vector_search/utils.py @@ -545,7 +545,7 @@ def _process_resource_embeddings(serialized_resources): point_id = vector_point_id(vector_point_key(doc)) stored_point = stored_points.get(point_id) if not should_generate_resource_embeddings(doc, stored_point=stored_point): - update_learning_resource_payload(doc, stored_point=stored_point) + update_learning_resource_payload(doc) continue metadata.append(doc) ids.append(point_id) @@ -562,20 +562,10 @@ def _process_resource_embeddings(serialized_resources): return None -def update_learning_resource_payload(serialized_document, stored_point=_UNSET): +def update_learning_resource_payload(serialized_document): """ Refresh a resource's Qdrant payload without re-embedding. - - When the caller supplies the already-retrieved stored point and its payload - matches, the write is skipped -- avoids a Qdrant write per resource on every - metadata sweep even when nothing changed. """ - if ( - stored_point is not _UNSET - and stored_point is not None - and stored_point.payload == serialized_document - ): - return point_id = vector_point_id(vector_point_key(serialized_document)) qdrant_client().overwrite_payload( collection_name=RESOURCES_COLLECTION_NAME, diff --git a/vector_search/utils_test.py b/vector_search/utils_test.py index c5607689c7..943579cbfb 100644 --- a/vector_search/utils_test.py +++ b/vector_search/utils_test.py @@ -2676,15 +2676,3 @@ def test_batch_retrieve_points_subbatches_at_cap(mocker): assert mock_client.return_value.retrieve.call_count == 3 for call in mock_client.return_value.retrieve.call_args_list: assert len(call.kwargs["ids"]) <= QDRANT_RETRIEVE_BATCH_SIZE - - -def test_update_learning_resource_payload_skips_write_when_stored_matches(mocker): - """Unchanged resource payload on a sweep → no overwrite_payload write.""" - doc = {"readable_id": "r1", "title": "t"} - stored = mocker.MagicMock() - stored.payload = dict(doc) - mock_client = mocker.patch("vector_search.utils.qdrant_client") - - update_learning_resource_payload(doc, stored_point=stored) - - mock_client.return_value.overwrite_payload.assert_not_called() From 7f03162485ba204665cc62f37e08b8e0d147f546 Mon Sep 17 00:00:00 2001 From: Matt Bertrand Date: Tue, 28 Jul 2026 17:20:50 -0400 Subject: [PATCH 4/6] enhancements --- learning_resources/tasks.py | 14 ++- learning_resources/tasks_test.py | 6 +- main/settings.py | 18 +++ main/settings_celery.py | 24 +++- vector_search/encoders/litellm.py | 13 ++- vector_search/encoders/litellm_test.py | 20 ++++ .../commands/generate_embeddings.py | 8 +- vector_search/tasks.py | 72 +++++++----- vector_search/tasks_test.py | 109 ++++++++++++------ vector_search/utils.py | 76 ++++++------ 10 files changed, 239 insertions(+), 121 deletions(-) diff --git a/learning_resources/tasks.py b/learning_resources/tasks.py index 6bd4b0a37e..c1eb75dd9a 100644 --- a/learning_resources/tasks.py +++ b/learning_resources/tasks.py @@ -699,11 +699,15 @@ def scrape_marketing_pages(self): course_tasks = [ marketing_page_for_resources.si(ids) - for ids in chunks(missing_course_ids, chunk_size=settings.QDRANT_CHUNK_SIZE) + for ids in chunks( + missing_course_ids, chunk_size=settings.MARKETING_PAGE_SCRAPE_CHUNK_SIZE + ) ] program_tasks = [ marketing_page_for_resources.si(ids) - for ids in chunks(sorted(program_ids), chunk_size=settings.QDRANT_CHUNK_SIZE) + for ids in chunks( + sorted(program_ids), chunk_size=settings.MARKETING_PAGE_SCRAPE_CHUNK_SIZE + ) ] if course_tasks and program_tasks: @@ -782,8 +786,10 @@ def marketing_page_for_resources(resource_ids): content_file_ids.append(content_file.id) if content_file.published: upsert_content_file.delay(content_file.id) - if content_file_ids: - generate_embeddings.delay(content_file_ids, CONTENT_FILE_TYPE, overwrite=True) + for ids in chunks( + content_file_ids, chunk_size=settings.QDRANT_CONTENT_FILE_CHUNK_SIZE + ): + generate_embeddings.delay(ids, CONTENT_FILE_TYPE, overwrite=True) @app.task( diff --git a/learning_resources/tasks_test.py b/learning_resources/tasks_test.py index bfd9d591a0..cb2ea32999 100644 --- a/learning_resources/tasks_test.py +++ b/learning_resources/tasks_test.py @@ -786,7 +786,7 @@ def test_scrape_marketing_pages(mocker, settings, mocked_celery): """Test that scrape_marketing_pages correctly identifies resources without marketing pages""" settings.EMBEDDINGS_EXTERNAL_FETCH_USE_WEBDRIVER = True - settings.QDRANT_CHUNK_SIZE = 2 + settings.MARKETING_PAGE_SCRAPE_CHUNK_SIZE = 2 course1 = models.LearningResource.objects.create( title="Course 1", @@ -842,7 +842,7 @@ def test_scrape_marketing_pages_orders_courses_before_programs( mocker, settings, mocked_celery ): """Courses are scraped in a group that runs before the programs group.""" - settings.QDRANT_CHUNK_SIZE = 10 + settings.MARKETING_PAGE_SCRAPE_CHUNK_SIZE = 10 course = models.LearningResource.objects.create( title="Course", url="https://example.com/course", @@ -891,7 +891,7 @@ def test_scrape_marketing_pages_queues_healable_programs( """A program that already has a page but is missing its children section (with a child course page available) is queued for re-scrape. """ - settings.QDRANT_CHUNK_SIZE = 10 + settings.MARKETING_PAGE_SCRAPE_CHUNK_SIZE = 10 course = models.LearningResource.objects.create( title="Course", url="https://example.com/course", diff --git a/main/settings.py b/main/settings.py index 22cc15b2d3..9f63f0425f 100644 --- a/main/settings.py +++ b/main/settings.py @@ -795,6 +795,14 @@ def get_all_config_keys(): default=25, ) +# Number of ids per remove_embeddings task. Separate from QDRANT_CHUNK_SIZE +# because removal issues one filter-based delete per id, and delete volume +# drives Qdrant's vacuum optimizer (CPU spikes). +QDRANT_DELETE_CHUNK_SIZE = get_int( + name="QDRANT_DELETE_CHUNK_SIZE", + default=10, +) + QDRANT_ENCODER = get_string( name="QDRANT_ENCODER", default="vector_search.encoders.gensim.GensimEncoder" ) @@ -839,9 +847,19 @@ def get_all_config_keys(): "EMBEDDINGS_EXTERNAL_FETCH_USE_WEBDRIVER", default=False ) WEBDRIVER_WAIT_SECONDS = get_int(name="WEBDRIVER_WAIT_SECONDS", default=10) + +# Resources scraped per marketing-page task. Each page costs two webdriver +# waits, so this bounds task runtime independently of the embedding chunk size. +MARKETING_PAGE_SCRAPE_CHUNK_SIZE = get_int( + name="MARKETING_PAGE_SCRAPE_CHUNK_SIZE", default=10 +) + LITELLM_TOKEN_ENCODING_NAME = get_string( name="LITELLM_TOKEN_ENCODING_NAME", default=None ) +# Documents per embedding request. 25 * the 8191-token per-document limit stays +# under OpenAI's 300k-token-per-request cap, which fails as a hard 400. +LITELLM_EMBEDDING_BATCH_SIZE = get_int(name="LITELLM_EMBEDDING_BATCH_SIZE", default=25) LITELLM_CUSTOM_PROVIDER = get_string(name="LITELLM_CUSTOM_PROVIDER", default="openai") LITELLM_API_BASE = get_string(name="LITELLM_API_BASE", default=None) diff --git a/main/settings_celery.py b/main/settings_celery.py index 026408608c..b459e6dc5c 100644 --- a/main/settings_celery.py +++ b/main/settings_celery.py @@ -25,12 +25,14 @@ # (the default here) that saturates memory. Keep results only long enough for # chord callbacks to consume them. CELERY_RESULT_EXPIRES = get_int("CELERY_RESULT_EXPIRES", 60 * 60) -# visibility_timeout must exceed the longest possible task runtime (inline -# summarization + embedding) plus max retry backoff (600s), or in-flight tasks -# get redelivered while still running, amplifying load during a backlog. +# visibility_timeout must exceed the longest possible task runtime (bounded by +# CELERY_EMBEDDINGS_TIME_LIMIT below) plus max retry backoff (600s), or in-flight +# tasks get redelivered while still running, amplifying load during a backlog. +# It is also the recovery delay for messages orphaned by a hard pod kill, so keep +# it only as high as the longest task needs. CELERY_BROKER_TRANSPORT_OPTIONS = { "visibility_timeout": get_int( - name="CELERY_BROKER_VISIBILITY_TIMEOUT", default=6 * 60 * 60 + name="CELERY_BROKER_VISIBILITY_TIMEOUT", default=2 * 60 * 60 ), } CELERY_BEAT_SCHEDULER = RedBeatScheduler @@ -229,5 +231,15 @@ CELERY_VECTOR_SEARCH_RATE_LIMIT = get_string( "CELERY_VECTOR_SEARCH_RATE_LIMIT", CELERY_RATE_LIMIT ) -# Per-worker execution rate for generate_embeddings; tunable without a deploy. -CELERY_EMBEDDINGS_RATE_LIMIT = get_string("CELERY_EMBEDDINGS_RATE_LIMIT", "200/m") +# Per-worker execution rate for generate_embeddings, counted in *tasks* per +# minute, so the item rate against Qdrant is this times the chunk size. Kept at +# 20/m so the QDRANT_CHUNK_SIZE=100 resource chunks hold the pre-chunking item +# rate (~2k/min/worker) instead of multiplying it. Changing it needs a worker +# restart; use celery's control.rate_limit for a live change. +CELERY_EMBEDDINGS_RATE_LIMIT = get_string("CELERY_EMBEDDINGS_RATE_LIMIT", "20/m") +# Bound generate_embeddings runtime so a wedged summarize/embed call cannot hold +# a message unacked for the whole visibility_timeout. Soft limit raises +# SoftTimeLimitExceeded (logged + failed like any other terminal error); the hard +# limit is the backstop if the task ignores it. +CELERY_EMBEDDINGS_SOFT_TIME_LIMIT = get_int("CELERY_EMBEDDINGS_SOFT_TIME_LIMIT", 1800) +CELERY_EMBEDDINGS_TIME_LIMIT = get_int("CELERY_EMBEDDINGS_TIME_LIMIT", 2400) diff --git a/vector_search/encoders/litellm.py b/vector_search/encoders/litellm.py index 89b0e5f9e5..017a6ee819 100644 --- a/vector_search/encoders/litellm.py +++ b/vector_search/encoders/litellm.py @@ -1,5 +1,6 @@ import logging import os +from itertools import batched from urllib.parse import urlparse import litellm @@ -44,7 +45,17 @@ def __init__(self, model_name): log.warning(msg) def embed_documents(self, documents): - return [result["embedding"] for result in self.get_embedding(documents)["data"]] + # Cap inputs per request: each document is truncated to the model's input + # limit upstream (8191 tokens for text-embedding-3-*), so a bounded count + # keeps a request under the provider's per-request token cap (OpenAI: + # 300k), which would otherwise fail as a non-transient 400. + embeddings = [] + for batch in batched(documents, settings.LITELLM_EMBEDDING_BATCH_SIZE): + embeddings.extend( + result["embedding"] + for result in self.get_embedding(list(batch))["data"] + ) + return embeddings def get_embedding(self, texts): if self.cache: diff --git a/vector_search/encoders/litellm_test.py b/vector_search/encoders/litellm_test.py index 94aec7b072..a1023aa178 100644 --- a/vector_search/encoders/litellm_test.py +++ b/vector_search/encoders/litellm_test.py @@ -33,6 +33,26 @@ def test_litellm_encoder_cache_enabled(mock_embedding): mock_embedding.assert_called_once_with(**expected_kwargs) +@patch("vector_search.encoders.litellm.embedding") +def test_litellm_encoder_batches_requests(mock_embedding, settings): + """ + Documents are split into requests of LITELLM_EMBEDDING_BATCH_SIZE: one + oversized request is a hard 400 from the provider, not a retryable error. + """ + settings.LITELLM_EMBEDDING_BATCH_SIZE = 2 + mock_embedding.return_value.to_dict.side_effect = lambda: { + "data": [{"embedding": [0.1]}] * len(mock_embedding.call_args.kwargs["input"]) + } + + embeddings = LiteLLMEncoder("test_model").embed_documents(["a", "b", "c"]) + + assert len(embeddings) == 3 + assert [call.kwargs["input"] for call in mock_embedding.call_args_list] == [ + ["a", "b"], + ["c"], + ] + + @patch("vector_search.encoders.litellm.embedding") def test_litellm_encoder_cache_disabled(mock_embedding): """ diff --git a/vector_search/management/commands/generate_embeddings.py b/vector_search/management/commands/generate_embeddings.py index 6076f3b046..a440481f63 100644 --- a/vector_search/management/commands/generate_embeddings.py +++ b/vector_search/management/commands/generate_embeddings.py @@ -98,14 +98,16 @@ def handle(self, *args, **options): # noqa: ARG002 f" Types to embed: {indexes_to_update}" ) - self.stdout.write("Waiting on task...") + self.stdout.write("Waiting on dispatch...") start = now_in_utc() error = task.get() if error: - msg = f"Geenerate embeddings errored: {error}" + msg = f"Generate embeddings errored: {error}" raise CommandError(msg) clear_views_cache() total_seconds = (now_in_utc() - start).total_seconds() self.stdout.write( - f"Embeddings generated and stored, took {total_seconds} seconds" + f"Embedding tasks dispatched in {total_seconds} seconds. Embedding runs " + "in the background on the embeddings queue; check worker logs or the " + "embeddings healthcheck for completion." ) diff --git a/vector_search/tasks.py b/vector_search/tasks.py index 916148856e..1db8a670c3 100644 --- a/vector_search/tasks.py +++ b/vector_search/tasks.py @@ -76,41 +76,42 @@ def _replace_with_chain(task, task_signatures): return task.replace(celery.chain(*task_signatures)) -def _replace_with_content_file_group( - task: celery.Task, content_file_ids: list[int], *, overwrite: bool +def _dispatch_content_file_chunks( + content_file_ids: list[int], *, overwrite: bool ) -> None: """ - Replace the parent with a group of content-file embedding chunks. + Dispatch content-file embedding chunks, fire and forget. - A group (not a chain) so chunks are not serialized behind one another and the - remaining work is not carried in every message; returns None when there is - nothing to embed. Failures surface via generate_embeddings' own logging and - the periodic embeddings_healthcheck, not a chain tail. + Chunks are published individually rather than via ``task.replace(group(...))``: + celery uplifts a replaced group into a chord, which keeps per-chunk result + bookkeeping in Redis even for ignore_result tasks. Callers are chain tails, + so nothing downstream needs to wait. Failures surface via + generate_embeddings' own logging/Sentry and the periodic + embeddings_healthcheck. """ - sigs = [ + _dispatch_signatures( generate_embeddings.si(ids, CONTENT_FILE_TYPE, overwrite=overwrite) for ids in chunks( content_file_ids, chunk_size=settings.QDRANT_CONTENT_FILE_CHUNK_SIZE ) - ] - if not sigs: - return None - return task.replace(celery.group(sigs)) + ) -def _dispatch_signatures(task_signatures) -> int: +def _dispatch_signatures(task_signatures) -> None: """ - Fire-and-forget dispatch of already-built chunk signatures (backfill paths). + Fire-and-forget dispatch of already-built chunk signatures. Publishes one message per chunk rather than a single catalog-wide group/chain, so a full-catalog backfill does not spike the broker with one huge message. - Nothing waits on the results (embedding tasks are ignore_result). + Nothing waits on the results (embedding tasks are ignore_result), so this + returns None: callers' return value stays "falsy on success, error string on + failure". """ count = 0 for sig in task_signatures: sig.apply_async() count += 1 - return count + log.info("Dispatched %d embedding chunk task(s)", count) def _queue_program_content_file_embedding_tasks(index_tasks, program_ids, overwrite): @@ -131,7 +132,7 @@ def _queue_program_content_file_embedding_tasks(index_tasks, program_ids, overwr generate_embeddings.si(ids, CONTENT_FILE_TYPE, overwrite) for ids in chunks( contentfile_ids, - chunk_size=settings.QDRANT_CHUNK_SIZE, + chunk_size=settings.QDRANT_CONTENT_FILE_CHUNK_SIZE, ) ] ) @@ -153,6 +154,8 @@ def _retry_countdown(retries: int) -> int: reject_on_worker_lost=True, max_retries=3, rate_limit=settings.CELERY_EMBEDDINGS_RATE_LIMIT, + soft_time_limit=settings.CELERY_EMBEDDINGS_SOFT_TIME_LIMIT, + time_limit=settings.CELERY_EMBEDDINGS_TIME_LIMIT, ignore_result=True, ) def generate_embeddings( @@ -160,6 +163,7 @@ def generate_embeddings( ids: list[int], resource_type: str, overwrite: bool, # noqa: FBT001 + failure_key: str | None = None, # noqa: ARG001 ) -> None: """ Generate learning resource embeddings and index in Qdrant. @@ -167,6 +171,10 @@ def generate_embeddings( Retries transient Qdrant/search errors with jittered backoff. On exhaustion or a non-transient error, logs and propagates the failure (surfaced via Sentry and the periodic embeddings_healthcheck). + + failure_key is accepted and ignored so messages published by the previous + release (which chained a finalize_embeddings tail) still run on a new worker. + Remove it, and finalize_embeddings below, one release after deploy. """ try: with wrap_retry_exception(*SEARCH_CONN_EXCEPTIONS): @@ -221,6 +229,15 @@ def remove_embeddings(ids, resource_type): raise +@app.task(ignore_result=True) +def finalize_embeddings(failure_key: str | None = None) -> None: + """ + No-op retained for one release so finalize_embeddings tails already queued by + the previous release do not fail as an unregistered task. Chunk failures now + surface via generate_embeddings' logging/Sentry. + """ + + @app.task def start_embed_resources(indexes, skip_content_files, overwrite): # noqa: C901 """ @@ -276,7 +293,7 @@ def start_embed_resources(indexes, skip_content_files, overwrite): # noqa: C901 generate_embeddings.si(ids, CONTENT_FILE_TYPE, overwrite) for ids in chunks( contentfiles, - chunk_size=settings.QDRANT_CHUNK_SIZE, + chunk_size=settings.QDRANT_CONTENT_FILE_CHUNK_SIZE, ) ] for resource_type in set(LEARNING_RESOURCE_TYPES) - {COURSE_TYPE}: @@ -381,7 +398,7 @@ def embed_learning_resources_by_id(ids, skip_content_files, overwrite): generate_embeddings.si(ids, CONTENT_FILE_TYPE, overwrite) for ids in chunks( content_ids, - chunk_size=settings.QDRANT_CHUNK_SIZE, + chunk_size=settings.QDRANT_CONTENT_FILE_CHUNK_SIZE, ) ] @@ -431,8 +448,8 @@ def embed_new_learning_resources(self): return self.replace(embed_tasks) -@app.task(bind=True, ignore_result=True) -def embed_new_content_files(self): +@app.task(ignore_result=True) +def embed_new_content_files(): """ Embed new content files from QDRANT_EMBEDDINGS_TASK_LOOKBACK_WINDOW minutes ago """ @@ -448,8 +465,7 @@ def embed_new_content_files(self): .exclude(learning_resource__published=False, learning_resource__test_mode=False) ) - return _replace_with_content_file_group( - self, + return _dispatch_content_file_chunks( list(new_content_files.values_list("id", flat=True)), overwrite=False, ) @@ -559,7 +575,7 @@ def is_stale(pid, checksum, meta): remove_qdrant_records(leftover_ids, CONTENT_FILE_TYPE) if not ids: return None - return _replace_with_content_file_group(self, ids, overwrite=True) + return _dispatch_content_file_chunks(ids, overwrite=True) @app.task(bind=True, ignore_result=True) @@ -572,7 +588,9 @@ def remove_run_content_files(self, run_id): ) tasks = [ remove_embeddings.si(ids, CONTENT_FILE_TYPE) - for ids in chunks(content_file_ids, chunk_size=settings.QDRANT_CHUNK_SIZE) + for ids in chunks( + content_file_ids, chunk_size=settings.QDRANT_DELETE_CHUNK_SIZE + ) ] return _replace_with_chain(self, tasks) @@ -589,7 +607,9 @@ def remove_unpublished_run_content_files(self, run_id): ) tasks = [ remove_embeddings.si(ids, CONTENT_FILE_TYPE) - for ids in chunks(content_file_ids, chunk_size=settings.QDRANT_CHUNK_SIZE) + for ids in chunks( + content_file_ids, chunk_size=settings.QDRANT_DELETE_CHUNK_SIZE + ) ] return _replace_with_chain(self, tasks) diff --git a/vector_search/tasks_test.py b/vector_search/tasks_test.py index 827ebeb5b6..c3133f24bc 100644 --- a/vector_search/tasks_test.py +++ b/vector_search/tasks_test.py @@ -104,6 +104,55 @@ def test_start_embed_resources(mocker, mocked_celery, index): assert mocked_celery.replace.call_count == 0 +def test_start_embed_resources_dispatches_one_message_per_chunk( + mocker, mocked_celery, settings +): + """ + Each chunk signature is published individually, and the task returns None so + the management command does not read a dispatch count as an error. + """ + settings.QDRANT_CHUNK_SIZE = 2 + LearningResourceFactory.create_batch(5, resource_type=PROGRAM_TYPE) + generate_embeddings_mock = mocker.patch( + "vector_search.tasks.generate_embeddings", autospec=True + ) + + assert ( + start_embed_resources([PROGRAM_TYPE], skip_content_files=True, overwrite=True) + is None + ) + + assert generate_embeddings_mock.si.call_count == 3 + assert generate_embeddings_mock.si.return_value.apply_async.call_count == 3 + + +def test_start_embed_resources_content_files_use_content_file_chunk_size( + mocker, mocked_celery, settings +): + """ + Content-file chunks in the backfill use QDRANT_CONTENT_FILE_CHUNK_SIZE: these + tasks summarize inline, so they must not inherit the larger resource chunk. + """ + settings.QDRANT_CHUNK_SIZE = 100 + settings.QDRANT_CONTENT_FILE_CHUNK_SIZE = 2 + mocker.patch("vector_search.tasks.load_course_blocklist", return_value=[]) + course = CourseFactory.create(etl_source=ETLSource.ocw.value) + ContentFileFactory.create_batch(5, run=course.learning_resource.runs.first()) + generate_embeddings_mock = mocker.patch( + "vector_search.tasks.generate_embeddings", autospec=True + ) + + start_embed_resources([COURSE_TYPE], skip_content_files=False, overwrite=True) + + content_file_calls = [ + call + for call in generate_embeddings_mock.si.call_args_list + if call.args[1] == CONTENT_FILE_TYPE + ] + assert len(content_file_calls) == 3 + assert all(len(call.args[0]) <= 2 for call in content_file_calls) + + @pytest.mark.parametrize( "index", ["course", "program"], @@ -269,8 +318,7 @@ def test_embed_new_content_files(mocker, mocked_celery): "vector_search.tasks.generate_embeddings", autospec=True ) - with pytest.raises(mocked_celery.replace_exception_class): - embed_new_content_files.delay() + embed_new_content_files.delay() embedded_ids = [ content_file_id @@ -282,13 +330,13 @@ def test_embed_new_content_files(mocker, mocked_celery): mock_call.kwargs.get("overwrite") is False for mock_call in generate_embeddings_mock.si.call_args_list ) - # content files now dispatch as a group of chunk signatures (no finalize tail) - group_args = list(mocked_celery.group.call_args.args[0]) - assert group_args == [ - generate_embeddings_mock.si.return_value - for _ in generate_embeddings_mock.si.call_args_list - ] - assert mocked_celery.replace.call_args[0][1] == mocked_celery.group.return_value + # content-file chunks are published individually: no group/chord bookkeeping + assert ( + generate_embeddings_mock.si.return_value.apply_async.call_count + == generate_embeddings_mock.si.call_count + ) + assert mocked_celery.group.call_count == 0 + assert mocked_celery.replace.call_count == 0 def test_remove_run_content_files(mocker, mocked_celery, settings): @@ -662,8 +710,7 @@ def test_embed_new_content_files_without_runs(mocker, mocked_celery): "vector_search.tasks.generate_embeddings", autospec=True ) - with pytest.raises(mocked_celery.replace_exception_class): - embed_new_content_files.delay() + embed_new_content_files.delay() embedded_ids = generate_embeddings_mock.si.call_args_list[0].args[0] for contentfile_id in content_files_without_run: assert contentfile_id in embedded_ids @@ -671,8 +718,8 @@ def test_embed_new_content_files_without_runs(mocker, mocked_celery): def test_embed_run_content_files(mocker, mocked_celery, settings): """ - embed_run_content_files should replace itself with embedding tasks for all - content files associated with the run. + embed_run_content_files should dispatch embedding tasks for all content files + associated with the run, one message per chunk. """ settings.QDRANT_CONTENT_FILE_CHUNK_SIZE = 2 run = LearningResourceRunFactory.create() @@ -685,8 +732,7 @@ def test_embed_run_content_files(mocker, mocked_celery, settings): "vector_search.tasks.generate_embeddings", autospec=True ) - with pytest.raises(mocked_celery.replace_exception_class): - embed_run_content_files.delay(run.id) + embed_run_content_files.delay(run.id) embedded_ids = [ content_file_id @@ -699,14 +745,11 @@ def test_embed_run_content_files(mocker, mocked_celery, settings): and mock_call.kwargs["overwrite"] is True for mock_call in generate_embeddings_mock.si.call_args_list ) - # content files dispatch as a group of chunk sigs (no finalize tail) - group_args = list(mocked_celery.group.call_args.args[0]) - assert group_args == [ - generate_embeddings_mock.si.return_value - for _ in generate_embeddings_mock.si.call_args_list - ] - assert mocked_celery.replace.call_count == 1 - assert mocked_celery.replace.call_args[0][1] == mocked_celery.group.return_value + # 3 files at chunk size 2 -> 2 chunks, each published on its own + assert generate_embeddings_mock.si.call_count == 2 + assert generate_embeddings_mock.si.return_value.apply_async.call_count == 2 + assert mocked_celery.group.call_count == 0 + assert mocked_celery.replace.call_count == 0 def test_embed_run_content_files_no_content_files(mocker, mocked_celery): @@ -735,7 +778,7 @@ def test_embed_run_content_files_no_files_returns_none(mocker, mocked_celery): def test_embed_run_content_files_skips_unpublished(mocker, mocked_celery, settings): """Unpublished files are never embedded.""" - settings.QDRANT_CHUNK_SIZE = 50 + settings.QDRANT_CONTENT_FILE_CHUNK_SIZE = 50 run = LearningResourceRunFactory.create() published = ContentFileFactory.create(run=run, published=True, content="aaa") ContentFileFactory.create(run=run, published=False, content="bbb") @@ -743,8 +786,7 @@ def test_embed_run_content_files_skips_unpublished(mocker, mocked_celery, settin "vector_search.tasks.generate_embeddings", autospec=True ) - with pytest.raises(mocked_celery.replace_exception_class): - embed_run_content_files.delay(run.id) + embed_run_content_files.delay(run.id) assert _embedded_content_file_ids(generate_embeddings_mock) == {published.id} @@ -754,7 +796,7 @@ def test_embed_run_content_files_skips_contentless(mocker, mocked_celery, settin Files without content never produce Qdrant points, so the pre-pass must not flag them as stale (they would otherwise be re-dispatched on every load). """ - settings.QDRANT_CHUNK_SIZE = 50 + settings.QDRANT_CONTENT_FILE_CHUNK_SIZE = 50 run = LearningResourceRunFactory.create() with_content = ContentFileFactory.create(run=run, published=True, content="aaa") ContentFileFactory.create(run=run, published=True, content="") @@ -763,8 +805,7 @@ def test_embed_run_content_files_skips_contentless(mocker, mocked_celery, settin "vector_search.tasks.generate_embeddings", autospec=True ) - with pytest.raises(mocked_celery.replace_exception_class): - embed_run_content_files.delay(run.id) + embed_run_content_files.delay(run.id) assert _embedded_content_file_ids(generate_embeddings_mock) == {with_content.id} @@ -858,7 +899,7 @@ def test_embed_run_content_files_pre_pass_skips_unchanged( unchanged file is skipped only if the task's pre-pass computes the same point id as the embed pipeline. """ - settings.QDRANT_CHUNK_SIZE = 50 + settings.QDRANT_CONTENT_FILE_CHUNK_SIZE = 50 run = LearningResourceRunFactory.create() # ContentFile.save() computes checksum from content unchanged = ContentFileFactory.create(run=run, published=True, content="aaa") @@ -876,8 +917,7 @@ def test_embed_run_content_files_pre_pass_skips_unchanged( "vector_search.tasks.generate_embeddings", autospec=True ) - with pytest.raises(mocked_celery.replace_exception_class): - embed_run_content_files.delay(run.id) + embed_run_content_files.delay(run.id) assert _embedded_content_file_ids(generate_embeddings_mock) == { stale.id, @@ -893,7 +933,7 @@ def test_embed_run_content_files_pre_pass_dispatches_metadata_only_change( A file with a matching checksum but drifted payload metadata (edited title, newly generated summary) is dispatched so its Qdrant payload gets refreshed. """ - settings.QDRANT_CHUNK_SIZE = 50 + settings.QDRANT_CONTENT_FILE_CHUNK_SIZE = 50 run = LearningResourceRunFactory.create() retitled = ContentFileFactory.create(run=run, published=True, content="aaa") summarized = ContentFileFactory.create(run=run, published=True, content="bbb") @@ -913,8 +953,7 @@ def test_embed_run_content_files_pre_pass_dispatches_metadata_only_change( "vector_search.tasks.generate_embeddings", autospec=True ) - with pytest.raises(mocked_celery.replace_exception_class): - embed_run_content_files.delay(run.id) + embed_run_content_files.delay(run.id) assert _embedded_content_file_ids(generate_embeddings_mock) == { retitled.id, diff --git a/vector_search/utils.py b/vector_search/utils.py index 57f47e75e3..e87fdcbcd5 100644 --- a/vector_search/utils.py +++ b/vector_search/utils.py @@ -980,6 +980,36 @@ def _summarize_content_files_for_embedding( return refreshed_docs +def _upsert_points(client, collection_name, points): + """ + Upsert points in QDRANT_POINT_UPLOAD_BATCH_SIZE batches. + + Batching bounds the request size: a dense+sparse point with a full resource + payload runs tens of KB, so an unbatched chunk can exceed the client timeout + or the server's request-size limit. + """ + batch = [] + for point in points: + batch.append(point) + if len(batch) >= settings.QDRANT_POINT_UPLOAD_BATCH_SIZE: + client.batch_update_points( + collection_name=collection_name, + update_operations=[ + models.UpsertOperation(upsert=models.PointsList(points=batch)), + ], + wait=False, + ) + batch = [] + if batch: + client.batch_update_points( + collection_name=collection_name, + update_operations=[ + models.UpsertOperation(upsert=models.PointsList(points=batch)), + ], + wait=False, + ) + + def embed_learning_resources(ids, resource_type, overwrite): # noqa: PLR0915, C901 """ Embed learning resources @@ -1087,40 +1117,10 @@ def process_batch(docs_batch): points_generator_iter = _generate_content_file_points( docs_batch, stored_payloads ) - points_upload_batch = [] - - for point in points_generator_iter: - points_upload_batch.append(point) - if len(points_upload_batch) >= settings.QDRANT_POINT_UPLOAD_BATCH_SIZE: - client.batch_update_points( - collection_name=collection_name, - update_operations=[ - models.UpsertOperation( - upsert=models.PointsList( - points=points_upload_batch, - ) - ), - ], - wait=False, - ) - points_upload_batch = [] + _upsert_points(client, collection_name, points_generator_iter) - if points_upload_batch: - client.batch_update_points( - collection_name=collection_name, - update_operations=[ - models.UpsertOperation( - upsert=models.PointsList( - points=points_upload_batch, - ) - ), - ], - wait=False, - ) - - # Explicit deletions to help GC + # Explicit deletion to help GC del points_generator_iter - del points_upload_batch # We don't delete docs_batch here because it's a reference passed in, # but the caller clears the list. @@ -1143,17 +1143,7 @@ def process_batch(docs_batch): points = None # Handled inside the loop if points: - client.batch_update_points( - collection_name=collection_name, - update_operations=[ - models.UpsertOperation( - upsert=models.PointsList( - points=points, - ) - ), - ], - wait=False, - ) + _upsert_points(client, collection_name, points) def _resource_vector_hits(search_result): From 52029cfaff78e29b376dc42576b9d5776c7f79f5 Mon Sep 17 00:00:00 2001 From: Matt Bertrand Date: Mon, 3 Aug 2026 12:30:08 -0400 Subject: [PATCH 5/6] Fix _dispatch_signatures docstring: publish failures raise, no error string Co-Authored-By: Claude Fable 5 --- vector_search/tasks.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/vector_search/tasks.py b/vector_search/tasks.py index 1db8a670c3..de5166b862 100644 --- a/vector_search/tasks.py +++ b/vector_search/tasks.py @@ -103,9 +103,9 @@ def _dispatch_signatures(task_signatures) -> None: Publishes one message per chunk rather than a single catalog-wide group/chain, so a full-catalog backfill does not spike the broker with one huge message. - Nothing waits on the results (embedding tasks are ignore_result), so this - returns None: callers' return value stays "falsy on success, error string on - failure". + Nothing waits on the results (embedding tasks are ignore_result). Always + returns None (falsy success for callers returning it directly); a publish + failure raises out of apply_async, failing the dispatching task loudly. """ count = 0 for sig in task_signatures: From bfac10c2ec18589dbb43e94fff46e2c430f28f56 Mon Sep 17 00:00:00 2001 From: Matt Bertrand Date: Mon, 3 Aug 2026 13:31:49 -0400 Subject: [PATCH 6/6] Simplify: drop dead single-retrieve fallbacks and redundant publish retry_policy should_generate_resource_embeddings / should_generate_content_embeddings now require the batch-retrieved stored_point (all production callers already pass it), removing the _UNSET sentinel, per-document fallback retrieves, and the now-unused _retrieve_content_file_point. try_with_retry_as_task relies on celery's default publish retry (task_publish_retry=True, 3 attempts) instead of an explicit near-identical retry_policy. Co-Authored-By: Claude Fable 5 --- learning_resources_search/plugins.py | 15 ++----- vector_search/utils.py | 59 +++++----------------------- vector_search/utils_test.py | 21 +++------- 3 files changed, 19 insertions(+), 76 deletions(-) diff --git a/learning_resources_search/plugins.py b/learning_resources_search/plugins.py index 59af690d00..c10b38295d 100644 --- a/learning_resources_search/plugins.py +++ b/learning_resources_search/plugins.py @@ -26,19 +26,12 @@ def try_with_retry_as_task(function, *args): Dispatch a task/signature to the broker with publish-time retry. Accepts a bare task object (plus positional args) or an already-built - signature (e.g. a chain). Publish-time retry absorbs transient broker - blips without a second manual publish that could double-dispatch. + signature (e.g. a chain). Celery's default publish retry + (task_publish_retry, 3 attempts) absorbs transient broker blips without + a second manual publish that could double-dispatch. """ signature = function.si(*args) if isinstance(function, Task) else function - signature.apply_async( - retry=True, - retry_policy={ - "max_retries": 3, - "interval_start": 0.2, - "interval_step": 0.5, - "interval_max": 2, - }, - ) + signature.apply_async() class SearchIndexPlugin: diff --git a/vector_search/utils.py b/vector_search/utils.py index e87fdcbcd5..4d5aff1b66 100644 --- a/vector_search/utils.py +++ b/vector_search/utils.py @@ -502,10 +502,6 @@ def _content_file_embedding_context(document): return document.get("content", "") -# Sentinel distinguishing "caller did not supply a stored point" from an -# explicit None (point genuinely absent from Qdrant). -_UNSET = object() - # Cap ids per Qdrant retrieve so large (backfill / healthcheck) batches stay # under the server's request-size limits. QDRANT_RETRIEVE_BATCH_SIZE = 256 @@ -627,21 +623,13 @@ def _set_payload(points, document, param_map, collection_name): ) -def should_generate_resource_embeddings(serialized_document, stored_point=_UNSET): +def should_generate_resource_embeddings(serialized_document, stored_point): """ Determine if we should generate embeddings for a learning resource. - Pass stored_point (the already-retrieved point, or None if absent) to reuse a - batched retrieve instead of issuing one retrieve per document. + stored_point is the already-retrieved Qdrant point (from a batched + retrieve), or None if absent. """ - if stored_point is _UNSET: - client = qdrant_client() - point_id = vector_point_id(vector_point_key(serialized_document)) - response = client.retrieve( - collection_name=RESOURCES_COLLECTION_NAME, - ids=[point_id], - ) - stored_point = response[0] if response else None if stored_point is not None: stored_embedding_content = _learning_resource_embedding_context( stored_point.payload @@ -654,26 +642,6 @@ def should_generate_resource_embeddings(serialized_document, stored_point=_UNSET return True -def _retrieve_content_file_point( - serialized_document: dict, point_id: str | None = None -): - client = qdrant_client() - if not point_id: - # we just need metadata from the first chunk - point_id = vector_point_id( - vector_point_key( - serialized_document, chunk_number=0, document_type="content_file" - ) - ) - response = client.retrieve( - collection_name=CONTENT_FILES_COLLECTION_NAME, - ids=[point_id], - ) - if len(response) > 0: - return response[0] - return None - - def _stored_content_payloads( point_ids: list[str], fields: tuple[str, ...] = ("checksum",) ) -> dict[str, dict]: @@ -699,23 +667,16 @@ def _stored_content_payloads( return stored -def should_generate_content_embeddings( - serialized_document: dict, point_id: str | None = None, stored_point=_UNSET -) -> bool: +def should_generate_content_embeddings(serialized_document: dict, stored_point) -> bool: """ Determine if we should generate embeddings for a content file. - Pass stored_point (the already-retrieved chunk-0 point, or None if absent) to - reuse a batched retrieve instead of issuing one retrieve per document. + stored_point is the already-retrieved chunk-0 Qdrant point (from a batched + retrieve), or None if absent. """ - point = ( - _retrieve_content_file_point(serialized_document, point_id=point_id) - if stored_point is _UNSET - else stored_point - ) - if not point: + if not stored_point: return True - qdrant_checksum = (point.payload or {}).get("checksum") + qdrant_checksum = (stored_point.payload or {}).get("checksum") return qdrant_checksum != serialized_document["checksum"] @@ -759,9 +720,7 @@ def _embed_course_metadata_as_contentfile(serialized_resources): vector_point_key(doc, document_type="course_information") ) if not should_generate_content_embeddings( - serialized_document, - document_point_id, - stored_point=stored_points.get(str(document_point_id)), + serialized_document, stored_points.get(str(document_point_id)) ): continue diff --git a/vector_search/utils_test.py b/vector_search/utils_test.py index 943579cbfb..62e557e86b 100644 --- a/vector_search/utils_test.py +++ b/vector_search/utils_test.py @@ -1006,18 +1006,15 @@ def test_should_generate_for_changed_resource(mocker): resource = LearningResourceFactory.create() serialized_resources = list(serialize_bulk_learning_resources([resource.id])) - mock_qdrant = mocker.MagicMock() fake_payload = { "title": "Different title", "description": serialized_resources[0]["description"], "full_description": serialized_resources[0]["full_description"], } mock_point = mocker.MagicMock() - # return record with different title + # stored record with different title mock_point.payload = fake_payload - mock_qdrant.retrieve.return_value = [mock_point] - mocker.patch("vector_search.utils.qdrant_client", return_value=mock_qdrant) - result = should_generate_resource_embeddings(serialized_resources[0]) + result = should_generate_resource_embeddings(serialized_resources[0], mock_point) assert result is True @@ -1111,13 +1108,10 @@ def test_should_generate_for_changed_content_file(mocker): content_file = ContentFileFactory.create(content="Test content") serialized_files = list(serialize_bulk_content_files([content_file.id])) - mock_qdrant = mocker.MagicMock() mock_point = mocker.MagicMock() - # return record with different checksum + # stored record with different checksum mock_point.payload = {"checksum": "different-checksum"} - mock_qdrant.retrieve.return_value = [mock_point] - mocker.patch("vector_search.utils.qdrant_client", return_value=mock_qdrant) - result = should_generate_content_embeddings(serialized_files[0]) + result = should_generate_content_embeddings(serialized_files[0], mock_point) assert result is True @@ -1204,13 +1198,10 @@ def test_should_not_generate_for_unchanged_content_file(mocker): content_file = ContentFileFactory.create(content="Test content") serialized_files = list(serialize_bulk_content_files([content_file.id])) - mock_qdrant = mocker.MagicMock() mock_point = mocker.MagicMock() - # return record with same checksum + # stored record with same checksum mock_point.payload = {"checksum": serialized_files[0]["checksum"]} - mock_qdrant.retrieve.return_value = [mock_point] - mocker.patch("vector_search.utils.qdrant_client", return_value=mock_qdrant) - result = should_generate_content_embeddings(serialized_files[0]) + result = should_generate_content_embeddings(serialized_files[0], mock_point) assert result is False