start_all.py
8.98 KB
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
"""
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()