"""课程工作台：声明式任务、单个计算任务、持久日志，不执行任意 shell。""" import json import os import re import shlex import signal import subprocess import sys import threading import time import uuid from pathlib import Path from .common import ROOT, write_json WORK = Path(os.environ.get("POETRY_WORKBENCH_DIR", ROOT / "artifacts/workbench")) FINAL = ROOT / "artifacts/runs/poet/best.pt" def integer(value, low, high, name): if isinstance(value, bool) or not re.fullmatch(r"-?\d+", str(value)): raise ValueError(f"{name}需要填写整数。") result = int(value) if not low <= result <= high: raise ValueError(f"{name}应在 {low}—{high} 之间。") return result def choice(value, options, name): if value not in options: raise ValueError(f"{name}请选择 {'、'.join(options)}。") return value def checkpoint(model): if model == "formal": return FINAL if model == "initial": return FINAL.parent/"initial.pt" if model == "stage6000": return FINAL.parent/"stage-6000.pt" if not re.fullmatch(r"wb-[a-f0-9]{12}", str(model)): raise ValueError("请选择工作台中的模型。") path = (WORK / "runs" / model / "best.pt").resolve() if not path.is_relative_to((WORK / "runs").resolve()) or not path.is_file(): raise ValueError("这个练习模型尚未保存，请先完成一次训练。") return path def build_command(action, params, job_id): """从已知功能和检查过的参数构造 argv；所有写入固定在练习目录。""" if not isinstance(params, dict): raise ValueError("参数应为对象。") params = dict(params) remote = bool(os.environ.get("POETRY_PUBLIC_ORIGIN")) if remote: if params.get("device", "auto") not in ("auto", "cpu"): raise ValueError("当前服务器使用 CPU；Mac 或 NVIDIA 加速请在对应电脑运行学习包。") if action == "train": if params.get("config", "tiny") != "tiny": raise ValueError("在线练习使用小模型；正式结构训练请下载学习包在本机运行。") params.setdefault("batch_size", 8) integer(params["batch_size"], 1, 8, "服务器每批数量") if action in ("train", "resume"): integer(params.get("steps", 50 if action == "train" else 75), 1, 500, "在线练习总步数") allowed = { "doctor": set(), "lesson": {"number", "kind"}, "materials": set(), "prepare": set(), "train": {"config", "steps", "batch_size", "device"}, "resume": {"run", "steps", "device"}, "write": {"start", "keywords", "form", "count", "temperature", "top_k", "seed", "model", "device", "json", "candidates"}, "evaluate": {"split", "batches", "model", "device"}, "inspect": {"start", "model", "device"}, "cli": {"stage", "start", "keywords", "form", "json", "device", "count"}, "help": {"command"}, "course": set(), "tests": set(), } if action not in allowed: raise ValueError("请选择课程中提供的操作。") if set(params) - allowed[action]: raise ValueError("出现了这个操作不支持的参数。") meta = {"action": action, "params": params} base = [sys.executable, "-u", "-m", "poetry_gpt"] device = choice(params.get("device", "auto"), ["auto", "cpu", "mps", "cuda"], "计算设备") if action == "doctor": argv = base + ["doctor"] elif action == "course": argv = base + ["course", "--status"] elif action == "help": cmd = choice(params.get("command", "all"), ["all", "write", "train", "evaluate", "course", "chat"], "功能") argv = base + ([] if cmd == "all" else [cmd]) + ["--help"] elif action == "lesson": number = integer(params.get("number", 1), 1, 26, "课号") kind = choice(params.get("kind", "experiment"), ["experiment", "solution", "starter", "exercise"], "实验类型") argv = base+["lesson",str(number),"--exercise"] if kind=="exercise" else [sys.executable, "-u", str(ROOT / f"lessons/{number:02}/{kind}.py")] elif action == "materials": argv = [sys.executable, "-u", str(ROOT / "lessons/03/inspect_material.py")] elif action == "prepare": target = WORK / "data" / job_id argv = base + ["prepare", "--output", str(target)] meta["data_directory"] = str(target) elif action in ("train", "resume"): default_steps = 50 if action == "train" else 75 steps = integer(params.get("steps", default_steps), 1, 18000, "总目标步数") if action == "train": name = "wb-" + uuid.uuid4().hex[:12] config = choice(params.get("config", "tiny"), ["tiny", "poet"], "模型规模") batch = integer(params.get("batch_size", 24), 1, 64, "每批数量") path = WORK / "runs" / name argv = base + ["train", "--config", str(ROOT/f"configs/{config}.json"), "--run-dir", str(path), "--steps", str(steps), "--batch-size", str(batch), "--device", device] else: name = params.get("run", "") if not re.fullmatch(r"wb-[a-f0-9]{12}", str(name)): raise ValueError("请先选择一次工作台训练。") path = (WORK / "runs" / name).resolve() if not path.is_relative_to((WORK / "runs").resolve()) or not (path/"latest.pt").is_file(): raise ValueError("没有可恢复的进度，先完成或停止一次短训练。") state = json.loads((path/"run.json").read_text()) if remote and (state.get("model", {}).get("width", 128) > 128 or state.get("training", {}).get("batch_size", 8) > 8): raise ValueError("这份进度超出在线练习规模，请下载后在本机恢复。") if steps <= state["step"]: raise ValueError(f"已完成 {state['step']} 步，总目标应当更大。") argv = base + ["train", "--run-dir", str(path), "--resume", str(path/"latest.pt"), "--steps", str(steps), "--device", device] meta.update(run=name, run_directory=str(path.relative_to(ROOT)) if path.is_relative_to(ROOT) else str(path), target_steps=steps) elif action in ("write", "evaluate", "inspect", "cli"): model = checkpoint(params.get("model", "formal")) if action == "inspect": start = str(params.get("start", "春江")) if not re.fullmatch(r"[\u3400-\u9fff]{1,8}", start): raise ValueError("观察前文请填写 1—8 个汉字。") argv = [sys.executable, "-u", "-m", "poetry_gpt.inspect_model", "--start", start, "--checkpoint", str(model), "--device", device] elif action == "evaluate": split = choice(params.get("split", "val"), ["val", "test"], "材料") batches = integer(params.get("batches", 10), 1, 60, "抽样批数") argv = base + ["evaluate", "--checkpoint", str(model), "--split", split, "--batches", str(batches), "--device", device] else: start, words = str(params.get("start", "")), str(params.get("keywords", "")) if len(start) > 12 or len(words) > 40: raise ValueError("请缩短开头或关键词。") form = choice(params.get("form", "five"), ["five", "seven"], "诗句外形") count = integer(params.get("count", 1), 0, 5, "生成数量") if action == "cli": stage = choice(str(params.get("stage", "1")), ["1", "2", "3"], "编写阶段") file = {"1":"01_arguments.py", "2":"02_write.py", "3":"03_output.py"}[stage] argv = [sys.executable, "-u", str(ROOT/"lessons/24/build_cli"/file), "--start", start, "--keywords", words, "--form", form] if stage != "1": argv += ["--device", device] if stage == "3": argv += ["--count", str(count), "--output", str(WORK/"jobs"/job_id/"poems.json")] if params.get("json", True): argv.append("--json") else: try: temp = float(params.get("temperature", 0.8)) except (ValueError,TypeError): raise ValueError("变化程度请填写数字。") if not 0.1 <= temp <= 2: raise ValueError("变化程度应在 0.1—2 之间。") argv = base + ["write", "--start", start, "--keywords", words, "--form", form, "--count", str(count), "--temperature", str(temp), "--top-k", str(integer(params.get("top_k",40),1,200,"候选字数")), "--candidates", str(integer(params.get("candidates",12),1,24,"候选诗数")), "--seed", str(integer(params.get("seed",2026),0,2**31-1,"随机种子")), "--checkpoint", str(model), "--device", device, "--output", str(WORK/"jobs"/job_id/"poems.json")] if params.get("json", False): argv.append("--json") else: argv = [sys.executable, "-u", "-m", "pytest", "-q", "tests/test_core.py"] # 仅用来展示命令；执行始终使用 argv 列表，绝不交给 shell。 shown = [str(x).replace(str(ROOT)+"/", "") for x in argv] shown[0] = ".venv/bin/python" meta["command"] = shlex.join(shown) return argv, meta class Jobs: def __init__(self): self.lock = threading.RLock() self.active = None self.process = None self.stopping = False self.items = {} (WORK/"jobs").mkdir(parents=True, exist_ok=True) for p in sorted((WORK/"jobs").glob("*/job.json"),key=lambda p:p.stat().st_mtime)[-100:]: try: row = json.loads(p.read_text()) if row["status"] in ("running", "stopping"): row.update(status="interrupted", note="服务重启前的任务，保留已写出的日志与模型进度。") write_json(p, row) self.items[row["id"]] = row except (ValueError, KeyError, OSError): continue def save(self, row): write_json(WORK/"jobs"/row["id"]/"job.json", row) def create(self, action, params): with self.lock: if self.active: raise RuntimeError("已有任务正在运行。等它完成，或先点“停止并保存”。") id = uuid.uuid4().hex argv, meta = build_command(action, params, id) row = {"id":id, **meta, "created":time.time(), "status":"running", "exit_code":None, "metrics":[]} folder = WORK/"jobs"/id;folder.mkdir(parents=True) (folder/"output.log").touch() self.items[id]=row; self.active=id; self.stopping=False; self.save(row) threading.Thread(target=self._run, args=(row,argv), daemon=True).start() return dict(row) def _run(self, row, argv): id=row["id"]; path=WORK/"jobs"/id/"output.log" env=os.environ.copy();env.update(PYTHONUNBUFFERED="1", PYTHONIOENCODING="utf-8", PYTHONPATH=str(ROOT)) try: process=subprocess.Popen(argv, cwd=ROOT, env=env, stdout=subprocess.PIPE, stderr=subprocess.STDOUT, stdin=subprocess.DEVNULL, text=True, encoding="utf-8", errors="replace", start_new_session=True) with self.lock: self.process=process if self.stopping: os.killpg(process.pid,signal.SIGINT) with path.open("a",encoding="utf-8") as stream: for line in process.stdout: stream.write(line);stream.flush() if len(line)>200000:continue try: data=json.loads(line) if isinstance(data,dict) and data.get("event") in ("baseline","train","validation"): with self.lock: row["metrics"].append(data) elif isinstance(data,dict) and data.get("event")=="model-inspection": with self.lock:row["inspection"]=data except ValueError:pass code=process.wait() with self.lock: row.update(exit_code=code,status="stopped" if self.stopping else "completed" if code==0 else "failed") except Exception as error: with path.open("a",encoding="utf-8") as stream:stream.write(f"无法运行：{error}\n") with self.lock:row.update(status="failed",exit_code=-1) finally: with self.lock: row["finished"]=time.time();self.save(row);self.active=None;self.process=None def get(self, id, offset=0): with self.lock: if id not in self.items: raise ValueError("找不到这个任务。") row=json.loads(json.dumps(self.items[id])); path=WORK/"jobs"/id/"output.log" offset=integer(offset,0,100_000_000,"日志位置") with path.open("rb") as stream: stream.seek(offset); text=stream.read(64000) # 一次响应末尾不能切断汉字的 UTF-8 编码。 while text: try:output=text.decode("utf-8");break except UnicodeDecodeError as error: if error.end==len(text):text=text[:error.start] else:output=text.decode("utf-8",errors="replace");break else:output="" row.update(output=output,offset=offset+len(text),more=path.stat().st_size>offset+len(text)) return row def stop(self, id): with self.lock: if self.active!=id: raise ValueError("这个任务已经结束。") self.stopping=True;self.items[id]["status"]="stopping";self.save(self.items[id]) process=self.process if process and process.poll() is None:os.killpg(process.pid,signal.SIGINT) if process: def finish(): try:process.wait(timeout=15) except subprocess.TimeoutExpired: if process.poll() is None:os.killpg(process.pid,signal.SIGTERM) threading.Thread(target=finish,daemon=True).start() def shutdown(self): if self.active: try:self.stop(self.active) except (ValueError,ProcessLookupError):pass if self.process: try:self.process.wait(timeout=16) except subprocess.TimeoutExpired:self.process.kill() def catalog(self): runs=[] for path in sorted((WORK/"runs").glob("wb-*/run.json"),key=lambda p:p.stat().st_mtime,reverse=True): try: row=json.loads(path.read_text());runs.append({"id":path.parent.name,**row}) except (ValueError,OSError):pass with self.lock: history=[{key:row.get(key) for key in ("id","action","status","created","exit_code","run","command")} for row in self.items.values()] return {"runs":runs,"jobs":sorted(history,key=lambda r:r["created"],reverse=True)[:30],"active":self.active, "formal_model":FINAL.is_file(),"data_ready":(ROOT/"artifacts/data/manifest.json").is_file()}