Adopt an existing job server socket instead of overwriting it

With multiple server replicas, each replica used to overwrite the job's
/run/dstack/server.sock on every open, so ownership churned between
replicas (they repeatedly stole the socket from each other) and every
overwrite briefly broke access. Probe the socket first and (re)create the
reverse forward only when it is missing or unreachable; otherwise adopt
the existing one. This keeps a single stable owner. If the owner replica
dies, another takes over on its next health check.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>

Author: Andrey Cheptsov

src/dstack/_internal/server/services/jobs/server_connection.py

@@ -91,17 +91,6 @@ async def open(self) -> None:
         try:
             remote_dir = shlex.quote(str(_REMOTE_SOCKET_PATH.parent))
             remote_socket = shlex.quote(str(_REMOTE_SOCKET_PATH))
-            # A new server owner replaces the stable path, making an orphaned forward unreachable.
-            await self._tunnel.aexec(
-                f"mkdir -p {remote_dir} && chmod 755 {remote_dir} && rm -f {remote_socket}"
-            )
-            server_socket = _get_server_socket()
-            self._tunnel.reverse_forwarded_sockets = [
-                SocketPair(
-                    local=server_socket,
-                    remote=UnixSocket(path=_REMOTE_SOCKET_PATH),
-                )
-            ]
             # Probe through the job socket itself: a socket path can remain after its listener
             # becomes unreachable.
             self._tunnel.forwarded_sockets = [
@@ -111,14 +100,34 @@ async def open(self) -> None:
                 )
             ]
             await self._tunnel.aopen()
-            # The socket carries no credentials. World access inside the isolated job container
-            # lets configurations using a non-root `user` reach it as well.
-            await self._tunnel.aexec(f"chmod 666 {remote_socket}")
+            # With multiple server replicas, do not overwrite a socket that another replica already
+            # serves: create the reverse forward only when the current one is missing or
+            # unreachable. This keeps a single stable owner and avoids ownership churn (replicas
+            # repeatedly stealing the socket) and the brief unavailability of overwriting a live one.
             if not await self._server_is_reachable():
-                raise SSHError(
-                    "dstack server is not reachable from the job"
-                    f" (forward target {server_socket.render()})"
+                # A new server owner replaces the stable path, making an orphaned forward
+                # unreachable.
+                await self._tunnel.aexec(
+                    f"mkdir -p {remote_dir} && chmod 755 {remote_dir} && rm -f {remote_socket}"
                 )
+                server_socket = _get_server_socket()
+                self._tunnel.reverse_forwarded_sockets = [
+                    SocketPair(
+                        local=server_socket,
+                        remote=UnixSocket(path=_REMOTE_SOCKET_PATH),
+                    )
+                ]
+                # The probe forward is already established; do not request it again.
+                self._tunnel.forwarded_sockets = []
+                await self._tunnel.aopen()
+                # The socket carries no credentials. World access inside the isolated job container
+                # lets configurations using a non-root `user` reach it as well.
+                await self._tunnel.aexec(f"chmod 666 {remote_socket}")
+                if not await self._server_is_reachable():
+                    raise SSHError(
+                        "dstack server is not reachable from the job"
+                        f" (forward target {server_socket.render()})"
+                    )
         except Exception:
             await self._tunnel.aclose()
             raise

src/tests/_internal/server/services/jobs/test_server_connection.py

@@ -20,6 +20,8 @@
 def tunnel_mock(tmp_path, monkeypatch: pytest.MonkeyPatch):
     monkeypatch.setattr(server_connection, "CONNECTIONS_DIR", tmp_path)
     tunnel = MagicMock()
+    tunnel.forwarded_sockets = []
+    tunnel.reverse_forwarded_sockets = []
     tunnel.acheck = AsyncMock(return_value=False)
     tunnel.aopen = AsyncMock()
     tunnel.aexec = AsyncMock(return_value="")
@@ -64,33 +66,66 @@ def test_get_server_socket_follows_server_bind_host(
 
 class TestJobServerConnection:
     @pytest.mark.asyncio
-    async def test_opens_private_reverse_socket(self, tunnel_mock):
+    async def test_becomes_owner_when_socket_unreachable(self, tunnel_mock):
         tunnel, tunnel_class = tunnel_mock
         job = Mock(id=uuid.uuid4())
         connection = JobServerConnection(job, job_runtime_data=None)
-        connection._server_is_reachable = AsyncMock(return_value=True)
+        # The probe after the local forward is unreachable (no healthy owner) -> become owner;
+        # the probe after the reverse forward confirms reachability.
+        connection._server_is_reachable = AsyncMock(side_effect=[False, True])
+
+        snapshots = []
+
+        async def record_aopen():
+            snapshots.append(
+                (list(tunnel.forwarded_sockets), list(tunnel.reverse_forwarded_sockets))
+            )
+
+        tunnel.aopen.side_effect = record_aopen
 
         await connection.open()
 
         tunnel_class.assert_called_once()
-        assert tunnel.aopen.await_count == 2
-        assert tunnel.reverse_forwarded_sockets == [
-            SocketPair(
-                local=IPSocket(host="127.0.0.1", port=server_connection.settings.SERVER_PORT),
-                remote=UnixSocket(path=server_connection._REMOTE_SOCKET_PATH),
-            )
+        probe_pair = SocketPair(
+            local=UnixSocket(path=connection._probe_socket_path),
+            remote=UnixSocket(path=server_connection._REMOTE_SOCKET_PATH),
+        )
+        reverse_pair = SocketPair(
+            local=IPSocket(host="127.0.0.1", port=server_connection.settings.SERVER_PORT),
+            remote=UnixSocket(path=server_connection._REMOTE_SOCKET_PATH),
+        )
+        # master (no forwards) -> add probe forward -> add reverse forward
+        assert snapshots == [
+            ([], []),
+            ([probe_pair], []),
+            ([], [reverse_pair]),
+        ]
+        commands = [call.args[0] for call in tunnel.aexec.await_args_list]
+        assert commands == [
+            "mkdir -p /run/dstack && chmod 755 /run/dstack && rm -f /run/dstack/server.sock",
+            "chmod 666 /run/dstack/server.sock",
         ]
+
+    @pytest.mark.asyncio
+    async def test_adopts_existing_healthy_socket(self, tunnel_mock):
+        tunnel, _ = tunnel_mock
+        job = Mock(id=uuid.uuid4())
+        connection = JobServerConnection(job, job_runtime_data=None)
+        # A healthy owner (another replica) already serves the socket.
+        connection._server_is_reachable = AsyncMock(return_value=True)
+
+        await connection.open()
+
+        # master + probe forward only; the reverse forward is not (re)created
+        assert tunnel.aopen.await_count == 2
+        assert tunnel.reverse_forwarded_sockets == []
+        tunnel.aexec.assert_not_awaited()
         assert tunnel.forwarded_sockets == [
             SocketPair(
                 local=UnixSocket(path=connection._probe_socket_path),
                 remote=UnixSocket(path=server_connection._REMOTE_SOCKET_PATH),
             )
         ]
-        commands = [call.args[0] for call in tunnel.aexec.await_args_list]
-        assert commands == [
-            "mkdir -p /run/dstack && chmod 755 /run/dstack && rm -f /run/dstack/server.sock",
-            "chmod 666 /run/dstack/server.sock",
-        ]
 
     @pytest.mark.asyncio
     async def test_reuses_live_tunnel_with_existing_socket(self, tunnel_mock):