448 lines
21 KiB
Python
448 lines
21 KiB
Python
"""Administrator-registered, content-bound local training sources (no client paths)."""
|
||
|
||
from __future__ import annotations
|
||
|
||
import hashlib
|
||
import json
|
||
import os
|
||
import re
|
||
import shutil
|
||
import stat
|
||
import subprocess
|
||
import tempfile
|
||
import threading
|
||
import time
|
||
from contextlib import contextmanager
|
||
from copy import deepcopy
|
||
from pathlib import Path
|
||
|
||
SOURCE_ID = re.compile(r"^[A-Za-z0-9_-]{1,64}$")
|
||
DIGEST = re.compile(r"^[0-9a-f]{64}$")
|
||
LIMITS = {
|
||
"checkpoint": 256 * 1024**2,
|
||
"onnx": 64 * 1024**2,
|
||
"env": 2 * 1024**2,
|
||
"agent": 128 * 1024,
|
||
}
|
||
TASKS = ["Unitree-Go2-Flat", "Unitree-Go2-ObstacleAvoidance"]
|
||
|
||
|
||
class SourceError(ValueError):
|
||
pass
|
||
|
||
|
||
def regular_bytes(path: Path, root: Path, limit: int) -> bytes:
|
||
"""Walk with dir_fd/O_NOFOLLOW, so symlink substitution cannot escape the root."""
|
||
if ".." in path.parts or not path.is_absolute():
|
||
raise SourceError("基础策略路径必须是允许根内的绝对路径，不能含 ..")
|
||
try:
|
||
parts = path.relative_to(root).parts
|
||
except ValueError as error:
|
||
raise SourceError("基础策略文件不在管理员配置的允许根内") from error
|
||
if not parts:
|
||
raise SourceError("基础策略必须是常规文件")
|
||
try:
|
||
fd = os.open(root, os.O_RDONLY | os.O_DIRECTORY | os.O_NOFOLLOW)
|
||
except OSError as error:
|
||
raise SourceError("基础策略快照目录缺失或无效，拒绝训练") from error
|
||
try:
|
||
for part in parts[:-1]:
|
||
next_fd = os.open(part, os.O_RDONLY | os.O_DIRECTORY | os.O_NOFOLLOW, dir_fd=fd)
|
||
os.close(fd)
|
||
fd = next_fd
|
||
leaf = os.open(parts[-1], os.O_RDONLY | os.O_NOFOLLOW | os.O_NONBLOCK, dir_fd=fd)
|
||
with os.fdopen(leaf, "rb") as stream:
|
||
info = os.fstat(stream.fileno())
|
||
if not stat.S_ISREG(info.st_mode) or info.st_size > limit:
|
||
raise SourceError("基础策略文件必须是常规文件且不能超过大小上限")
|
||
data = stream.read(limit + 1)
|
||
if len(data) > limit:
|
||
raise SourceError("基础策略文件超过大小上限")
|
||
return data
|
||
except OSError as error:
|
||
raise SourceError(
|
||
"基础策略文件缺失或含symlink；请提供.pt、policy.onnx和params配置常规文件"
|
||
) from error
|
||
finally:
|
||
os.close(fd)
|
||
|
||
|
||
class PretrainedSources:
|
||
def __init__(self, config_path: Path | None, store: Path, python: str, trainer_root: Path):
|
||
self.store = store.expanduser().resolve()
|
||
self.store.mkdir(parents=True, exist_ok=True, mode=0o700)
|
||
self.python = python
|
||
self.trainer_root = trainer_root.resolve()
|
||
self.entries: dict[str, dict] = {}
|
||
self._upload_lock = threading.Lock()
|
||
self._catalog_lock = threading.RLock()
|
||
self._restore_uploads()
|
||
if config_path is None:
|
||
return
|
||
try:
|
||
if config_path.stat().st_size > 128 * 1024:
|
||
raise SourceError("基础策略注册配置不能超过128KiB")
|
||
config = json.loads(config_path.read_text())
|
||
if not isinstance(config, dict) or set(config) != {"allowedRoots", "sources"}:
|
||
raise SourceError("注册配置只接受allowedRoots和sources")
|
||
if (
|
||
not isinstance(config["allowedRoots"], list)
|
||
or not isinstance(config["sources"], list)
|
||
or len(config["sources"]) > 32
|
||
):
|
||
raise SourceError("允许根/sources必须为数组，最多32个注册条目")
|
||
roots = [Path(p).expanduser().resolve(strict=True) for p in config["allowedRoots"]]
|
||
if not roots or any(not p.is_dir() for p in roots):
|
||
raise SourceError("请配置存在的本地允许根目录")
|
||
for entry in config["sources"]:
|
||
if not isinstance(entry, dict) or set(entry) != {
|
||
"id",
|
||
"label",
|
||
"checkpoint",
|
||
"onnx",
|
||
}:
|
||
raise SourceError("注册条目必须包含id/label/checkpoint/onnx")
|
||
key = entry["id"]
|
||
if not isinstance(key, str) or not SOURCE_ID.fullmatch(key) or key in self.entries:
|
||
raise SourceError("注册id无效或重复")
|
||
if not isinstance(entry["label"], str) or not 1 <= len(entry["label"]) <= 100:
|
||
raise SourceError("基础策略label必须为1–100字符")
|
||
record = {"id": key, "label": entry["label"], "ready": False, "compatibleTasks": []}
|
||
self.entries[key] = {"public": record}
|
||
try:
|
||
checkpoint, onnx = (
|
||
Path(entry["checkpoint"]).expanduser(),
|
||
Path(entry["onnx"]).expanduser(),
|
||
)
|
||
if not re.fullmatch(r"[A-Za-z0-9_.-]+\.pt", checkpoint.name):
|
||
raise SourceError("请选择ONNX对应的.pt训练checkpoint，ONNX不能直接续训")
|
||
files = {
|
||
"checkpoint": checkpoint,
|
||
"onnx": onnx,
|
||
"env": checkpoint.parent / "params/env.yaml",
|
||
"agent": checkpoint.parent / "params/agent.yaml",
|
||
}
|
||
data = {}
|
||
for name, path in files.items():
|
||
root = next((r for r in roots if path.is_relative_to(r)), None)
|
||
if root is None:
|
||
raise SourceError("注册文件不在允许根内")
|
||
data[name] = regular_bytes(path, root, LIMITS[name])
|
||
names = {
|
||
"checkpoint": checkpoint.name,
|
||
"onnx": "policy.onnx",
|
||
"env": "params/env.yaml",
|
||
"agent": "params/agent.yaml",
|
||
}
|
||
with tempfile.TemporaryDirectory(dir=self.store, prefix="import-") as temporary:
|
||
directory = Path(temporary)
|
||
for name, content in data.items():
|
||
destination = directory / names[name]
|
||
destination.parent.mkdir(exist_ok=True)
|
||
destination.write_bytes(content)
|
||
manifest = self._validate(directory, checkpoint.name, TASKS[0], None)
|
||
digest = manifest["source_id"]
|
||
bound = {
|
||
"sourceId": digest,
|
||
"registeredId": key,
|
||
"label": entry["label"],
|
||
"manifest": manifest,
|
||
}
|
||
destination = self.store / digest
|
||
if not destination.exists():
|
||
shutil.copytree(directory, destination)
|
||
for file in destination.rglob("*"):
|
||
file.chmod(0o555 if file.is_dir() else 0o444)
|
||
destination.chmod(0o555)
|
||
self.verify(bound)
|
||
record.update(
|
||
id=digest,
|
||
ready=True,
|
||
compatibleTasks=TASKS,
|
||
observationSizes=[47, 81, 97],
|
||
initialization=bound,
|
||
)
|
||
except (SourceError, OSError, ValueError, TypeError) as error:
|
||
record["error"] = f"{error}；请修正管理员注册配置并重启服务，不会退回随机初始化"
|
||
except (OSError, TypeError, ValueError) as error:
|
||
raise SourceError(f"基础策略注册配置无效：{error}") from error
|
||
|
||
def _restore_uploads(self):
|
||
# Only committed content directories are catalogued. Partial uploads are never sources.
|
||
for directory in self.store.iterdir():
|
||
if not DIGEST.fullmatch(directory.name) or not (directory / "upload.json").is_file():
|
||
continue
|
||
try:
|
||
manifest = json.loads(
|
||
regular_bytes(directory / "upload.json", self.store, 64 * 1024)
|
||
)
|
||
label = json.loads(regular_bytes(directory / "label.json", self.store, 1024))[
|
||
"label"
|
||
]
|
||
bound = {
|
||
"sourceId": directory.name,
|
||
"registeredId": directory.name,
|
||
"label": self._display_name(label),
|
||
"manifest": manifest,
|
||
}
|
||
if manifest["source_id"] != directory.name:
|
||
raise SourceError("上传来源身份损坏")
|
||
self.verify(bound)
|
||
self._publish_upload(bound)
|
||
except (SourceError, ValueError, KeyError, TypeError, OSError):
|
||
self.entries[directory.name] = {
|
||
"public": {
|
||
"id": directory.name,
|
||
"label": "失效的已上传策略",
|
||
"ready": False,
|
||
"compatibleTasks": [],
|
||
"error": "上传快照损坏，拒绝随机初始化；请重新上传原文件或恢复快照",
|
||
}
|
||
}
|
||
|
||
@staticmethod
|
||
def _display_name(name):
|
||
if not isinstance(name, str):
|
||
return "uploaded-policy"
|
||
name = name.replace("\\", "/").split("/")[-1]
|
||
return (
|
||
re.sub(r"[^\w .()\-]", "_", name, flags=re.UNICODE)[:100].strip(" .")
|
||
or "uploaded-policy"
|
||
)
|
||
|
||
def _publish_upload(self, bound):
|
||
record = {
|
||
"id": bound["sourceId"],
|
||
"label": bound["label"],
|
||
"ready": True,
|
||
"compatibleTasks": TASKS,
|
||
"observationSizes": [47, 81, 97],
|
||
"initialization": bound,
|
||
}
|
||
with self._catalog_lock:
|
||
self.entries[bound["sourceId"]] = {"public": record}
|
||
return deepcopy(record)
|
||
|
||
@contextmanager
|
||
def upload_slot(self, length):
|
||
# At most one receiving/decoding upload, 32 sources and 2GiB cumulative store.
|
||
if not self._upload_lock.acquire(blocking=False):
|
||
raise SourceError("已有文件正在上传/验证，请稍后重试")
|
||
try:
|
||
files = list(self.store.rglob("*"))
|
||
used = sum(p.lstat().st_size for p in files if not p.is_symlink() and p.is_file())
|
||
count = sum(1 for p in self.store.iterdir() if DIGEST.fullmatch(p.name))
|
||
if count >= 32 or used + length + 16 * 1024**2 > 2 * 1024**3:
|
||
raise SourceError("基础策略存储已达32个来源或2GiB上限，请由本机维护者释放空间")
|
||
with tempfile.TemporaryDirectory(dir=self.store, prefix="upload-") as temporary:
|
||
yield Path(temporary)
|
||
finally:
|
||
self._upload_lock.release()
|
||
|
||
def receive_upload(self, stream, length, fmt, template, display_name, *, set_timeout=None):
|
||
if template != "go2-legacy47-v1":
|
||
raise SourceError("必须明确确认go2-legacy47-v1观测与关节模板")
|
||
limit = {"pt": 256 * 1024**2, "onnx": 64 * 1024**2}.get(fmt)
|
||
if limit is None:
|
||
raise SourceError("仅支持单个.pt或.onnx，不接受ZIP/路径/配套目录")
|
||
if type(length) is not int or not 0 < length <= limit:
|
||
raise SourceError("上传文件为空或超过.pt 256MiB / ONNX 64MiB上限")
|
||
with self.upload_slot(length) as directory:
|
||
path = directory / f"upload.{fmt}"
|
||
deadline = time.monotonic() + 60
|
||
remaining = length
|
||
try:
|
||
with path.open("xb") as output:
|
||
while remaining:
|
||
timeout = deadline - time.monotonic()
|
||
if timeout <= 0:
|
||
raise SourceError("上传超时，请重试")
|
||
try:
|
||
if set_timeout is not None:
|
||
set_timeout(min(10, timeout))
|
||
# BufferedReader.read(n) may perform many recv calls whose socket
|
||
# timeouts reset with each trickled byte. read1 returns after one
|
||
# raw read, allowing the total deadline to be checked every time.
|
||
read_chunk = getattr(stream, "read1", stream.read)
|
||
chunk = read_chunk(min(1024 * 1024, remaining))
|
||
except OSError as error:
|
||
raise SourceError("上传超时/连接中断，临时文件已清理") from error
|
||
if not chunk or len(chunk) > remaining:
|
||
raise SourceError("上传连接中断或实际长度不符")
|
||
output.write(chunk)
|
||
remaining -= len(chunk)
|
||
finally:
|
||
if set_timeout is not None:
|
||
set_timeout(10)
|
||
request = {
|
||
"path": str(path),
|
||
"directory": str(directory),
|
||
"format": fmt,
|
||
"template": template,
|
||
}
|
||
env = {
|
||
**os.environ,
|
||
"CUDA_VISIBLE_DEVICES": "",
|
||
"OMP_NUM_THREADS": "1",
|
||
"OPENBLAS_NUM_THREADS": "1",
|
||
"MKL_NUM_THREADS": "1",
|
||
}
|
||
try:
|
||
result = subprocess.run(
|
||
[self.python, str(Path(__file__).parent / "rl/scripts/validate_upload.py")],
|
||
input=json.dumps(request),
|
||
capture_output=True,
|
||
text=True,
|
||
cwd=self.trainer_root,
|
||
env=env,
|
||
timeout=60,
|
||
)
|
||
manifest = json.loads(result.stdout)
|
||
except (OSError, subprocess.TimeoutExpired, ValueError) as error:
|
||
raise SourceError("模型验证超时/资源超限或训练Python依赖不可用") from error
|
||
if result.returncode:
|
||
raise SourceError(manifest.get("error", "不支持的模型文件"))
|
||
digest = manifest["source_id"]
|
||
if not DIGEST.fullmatch(digest):
|
||
raise SourceError("模型验证返回无效身份")
|
||
label = self._display_name(display_name)
|
||
(directory / "upload.json").write_text(json.dumps(manifest), encoding="utf-8")
|
||
(directory / "label.json").write_text(json.dumps({"label": label}), encoding="utf-8")
|
||
destination = self.store / digest
|
||
if destination.exists():
|
||
# Dedup never replaces a committed artifact or changes an old job binding.
|
||
old = next((r for r in self.catalog() if r["id"] == digest and r["ready"]), None)
|
||
if old is None:
|
||
raise SourceError("同ID旧快照已损坏，请恢复快照后重试；不会覆盖旧任务来源")
|
||
self.verify(old["initialization"])
|
||
return old
|
||
for file in directory.iterdir():
|
||
file.chmod(0o444)
|
||
directory.rename(destination)
|
||
destination.chmod(0o555)
|
||
bound = {
|
||
"sourceId": digest,
|
||
"registeredId": digest,
|
||
"label": label,
|
||
"manifest": manifest,
|
||
}
|
||
self.verify(bound)
|
||
return self._publish_upload(bound)
|
||
|
||
def _validate(self, directory: Path, checkpoint: str, task_id: str, task_config: dict | None):
|
||
command = [self.python, str(Path(__file__).parent / "rl/scripts/validate_pretrained.py")]
|
||
request = {
|
||
"directory": str(directory),
|
||
"checkpoint": checkpoint,
|
||
"taskId": task_id,
|
||
"taskConfig": task_config,
|
||
"uploaded": (directory / "upload.json").is_file(),
|
||
}
|
||
try:
|
||
result = subprocess.run(
|
||
command,
|
||
input=json.dumps(request),
|
||
cwd=self.trainer_root,
|
||
capture_output=True,
|
||
text=True,
|
||
timeout=60,
|
||
)
|
||
except (subprocess.TimeoutExpired, OSError) as error:
|
||
raise SourceError("基础策略验证器超时或不可用，请检查训练Python环境") from error
|
||
if result.returncode:
|
||
# Validator emits only an actionable error, never a traceback or file contents.
|
||
raise SourceError(
|
||
result.stderr.strip()[-1500:] or "基础策略验证器失败，请检查训练Python依赖"
|
||
)
|
||
try:
|
||
return json.loads(result.stdout)
|
||
except ValueError as error:
|
||
raise SourceError("基础策略验证器未返回有效身份描述") from error
|
||
|
||
def catalog(self):
|
||
with self._catalog_lock:
|
||
sources = {}
|
||
for entry in self.entries.values():
|
||
sources.setdefault(entry["public"]["id"], deepcopy(entry["public"]))
|
||
return list(sources.values())
|
||
|
||
def bind(self, source_id, task_id: str, task_config=None):
|
||
# Selection is content-addressed, so a stale UI cannot silently bind a newly
|
||
# registered model under the same administrator-friendly registration name.
|
||
entry = next(
|
||
(entry for entry in self.catalog() if entry["id"] == source_id),
|
||
None,
|
||
)
|
||
if not isinstance(source_id, str) or entry is None:
|
||
raise SourceError("基础策略内容ID不存在或已变化；请刷新服务连接，不接受文件路径")
|
||
if not entry["ready"]:
|
||
raise SourceError(entry["error"])
|
||
if task_id not in entry["compatibleTasks"]:
|
||
raise SourceError("该基础策略仅兼容Flat47/Obstacle81或97；Rough高度扫描不支持迁移")
|
||
bound = deepcopy(entry["initialization"])
|
||
directory = self.verify(bound)
|
||
manifest = self._validate(
|
||
directory, bound["manifest"]["artifacts"]["checkpoint"]["name"], task_id, task_config
|
||
)
|
||
if manifest["source_id"] != bound["sourceId"]:
|
||
raise SourceError("基础策略快照身份变化，拒绝初始化")
|
||
return bound
|
||
|
||
def verify(self, bound: dict) -> Path:
|
||
try:
|
||
digest = bound["sourceId"]
|
||
if not isinstance(digest, str) or not DIGEST.fullmatch(digest):
|
||
raise SourceError("基础策略快照身份无效")
|
||
directory = self.store / digest
|
||
artifacts = bound["manifest"]["artifacts"]
|
||
if "sourceFormat" in bound["manifest"]:
|
||
fmt = bound["manifest"]["sourceFormat"]
|
||
if fmt not in ("pt", "onnx"):
|
||
raise SourceError("上传格式无效")
|
||
stored = json.loads(regular_bytes(directory / "upload.json", self.store, 64 * 1024))
|
||
if stored != bound["manifest"]:
|
||
raise SourceError("上传manifest SHA绑定变化，拒绝初始化")
|
||
for name, relative, limit in (
|
||
(
|
||
"upload",
|
||
f"upload.{fmt}",
|
||
LIMITS["checkpoint"] if fmt == "pt" else LIMITS["onnx"],
|
||
),
|
||
("checkpoint", "actor.pt", 16 * 1024**2),
|
||
):
|
||
content = regular_bytes(directory / relative, self.store, limit)
|
||
if (
|
||
artifacts[name]["name"] != relative
|
||
or hashlib.sha256(content).hexdigest() != artifacts[name]["sha256"]
|
||
):
|
||
raise SourceError("上传快照SHA不匹配，拒绝训练")
|
||
return directory
|
||
for name, relative in {
|
||
"checkpoint": artifacts["checkpoint"]["name"],
|
||
"onnx": "policy.onnx",
|
||
"env": "params/env.yaml",
|
||
"agent": "params/agent.yaml",
|
||
}.items():
|
||
content = regular_bytes(directory / relative, self.store, LIMITS[name])
|
||
if hashlib.sha256(content).hexdigest() != artifacts[name]["sha256"]:
|
||
raise SourceError("基础策略快照SHA不匹配，拒绝训练；请恢复原快照")
|
||
return directory
|
||
except (KeyError, TypeError) as error:
|
||
raise SourceError("持久化基础策略描述损坏，拒绝训练") from error
|
||
|
||
def arguments(self, bound: dict) -> list[str]:
|
||
directory = self.verify(bound)
|
||
upload_args = (
|
||
["--pretrained-upload-manifest", str(directory / "upload.json")]
|
||
if "sourceFormat" in bound["manifest"]
|
||
else []
|
||
)
|
||
return upload_args + [
|
||
"--pretrained-checkpoint",
|
||
str(directory / bound["manifest"]["artifacts"]["checkpoint"]["name"]),
|
||
"--pretrained-allowed-roots",
|
||
json.dumps([str(directory)]),
|
||
"--pretrained-source-id",
|
||
bound["sourceId"],
|
||
]
|