fix: Target server exits automatically

Author: moehanabi

docs/advanced_features/remote_training.md

@@ -50,6 +50,19 @@ Request processing flow:
 - Retrieves model configuration (hidden_size, vocab_size, etc.) via POST `/setup`
 - Each training step sends a request via POST `/generate` and receives results via NCCL recv
 - Supports TP>1 training: only rank 0 sends requests, results are broadcast to other ranks
+- After the first successful connection, a background heartbeat is started; on `close()`, a best-effort `/disconnect` is sent
+
+### Client Lifecycle and Automatic Exit
+
+The target server tracks client activity and automatically shuts down after the client exits, preventing leftover GPU-occupying server processes after training completes:
+
+- After the client's first successful request or successful NCCL initialization, a background heartbeat thread is started, sending POST `/heartbeat` every 15 seconds by default
+- When the client exits normally, it sends a best-effort POST `/disconnect`; upon receiving it, the server immediately triggers shutdown
+- When the client exits abnormally, the server watchdog triggers shutdown after `--client-heartbeat-timeout` is exceeded (default 60 seconds)
+- The server only counts actual client API calls as active requests; `GET /health` and unrelated POSTs do not renew the watchdog timer
+- `--client-heartbeat-timeout 0` disables the server-side timeout watchdog, but `/disconnect` will still trigger automatic shutdown
+
+Since NCCL transport does not support safe disconnect and reconnect within the same server process, it is recommended to treat each target server process as a resource for a single training session: it automatically exits after training completes or the client disconnects, and a new instance is started for the next training run.
 
 ### NCCL Transport
 
@@ -198,6 +211,7 @@ export NCCL_IB_GID_INDEX=3             # RoCE GID index
 | `SPECFORGE_TOPK` | `0` | Server-side target_p top-k compression (`0` = full distribution) |
 | `SPECFORGE_TARGET_DTYPE` | `fp32` | target_p computation precision |
 | `SPECFORGE_GPU_ID` | auto | Specify GPU device ID |
+| `SPECFORGE_HEARTBEAT_INTERVAL` | `15` | Client heartbeat send interval (seconds; `<=0` means the heartbeat thread is not started) |
 
 ## 📊 Benchmark Results
 

scripts/launch_target_server.py

@@ -105,6 +105,12 @@ def parse_args():
         default=None,
         help="NCCL TCP rendezvous port for GPU-to-GPU data transfer (default: HTTP port + 100)",
     )
+    parser.add_argument(
+        "--client-heartbeat-timeout",
+        type=float,
+        default=60.0,
+        help="Seconds of client inactivity before treating as disconnected (0 to disable). Default: 60",
+    )
     return parser.parse_args()
 
 
@@ -154,6 +160,7 @@ def main():
         nccl_port=args.nccl_port if args.nccl_port else args.port + 100,
         host=args.host,
         attention_backend=args.attention_backend,
+        client_heartbeat_timeout=args.client_heartbeat_timeout,
     )
     server_app.load_model()
     logger.info("Model loaded successfully.")
@@ -169,16 +176,26 @@ def main():
             args.port,
             args.mode,
         )
+        logger.info(
+            "Client disconnect will shut down server (heartbeat timeout: %.0fs)",
+            args.client_heartbeat_timeout,
+        )
 
-        def shutdown(signum, frame):
-            logger.info("Received signal %d, shutting down...", signum)
+        def request_http_shutdown():
             # HTTPServer.shutdown() must be called from a different thread than
             # serve_forever(); signal handlers run on the main thread.
             threading.Thread(target=httpd.shutdown, daemon=True).start()
 
+        def shutdown(signum, frame):
+            logger.info("Received signal %d, shutting down...", signum)
+            request_http_shutdown()
+
         signal.signal(signal.SIGINT, shutdown)
         signal.signal(signal.SIGTERM, shutdown)
 
