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()
|