diff --git a/tests/sanity/check_license.py b/tests/sanity/check_license.py index e1cac117..994fadef 100644 --- a/tests/sanity/check_license.py +++ b/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, ] diff --git a/tests/test_mooncake_bootstrap.py b/tests/test_mooncake_bootstrap.py new file mode 100644 index 00000000..11b4dc0b --- /dev/null +++ b/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() diff --git a/transfer_queue/storage/bootstrap/mooncake_bootstrap.py b/transfer_queue/storage/bootstrap/mooncake_bootstrap.py index 6697222d..e95891f7 100644 --- a/transfer_queue/storage/bootstrap/mooncake_bootstrap.py +++ b/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()