+        # Register shutdown callback for client disconnect handling
+        server_app._shutdown_callback = request_http_shutdown
+
         try:
             httpd.serve_forever()
         except KeyboardInterrupt:

specforge/modeling/target/remote_target_client.py

@@ -250,6 +250,14 @@ def __init__(
         self._nccl_transport: Optional[NCCLTransport] = None
         self._nccl_init_attempted = False
         self._nccl_init_lock = threading.Lock()
+        # Heartbeat thread to keep server aware of client liveness
+        self._heartbeat_interval = float(
+            os.environ.get("SPECFORGE_HEARTBEAT_INTERVAL", "15")
+        )
+        self._heartbeat_stop = threading.Event()
+        self._heartbeat_thread: Optional[threading.Thread] = None
+        self._lifecycle_lock = threading.Lock()
+        self._closed = False
         atexit.register(self.close)
 
     def _get_nccl_port(self) -> int:
@@ -354,6 +362,8 @@ def _client_init():
                 return False
 
             logger.info("NCCL transport established successfully")
+            # Start heartbeat to keep server aware of client liveness
+            self._start_heartbeat()
             return True
 
     def _request(self, endpoint: str, payload: bytes) -> bytes:
@@ -370,6 +380,7 @@ def _request(self, endpoint: str, payload: bytes) -> bytes:
                     headers={"Content-Type": "application/octet-stream"},
                 )
                 resp.raise_for_status()
+                self._start_heartbeat()
                 return resp.content
             except (requests.ConnectionError, requests.Timeout) as exc:
                 last_exc = exc
@@ -413,6 +424,7 @@ def _request_transport(
                     url, data=payload, timeout=self.timeout, headers=headers
                 )
                 resp.raise_for_status()
+                self._start_heartbeat()
 
                 nccl_used = resp.headers.get(NCCL_HEADER) == "1"
 
@@ -442,11 +454,69 @@ def _request_transport(
                 ) from exc
 
     def close(self):
+        with self._lifecycle_lock:
+            if self._closed:
+                return
+            self._closed = True
+        # Stop heartbeat thread
+        self._stop_heartbeat()
+        # Notify server of disconnect (best-effort)
+        self._notify_disconnect()
+        # Destroy NCCL transport
         if self._nccl_transport is not None:
             self._nccl_transport.destroy()
             self._nccl_transport = None
         self._session.close()
 
+    def _notify_disconnect(self):
+        """Send disconnect notification to the server (best-effort)."""
+        try:
+            self._session.post(
+                f"{self.url}/disconnect",
+                data=b"",
+                timeout=5,
+                headers={"Content-Type": "application/octet-stream"},
+            )
+        except Exception:
+            # Best-effort: server may already be gone, network may be down
+            pass
+
+    def _start_heartbeat(self):
+        """Start background heartbeat thread."""
+        with self._lifecycle_lock:
+            if self._closed or self._heartbeat_thread is not None:
+                return
+            if self._heartbeat_interval <= 0:
+                return
+            self._heartbeat_stop.clear()
+            self._heartbeat_thread = threading.Thread(
+                target=self._heartbeat_loop, daemon=True
+            )
+            self._heartbeat_thread.start()
+
+    def _stop_heartbeat(self):
+        """Stop background heartbeat thread."""
+        with self._lifecycle_lock:
+            self._heartbeat_stop.set()
+            thread = self._heartbeat_thread
+            self._heartbeat_thread = None
+        if thread is not None:
+            thread.join(timeout=3)
+
+    def _heartbeat_loop(self):
+        """Periodically ping the server to indicate liveness."""
+        while not self._heartbeat_stop.wait(self._heartbeat_interval):
+            try:
+                self._session.post(
+                    f"{self.url}/heartbeat",
+                    data=b"",
+                    timeout=5,
+                    headers={"Content-Type": "application/octet-stream"},
+                )
+            except Exception:
+                # Server may be gone — stop sending heartbeats
+                break
+
     def __del__(self):
         try:
             self.close()