diff --git a/changelog.d/bulk-embeddings-pool.added.md b/changelog.d/bulk-embeddings-pool.added.md new file mode 100644 index 000000000..70b49c38e --- /dev/null +++ b/changelog.d/bulk-embeddings-pool.added.md @@ -0,0 +1,13 @@ +- **Separate "bulk" embeddings microservice pool for ingest.** `MicroserviceEmbedder` + gained an optional `embeddings_microservice_url_bulk` setting (on its `Settings` + dataclass in `opencontractserver/pipeline/embedders/sent_transformer_microservice.py`, + seeded from the `EMBEDDINGS_MICROSERVICE_URL_BULK` env var via + `migrate_pipeline_settings`). The ingest Celery tasks in + `opencontractserver/tasks/embeddings_task.py` tag their embed calls with + `use_bulk_pool=True`; `_get_service_config` routes those to the bulk URL when one is + configured, while search queries (untagged) stay on `embeddings_microservice_url`. + This isolates search-query latency from batch-ingest load with no change to the + embedder client, base class, or any query call site. Opt-in: when the bulk URL is + unset the flag is a no-op and ingest stays on the query pool, so single-pool + deployments are unaffected. The bulk URL is configured through the same + `PipelineSettings` singleton as the query URL — no separate configuration pathway. diff --git a/config/settings/base.py b/config/settings/base.py index 881d9b68a..7e78fb765 100644 --- a/config/settings/base.py +++ b/config/settings/base.py @@ -1168,6 +1168,13 @@ EMBEDDINGS_MICROSERVICE_URL = env( "EMBEDDINGS_MICROSERVICE_URL", default="http://vector-embedder:8000" ) +# Optional dedicated "bulk" pool for ingest embedding. This is the env binding +# that seeds the MicroserviceEmbedder ``embeddings_microservice_url_bulk`` +# setting (via ``migrate_pipeline_settings``); the value is consumed through the +# embedder's PipelineSettings singleton, NOT read directly here. When empty +# (the default), batch ingest stays on EMBEDDINGS_MICROSERVICE_URL — set it only +# to split ingest onto a separate pool while search queries stay warm. +EMBEDDINGS_MICROSERVICE_URL_BULK = env("EMBEDDINGS_MICROSERVICE_URL_BULK", default="") VECTOR_EMBEDDER_API_KEY = env("VECTOR_EMBEDDER_API_KEY", default="") # CLIP embedder configuration (768-dimensional vectors) CLIP_EMBEDDER_URL = env("CLIP_EMBEDDER_URL", default="http://vector-embedder:8000") diff --git a/docs/deployment/performance_tuning.md b/docs/deployment/performance_tuning.md index 1ae87c0b9..36b5ef162 100644 --- a/docs/deployment/performance_tuning.md +++ b/docs/deployment/performance_tuning.md @@ -187,6 +187,47 @@ gunicorn fleet). celery -A config.celery_app worker --concurrency=8 -Q ingest ``` +### Separate "bulk" embeddings pool for ingest + +Search queries embed one short string and need a *warm* microservice pod (a +cold start adds seconds to every query). Batch ingest embeds thousands of +strings and is happy to hit an autoscaled / scale-to-zero pool. Serving both +from one URL forces a compromise. + +`MicroserviceEmbedder` supports a dedicated bulk pool without adding a parallel +config pathway — the bulk URL is just another field on its `Settings` +singleton (`embeddings_microservice_url_bulk`, seeded from the +`EMBEDDINGS_MICROSERVICE_URL_BULK` env var via `migrate_pipeline_settings`, +alongside `embeddings_microservice_url`). The ingest Celery tasks in +`opencontractserver/tasks/embeddings_task.py` tag their embed calls with +`use_bulk_pool=True`; `MicroserviceEmbedder._get_service_config` then routes +those to the bulk URL while every (untagged) search query stays on +`embeddings_microservice_url`. When no bulk URL is configured the flag is a +no-op and ingest stays on the query pool, so single-pool deployments need no +change. + +```bash +# Point ingest at a separate autoscaled pool; queries stay on the warm pod. +EMBEDDINGS_MICROSERVICE_URL=http://vector-embedder:8000 # warm queries +EMBEDDINGS_MICROSERVICE_URL_BULK=http://vector-embedder-bulk:8000 # bulk ingest +# then reseed the MicroserviceEmbedder singleton so it picks up the new URL: +python manage.py migrate_pipeline_settings --component MicroserviceEmbedder --force +``` + +**Why `--force`:** `migrate_pipeline_settings` *preserves* existing +`PipelineSettings` values on re-run unless `--force` is given +(`migrate_pipeline_settings.py:233`). A deploy that ran the command once has +already persisted `embeddings_microservice_url_bulk: ""` (its default), so a +later plain `migrate_pipeline_settings` after setting the env var would keep the +empty value and silently leave ingest on the query pool. `--component +MicroserviceEmbedder` scopes the force to this one component so any hand-tuned +DB settings on other components are left untouched. + +**Scope (v1):** the bulk pool reuses the query pool's `vector_embedder_api_key` +and Cloud Run IAM auth — it is a separate URL, not a separately-credentialed +deployment. If the bulk pool needs its own API key, that is not configurable +yet. + ## What's still slow (open work) The 80% per-doc overhead breaks down roughly as: diff --git a/docs/sample_env_files/backend/local/.django b/docs/sample_env_files/backend/local/.django index 862ff2ff5..101dd3836 100644 --- a/docs/sample_env_files/backend/local/.django +++ b/docs/sample_env_files/backend/local/.django @@ -48,6 +48,12 @@ OPENAI_MODEL=gpt-4o # Microservice URLs # ------------------------------------------------------------------------------ EMBEDDINGS_MICROSERVICE_URL=http://vector-embedder:8000 +# Optional: dedicated pool for batch ingest embedding. Seeds the embedder's +# embeddings_microservice_url_bulk setting. After changing it, reseed with: +# python manage.py migrate_pipeline_settings --component MicroserviceEmbedder --force +# (--force is required because the command otherwise preserves the existing DB +# value). When unset, ingest uses EMBEDDINGS_MICROSERVICE_URL. +# EMBEDDINGS_MICROSERVICE_URL_BULK=http://vector-embedder-bulk:8000 DOCLING_PARSER_SERVICE_URL=http://docling-parser:8000/parse/ # Docling diff --git a/docs/sample_env_files/backend/production/.django b/docs/sample_env_files/backend/production/.django index 38f18eef3..00c4f5798 100644 --- a/docs/sample_env_files/backend/production/.django +++ b/docs/sample_env_files/backend/production/.django @@ -66,6 +66,13 @@ AUTH0_M2M_MANAGEMENT_GRANT_TYPE=client_credentials # bare metal, or managed service deployments, replace with the actual URLs. # ------------------------------------------------------------------------------ EMBEDDINGS_MICROSERVICE_URL=http://vector-embedder:8000 +# Optional: dedicated pool for batch ingest embedding. Seeds the embedder's +# embeddings_microservice_url_bulk setting. After changing it, reseed with: +# python manage.py migrate_pipeline_settings --component MicroserviceEmbedder --force +# (--force is required because the command otherwise preserves the existing DB +# value). When unset, ingest uses EMBEDDINGS_MICROSERVICE_URL. Set it to isolate +# search-query latency from ingest load. +# EMBEDDINGS_MICROSERVICE_URL_BULK=http://vector-embedder-bulk:8000 DOCLING_PARSER_SERVICE_URL=http://docling-parser:8000/parse/ # Docling diff --git a/opencontractserver/pipeline/embedders/sent_transformer_microservice.py b/opencontractserver/pipeline/embedders/sent_transformer_microservice.py index f0edfe54c..378dafc25 100644 --- a/opencontractserver/pipeline/embedders/sent_transformer_microservice.py +++ b/opencontractserver/pipeline/embedders/sent_transformer_microservice.py @@ -148,6 +148,23 @@ class Settings: ) }, ) + embeddings_microservice_url_bulk: str = field( + default="", + metadata={ + "pipeline_setting": PipelineSetting( + setting_type=SettingType.OPTIONAL, + required=False, + description=( + "Optional separate microservice URL for bulk ingest " + "embedding. When set, batch ingest (calls tagged with " + "use_bulk_pool=True) routes here while search queries " + "stay on embeddings_microservice_url; falls back to " + "embeddings_microservice_url when empty." + ), + env_var="EMBEDDINGS_MICROSERVICE_URL_BULK", + ) + }, + ) vector_embedder_api_key: str = field( default="", metadata={ @@ -191,9 +208,21 @@ def _get_service_config(self, all_kwargs: dict) -> tuple[str, dict]: """ s = self.settings if self.settings is not None else self.Settings() - service_url = all_kwargs.get( + query_url = all_kwargs.get( "embeddings_microservice_url", s.embeddings_microservice_url ) + # Ingest tasks tag their calls with ``use_bulk_pool=True``. When a + # dedicated bulk URL is configured we route those to it; otherwise (and + # for every query call, which never sets the flag) we stay on the + # always-warm query URL, so single-pool deployments are unaffected. The + # bulk URL lives in the same PipelineSettings singleton as the query URL + # rather than a separate config pathway. + bulk_url = all_kwargs.get( + "embeddings_microservice_url_bulk", s.embeddings_microservice_url_bulk + ) + use_bulk_pool = bool(all_kwargs.get("use_bulk_pool", False)) + service_url = bulk_url if (use_bulk_pool and bulk_url) else query_url + api_key = all_kwargs.get("vector_embedder_api_key", s.vector_embedder_api_key) use_cloud_run_iam_auth = bool( all_kwargs.get("use_cloud_run_iam_auth", s.use_cloud_run_iam_auth) diff --git a/opencontractserver/tasks/embeddings_task.py b/opencontractserver/tasks/embeddings_task.py index 3a8fe8116..9e2946372 100644 --- a/opencontractserver/tasks/embeddings_task.py +++ b/opencontractserver/tasks/embeddings_task.py @@ -69,7 +69,10 @@ def _create_text_embedding( f"with embedder {embedder_path} (text length={len(text)})" ) - vector = embedder.embed_text(text) + # Ingest routes through the dedicated bulk pool when one is configured + # (see MicroserviceEmbedder._get_service_config); embedders without a bulk + # URL ignore the flag and stay on their query URL. + vector = embedder.embed_text(text, use_bulk_pool=True) if vector is None: logger.error( @@ -592,7 +595,8 @@ def _batch_embed_items( def _embed_one(chunk): texts_only = [text for _, text in chunk] - return chunk, embedder.embed_texts_batch(texts_only) + # Ingest routes through the dedicated bulk pool when configured. + return chunk, embedder.embed_texts_batch(texts_only, use_bulk_pool=True) # Map future -> chunk index for logging/sub-batch numbering. # @@ -1050,7 +1054,8 @@ def _embed_relationship( embedder_path, len(text), ) - vector = embedder.embed_text(text) + # Ingest routes through the dedicated bulk pool when configured. + vector = embedder.embed_text(text, use_bulk_pool=True) if vector is None: logger.error( "Embedder %s returned None for relationship %s", diff --git a/opencontractserver/tests/test_batch_embedding.py b/opencontractserver/tests/test_batch_embedding.py index 916a8416a..c8236ad55 100644 --- a/opencontractserver/tests/test_batch_embedding.py +++ b/opencontractserver/tests/test_batch_embedding.py @@ -219,6 +219,35 @@ def test_basic_batch(self): self.assertEqual(result["failed"], 0) self.assertEqual(result["skipped"], 0) + def test_batch_tags_bulk_pool(self): + """Batch ingest tags embed_texts_batch with use_bulk_pool=True.""" + captured: dict = {} + + class RecordingBatchEmbedder(BaseEmbedder): + title = "Recording Batch" + description = "Test" + author = "Test" + dependencies = [] + vector_size = 384 + supported_file_types = [FileTypeEnum.TXT] + + def _embed_text_impl(self, text, **all_kwargs): + return [0.1] * self.vector_size + + def embed_texts_batch(self, texts, **kw): + captured.update(kw) + return [[0.1] * self.vector_size for _ in texts] + + annots = [_make_mock_annotation(1, "Hello world")] + result = self._make_result() + + _batch_embed_text_annotations( + annots, RecordingBatchEmbedder(), "test.RecordingBatchEmbedder", 50, result + ) + + self.assertEqual(result["succeeded"], 1) + self.assertTrue(captured.get("use_bulk_pool")) + def test_empty_text_skipped(self): """Annotations with empty/whitespace text are skipped.""" annots = [ @@ -711,10 +740,11 @@ def test_nan_values_handled_per_item(self, mock_post): class TestMicroserviceEmbedderSingleText(unittest.TestCase): """Test MicroserviceEmbedder._embed_text_impl and _get_service_config.""" - def _make_embedder(self, api_key="", use_cloud_run=False): + def _make_embedder(self, api_key="", use_cloud_run=False, bulk_url=""): embedder = MicroserviceEmbedder() embedder._settings = MicroserviceEmbedder.Settings( embeddings_microservice_url="http://test-service:8080", + embeddings_microservice_url_bulk=bulk_url, vector_embedder_api_key=api_key, use_cloud_run_iam_auth=use_cloud_run, ) @@ -835,6 +865,24 @@ def test_get_service_config_fallback_to_settings(self): url, headers = embedder._get_service_config({}) self.assertEqual(url, "http://test-service:8080") + def test_get_service_config_uses_bulk_pool_when_flagged(self): + """use_bulk_pool=True routes to the configured bulk URL (issue: ingest pool).""" + embedder = self._make_embedder(bulk_url="http://bulk-pool:9090") + url, _ = embedder._get_service_config({"use_bulk_pool": True}) + self.assertEqual(url, "http://bulk-pool:9090") + + def test_get_service_config_bulk_flag_without_bulk_url_falls_back(self): + """With the flag set but no bulk URL configured, stays on the query URL.""" + embedder = self._make_embedder(bulk_url="") + url, _ = embedder._get_service_config({"use_bulk_pool": True}) + self.assertEqual(url, "http://test-service:8080") + + def test_get_service_config_ignores_bulk_url_without_flag(self): + """A configured bulk URL is never used for query calls (no flag).""" + embedder = self._make_embedder(bulk_url="http://bulk-pool:9090") + url, _ = embedder._get_service_config({}) + self.assertEqual(url, "http://test-service:8080") + class TestCalculateEmbeddingsForAnnotationBatch(unittest.TestCase): """Integration tests for the calculate_embeddings_for_annotation_batch task. diff --git a/opencontractserver/tests/test_dual_embeddings.py b/opencontractserver/tests/test_dual_embeddings.py index b51322443..fd3d37e6a 100644 --- a/opencontractserver/tests/test_dual_embeddings.py +++ b/opencontractserver/tests/test_dual_embeddings.py @@ -30,7 +30,7 @@ class MockEmbedder: is_multimodal = False supports_images = False - def embed_text(self, text: str) -> list[float]: + def embed_text(self, text: str, **kwargs) -> list[float]: """Return a mock embedding vector.""" return [0.1] * 768 @@ -42,7 +42,7 @@ class MockCorpusEmbedder: is_multimodal = False supports_images = False - def embed_text(self, text: str) -> list[float]: + def embed_text(self, text: str, **kwargs) -> list[float]: """Return a different mock embedding vector.""" return [0.2] * 768 diff --git a/opencontractserver/tests/test_embeddings_task.py b/opencontractserver/tests/test_embeddings_task.py index 58dd96db6..b56ec346a 100644 --- a/opencontractserver/tests/test_embeddings_task.py +++ b/opencontractserver/tests/test_embeddings_task.py @@ -474,7 +474,9 @@ def test_calculate_embedding_for_annotation_text_with_explicit_embedder( mock_get_component.assert_called_with(explicit_embedder_path) # Verify embed_text was called - mock_embedder_instance.embed_text.assert_called_with("This is test text") + mock_embedder_instance.embed_text.assert_called_with( + "This is test text", use_bulk_pool=True + ) # The key test: verify that the explicit embedder_path was used mock_annot.add_embedding.assert_called_with(explicit_embedder_path, test_vector) @@ -553,7 +555,7 @@ def test_calculate_embedding_for_annotation_text_fallback_to_annotation_corpus( # Verify default embedder was called mock_get_default.assert_called_once() mock_default_embedder_instance.embed_text.assert_called_with( - "This is test text" + "This is test text", use_bulk_pool=True ) # Verify corpus was retrieved for dual embedding @@ -561,7 +563,9 @@ def test_calculate_embedding_for_annotation_text_fallback_to_annotation_corpus( # Verify corpus embedder was called for dual embedding mock_get_component.assert_called_with("corpus.embedder.path") - mock_corpus_embedder_instance.embed_text.assert_called_with("This is test text") + mock_corpus_embedder_instance.embed_text.assert_called_with( + "This is test text", use_bulk_pool=True + ) # Verify both embeddings were stored (default + corpus-specific) calls = mock_annot.add_embedding.call_args_list @@ -696,7 +700,9 @@ def test_annotation_with_images_non_multimodal_embedder_falls_back_to_text( ) # Should have called text embedding (fallback) - mock_embedder.embed_text.assert_called_once_with("Figure 1 caption") + mock_embedder.embed_text.assert_called_once_with( + "Figure 1 caption", use_bulk_pool=True + ) # Should have stored embedding mock_annot.add_embedding.assert_called_once_with( @@ -813,7 +819,9 @@ def test_annotation_multimodal_failure_falls_back_to_text( ) # Should have fallen back to text embedding - mock_embedder.embed_text.assert_called_once_with("Figure with error") + mock_embedder.embed_text.assert_called_once_with( + "Figure with error", use_bulk_pool=True + ) # Should have stored embedding mock_annot.add_embedding.assert_called_once_with( @@ -850,7 +858,9 @@ def test_annotation_text_only_modality_uses_text_embedding( ) # Should have called text embedding (no images to embed) - mock_embedder.embed_text.assert_called_once_with("Just plain text") + mock_embedder.embed_text.assert_called_once_with( + "Just plain text", use_bulk_pool=True + ) self.assertTrue(result) @@ -879,7 +889,9 @@ def test_annotation_no_modalities_defaults_to_text(self, mock_annotation_model): ) # Should default to text embedding - mock_embedder.embed_text.assert_called_once_with("No modalities set") + mock_embedder.embed_text.assert_called_once_with( + "No modalities set", use_bulk_pool=True + ) self.assertTrue(result) @@ -1553,7 +1565,7 @@ def test_successful_embedding_returns_true(self): ) self.assertTrue(result) - mock_embedder.embed_text.assert_called_once_with("HEAD\nT1") + mock_embedder.embed_text.assert_called_once_with("HEAD\nT1", use_bulk_pool=True) @patch( "opencontractserver.tasks.embeddings_task.synthesize_relationship_block_text" @@ -1572,7 +1584,9 @@ def test_synthesizes_text_when_not_precomputed(self, mock_synth): self.assertTrue(result) mock_synth.assert_called_once_with(mock_rel) - mock_embedder.embed_text.assert_called_once_with("synthesized") + mock_embedder.embed_text.assert_called_once_with( + "synthesized", use_bulk_pool=True + ) @patch( "opencontractserver.tasks.embeddings_task.synthesize_relationship_block_text" @@ -1595,7 +1609,9 @@ def test_precomputed_text_skips_synthesize(self, mock_synth): self.assertTrue(result) mock_synth.assert_not_called() - mock_embedder.embed_text.assert_called_once_with("precomputed") + mock_embedder.embed_text.assert_called_once_with( + "precomputed", use_bulk_pool=True + ) class TestCalculateEmbeddingsForRelationshipBatch(unittest.TestCase): @@ -1692,7 +1708,8 @@ def test_explicit_embedder_counts_outcomes( self.assertEqual(result["skipped"], 1) self.assertEqual(len(result["errors"]), 2) mock_embedder.embed_texts_batch.assert_called_once_with( - ["relationship text 1", "relationship text 2", "relationship text 3"] + ["relationship text 1", "relationship text 2", "relationship text 3"], + use_bulk_pool=True, ) rel1.add_embedding.assert_called_once_with("explicit.path", [0.1] * 384) rel2.add_embedding.assert_called_once_with("explicit.path", [0.2] * 384) @@ -1800,5 +1817,44 @@ def fake_dual(**kwargs): self.assertEqual(len(result["errors"]), 1) +class TestIngestBulkPoolRouting(unittest.TestCase): + """Ingest leaves tag their embedder calls with use_bulk_pool=True. + + The bulk-vs-query URL selection itself lives in + MicroserviceEmbedder._get_service_config (tested in test_batch_embedding.py); + these tests pin the ingest side of that contract — every text-ingest leaf + signals bulk intent so a configured bulk pool is actually used. + """ + + def test_create_text_embedding_tags_bulk_pool(self): + from opencontractserver.tasks.embeddings_task import _create_text_embedding + + mock_obj = MagicMock() + mock_embedder = MagicMock() + mock_embedder.embed_text.return_value = [0.1, 0.2] + + result = _create_text_embedding( + mock_obj, mock_embedder, "embedder.path", "hello", "document", 1 + ) + + self.assertTrue(result) + mock_embedder.embed_text.assert_called_once_with("hello", use_bulk_pool=True) + + def test_embed_relationship_tags_bulk_pool(self): + from opencontractserver.tasks.embeddings_task import _embed_relationship + + mock_rel = MagicMock() + mock_rel.id = 1 + mock_embedder = MagicMock() + mock_embedder.embed_text.return_value = [0.1, 0.2] + + result = _embed_relationship( + mock_rel, mock_embedder, "embedder.path", precomputed_text="HEAD\nT1" + ) + + self.assertTrue(result) + mock_embedder.embed_text.assert_called_once_with("HEAD\nT1", use_bulk_pool=True) + + if __name__ == "__main__": unittest.main()