257 lines
9.0 KiB
Python
257 lines
9.0 KiB
Python
"""
|
||
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()
|
||
|