177 lines
8.3 KiB
Python
177 lines
8.3 KiB
Python
|
"""Original-environment launcher, supervised by a per-run systemd service."""
|
||
|
|
import errno
|
||
|
|
import fcntl
|
||
|
|
import hashlib
|
||
|
|
import json
|
||
|
|
import os
|
||
|
|
from pathlib import Path
|
||
|
|
import pty
|
||
|
|
import pwd
|
||
|
|
import selectors
|
||
|
|
import signal
|
||
|
|
import stat
|
||
|
|
import struct
|
||
|
|
import subprocess
|
||
|
|
import sys
|
||
|
|
import termios
|
||
|
|
import time
|
||
|
|
from .artifacts import Collector
|
||
|
|
from .native_files import differences, outputs, preserve, runtime_snapshot
|
||
|
|
from .native_io import SecretFilter, load_environment
|
||
|
|
from .native_plans import build_plan, profile_for
|
||
|
|
from .settings import Settings
|
||
|
|
from .store import Store, now
|
||
|
|
|
||
|
|
|
||
|
|
def acquire_terminal():
|
||
|
|
# Popen has already created the session and duplicated the slave onto stdin.
|
||
|
|
# The original tmux/menu tools need a controlling terminal, not just isatty().
|
||
|
|
fcntl.ioctl(0, termios.TIOCSCTTY, 0)
|
||
|
|
|
||
|
|
|
||
|
|
def execute(settings, run):
|
||
|
|
script = {**run, "id": run["script_id"]}
|
||
|
|
plan = build_plan(settings, script, run_id=run["id"])
|
||
|
|
if plan["plan_token"] != run["execution"]["plan_token"] or not plan["ready"]:
|
||
|
|
raise ValueError("执行前环境或参数发生变化，未启动原测试")
|
||
|
|
if plan["impact"]["requires_confirmation"] and not run["execution"].get("impact_acknowledged"):
|
||
|
|
raise ValueError("缺少本次影响确认")
|
||
|
|
directory = settings.runs / run["id"]
|
||
|
|
(directory / "work").mkdir(mode=0o700, exist_ok=True)
|
||
|
|
env, secrets = load_environment(plan)
|
||
|
|
redactor = SecretFilter(secrets)
|
||
|
|
file_redactor = SecretFilter(secrets)
|
||
|
|
limit, written, truncated = settings.max_log, 0, False
|
||
|
|
def emit(data):
|
||
|
|
nonlocal written, truncated
|
||
|
|
remaining = limit - written
|
||
|
|
if len(data) > remaining:truncated = True
|
||
|
|
if remaining > 0:
|
||
|
|
sys.stdout.buffer.write(data[:remaining]);sys.stdout.buffer.flush()
|
||
|
|
written += min(len(data), remaining)
|
||
|
|
def message(text):emit(("[原环境] " + text + "\n").encode())
|
||
|
|
collector = Collector(settings, directory, lambda: outputs(plan["output_paths"]), redactor=file_redactor.clean,
|
||
|
|
stream_filter_factory=lambda: SecretFilter(secrets), file_limit=4096)
|
||
|
|
profile = profile_for(settings, script)
|
||
|
|
if profile.get("preserve_existing_outputs"):
|
||
|
|
message("正在备份原有同名结果，原测试尚未开始。")
|
||
|
|
count = preserve(collector.outputs, directory)
|
||
|
|
message(f"已保护 {count} 个原有结果文件。")
|
||
|
|
before = runtime_snapshot()
|
||
|
|
(directory / "runtime-before.json").write_text(json.dumps(before, ensure_ascii=False), encoding="utf-8")
|
||
|
|
stopped = False
|
||
|
|
def stop(signum, frame):
|
||
|
|
nonlocal stopped
|
||
|
|
stopped = True # systemd sends the same signal to the full cgroup and enforces the grace period.
|
||
|
|
signal.signal(signal.SIGTERM, stop)
|
||
|
|
signal.signal(signal.SIGINT, stop)
|
||
|
|
fifo = settings.native_socket.parent / ("input-" + run["id"])
|
||
|
|
master = slave = input_fd = None
|
||
|
|
selector = selectors.DefaultSelector()
|
||
|
|
process = None
|
||
|
|
code = 1
|
||
|
|
try:
|
||
|
|
account = pwd.getpwnam(plan["user"])
|
||
|
|
options = {"cwd": plan["cwd"], "env": env, "start_new_session": True,
|
||
|
|
"user": account.pw_uid, "group": account.pw_gid, "extra_groups": []}
|
||
|
|
if stopped or Store(settings).run(run["id"])["cancel_requested"]:
|
||
|
|
message("已在创建原入口前取消，原测试未启动。")
|
||
|
|
return 1
|
||
|
|
if plan["interactive"]:
|
||
|
|
master, slave = pty.openpty()
|
||
|
|
attributes = termios.tcgetattr(slave)
|
||
|
|
attributes[3] &= ~(termios.ECHO | termios.ECHONL)
|
||
|
|
termios.tcsetattr(slave, termios.TCSANOW, attributes)
|
||
|
|
fcntl.ioctl(slave, termios.TIOCSWINSZ, struct.pack("HHHH", 36, 120, 0, 0))
|
||
|
|
process = subprocess.Popen(plan["argv"], stdin=slave, stdout=slave, stderr=slave,
|
||
|
|
preexec_fn=acquire_terminal, **options)
|
||
|
|
os.close(slave);slave = None
|
||
|
|
os.mkfifo(fifo, 0o600)
|
||
|
|
input_fd = os.open(fifo, os.O_RDWR | os.O_NONBLOCK)
|
||
|
|
selector.register(input_fd, selectors.EVENT_READ, "input")
|
||
|
|
output_fd = master
|
||
|
|
else:
|
||
|
|
process = subprocess.Popen(plan["argv"], stdin=subprocess.DEVNULL, stdout=subprocess.PIPE, stderr=subprocess.STDOUT, **options)
|
||
|
|
output_fd = process.stdout.fileno()
|
||
|
|
os.set_blocking(output_fd, False)
|
||
|
|
selector.register(output_fd, selectors.EVENT_READ, "output")
|
||
|
|
message("已在原目录启动登记入口。取消仅处理本次宿主进程，Docker/tmux 或远端推理需结合收尾记录核对。")
|
||
|
|
last_follow = time.monotonic()
|
||
|
|
while selector.get_map():
|
||
|
|
for key, _ in selector.select(.3):
|
||
|
|
try:data = os.read(key.fd, 4096)
|
||
|
|
except OSError as exc:
|
||
|
|
if exc.errno in (errno.EIO, errno.EBADF):data = b""
|
||
|
|
elif exc.errno == errno.EAGAIN:continue
|
||
|
|
else:raise
|
||
|
|
if key.data == "input":
|
||
|
|
if data and process.poll() is None:
|
||
|
|
os.write(master, data)
|
||
|
|
elif data:emit(redactor.feed(data))
|
||
|
|
else:selector.unregister(key.fd)
|
||
|
|
if time.monotonic() - last_follow > 1:
|
||
|
|
collector.follow(lambda data: emit(file_redactor.clean(data)))
|
||
|
|
last_follow = time.monotonic()
|
||
|
|
if process.poll() is not None:
|
||
|
|
if input_fd is not None:
|
||
|
|
try:selector.unregister(input_fd)
|
||
|
|
except KeyError:pass
|
||
|
|
if not any(k.data == "output" for k in selector.get_map().values()):break
|
||
|
|
code = process.wait()
|
||
|
|
emit(redactor.finish())
|
||
|
|
finally:
|
||
|
|
selector.close()
|
||
|
|
if process and process.poll() is None:
|
||
|
|
os.killpg(process.pid, signal.SIGTERM)
|
||
|
|
try:process.wait(timeout=5)
|
||
|
|
except subprocess.TimeoutExpired:os.killpg(process.pid, signal.SIGKILL);process.wait(timeout=5)
|
||
|
|
for descriptor in (input_fd, master, slave):
|
||
|
|
if descriptor is not None:
|
||
|
|
try:os.close(descriptor)
|
||
|
|
except OSError:pass
|
||
|
|
if process and process.stdout:process.stdout.close()
|
||
|
|
if fifo.exists() and stat.S_ISFIFO(fifo.lstat().st_mode):fifo.unlink()
|
||
|
|
after = runtime_snapshot()
|
||
|
|
changes = differences(before, after)
|
||
|
|
if plan["interactive"] and code == 0 and changes["new_tmux_sessions"] and not stopped:
|
||
|
|
sessions = {item["id"] for item in changes["new_tmux_sessions"]}
|
||
|
|
message("原菜单已退出，观察到新增 tmux 会话；继续跟踪其结束，不把入口退出当作后台测试完成。")
|
||
|
|
while sessions and not stopped:
|
||
|
|
collector.follow(lambda data: emit(file_redactor.clean(data)))
|
||
|
|
time.sleep(2)
|
||
|
|
current = runtime_snapshot()
|
||
|
|
sessions &= {item["id"] for item in current["tmux_sessions"]}
|
||
|
|
after = runtime_snapshot()
|
||
|
|
changes = differences(before, after)
|
||
|
|
(directory / "work/runtime-changes.json").write_text(json.dumps(changes, ensure_ascii=False, indent=2), encoding="utf-8")
|
||
|
|
collector.finish_follow(lambda data: emit(file_redactor.clean(data)))
|
||
|
|
artifacts = collector.collect()
|
||
|
|
result = {"exit_code": code, "stopped": stopped, "finished_at": now(), "artifacts": artifacts,
|
||
|
|
"log_truncated": truncated, "archive_notices": list(dict.fromkeys(collector.notices))[:32],
|
||
|
|
"runtime_changes": changes, "command_preview": plan["command_preview"]}
|
||
|
|
tmp = directory / "native-result.tmp"
|
||
|
|
tmp.write_text(json.dumps(result, ensure_ascii=False), encoding="utf-8")
|
||
|
|
os.replace(tmp, directory / "native-result.json")
|
||
|
|
message(f"原入口进程退出，退出码 {code}。共享资源前后变化已记录；不代表原模型已恢复。")
|
||
|
|
return code
|
||
|
|
|
||
|
|
|
||
|
|
def main():
|
||
|
|
identity = sys.argv[1]
|
||
|
|
import re
|
||
|
|
if not re.fullmatch(r"[a-f0-9]{32}", identity):return 2
|
||
|
|
settings = Settings.from_env()
|
||
|
|
run = Store(settings).run(identity)
|
||
|
|
if run["state"] != "running" or run.get("execution", {}).get("mode") != "native":return 2
|
||
|
|
directory = settings.runs / identity
|
||
|
|
try:os.rename(directory / "native-permit", directory / "native-started")
|
||
|
|
except OSError:return 2
|
||
|
|
try:return execute(settings, run)
|
||
|
|
except Exception as exc:
|
||
|
|
print(f"[原环境] 启动或收尾失败：{type(exc).__name__}。未自动重试。", flush=True)
|
||
|
|
return 1
|
||
|
|
|
||
|
|
|
||
|
|
if __name__ == "__main__":
|
||
|
|
sys.exit(main())
|