server.py

# start_services.py
import subprocess
import time
import signal
import sys
import argparse
import os
from typing import Dict, List
# Configuration parameters
BASE_DIR = os.path.dirname(os.path.abspath(__file__))

eval_envs = {
    "N_GPUS": "8",
    "BASE_MODEL": "Qwen2.5-7B-Instruct_EVAL",
    "DATA_DIR": os.path.join(BASE_DIR, "v4data"),
    "ROLLOUT_TP_SIZE": "4",
    "EXPERIMENT_NAME": "medical",
    "PROJECT_NAME": "AgentLightningDebug",
    "VLLM_USE_V1": "1",
    "CUDA_VISIBLE_DEVICES": "0,1,2,3",
    "SAVE_DIR": os.path.join(BASE_DIR, "medical", "eval_saves"),
    "RUN_NAME": "grpo_two_step_search_scoretrain_v4",
    "SAVE_TOP_K": "5",
    "EVAL_STEP": "6",
    "EVAL_LABEL_FILE": os.path.join(BASE_DIR, "v4_data.jsonl")
}

train_envs = {
    "N_GPUS": "8",
    "BASE_MODEL": "Qwen2.5-7B-Instruct",
    "DATA_DIR": os.path.join(BASE_DIR, "v5data"),
    "ROLLOUT_TP_SIZE": "4",
    "EXPERIMENT_NAME": "medical",
    "PROJECT_NAME": "AgentLightningDebug",
    "VLLM_USE_V1": "1",
    "CUDA_VISIBLE_DEVICES": "0,1,2,3,4,5,6,7",
    "SAVE_DIR": os.path.join(BASE_DIR, "medical", "saves"),
    "RUN_NAME": "grpo_two_step_search_scoretrain_v5_full_7B_test",
    "SAVE_TOP_K": "5",
    "EVAL_LABEL_FILE": os.path.join(BASE_DIR, "v5_data.jsonl")
}

def pre_training_service():
    """预处理函数：可以在这里添加任何需要在启动训练服务前执行的操作"""
    # Example: if not os.path.exists(r"/root/Qwen2.5-7B-Instruct"):
    # Example: result = subprocess.run(["dl-hf-model", "Qwen2.5-7B-Instruct", "/root"], ...)
    # 注意：此处代码已替换为示例注释，请根据实际部署环境配置模型路径
    
    # 重启 Ray 集群
    print("🔄 正在重启 Ray 集群")
    result = subprocess.run(
        ["bash", os.path.join(BASE_DIR, "agent-lightning", "scripts", "restart_ray.sh")],
        check=True,
        stdout=subprocess.DEVNULL,
        stderr=sys.stderr
    )
    if result.returncode != 0:
        raise RuntimeError("❌ 重启 Ray 集群失败，请检查脚本或环境配置")
    print("✅ Ray 集群重启完成")


def pre_eval_service():
    # 将测试模型软连接到当前目录的Qwen2.5-7B-Instruct_EVAL,链接step xx 模型到 Qwen2.5-7B-Instruct_EVAL
    target_path = r"./Qwen2.5-7B-Instruct_EVAL"
    # os.unlink(target_path)
    # Example: source_path = rf"/root/{eval_envs['RUN_NAME']}/global_step_{eval_envs['EVAL_STEP']}/actor/huggingface/"
    # Example: os.symlink(src=source_path, dst=target_path)
    # print(f"✅ 已创建软连接: {source_path} -> {target_path}")

    # 重启 Ray 集群
    print("🔄 正在重启 Ray 集群")
    result = subprocess.run(
        ["bash", os.path.join(BASE_DIR, "agent-lightning", "scripts", "restart_ray.sh")],
        check=True,
        stdout=subprocess.DEVNULL,
        stderr=sys.stderr
    )
    if result.returncode != 0:
        raise RuntimeError("❌ 重启 Ray 集群失败，请检查脚本或环境配置")
    print("✅ Ray 集群重启完成")
    


SERVICE_DICT = {
    "search_agent": {
        "cwd": os.path.join(BASE_DIR, "medical"),
        "cmd": ["python", os.path.join(BASE_DIR, "medical", "search_agent.py")],
        "args": [],
        "log_file": os.path.join(BASE_DIR, "search_agent.log"),
        "pre_action": None
    },
    "retrieval_service": {
        "cwd": os.path.join(BASE_DIR, "agent-lightning"),
        "cmd": ["python", os.path.join(BASE_DIR, "agent-lightning", "search_server.py")],
        "args": [],
        "log_file": os.path.join(BASE_DIR, "retrieval_service.log"),
        "pre_action": None
    },
    "training_service": {
        "cwd": os.path.join(BASE_DIR, "medical"),
        "cmd": ["bash", "train.sh"],
        "args": [],
        "log_file": os.path.join(BASE_DIR, "training_service.log"),
        "pre_action": pre_training_service
    }
}

