diff --git a/bioengine/apps/manager.py b/bioengine/apps/manager.py index 783ca792..653e752b 100644 --- a/bioengine/apps/manager.py +++ b/bioengine/apps/manager.py @@ -28,6 +28,112 @@ ) +# A hypha_rpc service proxy is pinned to ONE client instance of the remote +# service. That instance can change client id at any time — the artifact manager +# re-registered under a new one with nothing restarted — and hypha_rpc's own +# websocket reconnect does not re-resolve cached proxies, so every later call +# addresses a client that no longer exists and hangs to timeout. Split by whether +# the call can have reached the server: a send-side failure provably did not, so +# retrying it is safe even for a write; a timeout or a mid-call disconnect may +# already have landed, so retrying those could double-execute a create or commit. +_PROXY_NEVER_SENT_MARKERS = ( + "failed to send the request", + "websocket reconnection timed out", +) +_PROXY_MAYBE_SENT_MARKERS = ( + "client disconnected", + "method call timed out", + "service not found", +) + +# Reads are safe to repeat even when the first attempt may already have run, so +# they are retried on either failure kind. This matters because the call that +# hung in the observed outage was a read — the write before it had already +# succeeded — so without this a stale handle still costs one failed deploy. +# Enumerated against the live service 2026-09-15 (50 methods); everything not +# listed (create, commit, edit, delete, publish, discard, put_file, vector and +# PR operations) re-raises instead. +_RETRY_SAFE_READS = frozenset( + {"read", "list", "search", "get_file", "read_file", "list_files"} +) +# ...but reads are not side-effect-free: ``silent`` defaults to False and +# ``read`` increments a view count. Only these three accept the parameter, so +# only these three can be silenced, and only on the retry — the first attempt +# should still count as a real view. +_SILENCEABLE_READS = frozenset({"read", "list", "get_file"}) + + +def _stale_proxy_kind(exc: BaseException) -> Optional[str]: + """``"never_sent"``, ``"maybe_sent"``, or ``None`` if not a stale-proxy failure.""" + message = str(exc).lower() + if any(marker in message for marker in _PROXY_NEVER_SENT_MARKERS): + return "never_sent" + if isinstance(exc, asyncio.TimeoutError) or any( + marker in message for marker in _PROXY_MAYBE_SENT_MARKERS + ): + return "maybe_sent" + return None + + +class _ReconnectingArtifactManager: + """The artifact-manager proxy, re-resolved when its client id goes stale. + + Wrapping at the point of resolution rather than at each call site is + deliberate: this object is handed to ``AppBuilder`` and to every + ``artifact_utils`` helper, so one wrapper covers all of them and no caller + has to remember the retry rule. + + Two axes decide whether the call is repeated against the fresh proxy: + + * a call that provably never left the process is always safe to repeat; + * a call that may already have executed is repeated only if it is a read, + because replaying a ``create`` or ``commit`` would double-execute it. + + Anything else re-raises, so the caller sees one failure rather than an + indefinite outage and the next call uses the refreshed handle. + """ + + def __init__(self, server: RemoteService, proxy: Any, logger: logging.Logger): + self._server = server + self._proxy = proxy + self._logger = logger + self._generation = 0 + self._lock = asyncio.Lock() + + async def _re_resolve(self, seen_generation: int) -> None: + async with self._lock: + # Concurrent callers all fail against the same dead proxy; only the + # first needs to replace it. + if self._generation != seen_generation: + return + self._proxy = await self._server.get_service("public/artifact-manager") + self._generation += 1 + self._logger.info("Re-resolved the artifact manager service proxy.") + + def __getattr__(self, name: str): + async def call(*args, **kwargs): + seen_generation = self._generation + try: + return await getattr(self._proxy, name)(*args, **kwargs) + except Exception as exc: + kind = _stale_proxy_kind(exc) + if kind is None: + raise + self._logger.warning( + f"Artifact manager call '{name}' failed against a stale " + f"service proxy: {exc}" + ) + await self._re_resolve(seen_generation) + if kind != "never_sent" and name not in _RETRY_SAFE_READS: + raise + if kind == "maybe_sent" and name in _SILENCEABLE_READS: + # The first attempt may already have counted this view. + kwargs = {**kwargs, "silent": True} + return await getattr(self._proxy, name)(*args, **kwargs) + + return call + + def _is_serve_controller_gone(exc: BaseException) -> bool: """Whether a ``serve.*`` call failed because the Serve controller is gone. @@ -1075,9 +1181,13 @@ async def complete_initialization( self.admin_users = admin_users try: - # Get artifact manager service - self.artifact_manager = await self.server.get_service( - "public/artifact-manager" + # Get artifact manager service. Wrapped because the proxy is cached + # for the process lifetime and the remote service can change client + # id without anything restarting. + self.artifact_manager = _ReconnectingArtifactManager( + server=self.server, + proxy=await self.server.get_service("public/artifact-manager"), + logger=self.logger, ) self.logger.info("Successfully connected to artifact manager.") except Exception as e: diff --git a/tests/apps/test_artifact_manager_reconnect.py b/tests/apps/test_artifact_manager_reconnect.py new file mode 100644 index 00000000..d23f7d2e --- /dev/null +++ b/tests/apps/test_artifact_manager_reconnect.py @@ -0,0 +1,306 @@ +"""The worker re-resolves its artifact-manager proxy instead of wedging forever. + +A hypha_rpc service proxy is pinned to one *client instance* of the remote +service. The artifact manager can re-register under a new client id with nothing +having restarted — observed at KTH, hypha-server at restartCount 0 and two days +old — and hypha_rpc's websocket reconnect does not re-resolve cached proxies. The +worker resolves this one once in ``initialize()`` and holds it for the process +lifetime, so every artifact-backed call (``deploy_app``, ``list_apps``, +``get_app_manifest``) then hung to timeout while every local call stayed green. +Nothing recovered it: ``auto_redeploy`` needs the artifact read that is broken, +and ``run_code`` executes in Ray tasks so it cannot reach the in-process handle. +model-runner was down ~62 minutes and only a pod restart cleared it (#0074). + +The retry rule is the delicate part, and it is not "retry transport errors": + +* a **send-side** failure provably never reached the server, so retrying it is + safe even for a write; +* a **timeout or mid-call disconnect** may already have executed, so retrying + would risk double-running a ``create`` or ``commit``. Those re-raise, and only + the handle is refreshed — one visible failure instead of an endless outage. +""" +from __future__ import annotations + +import asyncio +import logging + +import pytest + +from bioengine.apps.manager import ( + _ReconnectingArtifactManager, + _stale_proxy_kind, +) + +logger = logging.getLogger("test-artifact-manager") + + +class _Proxy: + """One resolved client instance of the artifact manager.""" + + def __init__(self, name: str, fails_with: BaseException | None = None): + self.name = name + self._fails_with = fails_with + self.calls: list[tuple] = [] + self.kwargs_seen: list[dict] = [] + + def _answer(self, method, artifact_id, kwargs): + self.calls.append((method, artifact_id)) + self.kwargs_seen.append(kwargs) + if self._fails_with is not None: + raise self._fails_with + return {"id": artifact_id, "served_by": self.name} + + async def read(self, artifact_id, **kwargs): + return self._answer("read", artifact_id, kwargs) + + async def read_file(self, artifact_id, **kwargs): + return self._answer("read_file", artifact_id, kwargs) + + async def commit(self, artifact_id, **kwargs): + return self._answer("commit", artifact_id, kwargs) + + +class _Server: + """Hands out a fresh proxy each time the service is resolved.""" + + def __init__(self, *proxies: _Proxy): + self._proxies = list(proxies) + self.resolutions = 0 + + async def get_service(self, service_id): + assert service_id == "public/artifact-manager" + self.resolutions += 1 + return self._proxies.pop(0) + + +def _wrap(server: _Server, proxy: _Proxy) -> _ReconnectingArtifactManager: + return _ReconnectingArtifactManager(server=server, proxy=proxy, logger=logger) + + +# ===== the classifier ===== + + +def test_send_side_failures_are_known_to_have_never_landed(): + assert ( + _stale_proxy_kind(RuntimeError("Failed to send the request when calling method")) + == "never_sent" + ) + assert ( + _stale_proxy_kind(RuntimeError("WebSocket reconnection timed out")) + == "never_sent" + ) + + +def test_failures_that_may_have_landed_are_classified_separately(): + assert _stale_proxy_kind(asyncio.TimeoutError()) == "maybe_sent" + assert ( + _stale_proxy_kind(ConnectionError("Client disconnected: ws/abc")) == "maybe_sent" + ) + assert ( + _stale_proxy_kind(RuntimeError("Method call timed out: ws/abc:services.x.read")) + == "maybe_sent" + ) + + +def test_ordinary_errors_are_not_stale_proxy_failures(): + assert _stale_proxy_kind(ValueError("artifact not found")) is None + assert _stale_proxy_kind(PermissionError("denied")) is None + + +# ===== pass-through ===== + + +@pytest.mark.asyncio +async def test_a_healthy_call_is_forwarded_untouched(): + proxy = _Proxy("first") + server = _Server() + am = _wrap(server, proxy) + + result = await am.read("ws/my-app") + + assert result["served_by"] == "first" + assert server.resolutions == 0, "a healthy call must not re-resolve" + + +@pytest.mark.asyncio +async def test_an_application_error_propagates_without_re_resolving(): + """Only transport failures mean the handle is stale.""" + proxy = _Proxy("first", fails_with=ValueError("artifact not found")) + server = _Server() + am = _wrap(server, proxy) + + with pytest.raises(ValueError): + await am.read("ws/missing") + + assert server.resolutions == 0 + + +# ===== the repair ===== + + +@pytest.mark.asyncio +async def test_a_send_side_failure_re_resolves_and_retries(): + dead = _Proxy("dead", fails_with=RuntimeError("Failed to send the request")) + fresh = _Proxy("fresh") + server = _Server(fresh) + am = _wrap(server, dead) + + result = await am.read("ws/my-app") + + assert server.resolutions == 1 + assert result["served_by"] == "fresh", "the retry must use the new proxy" + + +@pytest.mark.asyncio +async def test_a_timed_out_write_refreshes_the_handle_but_is_never_replayed(): + """The double-execute guard. + + A timeout cancels nothing at the far end, so a ``commit`` that timed out may + still be running. Replaying it would create a second version snapshot. + """ + dead = _Proxy("dead", fails_with=asyncio.TimeoutError()) + fresh = _Proxy("fresh") + server = _Server(fresh) + am = _wrap(server, dead) + + with pytest.raises(asyncio.TimeoutError): + await am.commit("ws/my-app") + + assert server.resolutions == 1, "the handle must still be refreshed" + assert fresh.calls == [], "a write that may have landed must not be replayed" + + +@pytest.mark.asyncio +async def test_a_timed_out_read_is_retried_against_the_fresh_proxy(): + """#0074's own signature, end to end. + + The call that hung in the outage was ``read`` — the ``upload_app`` write + before it had already succeeded, and ``deploy_app`` then failed at manifest + load. Repeating a read is safe however the first attempt failed, so this is + the case that takes the observed fault to zero failed calls rather than one. + """ + dead = _Proxy("dead", fails_with=asyncio.TimeoutError()) + fresh = _Proxy("fresh") + server = _Server(fresh) + am = _wrap(server, dead) + + result = await am.read("ws/my-app") + + assert server.resolutions == 1 + assert result["served_by"] == "fresh" + + +@pytest.mark.asyncio +async def test_a_retried_read_does_not_count_a_second_view(): + """Reads are repeatable but not side-effect-free. + + ``silent`` defaults to False and ``read`` increments a view count, so a + naive retry would double-count. The first attempt still counts normally; + only the retry is silenced. + """ + dead = _Proxy("dead", fails_with=asyncio.TimeoutError()) + fresh = _Proxy("fresh") + server = _Server(fresh) + am = _wrap(server, dead) + + await am.read("ws/my-app") + + assert fresh.kwargs_seen[-1].get("silent") is True + + +@pytest.mark.asyncio +async def test_a_send_side_read_retry_is_not_silenced(): + """A call that never reached the server counted nothing, so the retry is + the *first* real view and must be recorded as one.""" + dead = _Proxy("dead", fails_with=RuntimeError("Failed to send the request")) + fresh = _Proxy("fresh") + server = _Server(fresh) + am = _wrap(server, dead) + + await am.read("ws/my-app") + + assert "silent" not in fresh.kwargs_seen[-1] + + +@pytest.mark.asyncio +async def test_a_read_without_a_silent_parameter_is_retried_unsilenced(): + """``read_file`` and ``list_files`` do not accept ``silent``; passing it + would turn a recoverable timeout into a TypeError.""" + dead = _Proxy("dead", fails_with=asyncio.TimeoutError()) + fresh = _Proxy("fresh") + server = _Server(fresh) + am = _wrap(server, dead) + + await am.read_file("ws/my-app", file_path="manifest.yaml") + + assert "silent" not in fresh.kwargs_seen[-1] + + +@pytest.mark.asyncio +async def test_the_next_call_after_a_timed_out_write_succeeds(): + """The outage is one failed call, not an indefinite wedge.""" + dead = _Proxy("dead", fails_with=asyncio.TimeoutError()) + fresh = _Proxy("fresh") + server = _Server(fresh) + am = _wrap(server, dead) + + with pytest.raises(asyncio.TimeoutError): + await am.commit("ws/my-app") + + result = await am.read("ws/my-app") + assert result["served_by"] == "fresh" + assert server.resolutions == 1, "the handle was already replaced" + + +@pytest.mark.asyncio +async def test_concurrent_failures_re_resolve_the_service_once(): + """Every in-flight call fails against the same dead proxy. + + Without the generation guard each one would resolve its own replacement, + turning a single eviction into a burst of ``get_service`` calls. + """ + dead = _Proxy("dead", fails_with=RuntimeError("Failed to send the request")) + fresh = _Proxy("fresh") + server = _Server(fresh) + am = _wrap(server, dead) + + results = await asyncio.gather(*(am.read(f"ws/app-{i}") for i in range(5))) + + assert server.resolutions == 1 + assert all(r["served_by"] == "fresh" for r in results) + + +@pytest.mark.asyncio +async def test_a_failed_re_resolve_surfaces_rather_than_being_swallowed(): + """If the service genuinely cannot be resolved, say so.""" + + class _BrokenServer: + resolutions = 0 + + async def get_service(self, service_id): + raise RuntimeError("Hypha is returning 500") + + dead = _Proxy("dead", fails_with=RuntimeError("Failed to send the request")) + am = _wrap(_BrokenServer(), dead) + + with pytest.raises(RuntimeError, match="500"): + await am.read("ws/my-app") + + +# ===== the fix reaches the cached handle ===== + + +def test_the_manager_wraps_the_proxy_it_caches(): + """The wrapper is useless unless the manager actually installs it. + + This is the step that makes the fix reach ``AppBuilder`` and every + ``artifact_utils`` helper for free, since they are all handed this object. + """ + import inspect + + from bioengine.apps.manager import AppsManager + + source = inspect.getsource(AppsManager.complete_initialization) + + assert "_ReconnectingArtifactManager(" in source + assert "self.artifact_manager = _ReconnectingArtifactManager(" in source