gracefully handle pathway and milvusdb simultaneous start (#10472)

GitOrigin-RevId: 208e4833bb6b0b2419e691fbf7a9caffa929ab26

Author: zxqfd555

CHANGELOG.md

@@ -10,6 +10,7 @@ This project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.htm
 - `pw.io.postgres.write` now streams each batch into PostgreSQL through the binary `COPY` protocol instead of issuing one `INSERT` per row, giving a large throughput improvement (up to ~100x) on bulk writes. Both output modes use it: stream-of-changes copies straight into the target, while snapshot mode stages each batch in a temporary table and merges it with a single set-based upsert/delete.
 
 ### Fixed
+- `pw.io.milvus.write` no longer intermittently fails with a "server unavailable" / "connect failed" error when pointed at a local `.db` file. The embedded local Milvus server reports itself as started before it actually accepts connections, so under load the first connection could lose the race against the server coming up; the connector now retries the initial connection until the local server is ready.
 - Improved concurrent write handling in pw.io.sqlite.write for SQLite databases. Writes to the same database file now produce deterministic output in multi-worker and multi-table setups.
 - `pw.io.elasticsearch.write` no longer fails when a minibatch is big enough that its Elasticsearch `_bulk` request would exceed a server-side limit. The connector reads both the cluster's `http.max_content_length` (the `413 Request Entity Too Large` limit) and `indexing_pressure.memory.limit` (the `429 Too Many Requests` limit, which on a small-heap node trips well below 100 MB) at start-up, and splits the buffered documents across as many bulk requests as needed to stay under whichever is hit first — so large batches are still written in as few requests as possible instead of being rejected. (Both limits fall back to a conservative default if they cannot be read.)
 - `pw.io.elasticsearch.write` now retries transient bulk failures with backoff instead of failing the run on the first hiccup. A whole-request rejection or an individual document failing with `429`/`503` (back-pressure / temporary unavailability) is retried — resending only the documents the server reports as not yet applied, so a retry never duplicates data — while deterministic per-document failures (e.g. a type-mismatched value rejected with `400`) are now logged and skipped rather than silently dropped.

integration_tests/db_connectors/utils.py

@@ -847,15 +847,53 @@ def get_table_contents(
 MILVUS_VECTOR_DIM = 3
 
 
+def _is_milvus_transient_connect_error(e: Exception) -> bool:
+    """Whether ``e`` is milvus-lite's embedded-server connection race.
+
+    milvus-lite reports its local server as started as soon as the server
+    process is alive, *before* the server's local socket actually accepts
+    connections. The first client therefore races the socket coming up, and a
+    client that loses the race fails with a ``server unavailable`` /
+    ``connect failed`` error. Under heavy test parallelism the socket can take
+    a few seconds to appear, so this is retried in
+    :func:`_connect_milvus_client_with_retry`.
+    """
+    text = str(e)
+    return (
+        "Fail connecting to server" in text
+        or "server unavailable" in text
+        or "failed to connect to all addresses" in text
+        or "No such file or directory" in text
+    )
+
+
+def _connect_milvus_client_with_retry(uri: str, *, timeout: float = 30.0):
+    """Create a ``MilvusClient`` for ``uri``, retrying while a freshly started
+    milvus-lite embedded server brings its local socket up (see
+    :func:`_is_milvus_transient_connect_error`). Non-transient errors are raised
+    immediately."""
+    from pymilvus import MilvusClient
+
+    deadline = time.monotonic() + timeout
+    while True:
+        try:
+            return MilvusClient(uri)
+        except Exception as e:
+            if (
+                not _is_milvus_transient_connect_error(e)
+                or time.monotonic() >= deadline
+            ):
+                raise
+            time.sleep(0.25)
+
+
 def _make_milvus_client(uri: str):
     """Create a MilvusClient for the given URI.
 
     For local ``.db`` files, works around a pymilvus 2.6.x bug where the
     Unix-domain-socket address is not forwarded to the gRPC handler.  See
     ``pathway.io.milvus._make_client`` for the full explanation.
     """
-    from pymilvus import MilvusClient
-
     if uri.endswith(".db"):
         try:
             from milvus_lite.server_manager import server_manager_instance
@@ -865,10 +903,10 @@ def _make_milvus_client(uri: str):
                 raise RuntimeError(
                     f"milvus-lite failed to start a local server for: {uri}"
                 )
-            return MilvusClient(uds_uri)
+            return _connect_milvus_client_with_retry(uds_uri)
         except ImportError:
             pass
-    return MilvusClient(uri)
+    return _connect_milvus_client_with_retry(uri)
 
 
 def _is_milvus_transient_init_error(e: Exception) -> bool:

python/pathway/io/milvus/__init__.py

@@ -2,6 +2,7 @@
 
 from __future__ import annotations
 
+import time
 from typing import Iterable
 
 import numpy as np
@@ -52,6 +53,44 @@ def _prepare_row(row: dict) -> dict:
     return result
 
 
+def _is_milvus_transient_connect_error(e: Exception) -> bool:
+    """Whether ``e`` is a local milvus-lite embedded-server connection race.
+
+    milvus-lite reports its local server as started as soon as the server
+    process is alive, before the server's local socket actually accepts
+    connections. The first client therefore races the socket coming up, and a
+    client that loses the race fails with a ``server unavailable`` /
+    ``connect failed`` error. The socket can take a few seconds to appear when
+    many local databases start at once, so this is retried in
+    :func:`_connect_with_retry`.
+    """
+    text = str(e)
+    return (
+        "Fail connecting to server" in text
+        or "server unavailable" in text
+        or "failed to connect to all addresses" in text
+        or "No such file or directory" in text
+    )
+
+
+def _connect_with_retry(MilvusClient, uri: str, *, timeout: float = 30.0):
+    """Create a ``MilvusClient`` for ``uri``, retrying while a freshly started
+    local milvus-lite server brings its socket up (see
+    :func:`_is_milvus_transient_connect_error`). Non-transient errors are raised
+    immediately."""
+    deadline = time.monotonic() + timeout
+    while True:
+        try:
+            return MilvusClient(uri)
+        except Exception as e:
+            if (
+                not _is_milvus_transient_connect_error(e)
+                or time.monotonic() >= deadline
+            ):
+                raise
+            time.sleep(0.25)
+
+
 def _make_client(MilvusClient, uri: str):
     """Create a MilvusClient, working around a pymilvus 2.6.x bug.
 
@@ -63,6 +102,9 @@ def _make_client(MilvusClient, uri: str):
     connection hangs.  Passing ``address=`` as an explicit keyword argument
     causes it to flow through ``handler_kwargs`` directly into
     ``GrpcHandler.__init__``, where ``kwargs.get("address")`` picks it up.
+
+    A freshly started local server may not be accepting connections yet, so the
+    client is created through :func:`_connect_with_retry`.
     """
     if uri.endswith(".db"):
         with optional_imports("milvus"):
@@ -73,9 +115,9 @@ def _make_client(MilvusClient, uri: str):
                 raise RuntimeError(
                     f"milvus-lite failed to start a local server for: {uri}"
                 )
-            return MilvusClient(uds_uri)
+            return _connect_with_retry(MilvusClient, uds_uri)
 
-    return MilvusClient(uri)
+    return _connect_with_retry(MilvusClient, uri)
 
 
 @check_arg_types