Fix TCP PeerAgent advertised host

JimyMa · JimyMa · commit f861120c47f8 · 2026-06-04T21:45:23.000+08:00
diff --git a/dlslime/dlslime/peer_agent/_agent.py b/dlslime/dlslime/peer_agent/_agent.py
@@ -8,13 +8,16 @@
 from __future__ import annotations
 
 import inspect
+import ipaddress
 import json
 import os
+import socket
 import threading
 import time
 from concurrent.futures import ThreadPoolExecutor
 from dataclasses import dataclass
 from typing import Any, Dict, List, Optional, Set, Tuple, Union
+from urllib.parse import urlparse
 
 try:
     import httpx
@@ -407,6 +410,45 @@ def _local_address(self) -> str:
                 return str(address)
         return ""
 
+    def _ctrl_host(self) -> str:
+        ctrl_url = self.ctrl_url
+        if not ctrl_url.startswith(("http://", "https://")):
+            ctrl_url = f"http://{ctrl_url}"
+        return urlparse(ctrl_url).hostname or ""
+
+    @staticmethod
+    def _is_loopback_host(host: str) -> bool:
+        if host in {"localhost", "::1"}:
+            return True
+        try:
+            return ipaddress.ip_address(host).is_loopback
+        except ValueError:
+            return False
+
+    @staticmethod
+    def _local_ip_for_remote(remote_host: str) -> str:
+        if not remote_host:
+            return ""
+        try:
+            with socket.socket(socket.AF_INET, socket.SOCK_DGRAM) as sock:
+                sock.connect((remote_host, 80))
+                return sock.getsockname()[0]
+        except OSError:
+            return ""
+
+    def _resolve_tcp_local_host(self, local_host: Optional[str]) -> str:
+        if not local_host:
+            discovered = self._local_address()
+            if discovered and not self._is_loopback_host(discovered):
+                return discovered
+            ctrl_host = self._ctrl_host()
+            if ctrl_host and not self._is_loopback_host(ctrl_host):
+                routed = self._local_ip_for_remote(ctrl_host)
+                if routed:
+                    return routed
+            return discovered or "0.0.0.0"
+        return str(local_host)
+
     def _normalize_link_type(self, link_type: Optional[str]) -> str:
         if not link_type:
             return "UNKNOWN"
@@ -960,7 +1002,17 @@ def _ensure_connection_from_meta(
     ) -> DirectedConnection:
         """Create local connection state from a peer's directed request."""
         if str(meta.get("transport") or "rdma").lower() == "tcp":
