indexer: order-preserving concurrent embedding requests (#107)

app/AGENTS.md

@@ -12,7 +12,7 @@ The MCP server Databricks App: a FastMCP streamable-HTTP service exposing the in
 | `main.py` | FastMCP app: tool registration via `create_app()` factory, per-session `lifespan`, process-scoped engine singleton (`get_engine()`, `threading.Lock` + `atexit` dispose), `anyio` off-loop dispatch under a pool-sized `CapacityLimiter(5)`, `/health` (zero-DB) and `/ready` (`SELECT 1 FROM repos LIMIT 1` grant probe) |
 | `service.py` | Shared payload builders (`search_code_payload`, `list_repos_payload`, `get_file_payload`, `clamp_limit`): merges grep + symbol legs into the zoekt-parity envelope, base64url pagination cursors (`encode_cursor`/`decode_cursor`, `CursorError` never swallowed), `commit:` prefix resolution against `repo_branches`, permalink-branch selection |
 | `config.py` | `pydantic-settings` `Settings` with `CODE_SEARCH_` env prefix (timeouts, row limits, semantic tunables); unprefixed `LAKEBASE_ENDPOINT` via `validation_alias`; `SEMANTIC_EMBEDDING_DIM = 1024` single source of truth; `get_settings()` is `lru_cache`d once per process |
-| `embed.py` | `EmbedFn` seam (texts → unit-normalized 1024-dim vectors) and `databricks_embedder`: POSTs to the AI Gateway MLflow embeddings route via the SDK's raw API client; lazy SDK import; per-batch count check (`EmbeddingCountMismatchError`) and dim check (`EmbeddingDimMismatchError`) fail loudly instead of misaligning vectors |
+| `embed.py` | `EmbedFn` seam (texts → unit-normalized 1024-dim vectors) and `databricks_embedder`: POSTs to the AI Gateway MLflow embeddings route via the SDK's raw API client; lazy SDK import; per-batch count check (`EmbeddingCountMismatchError`) and dim check (`EmbeddingDimMismatchError`) fail loudly instead of misaligning vectors; batches dispatch through an order-preserving `ThreadPoolExecutor.map` (never `as_completed`) at `concurrency` (#107, indexer-only — the query path is one batch and stays serial) |
 | `app.yaml` | Databricks App runtime config: `ln -sf . app` symlink so `app.` imports resolve at the uploaded working-dir root; shell-form command so `DATABRICKS_APP_PORT` expands; sets only `LAKEBASE_ENDPOINT` |
 | `requirements.txt` | Deploy-time lockfile exported by `uv export --no-dev --no-hashes --no-emit-project` — regenerate, never hand-edit |
 | `__init__.py` | Empty package marker |

app/config.py

@@ -112,6 +112,14 @@ class Settings(BaseSettings):
     # which bounds file ingestion, not embedding-chunk granularity.
     semantic_chunk_max_tokens: int = 512
 
+    # In-flight embedding requests per worker (#107). The indexer clamps to 2 workers when
+    # semantic is on (indexer/repo_config.py:effective_workers), so total in-flight gateway
+    # requests are workers x concurrency: 2 x 4 = 8 at this default, 2 x 8 = 16 at the
+    # config.yaml-enforced ceiling of 8 -- both under the SDK's 20-connection pool
+    # (pool_block=True, so exceeding it would silently serialize rather than error). Setting
+    # this to 1 restores today's fully serial embed() and spawns no thread pool.
+    semantic_embedding_concurrency: int = 4
+
 
 @lru_cache(maxsize=1)
 def get_settings() -> Settings:

app/embed.py

@@ -18,6 +18,7 @@
 from __future__ import annotations
 
 from collections.abc import Callable, Sequence
+from concurrent.futures import ThreadPoolExecutor
 from typing import Any
 
 from app.config import SEMANTIC_EMBEDDING_DIM, Settings
@@ -50,13 +51,27 @@ def _assert_dims(vectors: Sequence[list[float]], dim: int) -> None:
 
 
 def _query_batch(
-    client: Any, endpoint: str, model: str, batch: list[str], *, max_retries: int
+    client: Any,
+    endpoint: str,
+    model: str,
+    batch: list[str],
+    *,
+    max_retries: int,
+    ordinal: int = 0,
+    offset: int = 0,
 ) -> list[list[float]]:
     """Query one batch, retrying up to ``max_retries`` times (small, bounded).
 
     ``endpoint`` is the gateway route path (e.g. ``/ai-gateway/mlflow/v1/embeddings``),
     POSTed via the SDK's raw API client; the response is OpenAI-shaped
     (``{"data": [{"embedding": [...]}, ...]}``).
+
+    ``ordinal``/``offset`` identify this batch's position in the caller's flat text
+    list (batch index and its starting position, respectively) purely so a count
+    mismatch can name the offending batch -- serially that batch was recoverable
+    from context (the one after the last success), but under concurrent dispatch
+    (:func:`databricks_embedder`) it is not. Both default to 0 so ``_query_batch``
+    stays directly callable in isolation.
     """
     last_exc: Exception | None = None
     for _attempt in range(max_retries + 1):
@@ -69,7 +84,8 @@ def _query_batch(
             # confusing IndexError (or silent corruption) much further downstream.
             if len(vectors) != len(batch):
                 raise EmbeddingCountMismatchError(
-                    f"embedder returned {len(vectors)} vectors for {len(batch)} texts"
+                    f"embedder returned {len(vectors)} vectors for {len(batch)} texts "
+                    f"(batch {ordinal}, texts[{offset}:{offset + len(batch)}])"
                 )
             return vectors
         except EmbeddingCountMismatchError:
@@ -89,6 +105,7 @@ def databricks_embedder(
     batch_size: int = 64,
     timeout: float = 20.0,
     max_retries: int = 2,
+    concurrency: int = 1,
 ) -> EmbedFn:
     """Build an :data:`EmbedFn` backed by the AI Gateway embeddings route ``endpoint``.
 
@@ -105,6 +122,19 @@ def databricks_embedder(
     client's concern. When omitted, the real ``WorkspaceClient`` is built with a
     ``Config`` carrying ``http_timeout_seconds=timeout`` (the raw API client has no
     per-call timeout).
+
+    ``concurrency`` dispatches up to that many batches at once via a
+    ``ThreadPoolExecutor``. It deliberately does NOT mirror this file's usual
+    default-mirroring convention (every other parameter here matches its
+    ``Settings`` twin): the concurrent path must be opt-in at the call site that
+    knows it is the indexer. ``get_embedder`` supplies the real value from
+    ``Settings.semantic_embedding_concurrency``; every direct caller and existing
+    unit test keeps today's serial semantics untouched. Do not "fix" this default
+    to match -- a higher default would make ``tests/unit/test_embed.py``'s
+    single-threaded fakes (e.g. ``_FakeApiClient.batches.append``) nondeterministic
+    under concurrent append. ``concurrency <= 1``, or a text list short enough to
+    produce only one batch, constructs no pool and spawns no thread -- the query
+    path (one text, one batch) is a strict no-op.
     """
     if client is None:
         from databricks.sdk import WorkspaceClient  # lazy: see module docstring
@@ -113,10 +143,33 @@ def databricks_embedder(
         client = WorkspaceClient(config=_SdkConfig(http_timeout_seconds=timeout))
 
     def embed(texts: list[str]) -> list[list[float]]:
-        vectors: list[list[float]] = []
-        for i in range(0, len(texts), batch_size):
-            batch = texts[i : i + batch_size]
-            vectors.extend(_query_batch(client, endpoint, model, batch, max_retries=max_retries))
+        batches = [texts[i : i + batch_size] for i in range(0, len(texts), batch_size)]
+        workers = max(1, min(concurrency, len(batches)))
+
+        def run_batch(indexed: tuple[int, list[str]]) -> list[list[float]]:
+            ordinal, batch = indexed
+            return _query_batch(
+                client,
+                endpoint,
+                model,
+                batch,
+                max_retries=max_retries,
+                ordinal=ordinal,
+                offset=ordinal * batch_size,
+            )
+
+        if workers == 1:
+            per_batch = [run_batch(item) for item in enumerate(batches)]
+        else:
+            # .map() yields in SUBMISSION order regardless of completion order --
+            # never as_completed(), which yields in completion order and would
+            # silently reorder vectors across files. See the module docstring's
+            # EmbeddingCountMismatchError note and tests/unit/test_embed.py's
+            # test_embed_module_never_uses_as_completed.
+            with ThreadPoolExecutor(max_workers=workers, thread_name_prefix="embed") as pool:
+                per_batch = list(pool.map(run_batch, enumerate(batches)))
+
+        vectors = [v for batch_vectors in per_batch for v in batch_vectors]
         _assert_dims(vectors, dim)
         return vectors
 
@@ -140,4 +193,5 @@ def get_embedder(cfg: Settings) -> EmbedFn:
         dim=cfg.semantic_embedding_dim,
         batch_size=cfg.semantic_embedding_batch_size,
         timeout=cfg.semantic_embedding_timeout_s,
+        concurrency=cfg.semantic_embedding_concurrency,
     )

config.yaml

@@ -94,10 +94,19 @@ connections:
 # `enabled: false` makes the job a true semantic no-op (no embedder built, no
 # chunking, the 2-worker memory clamp not applied) even if the env says enabled —
 # the fastest way to turn semantic off for the job alone.
+#
+# `embedding_concurrency` (#107) is in-flight embedding requests PER WORKER, sent
+# via a ThreadPoolExecutor that preserves submission order — vectors always come
+# back in the order their texts were sent, regardless of which request finishes
+# first. Total in-flight gateway requests for the job is workers x concurrency:
+# 2 x 4 = 8 at this default, 2 x 8 = 16 at the max of 8, both under the SDK's
+# 20-connection pool. Set to 1 to restore fully serial embedding (no thread pool
+# spawned at all) if you need to roll back.
 # semantic:
 #   enabled: true
 #   max_chunks_per_repo: 8000
 #   embedding_endpoint: /ai-gateway/mlflow/v1/embeddings
 #   embedding_model: system.ai.gte-large-en
 #   embedding_batch_size: 64
 #   embedding_timeout_s: 20.0
+#   embedding_concurrency: 4

docs/runbooks/indexing-parallelism.md

@@ -332,6 +332,21 @@ repo's chunks in memory (~0.5-0.8 GB per worker). The clamp is logged:
 INFO indexer.job [-]: semantic enabled: clamping index_concurrency 6 -> 2 (memory bound: ...)
 ```
 
+### Embedding concurrency (#107)
+
+`workers x concurrency` is the number that matters, not `concurrency` alone.
+Each of the (at most 2, semantic-clamped) workers dispatches up to
+`semantic.embedding_concurrency` embedding batches at once
+(`app/embed.py:databricks_embedder`, order-preserving `ThreadPoolExecutor.map`):
+2 x 4 = 8 in-flight gateway requests at the default, 2 x 8 = 16 at the
+config.yaml-enforced ceiling of 8 (the `CODE_SEARCH_SEMANTIC_EMBEDDING_CONCURRENCY`
+env var carries no ceiling, mirroring `semantic_embedding_batch_size`'s own
+unbounded env surface -- config.yaml is the job's real surface regardless), both
+under the SDK's 20-connection pool.
+`embedding_concurrency: 1` is the rollback switch — fully serial embedding, no
+thread pool spawned. See `docs/runbooks/semantic-enablement.md` §4 for the full
+in-flight/memory arithmetic and the 429 posture.
+
 ### The connection pool follows the workers
 
 Each worker holds exactly one connection, so the engine is built with

docs/runbooks/semantic-enablement.md

@@ -164,6 +164,46 @@ repo it names.
 `indexer/job.py` now applies to every index run by default (each worker materialises a
 whole repo's chunks) — see `docs/runbooks/indexing-parallelism.md`.
 
+**Concurrent embedding requests (#107):** each worker's `embed()` call
+(`app/embed.py`) dispatches up to `semantic.embedding_concurrency` batches at once
+via a `ThreadPoolExecutor`, using `.map()` — never `as_completed()` — so vectors
+always come back in submission order regardless of which request finishes first.
+Total in-flight gateway requests for the job is `effective_workers x concurrency`:
+2 x 4 = 8 at the default `embedding_concurrency: 4`, 2 x 8 = 16 at the config's
+`le=8` ceiling, both under the SDK's 20-connection pool
+(`HTTPAdapter(pool_connections=20, pool_maxsize=20, pool_block=True)` —
+`pool_block=True` means exceeding the pool **silently serializes** requests rather
+than raising, so staying under 20 is the load-bearing bound, not a nice-to-have).
+The only new per-in-flight-batch memory cost is transient request/response
+buffers (~3.5 MB each: ~2.1 MB parsed vectors + ~1.3 MB raw JSON response + a
+small request body) — ~28 MB at the default, ~56 MB at the ceiling, negligible
+beside the ~0.5–0.8 GB/worker baseline above. `embedding_concurrency: 1` restores
+today's fully serial embedding and spawns no thread pool at all (the rollback
+switch).
+
+429s from the AI Gateway are absorbed entirely by the `databricks-sdk`'s own
+`Retry-After`-honouring backoff (`_RetryAfterCustomizer`, defaulting to 1s when
+the header is absent) before `_query_batch`'s own bounded retry ever sees them —
+`app/embed.py` does not add a third retry layer. That absorption is invisible at
+the job's normal INFO log level: the SDK logs each throttle at DEBUG
+(`databricks.sdk.retries`). If you suspect throttling during a manual run, set
+`logging.getLogger("databricks.sdk.retries").setLevel(logging.DEBUG)` for that
+run only — never raise the root logger or the `databricks.sdk` parent logger,
+which would also re-enable a request/response body dump (the embedding request
+body is repo source code). This is an operator step for a one-off diagnostic
+run, never a code change (`tests/unit/test_job_redaction.py` tripwires
+`indexer/*.py` against exactly that).
+
+**Failure-path latency under concurrency:** when one batch raises, the pool's
+`__exit__` still waits for every other in-flight request in that worker's pool
+to finish before the exception propagates (there is no way to abort an
+in-flight HTTP call). A batch-0 failure that returned instantly under serial
+dispatch can now wait up to `concurrency - 1` requests' worth of time, each
+bounded by `semantic_embedding_timeout_s` (default 20s) plus the SDK's own
+retry budget. This is bounded and per-branch, not per-run: a sustained-outage
+branch still degrades to a core index without chunks (semantic is additive),
+just after a somewhat longer wait than serial dispatch's instant fail-fast.
+
 ## 5. Rollback note
 
 `0004`'s `downgrade()` drops the BM25/ANN indexes and the `chunks` table, but **does
@@ -192,6 +232,7 @@ semantic:
   embedding_model: system.ai.gte-large-en
   embedding_batch_size: 64
   embedding_timeout_s: 20.0
+  embedding_concurrency: 4         # -> Settings.semantic_embedding_concurrency (#107); 1 = serial
 ```
 
 Every field is optional; an omitted field falls through to the env value / default, and