fix: wait for Mooncake master exit before replacement

tests/sanity/check_license.py

@@ -18,11 +18,13 @@
 
 # Add license headers below
 license_head_huawei = "Copyright 2025 Huawei Technologies Co., Ltd. All Rights Reserved."
-license_head_tq = "Copyright 2025 The TransferQueue Team"
+license_head_tq_2025 = "Copyright 2025 The TransferQueue Team"
+license_head_tq_2026 = "Copyright 2026 The TransferQueue Team"
 
 license_headers = [
     license_head_huawei,
-    license_head_tq,
+    license_head_tq_2025,
+    license_head_tq_2026,
 ]
 
 

tests/test_mooncake_bootstrap.py

@@ -0,0 +1,84 @@
+# Copyright 2026 The TransferQueue Team
+# SPDX-License-Identifier: Apache-2.0
+
+"""A replacement master must wait for the previous listener to release its port."""
+
+import socket
+import subprocess
+import sys
+from pathlib import Path
+
+import psutil
+import pytest
+from omegaconf import OmegaConf
+
+from transfer_queue.storage.bootstrap import mooncake_bootstrap
+
+
+@pytest.mark.parametrize("non_child", [False, True], ids=["own-child", "other-parent"])
+@pytest.mark.parametrize("ignore_term", [False, True], ids=["delayed-exit", "forced-exit"])
+def test_replacement_waits_for_old_listener(monkeypatch, ignore_term, non_child):
+    if non_child and sys.platform != "linux":
+        pytest.skip("Non-child zombie regression requires Linux")
+    script = """
+import signal, socket, sys, time
+listener = socket.socket()
+listener.bind(("127.0.0.1", 0))
+listener.listen()
+def stop(*args):
+    time.sleep(0.3)
+    sys.exit(0)
+signal.signal(signal.SIGTERM, signal.SIG_IGN if sys.argv[1] == "True" else stop)
+print(listener.getsockname()[1], flush=True)
+while True:
+    time.sleep(1)
+"""
+    command = [sys.executable, "-u", "-c", script, str(ignore_term)]
+    if non_child:
+        owner_script = """
+import subprocess, sys
+master = subprocess.Popen(sys.argv[1:], stdout=subprocess.PIPE, text=True)
+try:
+    port = master.stdout.readline().strip()
+    print(master.pid, port, flush=True)
+    sys.stdin.readline()
+finally:
+    if master.poll() is None:
+        master.kill()
+    master.wait(timeout=5)
+    master.stdout.close()
+"""
+        command = [sys.executable, "-u", "-c", owner_script, *command]
+    old = subprocess.Popen(command, stdin=subprocess.PIPE, stdout=subprocess.PIPE, text=True)
+    try:
+        if non_child:
+            master_pid, port = map(int, old.stdout.readline().split())
+        else:
+            master_pid, port = old.pid, int(old.stdout.readline())
+        master = psutil.Process(master_pid)
+        monkeypatch.setattr(
+            mooncake_bootstrap.subprocess,
+            "run",
+            lambda *args, **kwargs: subprocess.CompletedProcess(args, 0, stdout=f"{master_pid}\n"),
+        )
+
+        def launch_replacement(*args, **kwargs):
+            if non_child:
+                assert old.poll() is None, "The master owner must stay alive without reaping its child"
+                assert master.status() == psutil.STATUS_ZOMBIE
+            else:
+                assert old.poll() is not None, "Replacement launched before the old master exited"
+            with socket.socket() as listener:
+                listener.bind(("127.0.0.1", port))
+            raise RuntimeError("Replacement boundary reached")
+
+        monkeypatch.setattr(mooncake_bootstrap.subprocess, "Popen", launch_replacement)
+        with pytest.raises(RuntimeError, match="Replacement boundary reached"):
+            config = OmegaConf.load(Path(__file__).resolve().parents[1] / "transfer_queue" / "config.yaml")
+            mooncake_bootstrap.initialize_mooncake_storage(config)
+    finally:
+        old.stdin.close()
+        if not non_child and old.poll() is None:
+            old.kill()
+        old.wait(timeout=5)
+        old.stdout.close()

transfer_queue/storage/bootstrap/mooncake_bootstrap.py

@@ -1,5 +1,5 @@
 # Copyright 2025 Huawei Technologies Co., Ltd. All Rights Reserved.
-# Copyright 2025 The TransferQueue Team
+# Copyright 2026 The TransferQueue Team
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.
@@ -18,6 +18,7 @@
 import time
 from urllib.parse import urlparse
 
+import psutil
 from omegaconf import DictConfig
 
 from transfer_queue.storage.bootstrap.provider import StorageBootstrapProvider
@@ -26,6 +27,24 @@
 logger = get_logger(__name__)
 
 
+def _wait_for_master_exit(masters: list[psutil.Process], timeout: float) -> list[psutil.Process]:
+    """Wait for masters to exit, including zombies owned by another process."""
+    deadline = time.monotonic() + timeout
+    while masters:
+        _, alive = psutil.wait_procs(masters, timeout=min(0.1, max(0, deadline - time.monotonic())))
+        masters = []
+        for master in alive:
+            try:
+                # Zombies have released their listeners; only their parent can reap them.
+                if master.status() != psutil.STATUS_ZOMBIE:
+                    masters.append(master)
+            except psutil.NoSuchProcess:
+                pass
+        if time.monotonic() >= deadline:
+            break
+    return masters
+
+
 @StorageBootstrapProvider.register_provider("MooncakeStore")
 def initialize_mooncake_storage(conf: DictConfig) -> subprocess.Popen | dict | None:
     """
@@ -58,11 +77,26 @@ def initialize_mooncake_storage(conf: DictConfig) -> subprocess.Popen | dict | N
         pids = check.stdout.strip().replace("\n", ", ")
         logger.info(f"Find existing mooncake_master (PID: {pids}), try to kill first...")
 
-        result = os.system('pkill -f "[m]ooncake_master"')
-        if result == 0:
-            logger.info("Successfully killed existing mooncake_master processes.")
-        else:
-            raise RuntimeError(f"Failed to kill existing mooncake_master processes (exit code: {result}).")
+        masters = []
+        for pid in check.stdout.split():
+            try:
+                master = psutil.Process(int(pid))
+                master.terminate()
+                masters.append(master)
+            except psutil.NoSuchProcess:
+                pass
+        # Sending SIGTERM is asynchronous; the new master must not race the old
+        # HTTP metadata listener for its port. Escalate only these captured PIDs.
+        alive = _wait_for_master_exit(masters, timeout=10)
+        for master in alive:
+            try:
+                master.kill()
+            except psutil.NoSuchProcess:
+                pass
+        alive = _wait_for_master_exit(alive, timeout=5)
+        if alive:
+            raise RuntimeError(f"Existing mooncake_master processes did not exit: {[p.pid for p in alive]}")
+        logger.info("Existing mooncake_master processes have exited.")
 
     # process metadata_server (normalize and validate)
     metadata_server_raw_address = str(conf.backend.MooncakeStore.metadata_server).strip()