Coverage for src/qdrant_loader/core/pipeline/workers/embedding_worker.py: 97%
101 statements
« prev ^ index » next coverage.py v7.15.0, created at 2026-07-20 10:15 +0000
1"""Embedding worker for processing chunks into embeddings."""
3import asyncio
4import gc
5from collections.abc import AsyncIterator
6from typing import Any
8import psutil
10from qdrant_loader.core.embedding.embedding_service import EmbeddingService
11from qdrant_loader.core.monitoring import prometheus_metrics
12from qdrant_loader.utils.logging import LoggingConfig
14from .base_worker import BaseWorker
16logger = LoggingConfig.get_logger(__name__)
19class EmbeddingWorker(BaseWorker):
20 """Handles chunk embedding with batching."""
22 def __init__(
23 self,
24 embedding_service: EmbeddingService,
25 max_workers: int = 4,
26 queue_size: int = 1000,
27 shutdown_event: asyncio.Event | None = None,
28 ):
29 super().__init__(max_workers, queue_size)
30 self.embedding_service = embedding_service
31 self.shutdown_event = shutdown_event or asyncio.Event()
33 async def process(self, chunks: list[Any]) -> list[tuple[Any, list[float] | None]]:
34 """Process a batch of chunks into embeddings.
36 The result is aligned 1:1 with ``chunks``: a chunk whose embedding
37 came back empty (invalid content was skipped by ``get_embeddings``)
38 is still included, paired with ``None``, so callers can account for
39 every chunk's fate instead of the failure silently disappearing.
41 Args:
42 chunks: List of chunks to embed
44 Returns:
45 List of (chunk, embedding) tuples; embedding is None on failure.
46 """
47 if not chunks:
48 return []
50 try:
51 logger.debug(f"EmbeddingWorker processing batch of {len(chunks)} items")
53 # Monitor memory usage
54 memory_percent = psutil.virtual_memory().percent
55 if memory_percent > 85:
56 logger.warning(
57 f"High memory usage detected: {memory_percent}%. Running garbage collection..."
58 )
59 gc.collect()
61 with prometheus_metrics.EMBEDDING_DURATION.time():
62 # Add timeout to prevent hanging and check for shutdown
63 embeddings = await asyncio.wait_for(
64 self.embedding_service.get_embeddings([c.content for c in chunks]),
65 timeout=300.0, # Increased to 5 minute timeout for large batches
66 )
68 # Check for shutdown before returning
69 if self.shutdown_event.is_set():
70 logger.debug("EmbeddingWorker skipping result due to shutdown")
71 return []
73 result: list[tuple[Any, list[float] | None]] = [
74 (chunk, emb if emb else None)
75 for chunk, emb in zip(chunks, embeddings, strict=False)
76 ]
77 skipped = sum(1 for _, emb in result if emb is None)
78 if skipped:
79 logger.warning(
80 f"Skipped {skipped} chunk(s) with empty embeddings, they will not be upserted"
81 )
82 logger.debug(f"EmbeddingWorker completed batch of {len(chunks)} items")
84 # Cleanup after large batches
85 if len(chunks) > 50:
86 gc.collect()
88 return result
90 except TimeoutError:
91 logger.error(
92 f"EmbeddingWorker timed out processing batch of {len(chunks)} items"
93 )
94 raise
95 except Exception as e:
96 logger.error(f"EmbeddingWorker error processing batch: {e}")
97 raise
99 async def _process_batch_guarded(
100 self, batch: list[Any], batch_index: int
101 ) -> list[tuple[Any, list[float] | None]]:
102 """Run process() for a batch under the shared concurrency semaphore.
104 Bounds how many embedding batches run concurrently to ``max_workers``
105 so that ``max_embed_workers`` actually governs concurrency, instead of
106 every batch running to completion before the next one starts.
108 If the whole batch raises (timeout, service error), every chunk in it
109 is still returned, paired with ``None``, instead of being dropped —
110 otherwise those chunks vanish from accounting entirely and their
111 parent documents can end up looking untouched rather than failed.
112 """
113 async with self.semaphore:
114 try:
115 logger.debug(
116 f"🔄 Processing embedding batch {batch_index} "
117 f"with {len(batch)} chunks..."
118 )
119 return await self.process(batch)
120 except Exception as e:
121 logger.error(f"EmbeddingWorker batch processing failed: {e}")
122 for chunk in batch:
123 logger.error(f"Embedding failed for chunk {chunk.id}: {e}")
124 return [(chunk, None) for chunk in batch]
126 async def process_chunks(
127 self, chunks: AsyncIterator[Any]
128 ) -> AsyncIterator[tuple[Any, list[float] | None]]:
129 """Process chunks into embeddings.
131 Batches are dispatched concurrently, but the number of batches
132 dispatched-and-not-yet-yielded is capped at ``max_workers``: once that
133 many are in flight, the oldest is awaited before a new one is created.
134 This bounds memory use (not just execution concurrency) when chunks
135 arrive far faster than embedding calls complete, while still yielding
136 results in submission order. Every dispatched chunk is yielded exactly
137 once, paired with ``None`` if it failed to embed, so downstream stages
138 can account for it instead of it silently vanishing.
140 Args:
141 chunks: AsyncIterator of chunks to process
143 Yields:
144 (chunk, embedding) tuples; embedding is None on failure.
145 """
146 logger.debug("EmbeddingWorker started")
147 logger.info(
148 f"🔄 Starting embedding generation (max_workers={self.max_workers})..."
149 )
150 batch_size = self.embedding_service.batch_size
151 batch: list[Any] = []
152 pending: list[asyncio.Task] = []
153 batch_index = 0
154 total_processed = 0
155 total_failed = 0
157 async def drain_oldest() -> list[tuple[Any, list[float] | None]]:
158 nonlocal total_processed, total_failed
159 results = await pending.pop(0)
160 if not results:
161 return []
163 succeeded = sum(1 for _, embedding in results if embedding is not None)
164 failed = len(results) - succeeded
165 total_processed += succeeded
166 total_failed += failed
167 logger.info(
168 f"🔗 Generated embeddings: {succeeded} items in batch"
169 + (f" ({failed} failed)" if failed else "")
170 + f", {total_processed} total processed"
171 )
172 return results
174 def dispatch(batch_to_dispatch: list[Any]) -> None:
175 nonlocal batch_index
176 batch_index += 1
177 pending.append(
178 asyncio.create_task(
179 self._process_batch_guarded(batch_to_dispatch, batch_index)
180 )
181 )
183 try:
184 async for chunk in chunks:
185 if self.shutdown_event.is_set():
186 logger.debug("EmbeddingWorker exiting due to shutdown")
187 break
189 batch.append(chunk)
191 # Dispatch batch when it reaches the desired size
192 if len(batch) >= batch_size:
193 batch_to_submit = batch
194 batch = []
196 # Keep at most max_workers batches in flight so memory
197 # use stays bounded, not just execution concurrency.
198 if len(pending) >= self.max_workers:
199 for result in await drain_oldest():
200 yield result
202 dispatch(batch_to_submit)
204 # Dispatch any remaining chunks in the final batch
205 if batch and not self.shutdown_event.is_set():
206 dispatch(batch)
208 while pending:
209 for result in await drain_oldest():
210 yield result
212 logger.info(
213 f"✅ Embedding completed: {total_processed} chunks processed, "
214 f"{total_failed} failed"
215 )
217 except asyncio.CancelledError:
218 logger.debug("EmbeddingWorker cancelled")
219 for task in pending:
220 if not task.done():
221 task.cancel()
222 raise
223 finally:
224 logger.debug("EmbeddingWorker exited")