class ServiceManager:
    def __init__(self):
        self.processes: Dict[str, subprocess.Popen] = {}
        self.log_files = {}  # 存储已打开的日志文件对象

    def _open_log_file(self, log_file_path: str):
        """确保日志文件所在目录存在，并打开文件"""
        # 提取目录
        log_dir = os.path.dirname(log_file_path)
        if log_dir:
            os.makedirs(log_dir, exist_ok=True)  # 确保目录存在
        f = open(log_file_path, 'w', encoding='utf-8')
        return f

    def start_service(self, args, name: str, cmd: List[str], cwd: str = None, log_file: str = None):
        print(f"🚀 启动服务: {name} | {' '.join(cmd)} (cwd: {cwd})")
        print(f"📝 日志将写入: {log_file}")

        # 打开日志文件
        try:
            logfile = self._open_log_file(log_file)
            self.log_files[name] = logfile  # 保存文件句柄用于后续关闭
        except Exception as e:
            raise RuntimeError(f"❌ 无法创建日志文件 {log_file}: {e}")
        if args.mode == "train":
            envs = {**os.environ, **train_envs}
        else:
            envs = {**os.environ, **eval_envs}
        # 启动进程
        proc = subprocess.Popen(
            cmd,
            cwd=cwd,
            stdout=logfile,
            stderr=logfile,
            bufsize=1,
            env={**os.environ, **envs} 
        )
        self.processes[name] = proc

        time.sleep(2)  # 等待启动
        if proc.poll() is not None:
            raise RuntimeError(f"❌ 服务 {name} 启动失败！检查日志: {log_file}")

        print(f"✅ {name} 已启动 (PID: {proc.pid})，日志: {log_file}")

    def stop_all(self):
        print("\n🛑 正在停止所有服务...")
        # 先停止所有进程
        for name, proc in self.processes.items():
            if proc.poll() is None:  # 进程仍在运行
                proc.send_signal(signal.SIGINT)
                try:
                    proc.wait(timeout=10)
                except subprocess.TimeoutExpired:
                    print(f"⚠️  {name} 未在10秒内退出，执行 kill")
                    proc.kill()
                print(f"🛑 {name} 已停止")

        # 再关闭所有日志文件
        for name, f in self.log_files.items():
            f.close()
            print(f"📄 日志已保存: {f.name}")

    def __enter__(self):
        return self

    def __exit__(self, exc_type, exc_value, traceback):
        self.stop_all()


def main():
    parser = argparse.ArgumentParser(description="启动服务管理器")
    parser.add_argument("--mode", type=str, default="train", choices=["train", "eval"],
                        help="运行模式，决定启动哪些服务")
    parser.add_argument("--services", nargs="+",
                        default=None,
                        help="要启动的服务列表")
    args = parser.parse_args()
    if not args.services:
        if args.mode == "train":
            args.services = ["search_agent", "retrieval_service", "training_service"]
        elif args.mode == "eval":
            args.services = ["search_agent", "retrieval_service", "eval_service"]
        else:
            print(f"❌ 未知模式: {args.mode}")
            sys.exit(1)
    else:
        args.services = args.services
    with ServiceManager() as sm:
        try:
            launched_services = []
            for svc in args.services:
                if svc not in SERVICE_DICT:
                    print(f"❌ 未知服务: {svc}")
                    continue
                # 执行预处理操作（如果有）
                pre_action = SERVICE_DICT[svc].get("pre_action")
                if pre_action:
                    try:
                        pre_action()
                    except Exception as e:
                        print(f"❌ 预处理操作失败: {e}")
                        continue
                service_info = SERVICE_DICT[svc]
                full_cmd = service_info["cmd"] + service_info["args"]
                cwd = service_info.get("cwd")
                log_file = service_info["log_file"]  # ✅ 完整的日志文件路径

                sm.start_service(
                    args,
                    name=svc,
                    cmd=full_cmd,
                    cwd=cwd,
                    log_file=log_file
                )
                launched_services.append(svc)

            if not launched_services:
                print("❌ 没有成功启动任何服务")
                sys.exit(1)

            print(f"✅ 已启动服务: {', '.join(launched_services)}")
            print("💡 按 Ctrl+C 可停止所有服务")

            # 保持主进程运行
            while True:
                time.sleep(1)

        except KeyboardInterrupt:
            print("\n👋 收到 Ctrl+C，正在退出...")
        except SystemExit:
            raise
        except Exception as e:
            print(f"❌ 服务启动失败: {e}")
            sys.exit(1)


if __name__ == "__main__":
    main()