-            local_key: ResourceKey = TcpResourceKey()
+            with self._connections_lock:
+                peer_connections = [
+                    conn
+                    for conn in self._connections.values()
+                    if conn.peer_alias == peer_alias
+                ]
+                if len(peer_connections) == 1:
+                    return peer_connections[0]
+            local_key: ResourceKey = TcpResourceKey(
+                host=self._resolve_tcp_local_host(None), port=0
+            )
             peer_key: ResourceKey = TcpResourceKey()
             return self._get_or_create_connection(
                 peer_alias,
@@ -1156,6 +1208,14 @@ def _mark_connection_connected(self, conn_id: str) -> None:
             self._connected_peers.add(conn_id)
             self._connected_peers_cond.notify_all()
 
+    def _mark_connection_failed(self, conn_id: str) -> None:
+        with self._connections_lock:
+            conn = self._connections.get(conn_id)
+            if conn is not None:
+                conn.mark_failed()
+        with self._connected_peers_lock:
+            self._connected_peers_cond.notify_all()
+
     def _has_notified_peer(self, conn_id: str) -> bool:
         """True iff we've already sent qp_ready for this connection."""
         with self._notified_peers_lock:
@@ -1183,16 +1243,18 @@ def connect_to(
         qp_num: Optional[int] = 1,
         min_bw: Optional[str] = None,
         transport: str = "rdma",
-        local_host: str = "0.0.0.0",
+        local_host: Optional[str] = None,
         local_port: int = 0,
     ) -> PeerConnection:
         """Start connecting to a peer and return a connection handle.
 
         ``transport='rdma'`` (default) selects the RDMA path, picking a NIC
         from the discovered local topology. ``transport='tcp'`` selects the
         TCP path: ``local_host``/``local_port`` configure the local bind
-        (port 0 = OS-assigned), and ``ib_port``/``qp_num``/``local_device``/
-        ``peer_device`` are ignored.
+        (port 0 = OS-assigned). If ``local_host`` is left as the default
+        ``None``, PeerAgent publishes its discovered local host address instead
+        so remote peers do not receive ``0.0.0.0`` in endpoint_info.
+        ``ib_port``/``qp_num``/``local_device``/``peer_device`` are ignored.
         """
         if not isinstance(peer_alias, str) or not peer_alias:
             raise TypeError("connect_to() requires a non-empty peer alias string")
@@ -1213,6 +1275,7 @@ def connect_to(
         t0 = time.perf_counter()
 
         if transport_norm == "tcp":
+            local_host = self._resolve_tcp_local_host(local_host)
             local_key: ResourceKey = TcpResourceKey(
                 host=local_host, port=int(local_port)
             )
@@ -1306,6 +1369,11 @@ def _wait_connected(self, conn_id: str, timeout_sec: float = 60.0) -> None:
                         f"+{(time.perf_counter() - t0) * 1000:.3f}ms"
                     )
                     return
+                with self._connections_lock:
+                    conn = self._connections.get(conn_id)
+                    state = conn.state if conn is not None else "unknown"
+                if state == "failed":
+                    raise RuntimeError(f"Connection {conn_id!r} failed")
                 remaining = deadline - time.monotonic()
                 if remaining <= 0:
                     raise TimeoutError(
diff --git a/dlslime/dlslime/peer_agent/_mailbox.py b/dlslime/dlslime/peer_agent/_mailbox.py
@@ -373,6 +373,15 @@ def _try_connect_peer_inner(
         # D. Complete RDMA handshake
         t_d = time.perf_counter()
         endpoint.connect(peer_qp_info)
+        if hasattr(endpoint, "is_connected") and not endpoint.is_connected():
+            logger.warning(
+                "StreamMailbox %s: endpoint.connect(%s) failed for endpoint_info=%s",
+                self._agent.alias,
+                peer,
+                peer_qp_info,
+            )
+            self._agent._mark_connection_failed(conn_id)
+            return
         # Stash for one-sided ops on transports (TCP) where remote MR info
         # rides on the endpoint_info JSON instead of a separate Redis record.
         conn.peer_endpoint_info = peer_qp_info
diff --git a/dlslime/examples/python/p2p_tcp_rc_read_peer_agent_two_process.py b/dlslime/examples/python/p2p_tcp_rc_read_peer_agent_two_process.py
@@ -0,0 +1,199 @@
+#!/usr/bin/env python3
+"""TCP one-sided read through two PeerAgent processes.
+
+The initiator reads bytes directly out of a memory region published by the
+target, then notifies the target so both processes can shut down cleanly.
+
+Prerequisites:
+    1. Start NanoCtrl:   cd NanoCtrl && cargo run --release
+    2. Redis must be reachable.
+    3. dlslime built with BUILD_TCP=ON (default).
+
+Usage:
+    python p2p_tcp_rc_read_peer_agent_two_process.py --role target
+    python p2p_tcp_rc_read_peer_agent_two_process.py --role initiator
+    python p2p_tcp_rc_read_peer_agent_two_process.py --role target --ctrl_address http://host:4479
+    python p2p_tcp_rc_read_peer_agent_two_process.py --role initiator --ctrl_address http://host:4479
+    python p2p_tcp_rc_read_peer_agent_two_process.py --role target --local_host 10.0.0.2
+    python p2p_tcp_rc_read_peer_agent_two_process.py --role initiator --local_host 10.0.0.3
+"""
+
+import argparse
+import ctypes
+from typing import Optional
+
+from dlslime import PeerAgent
+
+
+INITIATOR_ALIAS = "tcp_read_initiator"
+TARGET_ALIAS = "tcp_read_target"
+PAYLOAD = b"one-sided-read-via-peer-agent"
+DONE = b"done"
+
+
+def _buffer_from_bytes(data: bytes) -> ctypes.Array:
+    buf = ctypes.create_string_buffer(len(data))
+    ctypes.memmove(ctypes.addressof(buf), data, len(data))
+    return buf
+
+
+def run_initiator(
+    ctrl_url: str, local_host: Optional[str], local_port: int, connect_timeout: float
+) -> None:
+    agent = PeerAgent(ctrl_url=ctrl_url, alias=INITIATOR_ALIAS)
+    try:
+        buf_a = ctypes.create_string_buffer(64)
+        addr_a = ctypes.addressof(buf_a)
+
+        # Register the local MR before connect_to so endpoint_info() carries it
+        # when the rendezvous fires. Required for one-sided ops on TCP.
+        agent.register_memory_region("buf_a", addr_a, 0, 64)
+
+        if TARGET_ALIAS not in agent.list_agents():
+            raise RuntimeError(f"target agent {TARGET_ALIAS!r} is not running")
+
+        resolved_host = agent._resolve_tcp_local_host(local_host)
+        print(f"[initiator] TCP local host: {resolved_host}:{local_port}")
+        conn = agent.connect_to(
+            TARGET_ALIAS,
+            transport="tcp",
+            local_host=local_host,
+            local_port=local_port,
+        )
+        conn.wait(timeout=connect_timeout)
+        local_info = conn.endpoint.endpoint_info()
+        print(
+            f"[initiator] Connected over TCP: {INITIATOR_ALIAS} -> {TARGET_ALIAS} "
+            f"(local={local_info.get('host')}:{local_info.get('port')})"
+        )
+
+        done_buf = _buffer_from_bytes(DONE)
+
+        ep = conn.endpoint
+        if not ep.is_connected():
+            raise RuntimeError(
+                "TCP endpoint is not connected; make sure both processes publish "
+                "a peer-reachable --local_host"
+            )
+        peer_info = conn.peer_endpoint_info
+        assert peer_info is not None
+        remote_mr_info = peer_info.get("mr_info", {}).get("buf_b")
+        if remote_mr_info is None:
+            raise RuntimeError(
+                f"{TARGET_ALIAS!r} is connected but did not publish buf_b; "
+                "start the target process first"
+            )
+
+        h_local = ep.register_memory_region("buf_a_loc", addr_a, 0, 64)
+        h_remote = ep.register_remote_memory_region("buf_b_rem", remote_mr_info)
+
+        st = ep.read([(h_local, h_remote, 0, 0, len(PAYLOAD))]).wait()
+        assert st == 0, f"read failed: {st}"
+        assert (
+            bytes(buf_a[: len(PAYLOAD)]) == PAYLOAD
+        ), "initiator did not read bytes"
+        print(
+            "[initiator] target->initiator one-sided read = "
+            f"{bytes(buf_a[: len(PAYLOAD)])!r}  ok"
+        )
+
+        st = agent.send(
+            TARGET_ALIAS, (ctypes.addressof(done_buf), 0, len(DONE))
+        ).wait()
+        assert st == 0, f"done send failed: {st}"
+        print("[initiator] Sent completion notice.")
+    finally:
+        agent.shutdown()
+
+
+def run_target(
+    ctrl_url: str, local_host: Optional[str], local_port: int, connect_timeout: float
+) -> None:
+    agent = PeerAgent(ctrl_url=ctrl_url, alias=TARGET_ALIAS)
+    try:
+        buf_b = ctypes.create_string_buffer(64)
+        addr_b = ctypes.addressof(buf_b)
+
+        # Pre-fill the target buffer with the payload the initiator reads out.
+        ctypes.memmove(addr_b, PAYLOAD, len(PAYLOAD))
+
+        # Register MRs before connect_to so endpoint_info() carries them when
+        # the rendezvous fires. Required for one-sided ops on TCP.
+        agent.register_memory_region("buf_b", addr_b, 0, 64)
+
+        resolved_host = agent._resolve_tcp_local_host(local_host)
+        print(f"[target] TCP local host: {resolved_host}:{local_port}")
+        conn = agent.connect_to(
+            INITIATOR_ALIAS,
+            transport="tcp",
+            local_host=local_host,
+            local_port=local_port,
+        )
+        conn.wait(timeout=connect_timeout)
+        local_info = conn.endpoint.endpoint_info()
+        print(
+            f"[target] Connected over TCP: {TARGET_ALIAS} -> {INITIATOR_ALIAS} "
+            f"(local={local_info.get('host')}:{local_info.get('port')})"
+        )
+
+        done_buf = ctypes.create_string_buffer(len(DONE))
+
+        if not conn.endpoint.is_connected():
+            raise RuntimeError(
+                "TCP endpoint is not connected; make sure both processes publish "
+                "a peer-reachable --local_host"
+            )
+        print("[target] Published buffer; waiting for initiator completion.")
+
+        st = agent.recv(
+            INITIATOR_ALIAS, (ctypes.addressof(done_buf), 0, len(DONE))
+        ).wait()
+        assert st == 0, f"done recv failed: {st}"
+        assert bytes(done_buf[: len(DONE)]) == DONE, "initiator did not send done"
+        print("[target] Initiator completed read; shutting down.")
+    finally:
+        agent.shutdown()
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(
+        description="TCP one-sided read through two PeerAgent processes"
+    )
+    parser.add_argument(
+        "--role",
+        choices=("initiator", "target"),
+        required=True,
+        help="Process role to run",
+    )
+    parser.add_argument(
+        "--ctrl_address",
+        "--ctrl",
+        default="http://127.0.0.1:4479",
+        help="NanoCtrl URL",
+    )
+    parser.add_argument(
+        "--local_host",
+        default=None,
+        help="Local TCP host/IP to publish to the peer; defaults to discovered host",
+    )
+    parser.add_argument(
+        "--local_port",
+        type=int,
+        default=0,
+        help="Local TCP port to bind; 0 lets the OS choose",
+    )
+    parser.add_argument(
+        "--connect_timeout",
+        type=float,
+        default=60.0,
+        help="Seconds to wait for the peer-agent TCP connection",
+    )
+    args = parser.parse_args()
+    if args.role == "initiator":
+        run_initiator(
+            args.ctrl_address, args.local_host, args.local_port, args.connect_timeout
+        )
+    else:
+        run_target(
+            args.ctrl_address, args.local_host, args.local_port, args.connect_timeout
+        )
diff --git a/dlslime/tests/python/test_directed_connection_io.py b/dlslime/tests/python/test_directed_connection_io.py
