""" MinerU Tianshu - Unified Startup Script 天枢统一启动脚本 一键启动所有服务：API Server + LitServe Workers + Task Scheduler """ import subprocess import signal import sys import time import os from loguru import logger from pathlib import Path import argparse class TianshuLauncher: """天枢服务启动器""" def __init__( self, output_dir='/tmp/mineru_tianshu_output', api_port=8000, worker_port=9000, workers_per_device=1, devices='auto', accelerator='auto' ): self.output_dir = output_dir self.api_port = api_port self.worker_port = worker_port self.workers_per_device = workers_per_device self.devices = devices self.accelerator = accelerator self.processes = [] def start_services(self): """启动所有服务""" logger.info("=" * 70) logger.info("🚀 MinerU Tianshu - Starting All Services") logger.info("=" * 70) logger.info("天枢 - 企业级多GPU文档解析服务") logger.info("") try: # 1. 启动 API Server logger.info("📡 [1/3] Starting API Server...") env = os.environ.copy() env['API_PORT'] = str(self.api_port) api_proc = subprocess.Popen( [sys.executable, 'api_server.py'], cwd=Path(__file__).parent, env=env ) self.processes.append(('API Server', api_proc)) time.sleep(3) if api_proc.poll() is not None: logger.error("❌ API Server failed to start!") return False logger.info(f" ✅ API Server started (PID: {api_proc.pid})") logger.info(f" 📖 API Docs: http://localhost:{self.api_port}/docs") logger.info("") # 2. 启动 LitServe Worker Pool logger.info("⚙️ [2/3] Starting LitServe Worker Pool...") worker_cmd = [ sys.executable, 'litserve_worker.py', '--output-dir', self.output_dir, '--accelerator', self.accelerator, '--workers-per-device', str(self.workers_per_device), '--port', str(self.worker_port), '--devices', str(self.devices) if isinstance(self.devices, str) else ','.join(map(str, self.devices)) ] worker_proc = subprocess.Popen( worker_cmd, cwd=Path(__file__).parent ) self.processes.append(('LitServe Workers', worker_proc)) time.sleep(5) if worker_proc.poll() is not None: logger.error("❌ LitServe Workers failed to start!") return False logger.info(f" ✅ LitServe Workers started (PID: {worker_proc.pid})") logger.info(f" 🔌 Worker Port: {self.worker_port}") logger.info(f" 👷 Workers per Device: {self.workers_per_device}") logger.info("") # 3. 启动 Task Scheduler logger.info("🔄 [3/3] Starting Task Scheduler...") scheduler_cmd = [ sys.executable, 'task_scheduler.py', '--litserve-url', f'http://localhost:{self.worker_port}/predict', '--wait-for-workers' ] scheduler_proc = subprocess.Popen( scheduler_cmd, cwd=Path(__file__).parent ) self.processes.append(('Task Scheduler', scheduler_proc)) time.sleep(3) if scheduler_proc.poll() is not None: logger.error("❌ Task Scheduler failed to start!") return False logger.info(f" ✅ Task Scheduler started (PID: {scheduler_proc.pid})") logger.info("") # 启动成功 logger.info("=" * 70) logger.info("✅ All Services Started Successfully!") logger.info("=" * 70) logger.info("") logger.info("📚 Quick Start:") logger.info(f" • API Documentation: http://localhost:{self.api_port}/docs") logger.info(f" • Submit Task: POST http://localhost:{self.api_port}/api/v1/tasks/submit") logger.info(f" • Query Status: GET http://localhost:{self.api_port}/api/v1/tasks/{{task_id}}") logger.info(f" • Queue Stats: GET http://localhost:{self.api_port}/api/v1/queue/stats") logger.info("") logger.info("🔧 Service Details:") for name, proc in self.processes: logger.info(f" • {name:20s} PID: {proc.pid}") logger.info("") logger.info("⚠️ Press Ctrl+C to stop all services") logger.info("=" * 70) return True except Exception as e: logger.error(f"❌ Failed to start services: {e}") self.stop_services() return False def stop_services(self, signum=None, frame=None): """停止所有服务""" logger.info("") logger.info("=" * 70) logger.info("⏹️ Stopping All Services...") logger.info("=" * 70) for name, proc in self.processes: if proc.poll() is None: # 进程仍在运行 logger.info(f" Stopping {name} (PID: {proc.pid})...") proc.terminate() # 等待所有进程结束 for name, proc in self.processes: try: proc.wait(timeout=10) logger.info(f" ✅ {name} stopped") except subprocess.TimeoutExpired: logger.warning(f" ⚠️ {name} did not stop gracefully, forcing...") proc.kill() proc.wait() logger.info("=" * 70) logger.info("✅ All Services Stopped") logger.info("=" * 70) sys.exit(0) def wait(self): """等待所有服务""" try: while True: time.sleep(1) # 检查进程状态 for name, proc in self.processes: if proc.poll() is not None: logger.error(f"❌ {name} unexpectedly stopped!") self.stop_services() return except KeyboardInterrupt: self.stop_services() def main(): """主函数""" parser = argparse.ArgumentParser( description='MinerU Tianshu - 统一启动脚本', formatter_class=argparse.RawDescriptionHelpFormatter, epilog=""" 示例: # 使用默认配置启动（自动检测GPU） python start_all.py # 使用CPU模式 python start_all.py --accelerator cpu # 指定输出目录和端口 python start_all.py --output-dir /data/output --api-port 8080 # 每个GPU启动2个worker python start_all.py --accelerator cuda --workers-per-device 2 # 只使用指定的GPU python start_all.py --accelerator cuda --devices 0,1 """ ) parser.add_argument('--output-dir', type=str, default='/tmp/mineru_tianshu_output', help='输出目录 (默认: /tmp/mineru_tianshu_output)') parser.add_argument('--api-port', type=int, default=8000, help='API服务器端口 (默认: 8000)') parser.add_argument('--worker-port', type=int, default=9000, help='Worker服务器端口 (默认: 9000)') parser.add_argument('--accelerator', type=str, default='auto', choices=['auto', 'cuda', 'cpu', 'mps'], help='加速器类型 (默认: auto，自动检测)') parser.add_argument('--workers-per-device', type=int, default=1, help='每个GPU的worker数量 (默认: 1)') parser.add_argument('--devices', type=str, default='auto', help='使用的GPU设备，逗号分隔 (默认: auto，使用所有GPU)') args = parser.parse_args() # 处理 devices 参数 devices = args.devices if devices != 'auto': try: devices = [int(d) for d in devices.split(',')] except: logger.warning(f"Invalid devices format: {devices}, using 'auto'") devices = 'auto' # 创建启动器 launcher = TianshuLauncher( output_dir=args.output_dir, api_port=args.api_port, worker_port=args.worker_port, workers_per_device=args.workers_per_device, devices=devices, accelerator=args.accelerator ) # 设置信号处理 signal.signal(signal.SIGINT, launcher.stop_services) signal.signal(signal.SIGTERM, launcher.stop_services) # 启动服务 if launcher.start_services(): launcher.wait() else: sys.exit(1) if __name__ == '__main__': main()