287 lines
12 KiB
Python
287 lines
12 KiB
Python
"""Tray supervisor with authenticated local control and owned process cleanup."""
|
||
import argparse
|
||
import json
|
||
import logging
|
||
import multiprocessing as mp
|
||
import os
|
||
import secrets
|
||
import subprocess
|
||
import sys
|
||
import threading
|
||
import time
|
||
import webbrowser
|
||
from monitor_runtime.paths import data, resource, atomic_json
|
||
|
||
def worker(address, auth):
|
||
from multiprocessing.connection import Client
|
||
connection = Client(address, family='AF_PIPE', authkey=bytes.fromhex(auth))
|
||
try:
|
||
if connection.recv() != 'run':
|
||
return 1
|
||
from monitor_runtime.service import run
|
||
return run(connection)
|
||
finally:
|
||
connection.close()
|
||
|
||
def main():
|
||
parser = argparse.ArgumentParser()
|
||
parser.add_argument('--worker', nargs=2, metavar=('PIPE', 'AUTH'))
|
||
parser.add_argument('--stop', action='store_true')
|
||
parser.add_argument('--self-test', action='store_true')
|
||
parser.add_argument('--trust-model')
|
||
parser.add_argument('--yes', action='store_true')
|
||
args = parser.parse_args()
|
||
if os.name != 'nt':
|
||
raise RuntimeError('Monitor desktop requires Windows')
|
||
data('.runtime').mkdir(parents=True, exist_ok=True)
|
||
data('log').mkdir(parents=True, exist_ok=True)
|
||
logging.basicConfig(filename=str(data('log/desktop.log')), level=logging.INFO,
|
||
format='%(asctime)s %(levelname)s %(message)s')
|
||
if args.self_test:
|
||
from monitor_runtime.diagnostics import run
|
||
return run()
|
||
if args.worker:
|
||
return worker(*args.worker)
|
||
if args.trust_model:
|
||
return trust_model(args.trust_model, args.yes)
|
||
from app.services.lifecycle import ServiceLeaderLock
|
||
# Separate from the service lock; held for the entire supervisor lifetime.
|
||
lock = ServiceLeaderLock(data('.runtime/desktop.lock'))
|
||
if not lock.acquire():
|
||
return control_existing('stop' if args.stop else 'open')
|
||
if args.stop:
|
||
lock.release()
|
||
return 0
|
||
try:
|
||
return Supervisor().run()
|
||
finally:
|
||
lock.release()
|
||
|
||
def control_existing(command):
|
||
from multiprocessing.connection import Client
|
||
deadline = time.monotonic() + 10
|
||
while True:
|
||
try:
|
||
state = json.loads(data('.runtime/desktop.json').read_text(encoding='utf-8'))
|
||
with Client(state['pipe'], family='AF_PIPE', authkey=bytes.fromhex(state['auth'])) as connection:
|
||
connection.send(command)
|
||
response = connection.recv()
|
||
break
|
||
except (OSError, EOFError, json.JSONDecodeError):
|
||
if time.monotonic() >= deadline:
|
||
raise RuntimeError('已有实例的控制接口不可用，请查看日志')
|
||
time.sleep(.2)
|
||
if command == 'stop':
|
||
deadline = time.monotonic() + 40
|
||
from app.services.lifecycle import ServiceLeaderLock
|
||
while time.monotonic() < deadline:
|
||
lock = ServiceLeaderLock(data('.runtime/desktop.lock'))
|
||
if lock.acquire():
|
||
lock.release()
|
||
return 0
|
||
time.sleep(.3)
|
||
raise RuntimeError('Monitor 未能在 40 秒内退出，请查看日志')
|
||
return 0
|
||
|
||
def trust_model(filename, confirmed):
|
||
from pathlib import Path
|
||
from app.utils.ModelTrust import sha256_file
|
||
from monitor_runtime.windows import message
|
||
model = Path(filename).resolve()
|
||
if model.suffix.lower() != '.pt' or not model.is_file():
|
||
raise ValueError('请选择存在的 .pt 模型文件')
|
||
digest = sha256_file(model)
|
||
if not confirmed:
|
||
message('模型 SHA-256: ' + digest + '\n确认来源可信后使用 --yes 重新执行。')
|
||
return 2
|
||
path = data('.trusted-models.json')
|
||
doc = json.loads(path.read_text(encoding='utf-8')) if path.exists() else {'trusted_sha256': []}
|
||
doc['trusted_sha256'] = sorted(set(doc['trusted_sha256']) | {digest})
|
||
atomic_json(path, doc)
|
||
message('已登记可信模型 SHA-256：' + digest)
|
||
return 0
|
||
|
||
class Supervisor:
|
||
def __init__(self):
|
||
self.stop = threading.Event()
|
||
self.restart = threading.Event()
|
||
self.url = ''
|
||
self.status = '正在启动'
|
||
self.icon = None
|
||
self.child = None
|
||
self.pipe = None
|
||
self.job = None
|
||
|
||
def open_web(self, *_):
|
||
if self.url:
|
||
webbrowser.open(self.url)
|
||
|
||
def exit(self, *_):
|
||
self.stop.set()
|
||
|
||
def run(self):
|
||
import pystray
|
||
from PIL import Image
|
||
from multiprocessing.connection import Listener
|
||
auth = secrets.token_bytes(32)
|
||
address = r'\\.\pipe\MonitorControl-' + secrets.token_hex(16)
|
||
self.control = Listener(address, family='AF_PIPE', authkey=auth)
|
||
atomic_json(data('.runtime/desktop.json'), {'pipe': address, 'auth': auth.hex()})
|
||
threading.Thread(target=self.control_loop, daemon=True).start()
|
||
menu = pystray.Menu(
|
||
pystray.MenuItem('打开管理网页', self.open_web, default=True),
|
||
pystray.MenuItem(lambda _: self.status, lambda *_: None, enabled=False),
|
||
pystray.MenuItem('查看状态', lambda *_: self.show_status()),
|
||
pystray.MenuItem('打开日志目录', lambda *_: os.startfile(str(data('log')))),
|
||
pystray.MenuItem('重启服务', lambda *_: self.restart.set()),
|
||
pystray.MenuItem('退出', self.exit))
|
||
self.icon = pystray.Icon('Monitor', Image.open(resource('static/images/logo.png')), 'Monitor', menu)
|
||
threading.Thread(target=self.supervise, daemon=True).start()
|
||
try:
|
||
self.icon.run()
|
||
finally:
|
||
self.stop.set()
|
||
self.control.close()
|
||
return 0
|
||
|
||
def show_status(self):
|
||
from monitor_runtime.windows import message
|
||
message(self.status + '\n' + self.url.split('#')[0])
|
||
|
||
def control_loop(self):
|
||
while not self.stop.is_set():
|
||
try:
|
||
with self.control.accept() as conn:
|
||
command = conn.recv()
|
||
if command == 'stop':
|
||
self.stop.set()
|
||
elif command == 'open':
|
||
self.open_web()
|
||
conn.send('ok')
|
||
except (OSError, EOFError):
|
||
return
|
||
|
||
def start_child(self):
|
||
from multiprocessing.connection import Listener
|
||
from monitor_runtime.windows import Job
|
||
auth = secrets.token_bytes(32)
|
||
address = r'\\.\pipe\MonitorWorker-' + secrets.token_hex(16)
|
||
listener = Listener(address, family='AF_PIPE', authkey=auth)
|
||
args = [sys.executable]
|
||
if not getattr(sys, 'frozen', False):
|
||
args.append(str(resource('monitor_entry.py')))
|
||
args += ['--worker', address, auth.hex()]
|
||
env = os.environ.copy()
|
||
env['MONITOR_DESKTOP'] = '1'
|
||
env['MONITOR_SERVICE_MODE'] = 'disabled'
|
||
env['MONITOR_BOOTSTRAP_SERVICES'] = 'false'
|
||
env['MONITOR_DEBUG'] = 'false'
|
||
env['MONITOR_DATA_DIR'] = str(data(''))
|
||
self.job = Job()
|
||
log = open(data('log/service.log'), 'ab', buffering=0)
|
||
try:
|
||
self.child = subprocess.Popen(args, cwd=str(data('')), env=env, stdin=subprocess.DEVNULL,
|
||
stdout=log, stderr=log, creationflags=subprocess.CREATE_NO_WINDOW)
|
||
self.job.assign(self.child)
|
||
finally:
|
||
log.close()
|
||
# Child waits for run; Job assignment therefore precedes any child spawn.
|
||
accepted = []
|
||
def accept():
|
||
try:
|
||
accepted.append(listener.accept())
|
||
except OSError:
|
||
pass
|
||
thread = threading.Thread(target=accept, daemon=True)
|
||
thread.start()
|
||
deadline = time.monotonic() + 30
|
||
while not accepted:
|
||
if self.child.poll() is not None or time.monotonic() > deadline or self.stop.wait(.1):
|
||
listener.close()
|
||
raise RuntimeError('服务进程未能建立控制连接，请查看 service.log')
|
||
self.pipe = accepted[0]
|
||
listener.close()
|
||
self.pipe.send('run')
|
||
|
||
def stop_child(self):
|
||
if self.child and self.child.poll() is None:
|
||
try:
|
||
self.pipe.send('stop')
|
||
except (AttributeError, EOFError, OSError):
|
||
pass
|
||
try:
|
||
self.child.wait(timeout=25)
|
||
except subprocess.TimeoutExpired:
|
||
logging.error('Graceful shutdown timed out; closing owned job')
|
||
if self.pipe:
|
||
self.pipe.close()
|
||
self.pipe = None
|
||
if self.job:
|
||
self.job.close()
|
||
self.job = None
|
||
self.child = None
|
||
|
||
def supervise(self):
|
||
try:
|
||
while not self.stop.is_set():
|
||
self.restart.clear()
|
||
self.status = '正在启动'
|
||
self.start_child()
|
||
deadline = time.monotonic() + 180
|
||
last_heartbeat = time.monotonic()
|
||
opened = False
|
||
auto_restart = False
|
||
try:
|
||
while not self.stop.is_set() and not self.restart.is_set():
|
||
if self.pipe.poll(.3):
|
||
message = self.pipe.recv()
|
||
last_heartbeat = time.monotonic()
|
||
phase = message.get('phase')
|
||
if phase == 'failed':
|
||
raise RuntimeError(message.get('reason', '服务启动失败'))
|
||
self.status = {'ready': '服务运行中', 'setup': '等待首次初始化',
|
||
'license': '等待有效授权'}.get(phase, phase)
|
||
if message.get('url'):
|
||
self.url = message['url']
|
||
if not opened:
|
||
self.open_web()
|
||
opened = True
|
||
self.icon.update_menu()
|
||
code = self.child.poll()
|
||
if code is not None:
|
||
if code == 20:
|
||
auto_restart = True
|
||
break
|
||
raise RuntimeError('服务进程退出，代码 %s' % code)
|
||
if not opened and time.monotonic() > deadline:
|
||
raise RuntimeError('启动超过 180 秒，请查看日志')
|
||
if opened and time.monotonic() - last_heartbeat > 35:
|
||
raise RuntimeError('服务健康状态超时')
|
||
except EOFError:
|
||
code = self.child.wait(timeout=30)
|
||
if code == 20:
|
||
auto_restart = True
|
||
else:
|
||
raise RuntimeError('服务控制连接断开，代码 %s' % code)
|
||
finally:
|
||
self.stop_child()
|
||
if not auto_restart and not self.restart.is_set():
|
||
break
|
||
except Exception as exc:
|
||
logging.exception('Supervisor failure')
|
||
self.status = '启动失败：' + str(exc)
|
||
self.url = ''
|
||
self.icon.update_menu()
|
||
from monitor_runtime.windows import message
|
||
message(self.status + '\n日志：' + str(data('log')), True)
|
||
self.stop_child()
|
||
# Keep tray available; retries are explicit, never an infinite crash loop.
|
||
while not self.stop.wait(.3):
|
||
if self.restart.is_set():
|
||
self.supervise()
|
||
return
|
||
finally:
|
||
self.stop_child()
|
||
if self.stop.is_set():
|
||
self.icon.stop()
|