Reduce embedding task storms: broker visibility timeout, bounded dispatch, chunk sizing

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(

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",

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:

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)
 

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)

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:

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):
     """

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."
         )