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/learning_resources_search/plugins.py b/learning_resources_search/plugins.py index 0f3f5abfa6..c10b38295d 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,15 @@ 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). Celery's default publish retry + (task_publish_retry, 3 attempts) 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() class SearchIndexPlugin: diff --git a/main/settings.py b/main/settings.py index 378b0891ea..9f63f0425f 100644 --- a/main/settings.py +++ b/main/settings.py @@ -778,8 +778,28 @@ 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=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, +) + +# 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, ) @@ -827,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 62f11fce73..b459e6dc5c 100644 --- a/main/settings_celery.py +++ b/main/settings_celery.py @@ -25,6 +25,16 @@ # (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 (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=2 * 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 +231,15 @@ CELERY_VECTOR_SEARCH_RATE_LIMIT = get_string( "CELERY_VECTOR_SEARCH_RATE_LIMIT", CELERY_RATE_LIMIT ) +# 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 1d933c37c1..de5166b862 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,42 @@ def _replace_with_chain(task, task_signatures): return task.replace(celery.chain(*task_signatures)) -def _replace_with_finalized_chain( - task: celery.Task, content_file_ids: list[int], *, overwrite: bool +def _dispatch_content_file_chunks( + 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. + Dispatch content-file embedding chunks, fire and forget. + + 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. """ - failure_key = task.request.id - sigs = [ - generate_embeddings.si( - ids, CONTENT_FILE_TYPE, overwrite=overwrite, failure_key=failure_key + _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 ) - 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 _dispatch_signatures(task_signatures) -> None: + """ + 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). 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: + sig.apply_async() + count += 1 + log.info("Dispatched %d embedding chunk task(s)", count) def _queue_program_content_file_embedding_tasks(index_tasks, program_ids, overwrite): @@ -126,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, ) ] ) @@ -147,21 +153,28 @@ 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, + soft_time_limit=settings.CELERY_EMBEDDINGS_SOFT_TIME_LIMIT, + time_limit=settings.CELERY_EMBEDDINGS_TIME_LIMIT, + ignore_result=True, ) def generate_embeddings( self, ids: list[int], resource_type: str, overwrite: bool, # noqa: FBT001 - failure_key: str | None = None, + failure_key: str | None = None, # noqa: ARG001 ) -> 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). + + 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): @@ -171,6 +184,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 +195,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 +205,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): """ @@ -216,21 +229,17 @@ def remove_embeddings(ids, resource_type): raise -@app.task -def finalize_embeddings(failure_key: str) -> None: - """Chain tail: fail the parent task if any chunk recorded a failure.""" - cache = caches["redis"] - key = f"embed_errors:{failure_key}" - failures = cache.get(key, 0) - cache.delete(key) - if failures: - msg = f"{failures} embedding chunk(s) failed for {failure_key}" - log.error(msg) - raise RuntimeError(msg) +@app.task(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(bind=True) -def start_embed_resources(self, indexes, skip_content_files, overwrite): # noqa: C901 +@app.task +def start_embed_resources(indexes, skip_content_files, overwrite): # noqa: C901 """ Celery task to embed all learning resources for given indexes @@ -284,7 +293,7 @@ def start_embed_resources(self, indexes, skip_content_files, overwrite): # noqa 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}: @@ -329,13 +338,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 @@ -389,22 +398,21 @@ def embed_learning_resources_by_id(self, 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, ) ] 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,8 +448,8 @@ def embed_new_learning_resources(self): return self.replace(embed_tasks) -@app.task(bind=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 """ @@ -457,14 +465,13 @@ def embed_new_content_files(self): .exclude(learning_resource__published=False, learning_resource__test_mode=False) ) - return _replace_with_finalized_chain( - self, + return _dispatch_content_file_chunks( 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 +575,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 _dispatch_content_file_chunks(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 @@ -581,12 +588,14 @@ 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) -@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 @@ -598,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 e821d5bc16..c3133f24bc 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,65 @@ 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 + + +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( @@ -163,10 +200,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 +275,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 +317,26 @@ 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() + 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 + mock_call.kwargs.get("overwrite") is False + for mock_call in generate_embeddings_mock.si.call_args_list ) + # content-file chunks are published individually: no group/chord bookkeeping assert ( - finalize_embeddings_mock.si.call_args.args[0] - == generate_embeddings_mock.si.mock_calls[0].kwargs["failure_key"] - ) - chain_args = mocked_celery.chain.call_args.args - assert chain_args[:-1] == tuple( - generate_embeddings_mock.si.return_value - for _ in generate_embeddings_mock.si.mock_calls + generate_embeddings_mock.si.return_value.apply_async.call_count + == generate_embeddings_mock.si.call_count ) - assert chain_args[-1] == finalize_embeddings_mock.si.return_value + assert mocked_celery.group.call_count == 0 + assert mocked_celery.replace.call_count == 0 def test_remove_run_content_files(mocker, mocked_celery, settings): @@ -327,13 +359,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 +439,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 +486,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 +520,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 +554,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 +580,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 +615,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 +650,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 +669,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) @@ -690,19 +710,18 @@ 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() - embedded_ids = generate_embeddings_mock.si.mock_calls[0].args[0] + 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 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_CHUNK_SIZE = 2 + settings.QDRANT_CONTENT_FILE_CHUNK_SIZE = 2 run = LearningResourceRunFactory.create() content_file_ids = [ content_file.id @@ -712,37 +731,25 @@ 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) + 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 - ) - # 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 + for mock_call in generate_embeddings_mock.si.call_args_list ) - 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 + # 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): @@ -762,17 +769,16 @@ 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): """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") @@ -780,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} @@ -791,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="") @@ -800,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} @@ -895,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") @@ -913,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, @@ -930,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") @@ -950,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, @@ -998,6 +1000,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 @@ -1213,7 +1224,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 +1237,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 +1252,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 +1327,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..4d5aff1b66 100644 --- a/vector_search/utils.py +++ b/vector_search/utils.py @@ -502,6 +502,29 @@ def _content_file_embedding_context(document): return document.get("content", "") +# 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 +532,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): + 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) 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) @@ -593,20 +623,16 @@ 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): """ - Determine if we should generate embeddings for a learning resource + Determine if we should generate embeddings for a learning resource. + + stored_point is the already-retrieved Qdrant point (from a batched + retrieve), or None if absent. """ - 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 not None: stored_embedding_content = _learning_resource_embedding_context( - resource_payload + stored_point.payload ) current_embedding_content = _learning_resource_embedding_context( serialized_document @@ -616,26 +642,6 @@ def should_generate_resource_embeddings(serialized_document): 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]: @@ -661,16 +667,16 @@ def _stored_content_payloads( return stored -def should_generate_content_embeddings( - serialized_document: dict, point_id: str | None = None -) -> bool: +def should_generate_content_embeddings(serialized_document: dict, stored_point) -> bool: """ - Determine if we should generate embeddings for a content file + Determine if we should generate embeddings for a content file. + + 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 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"] @@ -688,6 +694,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 +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 + serialized_document, stored_points.get(str(document_point_id)) ): continue @@ -924,6 +939,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 @@ -1031,40 +1076,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. @@ -1087,17 +1102,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): diff --git a/vector_search/utils_test.py b/vector_search/utils_test.py index 37ff207cc7..62e557e86b 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, @@ -1004,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 @@ -1109,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 @@ -1202,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 @@ -1593,6 +1586,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 +2653,17 @@ 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