Fix: re-probe unhealthy ColPali endpoints after cooldown (#398)

Previously, when an endpoint failed once it was discarded from
healthy_endpoints and only re-added if ALL endpoints failed at
the same time. A single transient OOM on one of N endpoints could
silently halve sustained ingestion throughput until the worker
process restarted.

Now each endpoint is timestamped when marked unhealthy and
re-included after a 60s cooldown. If it's still failing, the
existing failure path will mark it unhealthy again — no change
to error semantics.

Co-authored-by: Claude Opus 4.7 (1M context) <noreply@anthropic.com>

Author: Adityav369

core/embedding/colpali_api_embedding_model.py

@@ -51,10 +51,27 @@ def __init__(self):
 
         # Track endpoint health for failover
         self.healthy_endpoints: set[str] = set(self.endpoints)
+        self._endpoint_unhealthy_since: Dict[str, float] = {}
+        self._unhealthy_recovery_seconds: float = 60.0
         self._endpoint_latencies: dict[str, float] = {}
         self.endpoint = self.endpoints[0]
         self._latest_ingest_metrics: Dict[str, float] = {}
 
+    def _recover_endpoints(self) -> None:
+        """Re-include endpoints whose unhealthy cooldown has elapsed."""
+        if not self._endpoint_unhealthy_since:
+            return
+        now = time.monotonic()
+        recovered = [
+            ep
+            for ep, marked_at in self._endpoint_unhealthy_since.items()
+            if now - marked_at >= self._unhealthy_recovery_seconds
+        ]
+        for ep in recovered:
+            self.healthy_endpoints.add(ep)
+            self._endpoint_unhealthy_since.pop(ep, None)
+            logger.info("Re-probing previously unhealthy ColPali endpoint: %s", ep)
+
     async def embed_for_ingestion(self, chunks: Union[Chunk, List[Chunk]]) -> List[MultiVector]:
         ingest_start = time.monotonic()
         # Normalize to list
@@ -131,6 +148,7 @@ async def _embed_inputs_distributed(
         if not indexed_inputs:
             return {}
 
+        self._recover_endpoints()
         # Use healthy endpoints, fall back to all if none healthy
         endpoints = list(self.healthy_endpoints) if self.healthy_endpoints else self.endpoints
         n_endpoints = len(endpoints)
@@ -166,6 +184,7 @@ async def _embed_inputs_distributed(
             elif isinstance(result, Exception):
                 logger.warning(f"Endpoint {endpoint} failed: {result}")
                 self.healthy_endpoints.discard(endpoint)
+                self._endpoint_unhealthy_since[endpoint] = time.monotonic()
                 failed_inputs.extend(batch)
             else:
                 merged.update(result)
@@ -182,6 +201,7 @@ async def _embed_inputs_distributed(
                 # All endpoints failed, reset health and raise
                 logger.error("All ColPali endpoints failed, resetting health status")
                 self.healthy_endpoints = set(self.endpoints)
+                self._endpoint_unhealthy_since.clear()
                 raise RuntimeError(
                     f"All {len(self.endpoints)} ColPali endpoints failed for {len(failed_inputs)} {input_type} inputs"
                 )
@@ -289,6 +309,7 @@ async def _call_api_endpoint(self, endpoint: str, inputs: List[str], input_type:
 
     async def embed_for_query(self, text: str) -> MultiVector:
         # Use first healthy endpoint for queries (single text, fast)
+        self._recover_endpoints()
         endpoint = next(iter(self.healthy_endpoints), self.endpoints[0])
         data = await self._call_api_endpoint(endpoint, [text], "text")
         if not data:
@@ -304,6 +325,7 @@ async def generate_embeddings(self, content: Union[str, Image]) -> np.ndarray:
         Returns:
             numpy array of embeddings.
         """
+        self._recover_endpoints()
         endpoint = next(iter(self.healthy_endpoints), self.endpoints[0])
 
         if isinstance(content, Image):