From 045bdaf72236cd47e0fd4b9a24be7d086235c03e Mon Sep 17 00:00:00 2001 From: IceRhymers Date: Sat, 25 Jul 2026 10:27:54 -0700 Subject: [PATCH] indexer: order-preserving concurrent embedding requests (#107) Batches now dispatch through a ThreadPoolExecutor.map (never as_completed, which yields in completion order), so vectors always return in submission order regardless of which request finishes first -- the whole point of the issue, since _precompute_chunk_writer re-slices the flat result positionally and a reorder would silently corrupt embeddings for the wrong file. - app/embed.py: databricks_embedder gains concurrency (default 1, a deliberate divergence from this file's default-mirroring convention so the concurrent path is opt-in only at the indexer call site); _query_batch gains ordinal/offset so a count mismatch names the offending batch (AC2). Existing retry semantics and the SDK's own 429/Retry-After handling are unchanged -- no third retry layer added. - app/config.py: semantic_embedding_concurrency = 4 (workers x concurrency = 8 in-flight by default, 16 at the config.yaml le=8 ceiling, both under the SDK's 20-connection pool). - indexer/repo_config.py: SemanticOverrides.embedding_concurrency (1..8) overlays the setting from config.yaml, the job's only reachable config surface. - scripts/measure_embedding_concurrency.py: offline scaling harness (no network); measured ~1x/2x/4x/8x at concurrency 1/2/4/8. - Docs: config.yaml, semantic-enablement.md, indexing-parallelism.md, app/AGENTS.md, indexer/AGENTS.md. - Tests: 24 cases in test_embed.py (submission-order proof plus a non-vacuity companion, AST-based as_completed tripwire, a structural queued-batch-cancellation bound, deterministic lowest-ordinal-failure reporting, the concurrency clamp's degenerate paths) and 3 sites in test_repo_config.py for the new config field. indexer/job.py is unchanged: it still calls embed_fn(all_texts) once and re-slices positionally. No INDEX_SEMANTICS_VERSION bump -- this changes dispatch mechanics, not extraction output. Refs #107 --- app/AGENTS.md | 2 +- app/config.py | 8 + app/embed.py | 66 +++++- config.yaml | 9 + docs/runbooks/indexing-parallelism.md | 15 ++ docs/runbooks/semantic-enablement.md | 41 ++++ indexer/AGENTS.md | 2 +- indexer/repo_config.py | 10 + scripts/measure_embedding_concurrency.py | 95 ++++++++ tests/unit/test_embed.py | 275 ++++++++++++++++++++++- tests/unit/test_repo_config.py | 12 +- 11 files changed, 521 insertions(+), 14 deletions(-) create mode 100644 scripts/measure_embedding_concurrency.py diff --git a/app/AGENTS.md b/app/AGENTS.md index 6b0867b..07e58b7 100644 --- a/app/AGENTS.md +++ b/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 | diff --git a/app/config.py b/app/config.py index e6bb63a..1c5322c 100644 --- a/app/config.py +++ b/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: diff --git a/app/embed.py b/app/embed.py index 2cff250..d65eec3 100644 --- a/app/embed.py +++ b/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, ) diff --git a/config.yaml b/config.yaml index 6094bfa..98e99a7 100644 --- a/config.yaml +++ b/config.yaml @@ -94,6 +94,14 @@ 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 @@ -101,3 +109,4 @@ connections: # embedding_model: system.ai.gte-large-en # embedding_batch_size: 64 # embedding_timeout_s: 20.0 +# embedding_concurrency: 4 diff --git a/docs/runbooks/indexing-parallelism.md b/docs/runbooks/indexing-parallelism.md index c7e99c7..f9ce380 100644 --- a/docs/runbooks/indexing-parallelism.md +++ b/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 diff --git a/docs/runbooks/semantic-enablement.md b/docs/runbooks/semantic-enablement.md index 51ba874..e554afb 100644 --- a/docs/runbooks/semantic-enablement.md +++ b/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 diff --git a/indexer/AGENTS.md b/indexer/AGENTS.md index 268b661..4477922 100644 --- a/indexer/AGENTS.md +++ b/indexer/AGENTS.md @@ -19,7 +19,7 @@ The serverless indexing job (`code-search-index`, wired as a `python_wheel_task` | `job.py` | Entry point and orchestration: `main()`/`run()`, `ThreadPoolExecutor` sized by `effective_workers` (unit of work = one repo, all branches sequential), batched `repo_branches` stamp read, per-branch skip-if-unchanged, per-branch `BranchOutcome` classification (`indexed`/`skipped`/`conflict`/`failed`), semantic precompute (`_precompute_chunk_writer`), `ContextVar`-based `[repo]` log attribution, and per-phase timing instrumentation (one `phase timing repo@branch: total=… resolve=… download=… parse=… embed=… db=… sweep=… other=…` INFO line per INDEXED branch, emitted from inside the worker so `[repo]` resolves; fixed unconditional field set; `_timed_items` charges lazy item production to `parse` so it is not fused into `db`; `resolve=`/`list=` on the per-repo `finished` line carry the repo-scoped costs that belong to no branch's total). Returns 1 if any branch failed; conflicts self-heal and do not fail the run. **Reconciliation checkpoint**: after every worker joins, `_decide_reconciliation` gates on zero failures/conflicts, every repo accounted for, and every repo's `discovery_complete`; on a pass, `_reconcile` reuses the pre-fan-out stamp snapshot to compute each repo's retired branches and the run's full desired repo set, then calls the injected `reconcile_retired_fn`/`reconcile_removed_fn` seams (default: the real `indexer.store` primitives) on one post-fan-out connection — retired branches first, repo purge second. `MAX_PURGE_SHRINK_FRACTION = 0.5` (hardcoded, no config knob) withholds ONLY the purge, as a logged incident signal, if it would remove a strict majority of currently stored repos; retired-branch cleanup on survivors still applies. `ReconcileProgress.committed_any` drives an honest "partially reconciled" vs "left stale" failure message — never the raw exception, only its phase/repo/type name. Never called from worker-thread code (`_index_one*`) — a source-level tripwire test enforces it. | | `languages.py` | Single source of truth shared by `parse.py` and `symbols.py`: `EXT_TO_LANG` (8 extensions -> tree-sitter language names), `SYMBOL_KINDS` (node type -> symbol kind per language), `EDGE_NODE_KINDS` (node type -> reference-edge kind per language, Python-only until #85), `MAX_FILE_BYTES` (1 MB), `SEMANTIC_CHUNK_MAX_CHARS` (2000, ~4 chars/token), and the frozen dataclasses `ParsedFile`, `Chunk`, `ExtractedSymbol`, `ExtractedEdge`, `FileExtraction`, `IndexCounts` (now carries `edges`). | | `parse.py` | `iter_source_files`: **no production caller since #106** — retained deliberately as the extraction ORACLE the streaming path (`ingest.py`) is pinned against by `tests/unit/test_ingest_parity.py`, and undeletable anyway without editing a `SEMANTICS_PATHS` module and forcing an `INDEX_SEMANTICS_VERSION` bump. Walk an extracted tree, yield every text file (unknown extensions kept with `lang=None`, since grep runs over all files); skips `.git/`, symlinks, files > `MAX_FILE_BYTES` (stat before read), NUL-sniffed binaries, and UTF-8 decode failures; strips surviving NULs (Postgres `text` rejects them). `iter_chunks`: deterministic line-aligned chunking, no overlap, no mid-line splits, 1-based inclusive line ranges. | -| `repo_config.py` | Pydantic schema for `config.yaml` (`RepoConfig` / `GitHubConnection` / `ExcludeRules`), `normalize_repo` (URL/SSH/bare -> canonical `org/repo`, GitHub hosts only), `parse_config` (pure) vs `read_workspace_config` (SDK I/O, wraps every failure in `ConfigError` with HTTP status so 404 "never synced" stays distinguishable from 403 "no permission"), and `effective_workers` (the semantic clamp). A connection with no `orgs`/`users`/`repos` fails validation. `index_concurrency`: 1-8, default 4 — a disk bound (0.5 GB peak per worker since #106), not a CPU one. `semantic_max_chunks_per_repo`: optional top-level `{"org/repo": N}` map overriding `app.config.Settings.semantic_max_chunks_per_repo` per repo; keys canonicalised via `normalize_repo` and rejected on collision post-casefold, values `>= 1`. Deliberately import-light: pydantic + PyYAML + stdlib only. | +| `repo_config.py` | Pydantic schema for `config.yaml` (`RepoConfig` / `GitHubConnection` / `ExcludeRules`), `normalize_repo` (URL/SSH/bare -> canonical `org/repo`, GitHub hosts only), `parse_config` (pure) vs `read_workspace_config` (SDK I/O, wraps every failure in `ConfigError` with HTTP status so 404 "never synced" stays distinguishable from 403 "no permission"), and `effective_workers` (the semantic clamp). A connection with no `orgs`/`users`/`repos` fails validation. `index_concurrency`: 1-8, default 4 — a disk bound (0.5 GB peak per worker since #106), not a CPU one. `semantic_max_chunks_per_repo`: optional top-level `{"org/repo": N}` map overriding `app.config.Settings.semantic_max_chunks_per_repo` per repo; keys canonicalised via `normalize_repo` and rejected on collision post-casefold, values `>= 1`. `SemanticOverrides.embedding_concurrency` (#107): optional `1..8`, overlays `Settings.semantic_embedding_concurrency` (in-flight embedding requests per worker; job-only, inert on the MCP/webui query path). Deliberately import-light: pydantic + PyYAML + stdlib only. | | `resolve.py` | `resolve_repos`: enumerate org/user selectors, apply `ExcludeRules` to enumerated repos only (explicit `repos` entries always win, unfiltered), dedup case-insensitively keeping first-seen spelling, union each connection's `branches:` globs into per-repo `RepoEntry.branch_globs`, then fail fast: `EmptyConfigError` on zero repos (indexing nothing must not exit 0), `RepoCeilingError` above `MAX_REPOS` (500, overridable via `--max_repos`). Also matches `RepoConfig.semantic_max_chunks_per_repo` case-insensitively onto each resolved repo's `RepoEntry.semantic_max_chunks` (`None` if unmatched); an override key matching no resolved repo logs a WARNING (typo guard). Enumerator params are a test seam only, not provider dispatch. | | `store.py` | `index_repo`: the single atomic unit of work for one (repo, branch) inside `with conn.begin():` — repos upsert, `repo_branches` CAS-baseline read, per-file classification (unchanged / membership-only / changed-new), membership sweep (strip this branch from unseen rows, delete rows with empty `branches`; skipped with a WARNING on an empty seen-set), then CAS stamp (raises `StaleIndexError` on mismatch, rolling everything back). `items` is `Iterable[tuple[ParsedFile, FileExtraction]]`. **Changed/new files are batched (#105)**, not written one at a time: `_flush_file_batch` accumulates up to `_BATCH_MAX_FILES` files or `_BATCH_MAX_CONTENT_BYTES` bytes of `pf.size` (whichever trips first, post-append — a single oversized file is never split), then issues ONE multi-row `files` upsert (`SET` uses `excluded.*` for every per-row column, NEVER a Python literal from one file — the trap that would collapse a batch's conflicting rows onto the last file's values), maps ids back from `RETURNING` by `(path, content_sha)` — never by row order, and RAISES on a missing id rather than warn-and-skip (a wrong id here would attach one file's symbols to another file's row) — then ONE `DELETE ... WHERE file_id = ANY(:ids)` + bulk insert each for `symbols`/`reference_edges` (`reference_edges` unconditional, even when a batch's `FileExtraction.edges` is all-empty, so stale rows never survive a re-index), then ONE `chunk_writer` call for the batch. An intra-batch duplicate `(path, content_sha)` is deduped first (last occurrence wins, one WARNING) — Postgres raises `ON CONFLICT DO UPDATE command cannot affect row a second time` and poisons the transaction otherwise. The sweep call site is wall-clocked into `indexer.timing`'s ambient `PhaseTimer` (`record("sweep", …)`) — never into `IndexCounts` (frozen, compared by value) and never through a new `index_repo` parameter (that signature is an injected seam); `store.py` emits **no** log record for it. Pure DML, no TEMP tables (job role has no TEMP privilege on Lakebase). `reconcile_retired_branches` and `reconcile_removed_repos`: the desired-state storage primitives `job.py`'s reconciliation checkpoint calls — each its own `conn.begin()`, repo-scoped membership subtraction (never delete-by-path) and an exact-match repo purge respectively; neither decides WHAT is retired/desired, only applies a caller-supplied set (see `job.py`'s row). | | `symbols.py` | `extract_file` (#84): one tree-sitter parse, one full-tree walk, emitting both `symbols` and `reference_edges` candidates (`FileExtraction`) in the same pass — nested definitions captured, named nodes only for symbols, kinds from `SYMBOL_KINDS`/`EDGE_NODE_KINDS`, 1-based lines. Stack entries carry the innermost named enclosing definition so edges attribute to it in O(1) with no second walk. `extract_symbols` is a thin wrapper (`extract_file(pf).symbols`) kept for the existing call sites. Python-only edge helpers (`_python_call_target`, `_python_import_edges`) are the seam #85 generalizes to other languages. Parser cache is per-thread (`threading.local`) as insurance against a future GIL-releasing `parse()`. | diff --git a/indexer/repo_config.py b/indexer/repo_config.py index 4d39a9f..d863771 100644 --- a/indexer/repo_config.py +++ b/indexer/repo_config.py @@ -150,6 +150,14 @@ class SemanticOverrides(BaseModel): block only moves the second operand. Setting both is coherent: raise the floor for the whole corpus here, spot-override the outliers in the map. + ``embedding_concurrency`` (#107) bounds in-flight embedding requests per + worker; ``databricks_embedder`` clamps it to the batch count, so a value + larger than the number of batches a repo produces degrades to that count + rather than spawning a thread per batch. The ``le=8`` ceiling here matches + the existing ``index_concurrency`` bounded-field pattern -- see + ``app.config.Settings.semantic_embedding_concurrency`` for the in-flight + arithmetic against the SDK's 20-connection pool. + Two ``Settings`` semantic knobs are deliberately absent, because exposing them would be a lie or a footgun: ``semantic_embedding_dim`` is pinned to ``SEMANTIC_EMBEDDING_DIM`` (the ``chunks.embedding`` column type and the 0004 @@ -180,6 +188,7 @@ class SemanticOverrides(BaseModel): embedding_model: str | None = Field(default=None, min_length=1) embedding_batch_size: int | None = Field(default=None, ge=1) embedding_timeout_s: float | None = Field(default=None, gt=0) + embedding_concurrency: int | None = Field(default=None, ge=1, le=8) @field_validator("embedding_endpoint") @classmethod @@ -232,6 +241,7 @@ def settings_overrides(self) -> dict[str, Any]: "embedding_model": "semantic_embedding_model", "embedding_batch_size": "semantic_embedding_batch_size", "embedding_timeout_s": "semantic_embedding_timeout_s", + "embedding_concurrency": "semantic_embedding_concurrency", } return { settings_field: value diff --git a/scripts/measure_embedding_concurrency.py b/scripts/measure_embedding_concurrency.py new file mode 100644 index 0000000..d72fa0d --- /dev/null +++ b/scripts/measure_embedding_concurrency.py @@ -0,0 +1,95 @@ +"""Measure embedding dispatch scaling under concurrency (#107, AC3 part 1). + +Offline companion to ``app.embed.databricks_embedder``. Drives the SAME dispatch +code path AC3 cares about (batches -> ``ThreadPoolExecutor.map`` -> flatten) +against a fake client whose ``do()`` sleeps a fixed per-batch latency -- no +workspace, no network, not run in CI. It measures that the dispatch actually +parallelizes; it says nothing about the real AI Gateway's throughput or 429 +behavior -- that is the other two halves of AC3's measurement protocol (a +real-repo run's ``embed=`` timing-line comparison, and a single named +``databricks.sdk.retries`` logger), which require a live workspace and are run +by hand, not by this script. See ``docs/runbooks/semantic-enablement.md`` §4. + +Usage: ``uv run python scripts/measure_embedding_concurrency.py +[--latency-s 0.05] [--num-batches 32] [--concurrencies 1,2,4,8]`` +""" + +from __future__ import annotations + +import argparse +import time +from collections.abc import Sequence +from typing import Any + +from app.embed import databricks_embedder + + +class _SleepingApiClient: + """Stands in for ``WorkspaceClient.api_client``: sleeps ``latency_s`` per + batch, then returns one dim-1 vector per text. Order/count correctness is + already pinned by ``tests/unit/test_embed.py``; this script only measures + wall clock.""" + + def __init__(self, latency_s: float) -> None: + self._latency_s = latency_s + + def do(self, method: str, path: str, *, body: dict[str, Any]) -> dict[str, Any]: + time.sleep(self._latency_s) + return {"data": [{"embedding": [0.0]} for _ in body["input"]]} + + +class _SleepingClient: + def __init__(self, latency_s: float) -> None: + self.api_client = _SleepingApiClient(latency_s) + + +def measure(*, num_batches: int, batch_size: int, latency_s: float, concurrency: int) -> float: + """Run one ``embed()`` call over a synthetic corpus of ``num_batches`` + batches and return the wall-clock seconds.""" + texts = [f"text-{i}" for i in range(num_batches * batch_size)] + client = _SleepingClient(latency_s) + embed = databricks_embedder( + "ep", "m", client=client, dim=1, batch_size=batch_size, concurrency=concurrency + ) + start = time.perf_counter() + vectors = embed(texts) + elapsed = time.perf_counter() - start + assert len(vectors) == len(texts) + return elapsed + + +def main(argv: Sequence[str] | None = None) -> int: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--latency-s", type=float, default=0.05, help="fake per-batch latency") + parser.add_argument("--num-batches", type=int, default=32, help="synthetic batch count") + parser.add_argument( + "--batch-size", + type=int, + default=1, + help="texts per batch (kept small so --num-batches controls the batch count directly)", + ) + parser.add_argument( + "--concurrencies", default="1,2,4,8", help="comma-separated concurrency levels" + ) + args = parser.parse_args(argv) + concurrencies = [int(c) for c in args.concurrencies.split(",")] + + print(f"{args.num_batches} batches x {args.latency_s}s fake latency each") + print(f"{'concurrency':>11s} {'wall_clock_s':>12s} {'speedup':>8s}") + baseline: float | None = None + for concurrency in concurrencies: + elapsed = measure( + num_batches=args.num_batches, + batch_size=args.batch_size, + latency_s=args.latency_s, + concurrency=concurrency, + ) + if baseline is None: + baseline = elapsed + speedup = baseline / elapsed if elapsed else float("inf") + print(f"{concurrency:>11d} {elapsed:>12.3f} {speedup:>7.2f}x") + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/tests/unit/test_embed.py b/tests/unit/test_embed.py index f228c58..98cc757 100644 --- a/tests/unit/test_embed.py +++ b/tests/unit/test_embed.py @@ -6,15 +6,21 @@ from __future__ import annotations +import ast +import inspect import sys +import threading +import time from typing import Any import pytest +from app.config import Settings from app.embed import ( EmbeddingCountMismatchError, EmbeddingDimMismatchError, databricks_embedder, + get_embedder, ) @@ -40,13 +46,26 @@ def __init__(self, vectors_fn: Any) -> None: @pytest.mark.unit def test_batching_splits_by_batch_size() -> None: + # concurrency=1 pinned explicitly here: this test's own submission-order + # assertion below would otherwise be the accidental tripwire for the + # default -- passing it explicitly means it no longer is, so + # test_default_concurrency_is_one below covers the default directly. client = _FakeClient(lambda texts: [[0.0, 0.0] for _ in texts]) - embed = databricks_embedder("ep", "m", client=client, dim=2, batch_size=2) + embed = databricks_embedder("ep", "m", client=client, dim=2, batch_size=2, concurrency=1) vectors = embed(["a", "b", "c", "d", "e"]) assert len(vectors) == 5 assert client.api_client.batches == [["a", "b"], ["c", "d"], ["e"]] +@pytest.mark.unit +def test_default_concurrency_is_one() -> None: + """The whole backward-compat guarantee rests on this default: every + non-indexer caller (app/search/semantic.py, every other test in this file) + relies on serial dispatch. Only get_embedder(cfg) is supposed to opt in + to concurrency>1 (see databricks_embedder's docstring).""" + assert inspect.signature(databricks_embedder).parameters["concurrency"].default == 1 + + @pytest.mark.unit def test_single_batch_when_under_batch_size() -> None: client = _FakeClient(lambda texts: [[0.0] for _ in texts]) @@ -121,9 +140,259 @@ def short(texts: list[str]) -> list[list[float]]: @pytest.mark.unit -def test_stub_path_never_imports_databricks_sdk() -> None: +@pytest.mark.parametrize("concurrency", [1, 4]) +def test_stub_path_never_imports_databricks_sdk(concurrency: int) -> None: + # batch_size=1 with two texts -> two batches, so concurrency=4 actually + # reaches the pool branch (concurrent.futures is stdlib; nothing new to + # import). At the module's default batch_size=64 this would be one batch + # and the extension would pin nothing at any concurrency (round-3 finding). sys.modules.pop("databricks.sdk", None) client = _FakeClient(lambda texts: [[0.0, 0.0] for _ in texts]) - embed = databricks_embedder("ep", "m", client=client, dim=2) + embed = databricks_embedder( + "ep", "m", client=client, dim=2, batch_size=1, concurrency=concurrency + ) embed(["hello", "world"]) assert "databricks.sdk" not in sys.modules + + +def _build_order_probe_client() -> tuple[Any, list[str]]: + """Fake client where batch "0" blocks until batch "3" (the last batch) + completes and releases it -- so batch 0 finishes LAST despite being + submitted first. ``completion_order`` records the FAKE's own completion + sequence, for T3 to prove this fixture is not passing vacuously. + """ + release = threading.Event() + completion_order: list[str] = [] + + class _OrderedApiClient: + def do(self, method: str, path: str, *, body: dict[str, Any]) -> dict[str, Any]: + assert method == "POST" + batch = list(body["input"]) + text = batch[0] + if text == "0": + assert release.wait(timeout=5), "the last batch never released batch 0" + completion_order.append(text) + if text == "3": + release.set() + return {"data": [{"embedding": [float(text)]}]} + + class _Client: + def __init__(self) -> None: + self.api_client = _OrderedApiClient() + + return _Client(), completion_order + + +@pytest.mark.unit +def test_concurrent_batches_return_in_submission_order() -> None: + """AC1: vectors come back in submission order even though batch 0 -- the + one that determines index 0 of the result -- is the LAST to complete.""" + client, _ = _build_order_probe_client() + embed = databricks_embedder("ep", "m", client=client, dim=1, batch_size=1, concurrency=4) + vectors = embed(["0", "1", "2", "3"]) + assert vectors == [[0.0], [1.0], [2.0], [3.0]] + + +@pytest.mark.unit +def test_out_of_order_completion_is_recorded_out_of_order() -> None: + """Proves test_concurrent_batches_return_in_submission_order isn't passing + vacuously: the fake's own completion order genuinely differs from submission + order (batch 0 finishes last, not first).""" + client, completion_order = _build_order_probe_client() + embed = databricks_embedder("ep", "m", client=client, dim=1, batch_size=1, concurrency=4) + embed(["0", "1", "2", "3"]) + assert completion_order != ["0", "1", "2", "3"] + assert completion_order[-1] == "0" + + +@pytest.mark.unit +def test_short_batch_under_concurrency_names_the_offending_batch() -> None: + """AC2: a mismatch in a LATE batch under concurrency still names it.""" + + def vectors_fn(batch: list[str]) -> list[list[float]]: + if batch == ["d"]: + return [] # 0 vectors for 1 text: short by one + return [[0.0] for _ in batch] + + client = _FakeClient(vectors_fn) + embed = databricks_embedder("ep", "m", client=client, dim=1, batch_size=1, concurrency=4) + with pytest.raises(EmbeddingCountMismatchError, match=r"0 vectors for 1 texts \(batch 3"): + embed(["a", "b", "c", "d"]) + + +@pytest.mark.unit +def test_offset_names_the_texts_slice_at_a_non_degenerate_batch_size() -> None: + """T4's ``batch_size=1`` fixture makes ``offset == ordinal`` trivially, so it + cannot catch a broken ``offset = ordinal * batch_size``. This uses a real + batch size (64) so the reported ``texts[192:256]`` only matches if the + multiplication is right.""" + + def vectors_fn(batch: list[str]) -> list[list[float]]: + if len(batch) == 64 and batch[0] == "192": + return [[0.0] for _ in batch[:-1]] # 63 vectors for 64 texts: short by one + return [[0.0] for _ in batch] + + client = _FakeClient(vectors_fn) + embed = databricks_embedder("ep", "m", client=client, dim=1, batch_size=64, concurrency=4) + texts = [str(i) for i in range(256)] + with pytest.raises( + EmbeddingCountMismatchError, + match=r"63 vectors for 64 texts \(batch 3, texts\[192:256\]\)", + ): + embed(texts) + + +@pytest.mark.unit +def test_first_offending_batch_in_submission_order_wins() -> None: + """A slow low-ordinal failure and a fast high-ordinal failure both occur; + `pool.map` must raise the lowest-ordinal one (deterministic error under + nondeterministic execution), not whichever failed first in wall clock.""" + + def vectors_fn(batch: list[str]) -> list[list[float]]: + text = batch[0] + if text == "1": + time.sleep(0.2) + raise RuntimeError("slow failure at ordinal 1") + if text == "6": + raise RuntimeError("fast failure at ordinal 6") + return [[0.0] for _ in batch] + + client = _FakeClient(vectors_fn) + embed = databricks_embedder( + "ep", "m", client=client, dim=1, batch_size=1, concurrency=8, max_retries=0 + ) + with pytest.raises(RuntimeError, match="slow failure at ordinal 1"): + embed([str(i) for i in range(8)]) + + +@pytest.mark.unit +def test_failure_cancels_queued_batches() -> None: + """Structural fixture, NOT a scheduler-derived one (see the plan's stall + sweep -- every wall-clock-derived constant fails somewhere). batch_size=1 + with 20 texts is load-bearing: at the module's default batch_size=64, 20 + texts is ONE batch, so this would take the serial path and pass vacuously + (`started == 1 <= 3`) with no pool in either implementation. + + Batch "0" raises immediately; every other batch parks on an Event that is + NEVER set. A parked worker structurally cannot pull the next queue item, so + `started` is fixed by the code path (concurrency slots occupied, plus one + more pulled by the freed failing worker before cancellation fires), not by + scheduling. Measured (per the plan) at exactly 3 with zero variance across + stall levels, while a submit()+serial-.result() refactor starts all 20. + """ + never = threading.Event() + started_lock = threading.Lock() + started = {"n": 0} + + def vectors_fn(batch: list[str]) -> list[list[float]]: + text = batch[0] + with started_lock: + started["n"] += 1 + if text == "0": + raise RuntimeError("endpoint down") + never.wait(timeout=0.5) # parks; never actually set -- bounds runtime only + return [[0.0] for _ in batch] + + client = _FakeClient(vectors_fn) + texts = [str(i) for i in range(20)] + embed = databricks_embedder( + "ep", "m", client=client, dim=1, batch_size=1, concurrency=2, max_retries=0 + ) + with pytest.raises(RuntimeError, match="endpoint down"): + embed(texts) + assert started["n"] <= 3 # concurrency (2) + 1, structural, not statistical + + +@pytest.mark.unit +@pytest.mark.parametrize("concurrency", [0, -1, 1]) +def test_concurrency_one_uses_the_calling_thread(concurrency: int) -> None: + """The kill switch is a real no-pool path, and the clamp handles bad + (<=0) values by degrading to serial rather than raising or spawning.""" + caller_thread = threading.current_thread() + seen_threads: list[threading.Thread] = [] + + def vectors_fn(batch: list[str]) -> list[list[float]]: + seen_threads.append(threading.current_thread()) + return [[0.0] for _ in batch] + + client = _FakeClient(vectors_fn) + embed = databricks_embedder( + "ep", "m", client=client, dim=1, batch_size=1, concurrency=concurrency + ) + embed(["a", "b", "c"]) + assert seen_threads == [caller_thread] * 3 + + +@pytest.mark.unit +def test_single_batch_never_spawns_a_pool() -> None: + """The query path (app/search/semantic.py: one text -> one batch) pays + nothing even at a high configured concurrency.""" + caller_thread = threading.current_thread() + seen_threads: list[threading.Thread] = [] + + def vectors_fn(batch: list[str]) -> list[list[float]]: + seen_threads.append(threading.current_thread()) + return [[0.0] for _ in batch] + + client = _FakeClient(vectors_fn) + embed = databricks_embedder("ep", "m", client=client, dim=1, concurrency=8) + embed(["only one text"]) + assert seen_threads == [caller_thread] + + +@pytest.mark.unit +def test_dim_mismatch_still_raises_under_concurrency() -> None: + """Aggregate _assert_dims still runs over the flattened result under the + pool path. batch_size=1 with two texts -> two batches, so this actually + exercises the pool branch rather than degrading to serial.""" + client = _FakeClient(lambda texts: [[0.0, 0.0, 0.0] for _ in texts]) # dim 3, expect 2 + embed = databricks_embedder("ep", "m", client=client, dim=2, batch_size=1, concurrency=4) + with pytest.raises(EmbeddingDimMismatchError): + embed(["a", "b"]) + + +@pytest.mark.unit +def test_embed_module_never_uses_as_completed() -> None: + """AST-based, not a substring check: a substring check would fail on this + very module's load-bearing explanatory comment about why as_completed must + never appear here (it yields in completion order, defeating AC1).""" + import app.embed as embed_module + + tree = ast.parse(inspect.getsource(embed_module)) + for node in ast.walk(tree): + if isinstance(node, ast.Name) and node.id == "as_completed": + pytest.fail("app.embed must never reference as_completed -- see module docstring") + if isinstance(node, ast.Attribute) and node.attr == "as_completed": + pytest.fail("app.embed must never reference as_completed -- see module docstring") + if isinstance(node, ast.ImportFrom) and any( + alias.name == "as_completed" for alias in node.names + ): + pytest.fail("app.embed must never import as_completed -- see module docstring") + + +@pytest.mark.unit +def test_get_embedder_threads_concurrency_from_settings(monkeypatch: pytest.MonkeyPatch) -> None: + captured: dict[str, Any] = {} + + def fake_databricks_embedder(endpoint: str, model: str, **kwargs: Any) -> Any: + captured.update(kwargs) + return lambda texts: [] + + monkeypatch.setattr("app.embed.databricks_embedder", fake_databricks_embedder) + cfg = Settings(semantic_embedding_endpoint="/ep", semantic_embedding_concurrency=7) + get_embedder(cfg) + assert captured["concurrency"] == 7 + + +@pytest.mark.unit +def test_empty_texts_under_concurrency() -> None: + calls = {"n": 0} + + def vectors_fn(texts: list[str]) -> list[list[float]]: + calls["n"] += 1 + return [[0.0] for _ in texts] + + client = _FakeClient(vectors_fn) + embed = databricks_embedder("ep", "m", client=client, dim=1, concurrency=8) + assert embed([]) == [] + assert calls["n"] == 0 diff --git a/tests/unit/test_repo_config.py b/tests/unit/test_repo_config.py index 51b88b2..e0d0eb4 100644 --- a/tests/unit/test_repo_config.py +++ b/tests/unit/test_repo_config.py @@ -323,7 +323,7 @@ def test_semantic_block_absent_is_a_noop() -> None: @pytest.mark.unit def test_semantic_block_full_maps_every_field_to_its_settings_name() -> None: - """A fully-populated block emits all six Settings-named keys, values intact.""" + """A fully-populated block emits all seven Settings-named keys, values intact.""" raw = ( b"version: 1\nconnections:\n - type: github\n users: [u]\n" b"semantic:\n" @@ -333,6 +333,7 @@ def test_semantic_block_full_maps_every_field_to_its_settings_name() -> None: b" embedding_model: acme.embed-v2\n" b" embedding_batch_size: 32\n" b" embedding_timeout_s: 15.0\n" + b" embedding_concurrency: 4\n" ) cfg = parse_config(raw, source="cfg") assert cfg.semantic.settings_overrides() == { @@ -342,6 +343,7 @@ def test_semantic_block_full_maps_every_field_to_its_settings_name() -> None: "semantic_embedding_model": "acme.embed-v2", "semantic_embedding_batch_size": 32, "semantic_embedding_timeout_s": 15.0, + "semantic_embedding_concurrency": 4, } @@ -371,6 +373,8 @@ def test_semantic_block_partial_emits_only_set_fields() -> None: ("embedding_timeout_s", b"-2.5"), ("embedding_endpoint", b'""'), ("embedding_model", b'""'), + ("embedding_concurrency", b"0"), + ("embedding_concurrency", b"9"), ], ) def test_semantic_block_rejects_out_of_bound_values(field: str, value: bytes) -> None: @@ -471,10 +475,11 @@ def test_settings_overrides_keys_are_real_settings_fields_and_types_survive() -> embedding_model="custom.model", embedding_batch_size=8, embedding_timeout_s=5.5, + embedding_concurrency=6, ) overrides = ov.settings_overrides() - # All six set -> all six emitted, and every key is a real Settings field. - assert len(overrides) == 6 + # All seven set -> all seven emitted, and every key is a real Settings field. + assert len(overrides) == 7 assert set(overrides) <= set(Settings.model_fields) # model_copy(update=) does NOT validate; re-validating the dumped model is what @@ -487,6 +492,7 @@ def test_settings_overrides_keys_are_real_settings_fields_and_types_survive() -> assert revalidated.semantic_embedding_model == "custom.model" assert revalidated.semantic_embedding_batch_size == 8 assert revalidated.semantic_embedding_timeout_s == 5.5 + assert revalidated.semantic_embedding_concurrency == 6 # --- parse failures -------------------------------------------------------