feat: initialize XAS test console with native execution
This commit is contained in:
AdamWen95
@@ -0,0 +1,270 @@
|
||||
"""Private dispatcher with a local Unix socket and one active test at a time."""
|
||||
import fcntl
|
||||
import hashlib
|
||||
import json
|
||||
import os
|
||||
from pathlib import Path
|
||||
import re
|
||||
import signal
|
||||
import shutil
|
||||
import subprocess
|
||||
import threading
|
||||
import time
|
||||
from .artifacts import Collector
|
||||
from .catalog import clean_source, safe_path, script_path, source_for, validate_config, SECRET, SERVICE, SMOKE_HASH
|
||||
from .settings import Settings
|
||||
from .store import Store, now
|
||||
from . import workspace
|
||||
|
||||
LABEL = "xas.console=dedicated"
|
||||
|
||||
|
||||
def docker(*args, timeout=20):
|
||||
result = subprocess.run(["docker", *args], capture_output=True, text=True, timeout=timeout)
|
||||
if result.returncode:
|
||||
# Arguments and secrets are never included in public error details.
|
||||
raise RuntimeError("Docker operation failed: " + args[0])
|
||||
return result.stdout.strip()
|
||||
|
||||
|
||||
class Worker:
|
||||
def __init__(self, settings, config):
|
||||
self.settings, self.config = settings, config
|
||||
self.store = Store(settings)
|
||||
self.stopping = False
|
||||
self.container = None
|
||||
|
||||
def stop(self, *args):
|
||||
self.stopping = True
|
||||
|
||||
def inspect_owned(self, container):
|
||||
info = json.loads(docker("inspect", container))[0]
|
||||
if info["Config"]["Labels"].get("xas.console") != "dedicated":
|
||||
raise RuntimeError("Refusing to manage a foreign container")
|
||||
return info
|
||||
|
||||
def remove(self, container):
|
||||
self.inspect_owned(container)
|
||||
docker("rm", "-f", container)
|
||||
|
||||
def recover(self):
|
||||
for container in docker("ps", "-aq", "--filter", "label=" + LABEL).splitlines():
|
||||
if container:
|
||||
self.remove(container)
|
||||
for mount in workspace.mounts(self.settings):
|
||||
workspace.remove(self.settings, mount)
|
||||
for run in self.store.runs(limit=10000):
|
||||
if run["state"] == "running":
|
||||
if run.get("execution", {}).get("mode") == "native":
|
||||
from .native_executor import NativeExecutor
|
||||
NativeExecutor(self).monitor(run)
|
||||
else:
|
||||
self.store.finish(run["id"], "interrupted", message="执行器重启，测试容器已停止；未自动补跑，远端请求结果可能未知")
|
||||
|
||||
def command(self, run, directory):
|
||||
if not re.fullmatch(r"[a-f0-9]{32}", run["id"]):
|
||||
raise ValueError("Invalid run identity")
|
||||
source = script_path(self.settings, run)
|
||||
raw = source.read_bytes()
|
||||
if hashlib.sha256(raw).hexdigest() != run["sha256"]:
|
||||
raise ValueError("脚本在排队后发生变化，本次未执行")
|
||||
if SECRET.search(clean_source(raw)):
|
||||
raise ValueError("脚本包含疑似凭据，未执行")
|
||||
if run["kind"] == "external" and SERVICE.search(clean_source(raw)):
|
||||
raise ValueError("涉及服务或设备管理的脚本尚未接入")
|
||||
config = validate_config(self.settings, run, run["config"])
|
||||
if run["language"] not in {"shell", "python"}:
|
||||
raise ValueError("Unsupported language")
|
||||
snapshot = directory / "source"
|
||||
snapshot.mkdir(mode=0o750)
|
||||
(snapshot / "script.txt").write_bytes(raw)
|
||||
(snapshot / "script.txt").chmod(0o444)
|
||||
work = directory / "work"
|
||||
work.mkdir(mode=0o770)
|
||||
uid = 0 if run["kind"] == "external" else 1000
|
||||
os.chown(work, uid, uid)
|
||||
workspace.create(self.settings, run["id"], uid)
|
||||
project = Path(self.config["project"])
|
||||
cmd = ["create", "--name", "adamwen-test-run-" + run["id"], "--label", LABEL,
|
||||
"--label", "xas.run=" + run["id"], "--runtime", "runc", "--network", "bridge",
|
||||
"--user", f"{uid}:{uid}", "--cap-drop", "ALL", "--security-opt", "no-new-privileges",
|
||||
"--pids-limit", "64", "--memory", "512m", "--cpus", "1", "--read-only",
|
||||
"--tmpfs", "/tmp:rw,nosuid,nodev,noexec,size=67108864", "--log-driver", "local",
|
||||
"--log-opt", "max-size=10m", "--log-opt", "max-file=2", "--interactive",
|
||||
"--env", "PYTHONUNBUFFERED=1", "--env", "PYTHONDONTWRITEBYTECODE=1",
|
||||
"--env", "XAS_RUN_DIR=/work", "--mount", f"type=bind,src={work},dst=/work",
|
||||
"--mount", f"type=bind,src={project / 'console'},dst=/opt/xas-console,readonly",
|
||||
"--mount", f"type=bind,src={snapshot / 'script.txt'},dst=/input/script{source.suffix},readonly"]
|
||||
cwd, execution_path = "/work", "/input/script" + source.suffix
|
||||
if run["kind"] == "external":
|
||||
root = source_for(self.settings, run["source_group"])
|
||||
cmd += ["--mount", f"type=bind,src={root},dst={root},readonly"]
|
||||
cwd = config["cwd"]
|
||||
# Preserve __file__ and sibling imports. Hash is rechecked immediately before launch.
|
||||
execution_path = str(source)
|
||||
cmd += ["--mount", f"type=bind,src={snapshot / 'script.txt'},dst={source},readonly"]
|
||||
for output in config["output_paths"]:
|
||||
target = Path(output)
|
||||
cmd += ["--mount", f"type=bind,src={target},dst={target}"]
|
||||
cmd += ["--workdir", cwd, "--entrypoint", "python3", self.config["runner_image"], "/opt/xas-console/job_main.py"]
|
||||
payload = {"language": run["language"], "script": execution_path, "args": config["args"], "profile": config["profile"]}
|
||||
if config["profile"] == "legacy_smoke":
|
||||
if run["sha256"] != SMOKE_HASH or config["args"] or run["id"] in self.config.get("forbidden_runs", []):
|
||||
raise ValueError("轻量模型测试必须使用已核对的原始 smoke.sh，不接受额外参数")
|
||||
credential = Path(self.config["smoke_credential"])
|
||||
pairs = [line.split("=", 1)[1].strip().strip("\"'") for line in credential.read_text().splitlines()
|
||||
if line.startswith("VLLM_API_KEY=")]
|
||||
if len(pairs) != 1 or not pairs[0]:
|
||||
raise ValueError("模型认证配置不可用")
|
||||
payload.update(key=pairs[0], upstream=self.config["smoke_upstream"])
|
||||
return cmd, payload
|
||||
|
||||
def execute(self, run):
|
||||
directory = safe_path(self.settings.runs / run["id"], self.settings.runs, must_exist=False)
|
||||
directory.mkdir(mode=0o750, exist_ok=False)
|
||||
log = (directory / "console.log").open("wb", buffering=0)
|
||||
log_lock = threading.Lock()
|
||||
overflow = threading.Event()
|
||||
process = None
|
||||
collector = None
|
||||
state, exit_code, message = "failed", None, "执行器发生错误"
|
||||
|
||||
def emit(data):
|
||||
with log_lock:
|
||||
remaining = self.settings.max_log - log.tell()
|
||||
if remaining <= len(data):
|
||||
overflow.set()
|
||||
if remaining > 0:
|
||||
log.write(data[:remaining])
|
||||
|
||||
try:
|
||||
if shutil.disk_usage(self.settings.root).free < 10 * 1024**3:
|
||||
raise ValueError("服务器可用磁盘低于 10 GiB，未启动测试")
|
||||
cmd, payload = self.command(run, directory)
|
||||
collector = Collector(self.settings, directory, run["config"]["output_paths"])
|
||||
self.container = docker(*cmd)
|
||||
process = subprocess.Popen(["docker", "start", "-ai", self.container], stdin=subprocess.PIPE,
|
||||
stdout=subprocess.PIPE, stderr=subprocess.STDOUT)
|
||||
process.stdin.write((json.dumps(payload) + "\n").encode())
|
||||
process.stdin.close()
|
||||
payload.clear()
|
||||
|
||||
def drain():
|
||||
while True:
|
||||
chunk = os.read(process.stdout.fileno(), 4096)
|
||||
if not chunk:
|
||||
break
|
||||
emit(chunk)
|
||||
reader = threading.Thread(target=drain, daemon=True)
|
||||
reader.start()
|
||||
started = time.monotonic()
|
||||
while process.poll() is None:
|
||||
self.store.heartbeat()
|
||||
collector.follow(emit)
|
||||
cancel = self.store.run(run["id"])["cancel_requested"]
|
||||
timed_out = time.monotonic() - started > run["config"]["timeout_seconds"]
|
||||
disk_full = shutil.disk_usage(self.settings.root).free < 10 * 1024**3 or collector.external_growth() > 128 * 1024**2
|
||||
if cancel or timed_out or overflow.is_set() or self.stopping or disk_full:
|
||||
self.inspect_owned(self.container)
|
||||
docker("kill", self.container)
|
||||
state = "cancelled" if cancel else ("timed_out" if timed_out else ("interrupted" if self.stopping else "failed"))
|
||||
message = "本次容器已停止；若已发送模型请求，远端推理是否结束需另行核对"
|
||||
if overflow.is_set():
|
||||
message = "日志达到 16 MiB 上限，本次测试已停止"
|
||||
if disk_full:
|
||||
message = "可用空间低于 10 GiB 或原目录输出新增超过 128 MiB，本次测试已停止；原文件保留"
|
||||
break
|
||||
time.sleep(0.5)
|
||||
process.wait(timeout=15)
|
||||
reader.join(timeout=5)
|
||||
if reader.is_alive():
|
||||
raise RuntimeError("Log reader did not stop")
|
||||
info = self.inspect_owned(self.container)
|
||||
exit_code = info["State"]["ExitCode"]
|
||||
if state == "failed" and message == "执行器发生错误":
|
||||
state = "succeeded" if exit_code == 0 else "failed"
|
||||
message = "脚本正常退出；模型质量请结合报告判断" if exit_code == 0 else f"脚本退出码 {exit_code}，请查看日志"
|
||||
if info["State"].get("OOMKilled"):
|
||||
message = "测试超过 512 MiB 内存限制"
|
||||
if info["State"].get("Error") and info["State"].get("StartedAt", "").startswith("0001-"):
|
||||
exit_code, state, message = None, "failed", "容器未成功启动，脚本未执行，请查看日志"
|
||||
except Exception as exc:
|
||||
message = str(exc) if isinstance(exc, ValueError) else f"执行器错误：{type(exc).__name__}；未自动重试"
|
||||
emit(("\n[控制台] " + message + "\n").encode())
|
||||
finally:
|
||||
if self.container:
|
||||
try:
|
||||
self.remove(self.container)
|
||||
self.container = None
|
||||
except Exception:
|
||||
state, message = "interrupted", "容器清理未确认，执行器停止接单，请管理员检查"
|
||||
self.stopping = True
|
||||
if process and process.poll() is None:
|
||||
process.terminate()
|
||||
process.wait(timeout=10)
|
||||
if collector:
|
||||
collector.finish_follow(emit)
|
||||
artifacts = collector.collect() if collector else []
|
||||
work = directory / "work"
|
||||
if work in workspace.mounts(self.settings):
|
||||
try:
|
||||
workspace.remove(self.settings, work)
|
||||
except Exception:
|
||||
state, message = "interrupted", "临时目录卸载未确认，执行器停止接单，请管理员检查"
|
||||
self.stopping = True
|
||||
if overflow.is_set() and state == "succeeded":
|
||||
state, message = "failed", "日志超过 16 MiB 上限，日志记录不完整，请查看已归档结果"
|
||||
notices = list(dict.fromkeys(collector.notices))[:32] if collector else []
|
||||
summary = {"run_id": run["id"], "script_name": run["script_name"], "state": state,
|
||||
"exit_code": exit_code, "sha256": run["sha256"], "version": run["version"],
|
||||
"synthetic": run["synthetic"], "message": message, "finished_at": now(),
|
||||
"archive_notices": notices, "log_truncated": overflow.is_set()}
|
||||
data = json.dumps(summary, ensure_ascii=False, indent=2).encode("utf-8")
|
||||
summary_path = directory / "summary.json"
|
||||
summary_path.write_bytes(data)
|
||||
artifact_id = hashlib.sha256(str(summary_path).encode()).hexdigest()[:24]
|
||||
(directory / "artifacts").mkdir(exist_ok=True)
|
||||
(directory / "artifacts" / artifact_id).write_bytes(data)
|
||||
artifacts.append({"id": artifact_id, "name": "summary.json", "size": len(data), "kind": "json"})
|
||||
for notice in notices:
|
||||
emit(("\n[归档] " + notice + "\n").encode())
|
||||
emit(("\n[控制台] " + message + "\n").encode())
|
||||
log.close()
|
||||
if process:
|
||||
process.stdout.close()
|
||||
self.store.finish(run["id"], state, exit_code=exit_code, message=message, artifacts=artifacts)
|
||||
|
||||
def loop(self):
|
||||
self.recover()
|
||||
while not self.stopping:
|
||||
self.store.heartbeat()
|
||||
run = self.store.claim()
|
||||
if run:
|
||||
if run.get("execution", {}).get("mode") == "native":
|
||||
from .native_executor import NativeExecutor
|
||||
NativeExecutor(self).execute(run)
|
||||
else:
|
||||
self.execute(run)
|
||||
else:
|
||||
time.sleep(1)
|
||||
|
||||
|
||||
def main():
|
||||
config_path = Path(os.environ["XAS_WORKER_CONFIG"])
|
||||
config = json.loads(config_path.read_text())
|
||||
settings = Settings.from_env()
|
||||
settings.prepare()
|
||||
lock = (settings.db.parent / "worker.lock").open("a")
|
||||
fcntl.flock(lock, fcntl.LOCK_EX | fcntl.LOCK_NB)
|
||||
worker = Worker(settings, config)
|
||||
from . import native_rpc
|
||||
rpc = native_rpc.start(settings, worker.store)
|
||||
signal.signal(signal.SIGTERM, worker.stop)
|
||||
signal.signal(signal.SIGINT, worker.stop)
|
||||
try:worker.loop()
|
||||
finally:
|
||||
rpc.shutdown();rpc.server_close()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
Reference in New Issue
Block a user