feat(compute): complete Origin job execution pipeline
This commit is contained in:
parent
faa2279f00
commit
3e0a719f80
10
DESIGN.md
10
DESIGN.md
|
|
@ -464,7 +464,15 @@ scheduled_jobs(§8.5) channel_bindings(§8.7,判别列+JSONB)
|
|||
|
||||
Node 通过 `Authorization: Bearer` 与 `X-Node-Id` 建立 `/v1/compute/nodes/connect` WebSocket。进程内 Connection Manager 保证同一节点单活，新连接关闭旧连接；`hello`/`heartbeat` 更新版本、容量、软件健康与最后在线时间。管理员禁用节点时先持久化禁用态，再关闭现有连接；断线收尾不得覆盖禁用态。当前单活只覆盖单 Web 进程，生产启用多实例前必须增加 Redis/PG fencing 或将 Node API 固定路由到单一控制面实例。
|
||||
|
||||
首批只落注册、认证、心跳、状态与禁用基础链路。`compute_jobs`、任务 offer/accept、Origin Worker、输入输出传输、重连对账和 Token 轮换属于后续垂直闭环，不以任意命令或脚本接口临时代替。
|
||||
第二阶段已增加 `compute_jobs` 账本与 `origin.plot@v1` 的 offer/accept 骨架。用户只能在本人 task 下以幂等键提交固定 schema；云端规范化请求并记录 SHA-256，按当前进程真实在线、能力匹配、健康且有空闲 slot 的 Node 创建短期 offer。Node 再次校验 schema、图形类型和输出格式，使用 write-through、flush 与原子 rename 先落本机任务目录，再回 `job_accept`；重复 job 只有 digest 一致才接受。过期或发送失败的 offer 回到队列，lease、Node 和 digest 不匹配的响应被拒绝。Node 接收后云端进入 `dispatched` 而非 `running`，并将 slot 降为 0；只有固定 Worker 真正启动后才进入 `origin_running`。
|
||||
|
||||
第三阶段补齐输入下载与恢复状态协议：`input_id` 固定为用户已有 artifact UUID，提交时快照文件名、大小和 SHA-256，只允许 CSV/XLSX/JSON 且不超过 100 MiB。Node 以自身 Bearer 身份访问任务绑定的只读下载端点，流式写入本 job 的 `input/`，同时限制声明大小并校验 SHA-256，完成后原子 rename；不暴露工作区路径。Node 会原子读取/补报 `terminal.json`，断线后云端把活动任务标记 `disconnected` 并保留 Node/lease，重连按 job、lease、digest 恢复下载或幂等补报终态，不自动重派。
|
||||
|
||||
第四阶段落地固定 Origin Worker：Node 仅从管理员安装的固定 Python 运行时启动随程序发布的 `worker.py`，参数只有本机 job 目录；请求不能指定脚本、解释器或文件路径。Worker 使用 `originpro` 生成 OPJU、PNG、SVG、PDF、plot spec 和 provenance，校验产物签名并原子写入终态；当前受控图形仅含 line、scatter、line_scatter 和双栏出版布局。进程内 pipeline 按 job 去重，并脱离单次 WebSocket 的取消令牌运行；连接中断只延迟状态/终态上报。Node 进程若在 Worker 启动后重启，则保守失败而不重复驱动 Origin，避免无法证明的双执行。
|
||||
|
||||
第五阶段完成输出上传与发布：Node 只按固定 manifest ID 逐项流式 PUT，并携带 Node、lease、request digest 与内容摘要；云端重新绑定任务身份，不信 Node 提供的路径或媒体类型。文件先进入用户根下隐藏暂存区，固定文件名、单文件/总大小和 SHA-256 全部验证后，目录级原子移动到 `<working_dir>/origin/<job_id>/`，再登记平台 artifact UUID 并写成功终态。重复 PUT、complete 和重连均按摘要幂等；部分上传不可见，只有完整集合才能发布。
|
||||
|
||||
后续仍需实现协作取消和 Token 轮换；不得以任意命令或脚本接口临时代替。当前单活与 offer 选择仍只覆盖单 Web 进程，生产启用多实例前必须增加 Redis/PG fencing 或固定路由到单一控制面实例。
|
||||
|
||||
---
|
||||
|
||||
|
|
|
|||
8
RUN.md
8
RUN.md
|
|
@ -1076,6 +1076,14 @@ windows-node/Zcbot.WindowsNode/bin/Debug/net10.0-windows/Zcbot.WindowsNode.exe e
|
|||
windows-node/Zcbot.WindowsNode/bin/Debug/net10.0-windows/Zcbot.WindowsNode.exe
|
||||
```
|
||||
|
||||
Origin Worker 使用独立固定 Python 运行时。先确认该交互式 Windows 账号已安装并可启动 Origin/OriginPro，再由管理员安装运行时（不要复用服务端 `.venv`）：
|
||||
|
||||
```powershell
|
||||
windows-node\install-origin-runtime.ps1 -BootstrapPython D:\programs\Python312\python.exe
|
||||
```
|
||||
|
||||
默认解释器为 `%ProgramData%\Zcbot\WindowsNode\runtimes\origin\python.exe`。如需放在其他受管目录，设置机器级 `ZCBOT_ORIGIN_PYTHON` 为绝对 `python.exe` 路径后重启 Node。运行时固定依赖见 `windows-node/origin-worker/requirements.txt`；任务请求无权选择解释器、脚本或路径。当前 Worker 支持 CSV/XLSX/JSON 输入，`line`、`scatter`、`line_scatter` 与 OPJU/PNG/SVG/PDF 输出。成功产物由 Node 流式上传，全部校验通过后发布到任务工作目录 `origin/<job_id>/`；上传中断会在重连时幂等续传。
|
||||
|
||||
注册配置写入 `%ProgramData%\Zcbot\WindowsNode\node.json`；Token 使用 DPAPI `LocalMachine` 加密，ACL 仅允许注册账号和 `SYSTEM`。应始终用同一专用 Windows 账号注册并运行 Node。当前 MVP 以该账号的登录后计划任务启动，不安装 Windows Service。
|
||||
|
||||
直接双击 EXE 启动托盘 UI：红点为未注册/身份失效，黄点为连接中，绿点为在线；双击托盘图标打开配置窗。原 CLI 注册入口继续保留，无 UI 模式使用 `Zcbot.WindowsNode.exe run --headless`。
|
||||
|
|
|
|||
|
|
@ -59,7 +59,7 @@ def register_published_artifacts(
|
|||
continue
|
||||
current_path = _rel(root, path)
|
||||
label = str(ref.get("label") or "")
|
||||
media_type = mimetypes.guess_type(path.name)[0]
|
||||
media_type = str(ref.get("media_type") or "") or mimetypes.guess_type(path.name)[0]
|
||||
size_bytes = path.stat().st_size
|
||||
content_sha256 = _hash_file(path)
|
||||
statement = pg_insert(Artifact).values(
|
||||
|
|
|
|||
|
|
@ -0,0 +1,638 @@
|
|||
"""Origin 受控计算任务的校验、幂等持久化和 offer 状态机。"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import re
|
||||
from datetime import datetime, timedelta, timezone
|
||||
from hashlib import sha256
|
||||
from uuid import UUID, uuid4
|
||||
|
||||
from sqlalchemy import select
|
||||
from sqlalchemy.exc import IntegrityError
|
||||
|
||||
from core.compute_nodes import ComputeNodeError, SUPPORTED_CAPABILITIES
|
||||
from core.storage.engine import session_scope
|
||||
from core.storage.models import Artifact, ComputeJob, ComputeNode, Task
|
||||
|
||||
OFFER_SECONDS = 60
|
||||
ALLOWED_PLOT_TYPES = frozenset(
|
||||
{"line", "scatter", "line_scatter"}
|
||||
)
|
||||
ALLOWED_OUTPUT_FORMATS = frozenset({"opju", "png", "svg", "pdf"})
|
||||
ALLOWED_INPUT_SUFFIXES = frozenset({".csv", ".xlsx", ".json"})
|
||||
MAX_INPUT_BYTES = 100 * 1024 * 1024
|
||||
MAX_OUTPUT_ARTIFACT_BYTES = 256 * 1024 * 1024
|
||||
MAX_OUTPUT_TOTAL_BYTES = 512 * 1024 * 1024
|
||||
OUTPUT_ARTIFACTS = {
|
||||
"project": ("project.opju", "application/x-origin-project", "opju"),
|
||||
"figure_png": ("figure.png", "image/png", "png"),
|
||||
"figure_svg": ("figure.svg", "image/svg+xml", "svg"),
|
||||
"figure_pdf": ("figure.pdf", "application/pdf", "pdf"),
|
||||
"plot_spec": ("plot-spec.json", "application/json", None),
|
||||
"provenance": ("provenance.json", "application/json", None),
|
||||
}
|
||||
|
||||
|
||||
def _has_only(value: dict, fields: set[str]) -> bool:
|
||||
return set(value).issubset(fields)
|
||||
|
||||
|
||||
def _canonical_request(request: dict) -> tuple[dict, str]:
|
||||
if not isinstance(request, dict) or set(request) != {"schema_version", "input", "plot", "output"}:
|
||||
raise ComputeNodeError("invalid origin plot request fields")
|
||||
if request.get("schema_version") != 1:
|
||||
raise ComputeNodeError("unsupported origin plot schema version")
|
||||
input_spec = request.get("input")
|
||||
plot = request.get("plot")
|
||||
output = request.get("output")
|
||||
if not all(isinstance(item, dict) for item in (input_spec, plot, output)):
|
||||
raise ComputeNodeError("origin plot request sections must be objects")
|
||||
if not _has_only(input_spec, {"input_id", "sheet"}):
|
||||
raise ComputeNodeError("unsupported origin input fields")
|
||||
try:
|
||||
UUID(str(input_spec.get("input_id") or ""))
|
||||
except ValueError as exc:
|
||||
raise ComputeNodeError("input.input_id must be an artifact UUID") from exc
|
||||
if "sheet" in input_spec and (
|
||||
not isinstance(input_spec["sheet"], str) or not 1 <= len(input_spec["sheet"]) <= 128
|
||||
):
|
||||
raise ComputeNodeError("input.sheet must be a string")
|
||||
if not _has_only(
|
||||
plot,
|
||||
{"type", "x", "y", "template", "title", "x_axis", "y_axis", "legend", "error_bars"},
|
||||
):
|
||||
raise ComputeNodeError("unsupported origin plot fields")
|
||||
if plot.get("type") not in ALLOWED_PLOT_TYPES:
|
||||
raise ComputeNodeError("unsupported origin plot type")
|
||||
if "title" in plot and (
|
||||
not isinstance(plot["title"], str) or len(plot["title"]) > 500
|
||||
):
|
||||
raise ComputeNodeError("plot.title must be a string")
|
||||
if plot.get("template", "publication_double_column") != "publication_double_column":
|
||||
raise ComputeNodeError("unsupported origin plot template")
|
||||
x_column = plot.get("x")
|
||||
y_columns = plot.get("y")
|
||||
if not isinstance(x_column, str) or not 1 <= len(x_column) <= 128:
|
||||
raise ComputeNodeError("plot.x must be a column name")
|
||||
if isinstance(y_columns, str):
|
||||
y_columns = [y_columns]
|
||||
if (
|
||||
not isinstance(y_columns, list)
|
||||
or not 1 <= len(y_columns) <= 16
|
||||
or len(y_columns) != len(set(y_columns))
|
||||
or any(not isinstance(item, str) or not 1 <= len(item) <= 128 for item in y_columns)
|
||||
):
|
||||
raise ComputeNodeError("plot.y must contain 1 to 16 unique column names")
|
||||
for axis_name in ("x_axis", "y_axis"):
|
||||
axis = plot.get(axis_name)
|
||||
if axis is not None and (
|
||||
not isinstance(axis, dict)
|
||||
or not _has_only(axis, {"title", "unit", "scale"})
|
||||
or axis.get("scale", "linear") != "linear"
|
||||
or any(
|
||||
name in axis and not isinstance(axis[name], str)
|
||||
for name in ("title", "unit")
|
||||
)
|
||||
):
|
||||
raise ComputeNodeError(f"invalid {axis_name}")
|
||||
legend = plot.get("legend")
|
||||
if legend is not None and (
|
||||
not isinstance(legend, dict)
|
||||
or not _has_only(legend, {"enabled", "position"})
|
||||
or ("enabled" in legend and not isinstance(legend["enabled"], bool))
|
||||
or legend.get("enabled", True) is not True
|
||||
or legend.get("position", "top_right") != "top_right"
|
||||
):
|
||||
raise ComputeNodeError("invalid plot.legend")
|
||||
if plot.get("error_bars") is not None:
|
||||
raise ComputeNodeError("error bars are not supported in origin.plot@v1")
|
||||
if not _has_only(output, {"formats", "dpi", "capture_screenshots", "record_video"}):
|
||||
raise ComputeNodeError("unsupported origin output fields")
|
||||
if any(
|
||||
name in output and not isinstance(output[name], bool)
|
||||
for name in ("capture_screenshots", "record_video")
|
||||
):
|
||||
raise ComputeNodeError("origin output capture flags must be boolean")
|
||||
if output.get("record_video", False):
|
||||
raise ComputeNodeError("origin video recording is not supported")
|
||||
dpi = output.get("dpi", 300)
|
||||
if not isinstance(dpi, int) or isinstance(dpi, bool) or not 72 <= dpi <= 1200:
|
||||
raise ComputeNodeError("output.dpi must be between 72 and 1200")
|
||||
formats = output.get("formats")
|
||||
if (
|
||||
not isinstance(formats, list)
|
||||
or not formats
|
||||
or len(formats) != len(set(formats))
|
||||
or any(item not in ALLOWED_OUTPUT_FORMATS for item in formats)
|
||||
):
|
||||
raise ComputeNodeError("output.formats contains unsupported values")
|
||||
encoded = json.dumps(request, ensure_ascii=False, sort_keys=True, separators=(",", ":"))
|
||||
if len(encoded.encode("utf-8")) > 256 * 1024:
|
||||
raise ComputeNodeError("origin plot request is too large")
|
||||
normalized = json.loads(encoded)
|
||||
return normalized, sha256(encoded.encode("utf-8")).hexdigest()
|
||||
|
||||
|
||||
def _job_dict(row: ComputeJob) -> dict:
|
||||
return {
|
||||
"job_id": str(row.job_id),
|
||||
"task_id": str(row.task_id),
|
||||
"capability": row.capability,
|
||||
"request_digest": row.request_digest,
|
||||
"node_id": str(row.node_id) if row.node_id else None,
|
||||
"status": row.status,
|
||||
"stage": row.stage,
|
||||
"progress": row.progress,
|
||||
"metrics": row.metrics,
|
||||
"error": row.error,
|
||||
"artifact_manifest": row.artifact_manifest,
|
||||
"created_at": row.created_at.isoformat() if row.created_at else None,
|
||||
"started_at": row.started_at.isoformat() if row.started_at else None,
|
||||
"terminal_at": row.terminal_at.isoformat() if row.terminal_at else None,
|
||||
}
|
||||
|
||||
|
||||
def create_job(
|
||||
user_id: UUID,
|
||||
task_id: UUID,
|
||||
*,
|
||||
idempotency_key: str,
|
||||
capability: str,
|
||||
request: dict,
|
||||
) -> tuple[dict, bool]:
|
||||
key = idempotency_key.strip()
|
||||
if not key or len(key) > 200:
|
||||
raise ComputeNodeError("idempotency_key must contain 1 to 200 characters")
|
||||
if capability not in SUPPORTED_CAPABILITIES:
|
||||
raise ComputeNodeError("unsupported capability")
|
||||
normalized, digest = _canonical_request(request)
|
||||
with session_scope() as session:
|
||||
task = session.execute(
|
||||
select(Task.task_id).where(Task.task_id == task_id, Task.user_id == user_id)
|
||||
).first()
|
||||
if task is None:
|
||||
raise ComputeNodeError("task not found")
|
||||
artifact_id = UUID(normalized["input"]["input_id"])
|
||||
artifact = session.execute(
|
||||
select(Artifact).where(
|
||||
Artifact.artifact_id == artifact_id,
|
||||
Artifact.user_id == user_id,
|
||||
Artifact.status == "active",
|
||||
)
|
||||
).scalar_one_or_none()
|
||||
if artifact is None:
|
||||
raise ComputeNodeError("input artifact not found")
|
||||
suffix = "." + artifact.current_path.rsplit(".", 1)[-1].lower() if "." in artifact.current_path else ""
|
||||
if suffix not in ALLOWED_INPUT_SUFFIXES:
|
||||
raise ComputeNodeError("input artifact type is not supported")
|
||||
if (
|
||||
artifact.size_bytes is None
|
||||
or artifact.size_bytes < 0
|
||||
or artifact.size_bytes > MAX_INPUT_BYTES
|
||||
or not artifact.content_sha256
|
||||
or len(artifact.content_sha256) != 64
|
||||
):
|
||||
raise ComputeNodeError("input artifact metadata is incomplete or too large")
|
||||
input_manifest = {
|
||||
"artifact_id": str(artifact.artifact_id),
|
||||
"filename": artifact.current_path.replace("\\", "/").rsplit("/", 1)[-1],
|
||||
"size_bytes": artifact.size_bytes,
|
||||
"sha256": artifact.content_sha256,
|
||||
}
|
||||
existing = session.execute(
|
||||
select(ComputeJob).where(
|
||||
ComputeJob.user_id == user_id,
|
||||
ComputeJob.idempotency_key == key,
|
||||
)
|
||||
).scalar_one_or_none()
|
||||
if existing is not None:
|
||||
if (
|
||||
existing.task_id != task_id
|
||||
or existing.capability != capability
|
||||
or existing.request_digest != digest
|
||||
):
|
||||
raise ComputeNodeError("idempotency key was already used for a different request")
|
||||
return _job_dict(existing), False
|
||||
row = ComputeJob(
|
||||
job_id=uuid4(),
|
||||
user_id=user_id,
|
||||
task_id=task_id,
|
||||
idempotency_key=key,
|
||||
capability=capability,
|
||||
request=normalized,
|
||||
request_digest=digest,
|
||||
input_manifest=input_manifest,
|
||||
status="queued",
|
||||
stage="",
|
||||
metrics={},
|
||||
error={},
|
||||
artifact_manifest=[],
|
||||
)
|
||||
try:
|
||||
with session.begin_nested():
|
||||
session.add(row)
|
||||
session.flush()
|
||||
return _job_dict(row), True
|
||||
except IntegrityError:
|
||||
existing = session.execute(
|
||||
select(ComputeJob).where(
|
||||
ComputeJob.user_id == user_id,
|
||||
ComputeJob.idempotency_key == key,
|
||||
)
|
||||
).scalar_one()
|
||||
if (
|
||||
existing.task_id != task_id
|
||||
or existing.capability != capability
|
||||
or existing.request_digest != digest
|
||||
):
|
||||
raise ComputeNodeError(
|
||||
"idempotency key was already used for a different request"
|
||||
)
|
||||
return _job_dict(existing), False
|
||||
|
||||
|
||||
def get_job(user_id: UUID, job_id: UUID) -> dict | None:
|
||||
with session_scope() as session:
|
||||
row = session.execute(
|
||||
select(ComputeJob).where(ComputeJob.job_id == job_id, ComputeJob.user_id == user_id)
|
||||
).scalar_one_or_none()
|
||||
return _job_dict(row) if row else None
|
||||
|
||||
|
||||
def offer_next_job(node_ids: set[UUID]) -> dict | None:
|
||||
"""从当前进程实际在线的节点中选择一个，为最早 queued job 创建短租约。"""
|
||||
if not node_ids:
|
||||
return None
|
||||
now = datetime.now(timezone.utc)
|
||||
with session_scope() as session:
|
||||
expired = session.execute(
|
||||
select(ComputeJob)
|
||||
.where(
|
||||
ComputeJob.status == "offered",
|
||||
ComputeJob.lease_expires_at <= now,
|
||||
)
|
||||
.with_for_update(skip_locked=True)
|
||||
).scalars()
|
||||
for item in expired:
|
||||
item.status = "queued"
|
||||
item.node_id = None
|
||||
item.lease_id = None
|
||||
item.lease_expires_at = None
|
||||
job = session.execute(
|
||||
select(ComputeJob)
|
||||
.where(ComputeJob.status == "queued")
|
||||
.order_by(ComputeJob.created_at, ComputeJob.job_id)
|
||||
.with_for_update(skip_locked=True)
|
||||
.limit(1)
|
||||
).scalar_one_or_none()
|
||||
if job is None:
|
||||
return None
|
||||
busy_node_ids = set(
|
||||
session.execute(
|
||||
select(ComputeJob.node_id).where(
|
||||
ComputeJob.node_id.is_not(None),
|
||||
ComputeJob.status.in_({"offered", "dispatched", "running"}),
|
||||
)
|
||||
).scalars()
|
||||
)
|
||||
nodes = session.execute(
|
||||
select(ComputeNode)
|
||||
.where(ComputeNode.node_id.in_(node_ids), ComputeNode.status == "online")
|
||||
.order_by(ComputeNode.last_seen_at.desc())
|
||||
).scalars()
|
||||
node = next(
|
||||
(
|
||||
item
|
||||
for item in nodes
|
||||
if item.node_id not in busy_node_ids
|
||||
and job.capability in item.capabilities
|
||||
and int((item.runtime or {}).get("available_slots") or 0) > 0
|
||||
),
|
||||
None,
|
||||
)
|
||||
if node is None:
|
||||
return None
|
||||
lease_id = uuid4()
|
||||
expires_at = now + timedelta(seconds=OFFER_SECONDS)
|
||||
job.node_id = node.node_id
|
||||
job.lease_id = lease_id
|
||||
job.lease_expires_at = expires_at
|
||||
job.status = "offered"
|
||||
return {
|
||||
"node_id": node.node_id,
|
||||
"payload": {
|
||||
"job_id": str(job.job_id),
|
||||
"lease_id": str(lease_id),
|
||||
"lease_expires_at": expires_at.isoformat(),
|
||||
"capability": job.capability,
|
||||
"request_digest": job.request_digest,
|
||||
"request": job.request,
|
||||
"input_transfer": {
|
||||
**job.input_manifest,
|
||||
"download_path": f"/v1/compute/jobs/{job.job_id}/input",
|
||||
},
|
||||
},
|
||||
}
|
||||
|
||||
|
||||
def get_job_input(node_id: UUID, job_id: UUID) -> dict | None:
|
||||
"""返回任务绑定的 artifact 定位信息；调用方仍需在 user_root 内安全解析。"""
|
||||
with session_scope() as session:
|
||||
job = session.execute(
|
||||
select(ComputeJob).where(
|
||||
ComputeJob.job_id == job_id,
|
||||
ComputeJob.node_id == node_id,
|
||||
ComputeJob.status.in_({"offered", "dispatched", "running", "disconnected"}),
|
||||
)
|
||||
).scalar_one_or_none()
|
||||
if job is None:
|
||||
return None
|
||||
artifact_id = UUID(job.input_manifest["artifact_id"])
|
||||
artifact = session.execute(
|
||||
select(Artifact).where(
|
||||
Artifact.artifact_id == artifact_id,
|
||||
Artifact.user_id == job.user_id,
|
||||
Artifact.status == "active",
|
||||
)
|
||||
).scalar_one_or_none()
|
||||
if artifact is None:
|
||||
return None
|
||||
return {
|
||||
"user_id": job.user_id,
|
||||
"current_path": artifact.current_path,
|
||||
**job.input_manifest,
|
||||
}
|
||||
|
||||
|
||||
def get_job_output_context(node_id: UUID, job_id: UUID, lease_id: UUID, digest: str) -> dict | None:
|
||||
"""返回 Node 输出上传上下文，不向 Node 暴露任何云端文件路径。"""
|
||||
with session_scope() as session:
|
||||
row = session.execute(
|
||||
select(ComputeJob, Task.working_dir)
|
||||
.join(Task, Task.task_id == ComputeJob.task_id)
|
||||
.where(ComputeJob.job_id == job_id)
|
||||
).one_or_none()
|
||||
if row is None:
|
||||
return None
|
||||
job, working_dir = row
|
||||
if (
|
||||
job.node_id != node_id
|
||||
or job.lease_id != lease_id
|
||||
or job.request_digest != digest
|
||||
or job.status not in {"dispatched", "running", "disconnected", "succeeded"}
|
||||
):
|
||||
return None
|
||||
return {
|
||||
"user_id": job.user_id,
|
||||
"task_id": job.task_id,
|
||||
"working_dir": working_dir,
|
||||
"request": job.request,
|
||||
"status": job.status,
|
||||
"artifact_manifest": job.artifact_manifest,
|
||||
}
|
||||
|
||||
|
||||
def validate_output_manifest(request: dict, manifest: object) -> list[dict]:
|
||||
if not isinstance(manifest, list):
|
||||
raise ComputeNodeError("job artifact manifest must be a list")
|
||||
requested_formats = set(request.get("output", {}).get("formats") or [])
|
||||
expected_ids = {"plot_spec", "provenance"}
|
||||
expected_ids.update(
|
||||
artifact_id
|
||||
for artifact_id, (_, _, output_format) in OUTPUT_ARTIFACTS.items()
|
||||
if output_format in requested_formats
|
||||
)
|
||||
if len(manifest) != len(expected_ids):
|
||||
raise ComputeNodeError("job artifact manifest is incomplete")
|
||||
normalized: list[dict] = []
|
||||
seen: set[str] = set()
|
||||
total = 0
|
||||
for raw in manifest:
|
||||
if not isinstance(raw, dict) or set(raw) != {
|
||||
"artifact_id", "filename", "media_type", "size_bytes", "sha256"
|
||||
}:
|
||||
raise ComputeNodeError("job artifact manifest entry is invalid")
|
||||
local_id = raw.get("artifact_id")
|
||||
if local_id not in expected_ids or local_id in seen:
|
||||
raise ComputeNodeError("job artifact manifest identity is invalid")
|
||||
filename, media_type, _ = OUTPUT_ARTIFACTS[local_id]
|
||||
size = raw.get("size_bytes")
|
||||
digest = raw.get("sha256")
|
||||
if raw.get("filename") != filename or raw.get("media_type") != media_type:
|
||||
raise ComputeNodeError("job artifact manifest metadata does not match its identity")
|
||||
if not isinstance(size, int) or isinstance(size, bool) or not 1 <= size <= MAX_OUTPUT_ARTIFACT_BYTES:
|
||||
raise ComputeNodeError("job output artifact size is invalid")
|
||||
if not isinstance(digest, str) or not re.fullmatch(r"[0-9a-f]{64}", digest):
|
||||
raise ComputeNodeError("job output artifact digest is invalid")
|
||||
total += size
|
||||
seen.add(local_id)
|
||||
normalized.append(dict(raw))
|
||||
if seen != expected_ids or total > MAX_OUTPUT_TOTAL_BYTES:
|
||||
raise ComputeNodeError("job artifact manifest is incomplete or too large")
|
||||
return normalized
|
||||
|
||||
|
||||
def abandon_offer(node_id: UUID, payload: dict) -> None:
|
||||
"""WebSocket 发送失败时只回滚仍属于该连接租约的 offer。"""
|
||||
try:
|
||||
job_id = UUID(str(payload.get("job_id", "")))
|
||||
lease_id = UUID(str(payload.get("lease_id", "")))
|
||||
except ValueError:
|
||||
return
|
||||
with session_scope() as session:
|
||||
job = session.execute(
|
||||
select(ComputeJob).where(ComputeJob.job_id == job_id).with_for_update()
|
||||
).scalar_one_or_none()
|
||||
if (
|
||||
job is not None
|
||||
and job.status == "offered"
|
||||
and job.node_id == node_id
|
||||
and job.lease_id == lease_id
|
||||
):
|
||||
job.status = "queued"
|
||||
job.node_id = None
|
||||
job.lease_id = None
|
||||
job.lease_expires_at = None
|
||||
|
||||
|
||||
def respond_to_offer(node_id: UUID, *, accepted: bool, payload: dict) -> None:
|
||||
try:
|
||||
job_id = UUID(str(payload.get("job_id", "")))
|
||||
lease_id = UUID(str(payload.get("lease_id", "")))
|
||||
except ValueError as exc:
|
||||
raise ComputeNodeError("invalid job offer response identity") from exc
|
||||
now = datetime.now(timezone.utc)
|
||||
with session_scope() as session:
|
||||
job = session.execute(
|
||||
select(ComputeJob).where(ComputeJob.job_id == job_id).with_for_update()
|
||||
).scalar_one_or_none()
|
||||
if job is None or job.node_id != node_id or job.lease_id != lease_id:
|
||||
raise ComputeNodeError("job offer is stale or does not belong to this node")
|
||||
if (
|
||||
accepted
|
||||
and job.status in {"dispatched", "running", "succeeded", "failed", "cancelled"}
|
||||
and payload.get("request_digest") == job.request_digest
|
||||
):
|
||||
return
|
||||
if job.status != "offered":
|
||||
raise ComputeNodeError("job offer is stale or does not belong to this node")
|
||||
if job.lease_expires_at is None or job.lease_expires_at <= now:
|
||||
job.status = "queued"
|
||||
job.node_id = None
|
||||
job.lease_id = None
|
||||
job.lease_expires_at = None
|
||||
raise ComputeNodeError("job offer has expired")
|
||||
if accepted:
|
||||
if payload.get("request_digest") != job.request_digest:
|
||||
raise ComputeNodeError("job request digest mismatch")
|
||||
job.status = "dispatched"
|
||||
job.stage = "accepted"
|
||||
job.error = {}
|
||||
else:
|
||||
job.status = "queued"
|
||||
job.node_id = None
|
||||
job.lease_id = None
|
||||
job.lease_expires_at = None
|
||||
job.error = {"code": "node_rejected", "detail": str(payload.get("reason") or "")[:500]}
|
||||
|
||||
|
||||
def update_job_state(node_id: UUID, payload: dict) -> None:
|
||||
job_id, lease_id, digest = _message_identity(payload)
|
||||
stage = str(payload.get("stage") or "")
|
||||
progress = payload.get("progress")
|
||||
metrics = payload.get("metrics") or {}
|
||||
if not stage or len(stage) > 100:
|
||||
raise ComputeNodeError("job stage is required")
|
||||
if not isinstance(progress, int) or isinstance(progress, bool) or not 0 <= progress <= 100:
|
||||
raise ComputeNodeError("job progress must be between 0 and 100")
|
||||
if not isinstance(metrics, dict) or len(json.dumps(metrics, ensure_ascii=False)) > 64 * 1024:
|
||||
raise ComputeNodeError("job metrics are invalid")
|
||||
now = datetime.now(timezone.utc)
|
||||
with session_scope() as session:
|
||||
job = session.execute(
|
||||
select(ComputeJob).where(ComputeJob.job_id == job_id).with_for_update()
|
||||
).scalar_one_or_none()
|
||||
_assert_job_message(job, node_id, lease_id, digest)
|
||||
if job.status in {"succeeded", "failed", "cancelled"}:
|
||||
return
|
||||
if not _can_accept_state(job.status):
|
||||
raise ComputeNodeError("job state cannot advance from its current status")
|
||||
job.status = (
|
||||
"dispatched"
|
||||
if stage in {"accepted", "waiting_input", "ready_to_run"}
|
||||
else "running"
|
||||
)
|
||||
job.stage = stage
|
||||
job.progress = progress
|
||||
job.metrics = metrics
|
||||
if job.status == "running" and job.started_at is None:
|
||||
job.started_at = now
|
||||
|
||||
|
||||
def record_job_terminal(node_id: UUID, payload: dict) -> None:
|
||||
job_id, lease_id, digest = _message_identity(payload)
|
||||
terminal_status = payload.get("status")
|
||||
if terminal_status not in {"succeeded", "failed", "cancelled"}:
|
||||
raise ComputeNodeError("invalid job terminal status")
|
||||
error = payload.get("error") or {}
|
||||
manifest = payload.get("artifact_manifest") or []
|
||||
if not isinstance(error, dict) or len(json.dumps(error, ensure_ascii=False)) > 64 * 1024:
|
||||
raise ComputeNodeError("job terminal error is invalid")
|
||||
if not isinstance(manifest, list) or len(json.dumps(manifest, ensure_ascii=False)) > 256 * 1024:
|
||||
raise ComputeNodeError("job artifact manifest is invalid")
|
||||
now = datetime.now(timezone.utc)
|
||||
with session_scope() as session:
|
||||
job = session.execute(
|
||||
select(ComputeJob).where(ComputeJob.job_id == job_id).with_for_update()
|
||||
).scalar_one_or_none()
|
||||
_assert_job_message(job, node_id, lease_id, digest)
|
||||
if terminal_status == "succeeded":
|
||||
expected = validate_output_manifest(
|
||||
job.request,
|
||||
[
|
||||
{
|
||||
"artifact_id": item.get("source_artifact_id"),
|
||||
"filename": item.get("filename"),
|
||||
"media_type": item.get("media_type"),
|
||||
"size_bytes": item.get("size_bytes"),
|
||||
"sha256": item.get("sha256"),
|
||||
}
|
||||
for item in manifest
|
||||
if isinstance(item, dict)
|
||||
],
|
||||
)
|
||||
if len(expected) != len(manifest) or any(
|
||||
not isinstance(item.get("artifact_id"), str)
|
||||
or not _is_uuid(item["artifact_id"])
|
||||
or not isinstance(item.get("path"), str)
|
||||
or not item["path"].startswith(f"origin/{job.job_id}/")
|
||||
for item in manifest
|
||||
):
|
||||
raise ComputeNodeError("successful job artifacts have not been published")
|
||||
if job.status in {"succeeded", "failed", "cancelled"}:
|
||||
if job.status != terminal_status:
|
||||
raise ComputeNodeError("job terminal status conflicts with existing terminal")
|
||||
return
|
||||
if job.status not in {"offered", "dispatched", "running", "disconnected"}:
|
||||
raise ComputeNodeError("job terminal cannot advance from its current status")
|
||||
job.status = terminal_status
|
||||
job.stage = "terminal"
|
||||
job.progress = 100 if terminal_status == "succeeded" else job.progress
|
||||
job.error = error
|
||||
job.artifact_manifest = manifest
|
||||
job.terminal_at = now
|
||||
|
||||
|
||||
def _is_uuid(value: str) -> bool:
|
||||
try:
|
||||
UUID(value)
|
||||
except ValueError:
|
||||
return False
|
||||
return True
|
||||
|
||||
|
||||
def mark_node_jobs_disconnected(node_id: UUID) -> None:
|
||||
"""连接丢失后保留 Node 归属和 lease，禁止任务被自动重派。"""
|
||||
with session_scope() as session:
|
||||
jobs = session.execute(
|
||||
select(ComputeJob)
|
||||
.where(
|
||||
ComputeJob.node_id == node_id,
|
||||
ComputeJob.status.in_({"dispatched", "running"}),
|
||||
)
|
||||
.with_for_update()
|
||||
).scalars()
|
||||
for job in jobs:
|
||||
job.status = "disconnected"
|
||||
|
||||
|
||||
def _can_accept_state(status: str) -> bool:
|
||||
return status in {"offered", "dispatched", "running", "disconnected"}
|
||||
|
||||
|
||||
def _message_identity(payload: dict) -> tuple[UUID, UUID, str]:
|
||||
try:
|
||||
job_id = UUID(str(payload.get("job_id", "")))
|
||||
lease_id = UUID(str(payload.get("lease_id", "")))
|
||||
except ValueError as exc:
|
||||
raise ComputeNodeError("invalid job message identity") from exc
|
||||
digest = str(payload.get("request_digest") or "")
|
||||
if len(digest) != 64:
|
||||
raise ComputeNodeError("invalid job request digest")
|
||||
return job_id, lease_id, digest
|
||||
|
||||
|
||||
def _assert_job_message(
|
||||
job: ComputeJob | None,
|
||||
node_id: UUID,
|
||||
lease_id: UUID,
|
||||
digest: str,
|
||||
) -> None:
|
||||
if (
|
||||
job is None
|
||||
or job.node_id != node_id
|
||||
or job.lease_id != lease_id
|
||||
or job.request_digest != digest
|
||||
):
|
||||
raise ComputeNodeError("job message does not belong to this node or lease")
|
||||
|
|
@ -465,6 +465,49 @@ class ComputeNode(Base):
|
|||
)
|
||||
|
||||
|
||||
class ComputeJob(Base):
|
||||
"""受控计算任务账本；请求只保存规范化参数和输入引用。"""
|
||||
|
||||
__tablename__ = "compute_jobs"
|
||||
__table_args__ = (
|
||||
UniqueConstraint("user_id", "idempotency_key", name="uq_compute_jobs_user_idempotency"),
|
||||
Index("ix_compute_jobs_status_created", "status", "created_at"),
|
||||
Index("ix_compute_jobs_node_status", "node_id", "status"),
|
||||
)
|
||||
|
||||
job_id: Mapped[UUID] = mapped_column(PG_UUID(as_uuid=True), primary_key=True, default=uuid4)
|
||||
user_id: Mapped[UUID] = mapped_column(
|
||||
PG_UUID(as_uuid=True), ForeignKey("users.user_id", ondelete="CASCADE"), nullable=False
|
||||
)
|
||||
task_id: Mapped[UUID] = mapped_column(
|
||||
PG_UUID(as_uuid=True), ForeignKey("tasks.task_id", ondelete="CASCADE"), nullable=False
|
||||
)
|
||||
idempotency_key: Mapped[str] = mapped_column(Text, nullable=False)
|
||||
capability: Mapped[str] = mapped_column(Text, nullable=False)
|
||||
request: Mapped[dict[str, Any]] = mapped_column(JSONB, nullable=False)
|
||||
request_digest: Mapped[str] = mapped_column(Text, nullable=False)
|
||||
input_manifest: Mapped[dict[str, Any]] = mapped_column(JSONB, nullable=False, default=dict)
|
||||
node_id: Mapped[Optional[UUID]] = mapped_column(
|
||||
PG_UUID(as_uuid=True), ForeignKey("compute_nodes.node_id", ondelete="SET NULL"), nullable=True
|
||||
)
|
||||
lease_id: Mapped[Optional[UUID]] = mapped_column(PG_UUID(as_uuid=True), nullable=True)
|
||||
lease_expires_at: Mapped[Optional[datetime]] = mapped_column(DateTime(timezone=True), nullable=True)
|
||||
status: Mapped[str] = mapped_column(Text, nullable=False, default="queued", server_default="queued")
|
||||
stage: Mapped[str] = mapped_column(Text, nullable=False, default="", server_default="")
|
||||
progress: Mapped[int] = mapped_column(Integer, nullable=False, default=0, server_default="0")
|
||||
metrics: Mapped[dict[str, Any]] = mapped_column(JSONB, nullable=False, default=dict)
|
||||
error: Mapped[dict[str, Any]] = mapped_column(JSONB, nullable=False, default=dict)
|
||||
artifact_manifest: Mapped[list[Any]] = mapped_column(JSONB, nullable=False, default=list)
|
||||
started_at: Mapped[Optional[datetime]] = mapped_column(DateTime(timezone=True), nullable=True)
|
||||
terminal_at: Mapped[Optional[datetime]] = mapped_column(DateTime(timezone=True), nullable=True)
|
||||
created_at: Mapped[datetime] = mapped_column(
|
||||
DateTime(timezone=True), server_default=func.now(), nullable=False
|
||||
)
|
||||
updated_at: Mapped[datetime] = mapped_column(
|
||||
DateTime(timezone=True), server_default=func.now(), onupdate=func.now(), nullable=False
|
||||
)
|
||||
|
||||
|
||||
class ExternalSystemDefinition(Base):
|
||||
"""管理员维护的可信外部系统目录；不含任何用户凭据。"""
|
||||
|
||||
|
|
|
|||
|
|
@ -0,0 +1,52 @@
|
|||
"""Add the Windows compute job ledger.
|
||||
|
||||
Revision ID: 0032
|
||||
Revises: 0031
|
||||
Create Date: 2026-08-13
|
||||
"""
|
||||
from collections.abc import Sequence
|
||||
|
||||
import sqlalchemy as sa
|
||||
from alembic import op
|
||||
from sqlalchemy.dialects import postgresql
|
||||
|
||||
revision: str = "0032"
|
||||
down_revision: str | None = "0031"
|
||||
branch_labels: str | Sequence[str] | None = None
|
||||
depends_on: str | Sequence[str] | None = None
|
||||
|
||||
|
||||
def upgrade() -> None:
|
||||
op.create_table(
|
||||
"compute_jobs",
|
||||
sa.Column("job_id", postgresql.UUID(as_uuid=True), primary_key=True),
|
||||
sa.Column("user_id", postgresql.UUID(as_uuid=True), sa.ForeignKey("users.user_id", ondelete="CASCADE"), nullable=False),
|
||||
sa.Column("task_id", postgresql.UUID(as_uuid=True), sa.ForeignKey("tasks.task_id", ondelete="CASCADE"), nullable=False),
|
||||
sa.Column("idempotency_key", sa.Text(), nullable=False),
|
||||
sa.Column("capability", sa.Text(), nullable=False),
|
||||
sa.Column("request", postgresql.JSONB(), nullable=False),
|
||||
sa.Column("request_digest", sa.Text(), nullable=False),
|
||||
sa.Column("input_manifest", postgresql.JSONB(), nullable=False),
|
||||
sa.Column("node_id", postgresql.UUID(as_uuid=True), sa.ForeignKey("compute_nodes.node_id", ondelete="SET NULL"), nullable=True),
|
||||
sa.Column("lease_id", postgresql.UUID(as_uuid=True), nullable=True),
|
||||
sa.Column("lease_expires_at", sa.DateTime(timezone=True), nullable=True),
|
||||
sa.Column("status", sa.Text(), server_default="queued", nullable=False),
|
||||
sa.Column("stage", sa.Text(), server_default="", nullable=False),
|
||||
sa.Column("progress", sa.Integer(), server_default="0", nullable=False),
|
||||
sa.Column("metrics", postgresql.JSONB(), nullable=False),
|
||||
sa.Column("error", postgresql.JSONB(), nullable=False),
|
||||
sa.Column("artifact_manifest", postgresql.JSONB(), nullable=False),
|
||||
sa.Column("started_at", sa.DateTime(timezone=True), nullable=True),
|
||||
sa.Column("terminal_at", sa.DateTime(timezone=True), nullable=True),
|
||||
sa.Column("created_at", sa.DateTime(timezone=True), server_default=sa.func.now(), nullable=False),
|
||||
sa.Column("updated_at", sa.DateTime(timezone=True), server_default=sa.func.now(), nullable=False),
|
||||
sa.UniqueConstraint("user_id", "idempotency_key", name="uq_compute_jobs_user_idempotency"),
|
||||
)
|
||||
op.create_index("ix_compute_jobs_status_created", "compute_jobs", ["status", "created_at"])
|
||||
op.create_index("ix_compute_jobs_node_status", "compute_jobs", ["node_id", "status"])
|
||||
|
||||
|
||||
def downgrade() -> None:
|
||||
op.drop_index("ix_compute_jobs_node_status", table_name="compute_jobs")
|
||||
op.drop_index("ix_compute_jobs_status_created", table_name="compute_jobs")
|
||||
op.drop_table("compute_jobs")
|
||||
|
|
@ -200,6 +200,8 @@ Node 断线且本地任务可能仍在执行时标记 `disconnected`，不得自
|
|||
|
||||
## 7. Origin 任务闭环
|
||||
|
||||
当前实现进度：云端任务账本、幂等提交、短期 offer、Node 本地原子保存与 accept/reject 已落地。输入以任务绑定的 artifact UUID 下载，Node 流式校验大小和 SHA-256 后原子保存。固定 Worker 使用管理员安装的隔离 Python 运行时与随程序发布的 `worker.py` 驱动 Origin，生成 OPJU、PNG、SVG、PDF、plot spec、provenance 和原子 `terminal.json`；运行不绑定单次 WebSocket，断线后继续执行。同一进程按 job 去重，Node 重启后不重复启动已留启动标记但无可信终态的任务。成功产物逐项流式上传到云端隐藏暂存区，云端复核任务身份、固定文件名、大小和 SHA-256 后，一次性发布到 `<working_dir>/origin/<job_id>/` 并登记平台 artifact UUID；Node 以 `upload-complete.json` 恢复中断上传。
|
||||
|
||||
```text
|
||||
用户上传 CSV/XLSX
|
||||
→ zcbot 生成受控 plot spec
|
||||
|
|
|
|||
|
|
@ -17,6 +17,15 @@ from core.compute_nodes import (
|
|||
_verify_secret,
|
||||
delete_node,
|
||||
)
|
||||
from core.compute_jobs import (
|
||||
_canonical_request,
|
||||
abandon_offer,
|
||||
mark_node_jobs_disconnected,
|
||||
record_job_terminal,
|
||||
respond_to_offer,
|
||||
update_job_state,
|
||||
validate_output_manifest,
|
||||
)
|
||||
from web.routers.compute_nodes import NodeConnectionManager, _bearer
|
||||
|
||||
|
||||
|
|
@ -98,6 +107,204 @@ class ComputeNodeMigrationTests(unittest.TestCase):
|
|||
self.assertIn("compute_nodes", rendered)
|
||||
self.assertIn("ix_compute_nodes_status", rendered)
|
||||
|
||||
def test_0032_upgrade_compiles_as_postgresql_ddl(self) -> None:
|
||||
statements: list[str] = []
|
||||
|
||||
def capture(sql, *multiparams, **params):
|
||||
statements.append(str(sql.compile(dialect=postgresql.dialect())))
|
||||
|
||||
engine = create_mock_engine("postgresql+psycopg://", capture)
|
||||
operations = Operations(MigrationContext.configure(engine.connect()))
|
||||
migration = importlib.import_module(
|
||||
"db.migrations.versions.20260813_1600_0032_compute_jobs"
|
||||
)
|
||||
with patch.object(migration, "op", operations):
|
||||
migration.upgrade()
|
||||
|
||||
rendered = "\n".join(statements)
|
||||
self.assertIn("compute_jobs", rendered)
|
||||
self.assertIn("uq_compute_jobs_user_idempotency", rendered)
|
||||
self.assertIn("ix_compute_jobs_status_created", rendered)
|
||||
|
||||
|
||||
class ComputeJobProtocolTests(unittest.TestCase):
|
||||
def test_origin_request_is_canonical_and_rejects_extra_fields(self) -> None:
|
||||
request = {
|
||||
"schema_version": 1,
|
||||
"input": {"input_id": str(uuid4()), "sheet": "Sheet1"},
|
||||
"plot": {"type": "line", "x": "x", "y": ["y"]},
|
||||
"output": {"formats": ["png", "opju"]},
|
||||
}
|
||||
normalized, digest = _canonical_request(request)
|
||||
self.assertEqual(normalized, request)
|
||||
self.assertEqual(len(digest), 64)
|
||||
with self.assertRaisesRegex(Exception, "invalid origin plot request fields"):
|
||||
_canonical_request({**request, "script": "anything"})
|
||||
with self.assertRaisesRegex(Exception, "unsupported origin plot fields"):
|
||||
_canonical_request({**request, "plot": {**request["plot"], "script": "anything"}})
|
||||
with self.assertRaisesRegex(Exception, "input.input_id must be an artifact UUID"):
|
||||
_canonical_request({**request, "input": {"input_id": "C:\\data.csv"}})
|
||||
|
||||
def test_origin_request_rejects_unimplemented_plot_semantics(self) -> None:
|
||||
request = {
|
||||
"schema_version": 1,
|
||||
"input": {"input_id": str(uuid4())},
|
||||
"plot": {"type": "scatter", "x": "time", "y": ["a", "b"]},
|
||||
"output": {"formats": ["png"], "dpi": 300},
|
||||
}
|
||||
for plot, message in (
|
||||
({**request["plot"], "template": "custom"}, "unsupported origin plot template"),
|
||||
({**request["plot"], "x_axis": {"scale": "log10"}}, "invalid x_axis"),
|
||||
({**request["plot"], "legend": {"enabled": False}}, "invalid plot.legend"),
|
||||
({**request["plot"], "y": ["a", "a"]}, "plot.y must contain"),
|
||||
):
|
||||
with self.subTest(message=message), self.assertRaisesRegex(Exception, message):
|
||||
_canonical_request({**request, "plot": plot})
|
||||
with self.assertRaisesRegex(Exception, "video recording is not supported"):
|
||||
_canonical_request(
|
||||
{**request, "output": {"formats": ["png"], "record_video": True}}
|
||||
)
|
||||
|
||||
def test_output_manifest_matches_exact_requested_formats(self) -> None:
|
||||
request = {"output": {"formats": ["opju", "png"]}}
|
||||
manifest = [
|
||||
{"artifact_id": "project", "filename": "project.opju", "media_type": "application/x-origin-project", "size_bytes": 10, "sha256": "a" * 64},
|
||||
{"artifact_id": "figure_png", "filename": "figure.png", "media_type": "image/png", "size_bytes": 20, "sha256": "b" * 64},
|
||||
{"artifact_id": "plot_spec", "filename": "plot-spec.json", "media_type": "application/json", "size_bytes": 30, "sha256": "c" * 64},
|
||||
{"artifact_id": "provenance", "filename": "provenance.json", "media_type": "application/json", "size_bytes": 40, "sha256": "d" * 64},
|
||||
]
|
||||
self.assertEqual(validate_output_manifest(request, manifest), manifest)
|
||||
with self.assertRaisesRegex(Exception, "incomplete"):
|
||||
validate_output_manifest(request, manifest[:-1])
|
||||
with self.assertRaisesRegex(Exception, "metadata"):
|
||||
validate_output_manifest(
|
||||
request,
|
||||
[{**manifest[0], "filename": "anything.opju"}, *manifest[1:]],
|
||||
)
|
||||
|
||||
@patch("core.compute_jobs.session_scope")
|
||||
def test_stale_offer_cannot_be_accepted_by_another_node(self, session_scope) -> None:
|
||||
session = session_scope.return_value.__enter__.return_value
|
||||
job = type("Job", (), {})()
|
||||
job.node_id = uuid4()
|
||||
job.lease_id = uuid4()
|
||||
job.status = "offered"
|
||||
session.execute.return_value.scalar_one_or_none.return_value = job
|
||||
with self.assertRaisesRegex(Exception, "stale or does not belong"):
|
||||
respond_to_offer(
|
||||
uuid4(),
|
||||
accepted=True,
|
||||
payload={"job_id": str(uuid4()), "lease_id": str(job.lease_id)},
|
||||
)
|
||||
|
||||
@patch("core.compute_jobs.session_scope")
|
||||
def test_failed_delivery_only_abandons_matching_offer(self, session_scope) -> None:
|
||||
session = session_scope.return_value.__enter__.return_value
|
||||
node_id = uuid4()
|
||||
lease_id = uuid4()
|
||||
job = type("Job", (), {})()
|
||||
job.node_id = node_id
|
||||
job.lease_id = lease_id
|
||||
job.status = "offered"
|
||||
session.execute.return_value.scalar_one_or_none.return_value = job
|
||||
abandon_offer(
|
||||
node_id,
|
||||
{"job_id": str(uuid4()), "lease_id": str(lease_id)},
|
||||
)
|
||||
self.assertEqual(job.status, "queued")
|
||||
self.assertIsNone(job.node_id)
|
||||
|
||||
def test_dispatcher_excludes_nodes_with_active_jobs(self) -> None:
|
||||
source = (
|
||||
Path(__file__).resolve().parents[1] / "core" / "compute_jobs.py"
|
||||
).read_text(encoding="utf-8")
|
||||
self.assertIn('{"offered", "dispatched", "running"}', source)
|
||||
self.assertIn("item.node_id not in busy_node_ids", source)
|
||||
|
||||
def test_input_download_rechecks_file_digest(self) -> None:
|
||||
source = (
|
||||
Path(__file__).resolve().parents[1]
|
||||
/ "web" / "routers" / "compute_nodes.py"
|
||||
).read_text(encoding="utf-8")
|
||||
self.assertIn("digest = sha256()", source)
|
||||
self.assertIn('digest.hexdigest() != item["sha256"]', source)
|
||||
|
||||
@patch("core.compute_jobs.session_scope")
|
||||
def test_job_state_restores_disconnected_job(self, session_scope) -> None:
|
||||
session = session_scope.return_value.__enter__.return_value
|
||||
node_id = uuid4()
|
||||
lease_id = uuid4()
|
||||
digest = "a" * 64
|
||||
job = type("Job", (), {})()
|
||||
job.node_id = node_id
|
||||
job.lease_id = lease_id
|
||||
job.request_digest = digest
|
||||
job.status = "disconnected"
|
||||
job.started_at = None
|
||||
session.execute.return_value.scalar_one_or_none.return_value = job
|
||||
update_job_state(node_id, {
|
||||
"job_id": str(uuid4()),
|
||||
"lease_id": str(lease_id),
|
||||
"request_digest": digest,
|
||||
"stage": "waiting_input",
|
||||
"progress": 0,
|
||||
"metrics": {},
|
||||
})
|
||||
self.assertEqual(job.status, "dispatched")
|
||||
self.assertEqual(job.stage, "waiting_input")
|
||||
|
||||
@patch("core.compute_jobs.session_scope")
|
||||
def test_ready_to_run_is_not_reported_as_running(self, session_scope) -> None:
|
||||
session = session_scope.return_value.__enter__.return_value
|
||||
node_id = uuid4()
|
||||
lease_id = uuid4()
|
||||
digest = "c" * 64
|
||||
job = type("Job", (), {})()
|
||||
job.node_id = node_id
|
||||
job.lease_id = lease_id
|
||||
job.request_digest = digest
|
||||
job.status = "dispatched"
|
||||
job.started_at = None
|
||||
session.execute.return_value.scalar_one_or_none.return_value = job
|
||||
update_job_state(node_id, {
|
||||
"job_id": str(uuid4()), "lease_id": str(lease_id),
|
||||
"request_digest": digest, "stage": "ready_to_run",
|
||||
"progress": 5, "metrics": {"input_bytes": 10},
|
||||
})
|
||||
self.assertEqual(job.status, "dispatched")
|
||||
|
||||
@patch("core.compute_jobs.session_scope")
|
||||
def test_terminal_replay_is_idempotent(self, session_scope) -> None:
|
||||
session = session_scope.return_value.__enter__.return_value
|
||||
node_id = uuid4()
|
||||
lease_id = uuid4()
|
||||
digest = "b" * 64
|
||||
job = type("Job", (), {})()
|
||||
job.node_id = node_id
|
||||
job.lease_id = lease_id
|
||||
job.request_digest = digest
|
||||
job.status = "failed"
|
||||
session.execute.return_value.scalar_one_or_none.return_value = job
|
||||
record_job_terminal(node_id, {
|
||||
"job_id": str(uuid4()),
|
||||
"lease_id": str(lease_id),
|
||||
"request_digest": digest,
|
||||
"status": "failed",
|
||||
"error": {"code": "TEST"},
|
||||
"artifact_manifest": [],
|
||||
})
|
||||
self.assertEqual(job.status, "failed")
|
||||
|
||||
@patch("core.compute_jobs.session_scope")
|
||||
def test_disconnect_does_not_requeue_active_jobs(self, session_scope) -> None:
|
||||
session = session_scope.return_value.__enter__.return_value
|
||||
first = type("Job", (), {"status": "running"})()
|
||||
second = type("Job", (), {"status": "dispatched"})()
|
||||
session.execute.return_value.scalars.return_value = [first, second]
|
||||
mark_node_jobs_disconnected(uuid4())
|
||||
self.assertEqual(first.status, "disconnected")
|
||||
self.assertEqual(second.status, "disconnected")
|
||||
|
||||
|
||||
class ComputeNodeDeleteTests(unittest.TestCase):
|
||||
@patch("core.compute_nodes.session_scope")
|
||||
|
|
|
|||
|
|
@ -0,0 +1,65 @@
|
|||
from __future__ import annotations
|
||||
|
||||
import hashlib
|
||||
import tempfile
|
||||
import unittest
|
||||
from pathlib import Path
|
||||
from unittest.mock import patch
|
||||
from uuid import uuid4
|
||||
|
||||
from web.routers.compute_nodes import _publish_compute_outputs
|
||||
|
||||
|
||||
class ComputeOutputPublishTests(unittest.TestCase):
|
||||
def test_complete_set_moves_atomically_and_can_be_replayed(self) -> None:
|
||||
with tempfile.TemporaryDirectory() as directory:
|
||||
root = Path(directory)
|
||||
job_id = uuid4()
|
||||
working_dir = root / "research"
|
||||
staging = root / ".zcbot_compute_staging" / str(job_id)
|
||||
staging.mkdir(parents=True)
|
||||
working_dir.mkdir()
|
||||
content = b"origin-result"
|
||||
(staging / "figure.png").write_bytes(content)
|
||||
manifest = [{
|
||||
"artifact_id": "figure_png",
|
||||
"filename": "figure.png",
|
||||
"media_type": "image/png",
|
||||
"size_bytes": len(content),
|
||||
"sha256": hashlib.sha256(content).hexdigest(),
|
||||
}]
|
||||
context = {
|
||||
"user_id": uuid4(),
|
||||
"task_id": uuid4(),
|
||||
"working_dir": "research",
|
||||
}
|
||||
|
||||
def register(**kwargs):
|
||||
return tuple({
|
||||
"version": 2,
|
||||
"scope": "working_dir",
|
||||
"path": ref["path"],
|
||||
"label": ref["label"],
|
||||
"artifact_id": str(uuid4()),
|
||||
} for ref in kwargs["refs"])
|
||||
|
||||
with (
|
||||
patch("web.routers.compute_nodes.load_user_root", return_value=root),
|
||||
patch(
|
||||
"web.routers.compute_nodes.register_published_artifacts",
|
||||
side_effect=register,
|
||||
),
|
||||
):
|
||||
first = _publish_compute_outputs(job_id, context, manifest)
|
||||
second = _publish_compute_outputs(job_id, context, manifest)
|
||||
|
||||
published = working_dir / "origin" / str(job_id) / "figure.png"
|
||||
self.assertEqual(published.read_bytes(), content)
|
||||
self.assertFalse(staging.exists())
|
||||
self.assertEqual(first[0]["source_artifact_id"], "figure_png")
|
||||
self.assertEqual(first[0]["path"], f"origin/{job_id}/figure.png")
|
||||
self.assertEqual(second[0]["source_artifact_id"], "figure_png")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
|
|
@ -0,0 +1,46 @@
|
|||
from __future__ import annotations
|
||||
|
||||
import importlib.util
|
||||
import json
|
||||
import tempfile
|
||||
import unittest
|
||||
from pathlib import Path
|
||||
|
||||
|
||||
WORKER_PATH = (
|
||||
Path(__file__).resolve().parents[1] / "windows-node" / "origin-worker" / "worker.py"
|
||||
)
|
||||
SPEC = importlib.util.spec_from_file_location("zcbot_origin_worker", WORKER_PATH)
|
||||
assert SPEC and SPEC.loader
|
||||
worker = importlib.util.module_from_spec(SPEC)
|
||||
SPEC.loader.exec_module(worker)
|
||||
|
||||
|
||||
class OriginWorkerUnitTests(unittest.TestCase):
|
||||
def test_csv_and_json_inputs_are_read_without_origin(self) -> None:
|
||||
with tempfile.TemporaryDirectory() as directory:
|
||||
root = Path(directory)
|
||||
csv_path = root / "input.csv"
|
||||
csv_path.write_text("x,y\n1,2\n3,4\n", encoding="utf-8")
|
||||
self.assertEqual(worker._read_rows(csv_path, None), (["x", "y"], [["1", "2"], ["3", "4"]]))
|
||||
|
||||
json_path = root / "input.json"
|
||||
json_path.write_text(json.dumps([{"x": 1, "y": 2}, {"x": 3, "y": 4}]), encoding="utf-8")
|
||||
self.assertEqual(worker._read_rows(json_path, None), (["x", "y"], [[1, 2], [3, 4]]))
|
||||
|
||||
def test_manifest_uses_stable_id_and_streaming_digest(self) -> None:
|
||||
with tempfile.TemporaryDirectory() as directory:
|
||||
path = Path(directory) / "plot-spec.json"
|
||||
path.write_text("{}", encoding="utf-8")
|
||||
manifest = worker._manifest(path, "application/json")
|
||||
self.assertEqual(manifest["artifact_id"], "plot_spec")
|
||||
self.assertEqual(manifest["sha256"], worker._file_sha256(path))
|
||||
self.assertEqual(manifest["size_bytes"], 2)
|
||||
|
||||
def test_axis_title_includes_units(self) -> None:
|
||||
self.assertEqual(worker._axis_title({"title": "Stress", "unit": "MPa"}, "Y"), "Stress (MPa)")
|
||||
self.assertEqual(worker._axis_title(None, "Time"), "Time")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
|
|
@ -35,7 +35,11 @@ class WindowsNodeSourceTests(unittest.TestCase):
|
|||
self.assertIn(marker, source)
|
||||
|
||||
def test_node_does_not_expose_arbitrary_execution_primitives(self) -> None:
|
||||
source = "\n".join(path.read_text(encoding="utf-8") for path in PROJECT.glob("*.cs"))
|
||||
source = "\n".join(
|
||||
path.read_text(encoding="utf-8")
|
||||
for path in PROJECT.glob("*.cs")
|
||||
if path.name != "OriginWorkerRunner.cs"
|
||||
)
|
||||
for forbidden in ("Process.Start", "cmd.exe", "powershell.exe", "LabTalk"):
|
||||
self.assertNotIn(forbidden, source)
|
||||
|
||||
|
|
@ -102,8 +106,11 @@ class WindowsNodeSourceTests(unittest.TestCase):
|
|||
self.assertIn('AutomationProgId = @"Origin.ApplicationSI\\CLSID"', probe)
|
||||
self.assertIn("RegistryHive.LocalMachine", probe)
|
||||
self.assertIn("RegistryHive.CurrentUser", probe)
|
||||
self.assertIn('new("OriginPro", version, "0.1.0", health, detail)', probe)
|
||||
self.assertIn('available_slots = origin.Health == "ready" ? 1 : 0', connection)
|
||||
self.assertIn('new("OriginPro", version, "0.2.0", health, detail)', probe)
|
||||
self.assertIn(
|
||||
'available_slots = origin.Health == "ready" && !jobInbox.HasPendingJobs ? 1 : 0',
|
||||
connection,
|
||||
)
|
||||
self.assertNotIn("CreateInstance", probe)
|
||||
self.assertNotIn("Process.Start", probe)
|
||||
for marker in (
|
||||
|
|
@ -114,6 +121,74 @@ class WindowsNodeSourceTests(unittest.TestCase):
|
|||
):
|
||||
self.assertIn(marker, connection)
|
||||
|
||||
def test_job_offer_is_persisted_before_acceptance(self) -> None:
|
||||
inbox = (PROJECT / "JobInboxStore.cs").read_text(encoding="utf-8")
|
||||
connection = (PROJECT / "NodeConnectionLoop.cs").read_text(encoding="utf-8")
|
||||
self.assertIn('capabilityValue.GetString() != "origin.plot@v1"', inbox)
|
||||
self.assertIn("PlotTypes.Contains", inbox)
|
||||
self.assertIn("OutputFormats.Contains", inbox)
|
||||
self.assertIn("FileOptions.WriteThrough", inbox)
|
||||
self.assertIn("stream.Flush(flushToDisk: true)", inbox)
|
||||
new_record = inbox.split("var record =", 1)[1].split("private static JsonElement?", 1)[0]
|
||||
self.assertLess(
|
||||
new_record.index("AtomicWrite(requestPath, record"),
|
||||
new_record.index("JobOfferResult.Accept"),
|
||||
)
|
||||
self.assertIn('offerResult.Accepted ? "job_accept" : "job_reject"', connection)
|
||||
self.assertIn("sendLock.WaitAsync", connection)
|
||||
self.assertIn("!jobInbox.HasPendingJobs ? 1 : 0", connection)
|
||||
self.assertIn("ReportRecoverableJobsAsync", connection)
|
||||
self.assertIn("ConcurrentDictionary<Guid, Task> jobPipelines", connection)
|
||||
self.assertIn("StartJobPipeline(socket, acceptedJob)", connection)
|
||||
self.assertIn("inputDownloader.DownloadAsync(job, CancellationToken.None)", connection)
|
||||
self.assertIn('stage = "uploading_outputs"', connection)
|
||||
self.assertIn("&& !job.UploadComplete", connection)
|
||||
self.assertIn("StartJobPipeline(socket, job)", connection)
|
||||
self.assertIn('stage = "waiting_input"', connection)
|
||||
self.assertIn('Path.Combine(jobDirectory, "terminal.json")', inbox)
|
||||
self.assertIn("AtomicWrite(requestPath, updated, overwrite: true)", inbox)
|
||||
|
||||
downloader = (PROJECT / "JobInputDownloader.cs").read_text(encoding="utf-8")
|
||||
self.assertIn('new AuthenticationHeaderValue("Bearer", config.NodeToken)', downloader)
|
||||
self.assertIn('DefaultRequestHeaders.Add("X-Node-Id"', downloader)
|
||||
self.assertIn("HttpCompletionOption.ResponseHeadersRead", downloader)
|
||||
self.assertIn("IncrementalHash.CreateHash", downloader)
|
||||
self.assertIn("total > expectedSize", downloader)
|
||||
self.assertIn("File.Move(temporaryPath, destination, overwrite: false)", downloader)
|
||||
self.assertNotIn("Process.Start", downloader)
|
||||
|
||||
def test_origin_worker_launch_is_fixed_and_terminal_driven(self) -> None:
|
||||
runner = (PROJECT / "OriginWorkerRunner.cs").read_text(encoding="utf-8")
|
||||
connection = (PROJECT / "NodeConnectionLoop.cs").read_text(encoding="utf-8")
|
||||
project = (PROJECT / "Zcbot.WindowsNode.csproj").read_text(encoding="utf-8")
|
||||
worker = (ROOT / "origin-worker" / "worker.py").read_text(encoding="utf-8")
|
||||
self.assertIn('Environment.GetEnvironmentVariable("ZCBOT_ORIGIN_PYTHON")', runner)
|
||||
self.assertIn('Path.Combine(paths.RootDirectory, "runtimes", "origin", "python.exe")', runner)
|
||||
self.assertIn("UseShellExecute = false", runner)
|
||||
self.assertIn("startInfo.ArgumentList.Add(workerScript)", runner)
|
||||
self.assertIn("startInfo.ArgumentList.Add(jobDirectory)", runner)
|
||||
self.assertIn('Path.Combine(jobDirectory, "terminal.json")', runner)
|
||||
self.assertIn('"NODE_RESTARTED_DURING_JOB"', runner)
|
||||
self.assertNotIn("RunAsync(RecoverableJob job, CancellationToken", runner)
|
||||
self.assertIn("origin-worker\\worker.py", project)
|
||||
self.assertIn("if op.oext:", worker)
|
||||
self.assertIn("op.exit()", worker)
|
||||
self.assertIn("op.new_graph", worker)
|
||||
self.assertIn("layer.add_plot", worker)
|
||||
self.assertIn("op.save", worker)
|
||||
self.assertIn("graph.save_fig", worker)
|
||||
self.assertIn('_atomic_json(job_dir / "terminal.json"', worker)
|
||||
for forbidden in ("subprocess", "eval(", "exec(", "os.system"):
|
||||
self.assertNotIn(forbidden, worker)
|
||||
|
||||
uploader = (PROJECT / "JobOutputUploader.cs").read_text(encoding="utf-8")
|
||||
self.assertIn('new AuthenticationHeaderValue("Bearer", config.NodeToken)', uploader)
|
||||
self.assertIn('DefaultRequestHeaders.Add("X-Node-Id"', uploader)
|
||||
self.assertIn('DefaultRequestHeaders.Add("X-Lease-Id"', uploader)
|
||||
self.assertIn("SHA256.HashDataAsync", uploader)
|
||||
self.assertIn("upload-complete.json", connection + uploader)
|
||||
self.assertNotIn("Process.Start", uploader)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
|
|
|
|||
|
|
@ -206,7 +206,9 @@ def create_app() -> FastAPI:
|
|||
register_asr_routes(app, require_user=require_user, auth_cfg=auth_cfg)
|
||||
register_task_routes(app, require_user=require_user)
|
||||
register_message_routes(app, require_user=require_user)
|
||||
register_compute_node_routes(app, require_admin=require_admin)
|
||||
register_compute_node_routes(
|
||||
app, require_user=require_user, require_admin=require_admin
|
||||
)
|
||||
|
||||
# ───────────── 管理后台(admin-only)─────────────
|
||||
register_admin_routes(app, require_admin)
|
||||
|
|
|
|||
|
|
@ -3,9 +3,13 @@
|
|||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import os
|
||||
from hashlib import sha256
|
||||
from pathlib import Path
|
||||
from uuid import UUID
|
||||
|
||||
from fastapi import Depends, HTTPException, WebSocket, WebSocketDisconnect, status
|
||||
from fastapi import Depends, Header, HTTPException, Request, WebSocket, WebSocketDisconnect, status
|
||||
from fastapi.responses import FileResponse
|
||||
|
||||
from core.compute_nodes import (
|
||||
ComputeNodeError,
|
||||
|
|
@ -18,22 +22,43 @@ from core.compute_nodes import (
|
|||
set_node_disabled,
|
||||
update_node_runtime,
|
||||
)
|
||||
from core.compute_jobs import (
|
||||
MAX_OUTPUT_ARTIFACT_BYTES,
|
||||
MAX_OUTPUT_TOTAL_BYTES,
|
||||
OUTPUT_ARTIFACTS,
|
||||
abandon_offer,
|
||||
create_job,
|
||||
get_job,
|
||||
get_job_input,
|
||||
get_job_output_context,
|
||||
mark_node_jobs_disconnected,
|
||||
offer_next_job,
|
||||
record_job_terminal,
|
||||
respond_to_offer,
|
||||
update_job_state,
|
||||
validate_output_manifest,
|
||||
)
|
||||
from core.artifact_lifecycle import register_published_artifacts
|
||||
from web.schemas import (
|
||||
ComputeEnrollmentCreateRequest,
|
||||
ComputeJobCreateRequest,
|
||||
ComputeNodeDisableRequest,
|
||||
ComputeNodeEnrollRequest,
|
||||
)
|
||||
from web.userfiles import load_user_root, safe_join
|
||||
|
||||
|
||||
class NodeConnectionManager:
|
||||
def __init__(self) -> None:
|
||||
self._connections: dict[UUID, WebSocket] = {}
|
||||
self._send_locks: dict[UUID, asyncio.Lock] = {}
|
||||
self._lock = asyncio.Lock()
|
||||
|
||||
async def activate(self, node_id: UUID, websocket: WebSocket) -> None:
|
||||
async with self._lock:
|
||||
old = self._connections.get(node_id)
|
||||
self._connections[node_id] = websocket
|
||||
self._send_locks.setdefault(node_id, asyncio.Lock())
|
||||
if old is not None and old is not websocket:
|
||||
await old.close(code=4001, reason="replaced by a newer connection")
|
||||
|
||||
|
|
@ -41,15 +66,41 @@ class NodeConnectionManager:
|
|||
async with self._lock:
|
||||
if self._connections.get(node_id) is websocket:
|
||||
self._connections.pop(node_id, None)
|
||||
self._send_locks.pop(node_id, None)
|
||||
return True
|
||||
return False
|
||||
|
||||
async def close(self, node_id: UUID) -> None:
|
||||
async with self._lock:
|
||||
websocket = self._connections.pop(node_id, None)
|
||||
self._send_locks.pop(node_id, None)
|
||||
if websocket is not None:
|
||||
await websocket.close(code=4003, reason="node disabled")
|
||||
|
||||
async def node_ids(self) -> set[UUID]:
|
||||
async with self._lock:
|
||||
return set(self._connections)
|
||||
|
||||
async def send(self, node_id: UUID, message: dict) -> bool:
|
||||
async with self._lock:
|
||||
websocket = self._connections.get(node_id)
|
||||
send_lock = self._send_locks.get(node_id)
|
||||
if websocket is None or send_lock is None:
|
||||
return False
|
||||
async with send_lock:
|
||||
await websocket.send_json(message)
|
||||
return True
|
||||
|
||||
async def send_on(self, node_id: UUID, websocket: WebSocket, message: dict) -> bool:
|
||||
async with self._lock:
|
||||
current = self._connections.get(node_id)
|
||||
send_lock = self._send_locks.get(node_id)
|
||||
if current is not websocket or send_lock is None:
|
||||
return False
|
||||
async with send_lock:
|
||||
await websocket.send_json(message)
|
||||
return True
|
||||
|
||||
|
||||
node_connections = NodeConnectionManager()
|
||||
|
||||
|
|
@ -61,7 +112,93 @@ def _bearer(authorization: str | None) -> str:
|
|||
return token
|
||||
|
||||
|
||||
def register_compute_node_routes(app, *, require_admin) -> None:
|
||||
def _authenticate_output_request(
|
||||
job_id: UUID,
|
||||
authorization: str | None,
|
||||
x_node_id: str,
|
||||
x_lease_id: str,
|
||||
x_request_digest: str,
|
||||
) -> tuple[UUID, UUID, dict]:
|
||||
try:
|
||||
node_id = UUID(x_node_id)
|
||||
lease_id = UUID(x_lease_id)
|
||||
authenticate_node(node_id, _bearer(authorization))
|
||||
except (ValueError, ComputeNodeError) as exc:
|
||||
raise HTTPException(401, "invalid node credentials or job identity") from exc
|
||||
context = get_job_output_context(node_id, job_id, lease_id, x_request_digest)
|
||||
if context is None:
|
||||
raise HTTPException(404, "compute job output target not found")
|
||||
return node_id, lease_id, context
|
||||
|
||||
|
||||
def _hash_file(path: Path) -> str:
|
||||
digest = sha256()
|
||||
with path.open("rb") as handle:
|
||||
for chunk in iter(lambda: handle.read(1024 * 1024), b""):
|
||||
digest.update(chunk)
|
||||
return digest.hexdigest()
|
||||
|
||||
|
||||
def _reject_symlink_path(root: Path, target: Path) -> None:
|
||||
root = root.resolve()
|
||||
current = root
|
||||
for part in target.relative_to(root).parts:
|
||||
current = current / part
|
||||
if current.is_symlink():
|
||||
raise HTTPException(409, "compute output path contains a symbolic link")
|
||||
|
||||
|
||||
def _publish_compute_outputs(job_id: UUID, context: dict, manifest: list[dict]) -> list[dict]:
|
||||
root = load_user_root(context["user_id"])
|
||||
working_dir = safe_join(root, context["working_dir"])
|
||||
staging = safe_join(root, f".zcbot_compute_staging/{job_id}")
|
||||
relative_output = Path("origin") / str(job_id)
|
||||
destination = safe_join(working_dir, relative_output.as_posix())
|
||||
source = staging if staging.is_dir() else destination
|
||||
_reject_symlink_path(root, source)
|
||||
_reject_symlink_path(root, destination)
|
||||
for item in manifest:
|
||||
path = source / item["filename"]
|
||||
if (
|
||||
not path.is_file()
|
||||
or path.stat().st_size != item["size_bytes"]
|
||||
or _hash_file(path) != item["sha256"]
|
||||
):
|
||||
raise ComputeNodeError(f"uploaded artifact is missing or invalid: {item['artifact_id']}")
|
||||
if source == staging:
|
||||
destination.parent.mkdir(parents=True, exist_ok=True)
|
||||
if destination.exists():
|
||||
raise ComputeNodeError("compute output destination already exists unexpectedly")
|
||||
os.replace(staging, destination)
|
||||
try:
|
||||
staging.parent.rmdir()
|
||||
except OSError:
|
||||
pass
|
||||
refs = tuple({
|
||||
"path": (relative_output / item["filename"]).as_posix(),
|
||||
"label": item["filename"],
|
||||
"media_type": item["media_type"],
|
||||
} for item in manifest)
|
||||
published_refs = register_published_artifacts(
|
||||
user_id=context["user_id"],
|
||||
task_id=context["task_id"],
|
||||
user_root=root,
|
||||
working_dir=working_dir,
|
||||
refs=refs,
|
||||
)
|
||||
refs_by_path = {item["path"]: item for item in published_refs}
|
||||
return [
|
||||
{
|
||||
**item,
|
||||
"source_artifact_id": item["artifact_id"],
|
||||
"artifact_id": refs_by_path[(relative_output / item["filename"]).as_posix()]["artifact_id"],
|
||||
"path": (relative_output / item["filename"]).as_posix(),
|
||||
}
|
||||
for item in manifest
|
||||
]
|
||||
|
||||
|
||||
def register_compute_node_routes(app, *, require_user, require_admin) -> None:
|
||||
@app.post(
|
||||
"/v1/compute/nodes/enroll",
|
||||
tags=["compute-nodes"],
|
||||
|
|
@ -73,6 +210,147 @@ def register_compute_node_routes(app, *, require_admin) -> None:
|
|||
except ComputeNodeError as exc:
|
||||
raise HTTPException(400, str(exc)) from exc
|
||||
|
||||
@app.get("/v1/compute/jobs/{job_id}/input", tags=["compute-nodes"])
|
||||
def download_compute_job_input(
|
||||
job_id: UUID,
|
||||
authorization: str | None = Header(default=None),
|
||||
x_node_id: str = Header(default=""),
|
||||
):
|
||||
try:
|
||||
node_id = UUID(x_node_id)
|
||||
authenticate_node(node_id, _bearer(authorization))
|
||||
except (ValueError, ComputeNodeError) as exc:
|
||||
raise HTTPException(401, "invalid node credentials") from exc
|
||||
item = get_job_input(node_id, job_id)
|
||||
if item is None:
|
||||
raise HTTPException(404, "compute job input not found")
|
||||
target = safe_join(load_user_root(item["user_id"]), item["current_path"])
|
||||
if not target.is_file():
|
||||
raise HTTPException(404, "compute job input file not found")
|
||||
stat = target.stat()
|
||||
if stat.st_size != item["size_bytes"]:
|
||||
raise HTTPException(409, "compute job input changed after submission")
|
||||
digest = sha256()
|
||||
with target.open("rb") as handle:
|
||||
for chunk in iter(lambda: handle.read(1024 * 1024), b""):
|
||||
digest.update(chunk)
|
||||
if digest.hexdigest() != item["sha256"]:
|
||||
raise HTTPException(409, "compute job input changed after submission")
|
||||
return FileResponse(
|
||||
path=str(target),
|
||||
filename=item["filename"],
|
||||
media_type="application/octet-stream",
|
||||
headers={
|
||||
"Cache-Control": "no-store",
|
||||
"X-Content-SHA256": item["sha256"],
|
||||
},
|
||||
)
|
||||
|
||||
@app.put(
|
||||
"/v1/compute/jobs/{job_id}/outputs/{artifact_id}",
|
||||
tags=["compute-nodes"],
|
||||
status_code=status.HTTP_204_NO_CONTENT,
|
||||
)
|
||||
async def upload_compute_job_output(
|
||||
job_id: UUID,
|
||||
artifact_id: str,
|
||||
request: Request,
|
||||
authorization: str | None = Header(default=None),
|
||||
x_node_id: str = Header(default=""),
|
||||
x_lease_id: str = Header(default=""),
|
||||
x_request_digest: str = Header(default=""),
|
||||
x_content_sha256: str = Header(default=""),
|
||||
x_content_length: int = Header(default=-1),
|
||||
):
|
||||
_, _, context = await asyncio.to_thread(
|
||||
_authenticate_output_request,
|
||||
job_id, authorization, x_node_id, x_lease_id, x_request_digest,
|
||||
)
|
||||
metadata = OUTPUT_ARTIFACTS.get(artifact_id)
|
||||
if metadata is None:
|
||||
raise HTTPException(400, "unsupported output artifact identity")
|
||||
filename, _, output_format = metadata
|
||||
requested_formats = set(context["request"].get("output", {}).get("formats") or [])
|
||||
if output_format is not None and output_format not in requested_formats:
|
||||
raise HTTPException(400, "output artifact was not requested")
|
||||
if not 1 <= x_content_length <= MAX_OUTPUT_ARTIFACT_BYTES:
|
||||
raise HTTPException(400, "output artifact size is invalid")
|
||||
if len(x_content_sha256) != 64 or any(c not in "0123456789abcdef" for c in x_content_sha256):
|
||||
raise HTTPException(400, "output artifact digest is invalid")
|
||||
root = load_user_root(context["user_id"])
|
||||
published = safe_join(
|
||||
safe_join(root, context["working_dir"]),
|
||||
f"origin/{job_id}/{filename}",
|
||||
)
|
||||
if published.is_file():
|
||||
if published.stat().st_size == x_content_length and _hash_file(published) == x_content_sha256:
|
||||
return None
|
||||
raise HTTPException(409, "published output conflicts with uploaded artifact")
|
||||
staging = safe_join(root, f".zcbot_compute_staging/{job_id}")
|
||||
_reject_symlink_path(root, staging)
|
||||
staging.mkdir(parents=True, exist_ok=True)
|
||||
destination = staging / filename
|
||||
if destination.is_file():
|
||||
if destination.stat().st_size == x_content_length and _hash_file(destination) == x_content_sha256:
|
||||
return None
|
||||
raise HTTPException(409, "uploaded output conflicts with existing staging file")
|
||||
staged_total = sum(
|
||||
item.stat().st_size for item in staging.iterdir() if item.is_file()
|
||||
)
|
||||
if staged_total + x_content_length > MAX_OUTPUT_TOTAL_BYTES:
|
||||
raise HTTPException(413, "compute job outputs exceed the total size limit")
|
||||
temporary = destination.with_name(destination.name + ".tmp-" + os.urandom(8).hex())
|
||||
digest = sha256()
|
||||
total = 0
|
||||
try:
|
||||
with temporary.open("xb") as handle:
|
||||
async for chunk in request.stream():
|
||||
total += len(chunk)
|
||||
if total > x_content_length or total > MAX_OUTPUT_ARTIFACT_BYTES:
|
||||
raise HTTPException(413, "output artifact exceeded declared size")
|
||||
digest.update(chunk)
|
||||
handle.write(chunk)
|
||||
handle.flush()
|
||||
os.fsync(handle.fileno())
|
||||
if total != x_content_length or digest.hexdigest() != x_content_sha256:
|
||||
raise HTTPException(400, "output artifact did not match declared metadata")
|
||||
os.replace(temporary, destination)
|
||||
finally:
|
||||
temporary.unlink(missing_ok=True)
|
||||
return None
|
||||
|
||||
@app.post("/v1/compute/jobs/{job_id}/outputs/complete", tags=["compute-nodes"])
|
||||
async def complete_compute_job_outputs(
|
||||
job_id: UUID,
|
||||
request: Request,
|
||||
authorization: str | None = Header(default=None),
|
||||
x_node_id: str = Header(default=""),
|
||||
x_lease_id: str = Header(default=""),
|
||||
x_request_digest: str = Header(default=""),
|
||||
):
|
||||
node_id, lease_id, context = await asyncio.to_thread(
|
||||
_authenticate_output_request,
|
||||
job_id, authorization, x_node_id, x_lease_id, x_request_digest,
|
||||
)
|
||||
body = await request.json()
|
||||
if not isinstance(body, dict):
|
||||
raise HTTPException(400, "output completion body must be an object")
|
||||
try:
|
||||
manifest = validate_output_manifest(context["request"], body.get("artifact_manifest"))
|
||||
published = await asyncio.to_thread(_publish_compute_outputs, job_id, context, manifest)
|
||||
terminal = {
|
||||
"job_id": str(job_id),
|
||||
"lease_id": str(lease_id),
|
||||
"request_digest": x_request_digest,
|
||||
"status": "succeeded",
|
||||
"error": {},
|
||||
"artifact_manifest": published,
|
||||
}
|
||||
await asyncio.to_thread(record_job_terminal, node_id, terminal)
|
||||
except (ComputeNodeError, KeyError, TypeError) as exc:
|
||||
raise HTTPException(409, str(exc)) from exc
|
||||
return {"status": "succeeded", "artifact_manifest": published}
|
||||
|
||||
@app.websocket("/v1/compute/nodes/connect")
|
||||
async def node_connect(websocket: WebSocket):
|
||||
try:
|
||||
|
|
@ -89,15 +367,53 @@ def register_compute_node_routes(app, *, require_admin) -> None:
|
|||
await websocket.accept()
|
||||
await node_connections.activate(node_id, websocket)
|
||||
try:
|
||||
await websocket.send_json({"type": "connected", "heartbeat_seconds": 15})
|
||||
await node_connections.send_on(
|
||||
node_id,
|
||||
websocket,
|
||||
{"type": "connected", "heartbeat_seconds": 15},
|
||||
)
|
||||
while True:
|
||||
message = await websocket.receive_json()
|
||||
message_type = message.get("type")
|
||||
payload = message.get("payload") or {}
|
||||
if message_type not in {"hello", "heartbeat"} or not isinstance(
|
||||
payload, dict
|
||||
):
|
||||
await websocket.send_json(
|
||||
if not isinstance(payload, dict):
|
||||
await node_connections.send_on(node_id, websocket,
|
||||
{"type": "error", "code": "unsupported_message"}
|
||||
)
|
||||
continue
|
||||
if message_type in {"job_accept", "job_reject"}:
|
||||
await asyncio.to_thread(
|
||||
respond_to_offer,
|
||||
node_id,
|
||||
accepted=message_type == "job_accept",
|
||||
payload=payload,
|
||||
)
|
||||
await node_connections.send_on(node_id, websocket,
|
||||
{"type": "ack", "message_id": message.get("message_id")}
|
||||
)
|
||||
continue
|
||||
if message_type in {"job_state", "job_terminal"}:
|
||||
if message_type == "job_terminal" and payload.get("status") == "succeeded":
|
||||
await node_connections.send_on(node_id, websocket,
|
||||
{
|
||||
"type": "error",
|
||||
"code": "outputs_not_published",
|
||||
"message_id": message.get("message_id"),
|
||||
}
|
||||
)
|
||||
continue
|
||||
handler = (
|
||||
update_job_state
|
||||
if message_type == "job_state"
|
||||
else record_job_terminal
|
||||
)
|
||||
await asyncio.to_thread(handler, node_id, payload)
|
||||
await node_connections.send_on(node_id, websocket,
|
||||
{"type": "ack", "message_id": message.get("message_id")}
|
||||
)
|
||||
continue
|
||||
if message_type not in {"hello", "heartbeat"}:
|
||||
await node_connections.send_on(node_id, websocket,
|
||||
{"type": "error", "code": "unsupported_message"}
|
||||
)
|
||||
continue
|
||||
|
|
@ -112,14 +428,63 @@ def register_compute_node_routes(app, *, require_admin) -> None:
|
|||
status="online",
|
||||
runtime=payload,
|
||||
)
|
||||
await websocket.send_json(
|
||||
await node_connections.send_on(node_id, websocket,
|
||||
{"type": "ack", "message_id": message.get("message_id")}
|
||||
)
|
||||
offer = await asyncio.to_thread(
|
||||
offer_next_job, await node_connections.node_ids()
|
||||
)
|
||||
if offer is not None:
|
||||
delivered = await node_connections.send(
|
||||
offer["node_id"],
|
||||
{"type": "job_offer", "payload": offer["payload"]},
|
||||
)
|
||||
if not delivered:
|
||||
await asyncio.to_thread(
|
||||
abandon_offer, offer["node_id"], offer["payload"]
|
||||
)
|
||||
except (ComputeNodeError, WebSocketDisconnect, RuntimeError, ValueError):
|
||||
pass
|
||||
finally:
|
||||
if await node_connections.remove(node_id, websocket):
|
||||
await asyncio.to_thread(mark_node_offline, node_id)
|
||||
await asyncio.to_thread(mark_node_jobs_disconnected, node_id)
|
||||
|
||||
@app.post("/v1/tasks/{task_id}/compute-jobs", tags=["compute-jobs"])
|
||||
async def submit_compute_job(
|
||||
task_id: UUID,
|
||||
body: ComputeJobCreateRequest,
|
||||
user_id: UUID = Depends(require_user), # noqa: B008
|
||||
):
|
||||
try:
|
||||
job, created = await asyncio.to_thread(
|
||||
create_job, user_id, task_id, **body.model_dump()
|
||||
)
|
||||
offer = await asyncio.to_thread(
|
||||
offer_next_job, await node_connections.node_ids()
|
||||
)
|
||||
if offer is not None:
|
||||
delivered = await node_connections.send(
|
||||
offer["node_id"], {"type": "job_offer", "payload": offer["payload"]}
|
||||
)
|
||||
if not delivered:
|
||||
await asyncio.to_thread(
|
||||
abandon_offer, offer["node_id"], offer["payload"]
|
||||
)
|
||||
return {**job, "created": created}
|
||||
except ComputeNodeError as exc:
|
||||
detail = str(exc)
|
||||
raise HTTPException(404 if detail == "task not found" else 400, detail) from exc
|
||||
|
||||
@app.get("/v1/compute-jobs/{job_id}", tags=["compute-jobs"])
|
||||
def read_compute_job(
|
||||
job_id: UUID,
|
||||
user_id: UUID = Depends(require_user), # noqa: B008
|
||||
):
|
||||
job = get_job(user_id, job_id)
|
||||
if job is None:
|
||||
raise HTTPException(404, "compute job not found")
|
||||
return job
|
||||
|
||||
@app.post("/v1/admin/compute-node-enrollments", tags=["admin"])
|
||||
def admin_create_compute_enrollment(
|
||||
|
|
|
|||
|
|
@ -133,3 +133,9 @@ class ComputeNodeEnrollRequest(BaseModel):
|
|||
|
||||
class ComputeNodeDisableRequest(BaseModel):
|
||||
disabled: bool = True
|
||||
|
||||
|
||||
class ComputeJobCreateRequest(BaseModel):
|
||||
idempotency_key: str
|
||||
capability: str = "origin.plot@v1"
|
||||
request: dict = Field(default_factory=dict)
|
||||
|
|
|
|||
|
|
@ -2,7 +2,15 @@
|
|||
|
||||
内网 MVP 的 Windows 执行节点，目标运行环境为 Windows 11 Enterprise + .NET 10 SDK 10.0.303。仓库根目录 `global.json` 固定 SDK patch；客户端只使用 .NET Windows Desktop Framework，不依赖第三方 NuGet 包。
|
||||
|
||||
当前实现托盘状态角标、小型配置窗口、注册、DPAPI/ACL 配置保存、WebSocket `hello`/心跳和退避重连，并只读探测 Origin/OriginPro 安装版本、COM 自动化组件与桌面会话状态。尚未实现 `compute_jobs`、Origin Worker、任务目录和产物上传。
|
||||
当前实现托盘状态角标、小型配置窗口、注册、DPAPI/ACL 配置保存、WebSocket `hello`/心跳和退避重连，并只读探测 Origin/OriginPro 安装版本、COM 自动化组件与桌面会话状态。Node 可以接收受控的 `origin.plot@v1` offer，在本机任务目录原子保存请求后回报 accept/reject；随后以 Node 身份流式下载任务绑定的 CSV/XLSX/JSON，校验大小与 SHA-256 后原子保存。固定 Origin Worker 独立于单次 WebSocket 执行，断线不终止已启动绘图。成功产物按 manifest 逐项流式上传并由云端复核大小与 SHA-256，全部完成后原子发布到任务工作目录的 `origin/<job_id>/`；中断后按本地 `upload-complete.json` 幂等续传。
|
||||
|
||||
固定 Origin Worker 已支持 `line`、`scatter` 和 `line_scatter`，生成 OPJU、PNG、SVG、PDF、plot spec、provenance 与原子 `terminal.json`。运行时独立于 zcbot 服务端 Python，管理员执行：
|
||||
|
||||
```powershell
|
||||
.\install-origin-runtime.ps1 -BootstrapPython D:\programs\Python312\python.exe
|
||||
```
|
||||
|
||||
默认安装到 `%ProgramData%\Zcbot\WindowsNode\runtimes\origin\python.exe`；也可由管理员用绝对路径环境变量 `ZCBOT_ORIGIN_PYTHON` 指向固定解释器。任务请求不能指定解释器、脚本或路径。
|
||||
|
||||
注册和运行必须使用同一专用 Windows 账号。MVP 通过该账号的登录后计划任务自动启动，不以 Windows Service 在 Session 0 运行。
|
||||
|
||||
|
|
|
|||
|
|
@ -0,0 +1,346 @@
|
|||
using System.Text;
|
||||
using System.Text.Json;
|
||||
|
||||
namespace Zcbot.WindowsNode;
|
||||
|
||||
internal sealed class JobInboxStore(string jobsDirectory)
|
||||
{
|
||||
private static readonly JsonSerializerOptions JsonOptions = new() { WriteIndented = true };
|
||||
private static readonly HashSet<string> PlotTypes =
|
||||
["line", "scatter", "line_scatter"];
|
||||
private static readonly HashSet<string> OutputFormats = ["opju", "png", "svg", "pdf"];
|
||||
|
||||
internal bool HasPendingJobs => Directory.Exists(jobsDirectory)
|
||||
&& ReadRecoverableJobs().Any(item =>
|
||||
item.Terminal is null
|
||||
|| item.Terminal.Value.GetProperty("status").GetString() == "succeeded"
|
||||
&& !item.UploadComplete);
|
||||
|
||||
internal IReadOnlyList<RecoverableJob> ReadRecoverableJobs()
|
||||
{
|
||||
if (!Directory.Exists(jobsDirectory))
|
||||
{
|
||||
return [];
|
||||
}
|
||||
var jobs = new List<RecoverableJob>();
|
||||
foreach (var requestPath in Directory.EnumerateFiles(
|
||||
jobsDirectory, "request.json", SearchOption.AllDirectories))
|
||||
{
|
||||
try
|
||||
{
|
||||
using var request = JsonDocument.Parse(File.ReadAllBytes(requestPath));
|
||||
var root = request.RootElement;
|
||||
if (!TryReadGuid(root, "job_id", out var jobId)
|
||||
|| !TryReadGuid(root, "lease_id", out var leaseId)
|
||||
|| !root.TryGetProperty("request_digest", out var digestValue)
|
||||
|| digestValue.GetString() is not { Length: 64 } requestDigest)
|
||||
{
|
||||
continue;
|
||||
}
|
||||
var jobDirectory = Directory.GetParent(Directory.GetParent(requestPath)!.FullName)!.FullName;
|
||||
jobs.Add(new RecoverableJob(
|
||||
jobId,
|
||||
leaseId,
|
||||
requestDigest,
|
||||
root.TryGetProperty("input_transfer", out var transfer)
|
||||
? transfer.Clone() : null,
|
||||
ReadTerminal(Path.Combine(jobDirectory, "terminal.json")),
|
||||
File.Exists(Path.Combine(jobDirectory, "upload-complete.json"))));
|
||||
}
|
||||
catch (Exception exception) when (
|
||||
exception is JsonException or IOException or UnauthorizedAccessException)
|
||||
{
|
||||
}
|
||||
}
|
||||
return jobs;
|
||||
}
|
||||
|
||||
internal JobOfferResult Accept(JsonElement payload)
|
||||
{
|
||||
if (!TryReadGuid(payload, "job_id", out var jobId)
|
||||
|| !TryReadGuid(payload, "lease_id", out var leaseId)
|
||||
|| !payload.TryGetProperty("request_digest", out var digestValue)
|
||||
|| digestValue.ValueKind != JsonValueKind.String
|
||||
|| digestValue.GetString() is not { Length: 64 } requestDigest
|
||||
|| !payload.TryGetProperty("capability", out var capabilityValue)
|
||||
|| capabilityValue.GetString() != "origin.plot@v1"
|
||||
|| !payload.TryGetProperty("request", out var request)
|
||||
|| request.ValueKind != JsonValueKind.Object
|
||||
|| !payload.TryGetProperty("input_transfer", out var inputTransfer)
|
||||
|| !IsValidInputTransfer(inputTransfer))
|
||||
{
|
||||
return JobOfferResult.Reject("invalid_offer");
|
||||
}
|
||||
if (!IsValidRequest(request))
|
||||
{
|
||||
return JobOfferResult.Reject("unsupported_request");
|
||||
}
|
||||
|
||||
var directory = Path.Combine(jobsDirectory, jobId.ToString("D"));
|
||||
var requestDirectory = Path.Combine(directory, "request");
|
||||
var requestPath = Path.Combine(requestDirectory, "request.json");
|
||||
Directory.CreateDirectory(requestDirectory);
|
||||
if (File.Exists(requestPath))
|
||||
{
|
||||
try
|
||||
{
|
||||
using var existing = JsonDocument.Parse(File.ReadAllBytes(requestPath));
|
||||
var root = existing.RootElement;
|
||||
var sameDigest = root.TryGetProperty("request_digest", out var existingDigest)
|
||||
&& existingDigest.GetString() == requestDigest;
|
||||
if (!sameDigest)
|
||||
{
|
||||
return JobOfferResult.Reject("job_digest_conflict");
|
||||
}
|
||||
if (!TryReadGuid(root, "lease_id", out var existingLease)
|
||||
|| existingLease != leaseId)
|
||||
{
|
||||
var updated = JsonSerializer.SerializeToUtf8Bytes(new
|
||||
{
|
||||
job_id = jobId,
|
||||
lease_id = leaseId,
|
||||
request_digest = requestDigest,
|
||||
capability = "origin.plot@v1",
|
||||
accepted_at = DateTimeOffset.UtcNow,
|
||||
request,
|
||||
input_transfer = inputTransfer,
|
||||
}, JsonOptions);
|
||||
AtomicWrite(requestPath, updated, overwrite: true);
|
||||
}
|
||||
return JobOfferResult.Accept(jobId, leaseId, requestDigest);
|
||||
}
|
||||
catch (JsonException)
|
||||
{
|
||||
return JobOfferResult.Reject("local_job_record_invalid");
|
||||
}
|
||||
}
|
||||
|
||||
var record = JsonSerializer.SerializeToUtf8Bytes(new
|
||||
{
|
||||
job_id = jobId,
|
||||
lease_id = leaseId,
|
||||
request_digest = requestDigest,
|
||||
capability = "origin.plot@v1",
|
||||
accepted_at = DateTimeOffset.UtcNow,
|
||||
request,
|
||||
input_transfer = inputTransfer,
|
||||
}, JsonOptions);
|
||||
try
|
||||
{
|
||||
AtomicWrite(requestPath, record, overwrite: false);
|
||||
return JobOfferResult.Accept(jobId, leaseId, requestDigest);
|
||||
}
|
||||
catch (IOException)
|
||||
{
|
||||
return JobOfferResult.Reject("local_job_persist_failed");
|
||||
}
|
||||
}
|
||||
|
||||
private static JsonElement? ReadTerminal(string path)
|
||||
{
|
||||
if (!File.Exists(path))
|
||||
{
|
||||
return null;
|
||||
}
|
||||
using var document = JsonDocument.Parse(File.ReadAllBytes(path));
|
||||
return document.RootElement.Clone();
|
||||
}
|
||||
|
||||
private static void AtomicWrite(string path, byte[] content, bool overwrite)
|
||||
{
|
||||
Directory.CreateDirectory(Path.GetDirectoryName(path)!);
|
||||
var temporaryPath = path + ".tmp-" + Guid.NewGuid().ToString("N");
|
||||
try
|
||||
{
|
||||
using (var stream = new FileStream(
|
||||
temporaryPath, FileMode.CreateNew, FileAccess.Write, FileShare.None,
|
||||
bufferSize: 4096, FileOptions.WriteThrough))
|
||||
{
|
||||
stream.Write(content);
|
||||
stream.Flush(flushToDisk: true);
|
||||
}
|
||||
File.Move(temporaryPath, path, overwrite);
|
||||
}
|
||||
finally
|
||||
{
|
||||
if (File.Exists(temporaryPath)) File.Delete(temporaryPath);
|
||||
}
|
||||
}
|
||||
|
||||
private static bool IsValidRequest(JsonElement request) =>
|
||||
HasOnlyProperties(request, "schema_version", "input", "plot", "output")
|
||||
&& request.TryGetProperty("schema_version", out var schemaVersion)
|
||||
&& schemaVersion.TryGetInt32(out var version)
|
||||
&& version == 1
|
||||
&& request.TryGetProperty("input", out var input)
|
||||
&& input.ValueKind == JsonValueKind.Object
|
||||
&& HasOnlyProperties(input, "input_id", "sheet")
|
||||
&& input.TryGetProperty("input_id", out var inputId)
|
||||
&& inputId.ValueKind == JsonValueKind.String
|
||||
&& !string.IsNullOrWhiteSpace(inputId.GetString())
|
||||
&& request.TryGetProperty("plot", out var plot)
|
||||
&& plot.ValueKind == JsonValueKind.Object
|
||||
&& HasOnlyProperties(
|
||||
plot, "type", "x", "y", "template", "title", "x_axis", "y_axis", "legend", "error_bars")
|
||||
&& plot.TryGetProperty("type", out var plotType)
|
||||
&& plotType.ValueKind == JsonValueKind.String
|
||||
&& PlotTypes.Contains(plotType.GetString() ?? "")
|
||||
&& (!plot.TryGetProperty("title", out var title)
|
||||
|| title.ValueKind == JsonValueKind.String && title.GetString()!.Length <= 500)
|
||||
&& plot.TryGetProperty("x", out var x)
|
||||
&& IsColumnName(x)
|
||||
&& plot.TryGetProperty("y", out var y)
|
||||
&& IsValidYColumns(y)
|
||||
&& (!plot.TryGetProperty("template", out var template)
|
||||
|| template.GetString() == "publication_double_column")
|
||||
&& IsValidAxis(plot, "x_axis")
|
||||
&& IsValidAxis(plot, "y_axis")
|
||||
&& IsValidLegend(plot)
|
||||
&& !plot.TryGetProperty("error_bars", out _)
|
||||
&& request.TryGetProperty("output", out var output)
|
||||
&& output.ValueKind == JsonValueKind.Object
|
||||
&& HasOnlyProperties(output, "formats", "dpi", "capture_screenshots", "record_video")
|
||||
&& output.TryGetProperty("formats", out var formats)
|
||||
&& formats.ValueKind == JsonValueKind.Array
|
||||
&& formats.GetArrayLength() > 0
|
||||
&& IsValidFormats(formats)
|
||||
&& (!output.TryGetProperty("dpi", out var dpi)
|
||||
|| dpi.TryGetInt32(out var dpiValue) && dpiValue is >= 72 and <= 1200)
|
||||
&& IsOptionalBoolean(output, "capture_screenshots")
|
||||
&& IsOptionalBoolean(output, "record_video")
|
||||
&& (!output.TryGetProperty("record_video", out var recordVideo)
|
||||
|| recordVideo.ValueKind == JsonValueKind.False);
|
||||
|
||||
private static bool IsColumnName(JsonElement value) =>
|
||||
value.ValueKind == JsonValueKind.String
|
||||
&& value.GetString() is { Length: >= 1 and <= 128 };
|
||||
|
||||
private static bool IsValidFormats(JsonElement formats)
|
||||
{
|
||||
var values = formats.EnumerateArray().ToArray();
|
||||
return values.All(item =>
|
||||
item.ValueKind == JsonValueKind.String
|
||||
&& OutputFormats.Contains(item.GetString() ?? ""))
|
||||
&& values.Select(item => item.GetString()).Distinct(StringComparer.Ordinal).Count()
|
||||
== values.Length;
|
||||
}
|
||||
|
||||
private static bool IsValidYColumns(JsonElement value)
|
||||
{
|
||||
if (IsColumnName(value)) return true;
|
||||
if (value.ValueKind != JsonValueKind.Array
|
||||
|| value.GetArrayLength() is < 1 or > 16)
|
||||
{
|
||||
return false;
|
||||
}
|
||||
var names = value.EnumerateArray().Select(item => item.GetString()).ToArray();
|
||||
return value.EnumerateArray().All(IsColumnName)
|
||||
&& names.Distinct(StringComparer.Ordinal).Count() == names.Length;
|
||||
}
|
||||
|
||||
private static bool IsValidAxis(JsonElement plot, string name)
|
||||
{
|
||||
if (!plot.TryGetProperty(name, out var axis)) return true;
|
||||
return axis.ValueKind == JsonValueKind.Object
|
||||
&& HasOnlyProperties(axis, "title", "unit", "scale")
|
||||
&& (!axis.TryGetProperty("title", out var title) || title.ValueKind == JsonValueKind.String)
|
||||
&& (!axis.TryGetProperty("unit", out var unit) || unit.ValueKind == JsonValueKind.String)
|
||||
&& (!axis.TryGetProperty("scale", out var scale) || scale.GetString() == "linear");
|
||||
}
|
||||
|
||||
private static bool IsValidLegend(JsonElement plot)
|
||||
{
|
||||
if (!plot.TryGetProperty("legend", out var legend)) return true;
|
||||
return legend.ValueKind == JsonValueKind.Object
|
||||
&& HasOnlyProperties(legend, "enabled", "position")
|
||||
&& (!legend.TryGetProperty("enabled", out var enabled)
|
||||
|| enabled.ValueKind == JsonValueKind.True)
|
||||
&& (!legend.TryGetProperty("position", out var position)
|
||||
|| position.GetString() == "top_right");
|
||||
}
|
||||
|
||||
private static bool IsOptionalBoolean(JsonElement value, string name) =>
|
||||
!value.TryGetProperty(name, out var property)
|
||||
|| property.ValueKind is JsonValueKind.True or JsonValueKind.False;
|
||||
|
||||
private static bool IsValidInputTransfer(JsonElement transfer) =>
|
||||
transfer.ValueKind == JsonValueKind.Object
|
||||
&& HasOnlyProperties(transfer, "artifact_id", "filename", "size_bytes", "sha256", "download_path")
|
||||
&& transfer.TryGetProperty("artifact_id", out var artifactId)
|
||||
&& Guid.TryParse(artifactId.GetString(), out _)
|
||||
&& transfer.TryGetProperty("filename", out var filename)
|
||||
&& filename.ValueKind == JsonValueKind.String
|
||||
&& Path.GetFileName(filename.GetString()) == filename.GetString()
|
||||
&& transfer.TryGetProperty("size_bytes", out var size)
|
||||
&& size.TryGetInt64(out var sizeBytes)
|
||||
&& sizeBytes is >= 0 and <= 104_857_600
|
||||
&& transfer.TryGetProperty("sha256", out var sha)
|
||||
&& sha.GetString() is { Length: 64 }
|
||||
&& transfer.TryGetProperty("download_path", out var downloadPath)
|
||||
&& downloadPath.GetString()?.StartsWith("/v1/compute/jobs/", StringComparison.Ordinal) == true
|
||||
&& downloadPath.GetString()?.EndsWith("/input", StringComparison.Ordinal) == true;
|
||||
|
||||
internal string InputPath(RecoverableJob job)
|
||||
{
|
||||
if (job.InputTransfer is not JsonElement transfer
|
||||
|| !IsValidInputTransfer(transfer))
|
||||
{
|
||||
throw new InvalidDataException("Stored input transfer is invalid.");
|
||||
}
|
||||
var filename = transfer.GetProperty("filename").GetString()!;
|
||||
return Path.Combine(jobsDirectory, job.JobId.ToString("D"), "input", filename);
|
||||
}
|
||||
|
||||
internal void WriteTerminal(
|
||||
RecoverableJob job, string status, string code, string detail)
|
||||
{
|
||||
var path = Path.Combine(jobsDirectory, job.JobId.ToString("D"), "terminal.json");
|
||||
if (File.Exists(path))
|
||||
{
|
||||
return;
|
||||
}
|
||||
var content = JsonSerializer.SerializeToUtf8Bytes(new
|
||||
{
|
||||
job_id = job.JobId,
|
||||
lease_id = job.LeaseId,
|
||||
request_digest = job.RequestDigest,
|
||||
status,
|
||||
error = new { code, detail },
|
||||
artifact_manifest = Array.Empty<object>(),
|
||||
terminal_at = DateTimeOffset.UtcNow,
|
||||
}, JsonOptions);
|
||||
AtomicWrite(path, content, overwrite: false);
|
||||
}
|
||||
|
||||
private static bool HasOnlyProperties(JsonElement value, params string[] allowed)
|
||||
{
|
||||
var names = new HashSet<string>(allowed, StringComparer.Ordinal);
|
||||
return value.EnumerateObject().All(item => names.Contains(item.Name));
|
||||
}
|
||||
|
||||
private static bool TryReadGuid(JsonElement payload, string name, out Guid value)
|
||||
{
|
||||
value = Guid.Empty;
|
||||
return payload.TryGetProperty(name, out var property)
|
||||
&& property.ValueKind == JsonValueKind.String
|
||||
&& Guid.TryParse(property.GetString(), out value);
|
||||
}
|
||||
}
|
||||
|
||||
internal sealed record RecoverableJob(
|
||||
Guid JobId,
|
||||
Guid LeaseId,
|
||||
string RequestDigest,
|
||||
JsonElement? InputTransfer,
|
||||
JsonElement? Terminal,
|
||||
bool UploadComplete);
|
||||
|
||||
internal sealed record JobOfferResult(
|
||||
bool Accepted, Guid JobId, Guid LeaseId, string RequestDigest, string Reason)
|
||||
{
|
||||
internal static JobOfferResult Accept(Guid jobId, Guid leaseId, string requestDigest) =>
|
||||
new(true, jobId, leaseId, requestDigest, "");
|
||||
|
||||
internal static JobOfferResult Reject(string reason) =>
|
||||
new(false, Guid.Empty, Guid.Empty, "", reason);
|
||||
}
|
||||
|
|
@ -0,0 +1,99 @@
|
|||
using System.Net.Http.Headers;
|
||||
using System.Security.Cryptography;
|
||||
using System.Text.Json;
|
||||
|
||||
namespace Zcbot.WindowsNode;
|
||||
|
||||
internal sealed class JobInputDownloader(NodeConfig config, JobInboxStore inbox)
|
||||
{
|
||||
internal async Task DownloadAsync(RecoverableJob job, CancellationToken cancellationToken)
|
||||
{
|
||||
if (job.InputTransfer is not JsonElement transfer)
|
||||
{
|
||||
throw new InvalidDataException("Job input transfer is missing.");
|
||||
}
|
||||
var downloadPath = transfer.GetProperty("download_path").GetString()!;
|
||||
if (!downloadPath.StartsWith("/v1/compute/jobs/", StringComparison.Ordinal)
|
||||
|| !downloadPath.EndsWith("/input", StringComparison.Ordinal)
|
||||
|| !Uri.TryCreate(downloadPath, UriKind.Relative, out var relativeUri))
|
||||
{
|
||||
throw new InvalidDataException("Job input download path is invalid.");
|
||||
}
|
||||
var expectedSize = transfer.GetProperty("size_bytes").GetInt64();
|
||||
var expectedSha256 = transfer.GetProperty("sha256").GetString()!;
|
||||
var destination = inbox.InputPath(job);
|
||||
if (File.Exists(destination))
|
||||
{
|
||||
await VerifyExistingAsync(destination, expectedSize, expectedSha256, cancellationToken);
|
||||
return;
|
||||
}
|
||||
|
||||
Directory.CreateDirectory(Path.GetDirectoryName(destination)!);
|
||||
var temporaryPath = destination + ".tmp-" + Guid.NewGuid().ToString("N");
|
||||
using var client = new HttpClient { BaseAddress = config.ServerUrl };
|
||||
client.DefaultRequestHeaders.Authorization =
|
||||
new AuthenticationHeaderValue("Bearer", config.NodeToken);
|
||||
client.DefaultRequestHeaders.Add("X-Node-Id", config.NodeId.ToString());
|
||||
try
|
||||
{
|
||||
using var response = await client.GetAsync(
|
||||
relativeUri, HttpCompletionOption.ResponseHeadersRead, cancellationToken);
|
||||
response.EnsureSuccessStatusCode();
|
||||
if (response.Content.Headers.ContentLength is long contentLength
|
||||
&& contentLength != expectedSize)
|
||||
{
|
||||
throw new InvalidDataException("Job input size header does not match the manifest.");
|
||||
}
|
||||
await using var source = await response.Content.ReadAsStreamAsync(cancellationToken);
|
||||
await using var target = new FileStream(
|
||||
temporaryPath, FileMode.CreateNew, FileAccess.Write, FileShare.None,
|
||||
bufferSize: 64 * 1024, FileOptions.Asynchronous | FileOptions.WriteThrough);
|
||||
using var hash = IncrementalHash.CreateHash(HashAlgorithmName.SHA256);
|
||||
var buffer = new byte[64 * 1024];
|
||||
long total = 0;
|
||||
while (true)
|
||||
{
|
||||
var count = await source.ReadAsync(buffer, cancellationToken);
|
||||
if (count == 0) break;
|
||||
total += count;
|
||||
if (total > expectedSize)
|
||||
{
|
||||
throw new InvalidDataException("Job input exceeded its declared size.");
|
||||
}
|
||||
hash.AppendData(buffer, 0, count);
|
||||
await target.WriteAsync(buffer.AsMemory(0, count), cancellationToken);
|
||||
}
|
||||
await target.FlushAsync(cancellationToken);
|
||||
target.Flush(flushToDisk: true);
|
||||
var actualSha256 = Convert.ToHexString(hash.GetHashAndReset()).ToLowerInvariant();
|
||||
if (total != expectedSize || actualSha256 != expectedSha256)
|
||||
{
|
||||
throw new InvalidDataException("Job input digest does not match the manifest.");
|
||||
}
|
||||
target.Close();
|
||||
File.Move(temporaryPath, destination, overwrite: false);
|
||||
}
|
||||
finally
|
||||
{
|
||||
if (File.Exists(temporaryPath)) File.Delete(temporaryPath);
|
||||
}
|
||||
}
|
||||
|
||||
private static async Task VerifyExistingAsync(
|
||||
string path, long expectedSize, string expectedSha256, CancellationToken cancellationToken)
|
||||
{
|
||||
var info = new FileInfo(path);
|
||||
if (info.Length != expectedSize)
|
||||
{
|
||||
throw new InvalidDataException("Existing job input size does not match the manifest.");
|
||||
}
|
||||
await using var stream = new FileStream(
|
||||
path, FileMode.Open, FileAccess.Read, FileShare.Read, 64 * 1024, FileOptions.Asynchronous);
|
||||
var digest = Convert.ToHexString(
|
||||
await SHA256.HashDataAsync(stream, cancellationToken)).ToLowerInvariant();
|
||||
if (digest != expectedSha256)
|
||||
{
|
||||
throw new InvalidDataException("Existing job input digest does not match the manifest.");
|
||||
}
|
||||
}
|
||||
}
|
||||
|
|
@ -0,0 +1,100 @@
|
|||
using System.Net.Http.Headers;
|
||||
using System.Security.Cryptography;
|
||||
using System.Text;
|
||||
using System.Text.Json;
|
||||
|
||||
namespace Zcbot.WindowsNode;
|
||||
|
||||
internal sealed class JobOutputUploader(NodeConfig config)
|
||||
{
|
||||
private static readonly JsonSerializerOptions JsonOptions = new() { WriteIndented = true };
|
||||
|
||||
internal async Task UploadAsync(RecoverableJob job)
|
||||
{
|
||||
var jobDirectory = Path.Combine(
|
||||
NodePaths.ForCurrentMachine().JobsDirectory, job.JobId.ToString("D"));
|
||||
var completionPath = Path.Combine(jobDirectory, "upload-complete.json");
|
||||
if (File.Exists(completionPath)) return;
|
||||
var terminalPath = Path.Combine(jobDirectory, "terminal.json");
|
||||
using var terminal = JsonDocument.Parse(await File.ReadAllBytesAsync(terminalPath));
|
||||
if (terminal.RootElement.GetProperty("status").GetString() != "succeeded") return;
|
||||
var manifest = terminal.RootElement.GetProperty("artifact_manifest").Clone();
|
||||
|
||||
using var client = new HttpClient { BaseAddress = config.ServerUrl };
|
||||
client.DefaultRequestHeaders.Authorization =
|
||||
new AuthenticationHeaderValue("Bearer", config.NodeToken);
|
||||
client.DefaultRequestHeaders.Add("X-Node-Id", config.NodeId.ToString());
|
||||
client.DefaultRequestHeaders.Add("X-Lease-Id", job.LeaseId.ToString());
|
||||
client.DefaultRequestHeaders.Add("X-Request-Digest", job.RequestDigest);
|
||||
foreach (var artifact in manifest.EnumerateArray())
|
||||
{
|
||||
var localId = artifact.GetProperty("artifact_id").GetString()!;
|
||||
var filename = artifact.GetProperty("filename").GetString()!;
|
||||
var expectedSize = artifact.GetProperty("size_bytes").GetInt64();
|
||||
var expectedDigest = artifact.GetProperty("sha256").GetString()!;
|
||||
var path = Path.Combine(jobDirectory, "output", filename);
|
||||
var info = new FileInfo(path);
|
||||
if (!info.Exists || info.Length != expectedSize)
|
||||
{
|
||||
throw new InvalidDataException($"Output artifact is missing or changed: {localId}.");
|
||||
}
|
||||
await using (var verify = new FileStream(
|
||||
path, FileMode.Open, FileAccess.Read, FileShare.Read, 64 * 1024,
|
||||
FileOptions.Asynchronous | FileOptions.SequentialScan))
|
||||
{
|
||||
var digest = Convert.ToHexString(
|
||||
await SHA256.HashDataAsync(verify)).ToLowerInvariant();
|
||||
if (digest != expectedDigest)
|
||||
{
|
||||
throw new InvalidDataException($"Output artifact digest changed: {localId}.");
|
||||
}
|
||||
}
|
||||
await using var stream = new FileStream(
|
||||
path, FileMode.Open, FileAccess.Read, FileShare.Read, 64 * 1024,
|
||||
FileOptions.Asynchronous | FileOptions.SequentialScan);
|
||||
using var content = new StreamContent(stream);
|
||||
content.Headers.ContentType = new MediaTypeHeaderValue("application/octet-stream");
|
||||
content.Headers.ContentLength = expectedSize;
|
||||
content.Headers.Add("X-Content-SHA256", expectedDigest);
|
||||
content.Headers.Add("X-Content-Length", expectedSize.ToString());
|
||||
using var response = await client.PutAsync(
|
||||
$"/v1/compute/jobs/{job.JobId:D}/outputs/{Uri.EscapeDataString(localId)}",
|
||||
content);
|
||||
response.EnsureSuccessStatusCode();
|
||||
}
|
||||
|
||||
using var completeContent = new StringContent(
|
||||
JsonSerializer.Serialize(new { artifact_manifest = manifest }),
|
||||
Encoding.UTF8,
|
||||
"application/json");
|
||||
using var completeResponse = await client.PostAsync(
|
||||
$"/v1/compute/jobs/{job.JobId:D}/outputs/complete", completeContent);
|
||||
completeResponse.EnsureSuccessStatusCode();
|
||||
var responseBody = await completeResponse.Content.ReadAsByteArrayAsync();
|
||||
AtomicWrite(completionPath, responseBody);
|
||||
}
|
||||
|
||||
private static void AtomicWrite(string path, byte[] responseBody)
|
||||
{
|
||||
using var response = JsonDocument.Parse(responseBody);
|
||||
var content = JsonSerializer.SerializeToUtf8Bytes(new
|
||||
{
|
||||
completed_at = DateTimeOffset.UtcNow,
|
||||
response = response.RootElement,
|
||||
}, JsonOptions);
|
||||
var temporary = path + ".tmp-" + Guid.NewGuid().ToString("N");
|
||||
try
|
||||
{
|
||||
using var stream = new FileStream(
|
||||
temporary, FileMode.CreateNew, FileAccess.Write, FileShare.None,
|
||||
4096, FileOptions.WriteThrough);
|
||||
stream.Write(content);
|
||||
stream.Flush(flushToDisk: true);
|
||||
File.Move(temporary, path, overwrite: false);
|
||||
}
|
||||
finally
|
||||
{
|
||||
if (File.Exists(temporary)) File.Delete(temporary);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
|
@ -4,11 +4,20 @@ using System.Reflection;
|
|||
using System.Runtime.InteropServices;
|
||||
using System.Text;
|
||||
using System.Text.Json;
|
||||
using System.Collections.Concurrent;
|
||||
|
||||
namespace Zcbot.WindowsNode;
|
||||
|
||||
internal sealed class NodeConnectionLoop(NodeConfig config, Action<NodeStatus>? statusChanged = null)
|
||||
{
|
||||
private readonly SemaphoreSlim sendLock = new(1, 1);
|
||||
private readonly JobInboxStore jobInbox = new(NodePaths.ForCurrentMachine().JobsDirectory);
|
||||
private readonly JobInputDownloader inputDownloader = new(
|
||||
config, new JobInboxStore(NodePaths.ForCurrentMachine().JobsDirectory));
|
||||
private readonly OriginWorkerRunner workerRunner = new(
|
||||
new JobInboxStore(NodePaths.ForCurrentMachine().JobsDirectory));
|
||||
private readonly JobOutputUploader outputUploader = new(config);
|
||||
private readonly ConcurrentDictionary<Guid, Task> jobPipelines = new();
|
||||
private static readonly TimeSpan[] Backoff =
|
||||
[
|
||||
TimeSpan.FromSeconds(1),
|
||||
|
|
@ -91,6 +100,7 @@ internal sealed class NodeConnectionLoop(NodeConfig config, Action<NodeStatus>?
|
|||
Report(NodeState.Online, "已连接");
|
||||
|
||||
await SendAsync(socket, "hello", RuntimePayload(), cancellationToken);
|
||||
await ReportRecoverableJobsAsync(socket, cancellationToken);
|
||||
using var heartbeatStop = CancellationTokenSource.CreateLinkedTokenSource(cancellationToken);
|
||||
var heartbeat = HeartbeatLoopAsync(socket, heartbeatStop.Token);
|
||||
try
|
||||
|
|
@ -116,10 +126,52 @@ internal sealed class NodeConnectionLoop(NodeConfig config, Action<NodeStatus>?
|
|||
while (await timer.WaitForNextTickAsync(cancellationToken))
|
||||
{
|
||||
await SendAsync(socket, "heartbeat", RuntimePayload(), cancellationToken);
|
||||
foreach (var job in jobInbox.ReadRecoverableJobs())
|
||||
{
|
||||
if (job.Terminal is JsonElement terminal
|
||||
&& terminal.GetProperty("status").GetString() == "succeeded"
|
||||
&& !job.UploadComplete)
|
||||
{
|
||||
StartJobPipeline(socket, job);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
private static async Task ReceiveLoopAsync(
|
||||
private async Task ReportRecoverableJobsAsync(
|
||||
ClientWebSocket socket, CancellationToken cancellationToken)
|
||||
{
|
||||
foreach (var job in jobInbox.ReadRecoverableJobs())
|
||||
{
|
||||
if (job.Terminal is JsonElement terminal)
|
||||
{
|
||||
if (terminal.GetProperty("status").GetString() == "succeeded")
|
||||
{
|
||||
if (!job.UploadComplete)
|
||||
{
|
||||
StartJobPipeline(socket, job);
|
||||
}
|
||||
}
|
||||
else
|
||||
{
|
||||
await SendAsync(socket, "job_terminal", terminal, cancellationToken);
|
||||
}
|
||||
continue;
|
||||
}
|
||||
await SendAsync(socket, "job_state", new
|
||||
{
|
||||
job_id = job.JobId,
|
||||
lease_id = job.LeaseId,
|
||||
request_digest = job.RequestDigest,
|
||||
stage = "waiting_input",
|
||||
progress = 0,
|
||||
metrics = new { },
|
||||
}, cancellationToken);
|
||||
StartJobPipeline(socket, job);
|
||||
}
|
||||
}
|
||||
|
||||
private async Task ReceiveLoopAsync(
|
||||
ClientWebSocket socket, CancellationToken cancellationToken)
|
||||
{
|
||||
var buffer = new byte[16 * 1024];
|
||||
|
|
@ -160,11 +212,179 @@ internal sealed class NodeConnectionLoop(NodeConfig config, Action<NodeStatus>?
|
|||
if (document.RootElement.TryGetProperty("type", out var type))
|
||||
{
|
||||
Console.WriteLine($"[INFO] Server message: {type.GetString()}.");
|
||||
if (type.GetString() == "job_offer"
|
||||
&& document.RootElement.TryGetProperty("payload", out var payload))
|
||||
{
|
||||
var offerResult = jobInbox.Accept(payload);
|
||||
await SendAsync(
|
||||
socket,
|
||||
offerResult.Accepted ? "job_accept" : "job_reject",
|
||||
offerResult.Accepted
|
||||
? new
|
||||
{
|
||||
job_id = offerResult.JobId,
|
||||
lease_id = offerResult.LeaseId,
|
||||
request_digest = offerResult.RequestDigest,
|
||||
}
|
||||
: new
|
||||
{
|
||||
job_id = payload.TryGetProperty("job_id", out var jobId)
|
||||
? jobId.GetString() : "",
|
||||
lease_id = payload.TryGetProperty("lease_id", out var leaseId)
|
||||
? leaseId.GetString() : "",
|
||||
reason = offerResult.Reason,
|
||||
},
|
||||
cancellationToken);
|
||||
if (offerResult.Accepted)
|
||||
{
|
||||
await SendAsync(socket, "job_state", new
|
||||
{
|
||||
job_id = offerResult.JobId,
|
||||
lease_id = offerResult.LeaseId,
|
||||
request_digest = offerResult.RequestDigest,
|
||||
stage = "waiting_input",
|
||||
progress = 0,
|
||||
metrics = new { },
|
||||
}, cancellationToken);
|
||||
var acceptedJob = jobInbox.ReadRecoverableJobs()
|
||||
.Single(item => item.JobId == offerResult.JobId);
|
||||
StartJobPipeline(socket, acceptedJob);
|
||||
}
|
||||
}
|
||||
}
|
||||
message.SetLength(0);
|
||||
}
|
||||
}
|
||||
|
||||
private void StartJobPipeline(ClientWebSocket socket, RecoverableJob job)
|
||||
{
|
||||
var completion = new TaskCompletionSource(
|
||||
TaskCreationOptions.RunContinuationsAsynchronously);
|
||||
if (jobPipelines.TryAdd(job.JobId, completion.Task))
|
||||
{
|
||||
_ = RunJobPipelineAndReleaseAsync(socket, job, completion);
|
||||
}
|
||||
}
|
||||
|
||||
private async Task RunJobPipelineAndReleaseAsync(
|
||||
ClientWebSocket socket, RecoverableJob job, TaskCompletionSource completion)
|
||||
{
|
||||
try
|
||||
{
|
||||
await RunJobPipelineAsync(socket, job);
|
||||
}
|
||||
finally
|
||||
{
|
||||
completion.TrySetResult();
|
||||
jobPipelines.TryRemove(job.JobId, out _);
|
||||
}
|
||||
}
|
||||
|
||||
private async Task RunJobPipelineAsync(ClientWebSocket socket, RecoverableJob job)
|
||||
{
|
||||
try
|
||||
{
|
||||
if (job.Terminal is null)
|
||||
{
|
||||
await inputDownloader.DownloadAsync(job, CancellationToken.None);
|
||||
await TrySendAsync(socket, "job_state", new
|
||||
{
|
||||
job_id = job.JobId,
|
||||
lease_id = job.LeaseId,
|
||||
request_digest = job.RequestDigest,
|
||||
stage = "ready_to_run",
|
||||
progress = 5,
|
||||
metrics = new { input_bytes = job.InputTransfer?.GetProperty("size_bytes").GetInt64() },
|
||||
});
|
||||
await TrySendAsync(socket, "job_state", new
|
||||
{
|
||||
job_id = job.JobId,
|
||||
lease_id = job.LeaseId,
|
||||
request_digest = job.RequestDigest,
|
||||
stage = "origin_running",
|
||||
progress = 10,
|
||||
metrics = new { },
|
||||
});
|
||||
await workerRunner.RunAsync(job);
|
||||
}
|
||||
var refreshed = jobInbox.ReadRecoverableJobs()
|
||||
.Single(item => item.JobId == job.JobId);
|
||||
if (refreshed.Terminal is not JsonElement terminal)
|
||||
{
|
||||
throw new InvalidDataException("Origin worker did not create a terminal record.");
|
||||
}
|
||||
if (terminal.GetProperty("status").GetString() != "succeeded")
|
||||
{
|
||||
await TrySendAsync(socket, "job_terminal", terminal);
|
||||
return;
|
||||
}
|
||||
await TrySendAsync(socket, "job_state", new
|
||||
{
|
||||
job_id = job.JobId,
|
||||
lease_id = job.LeaseId,
|
||||
request_digest = job.RequestDigest,
|
||||
stage = "uploading_outputs",
|
||||
progress = 90,
|
||||
metrics = new { },
|
||||
});
|
||||
await outputUploader.UploadAsync(refreshed);
|
||||
}
|
||||
catch (Exception exception) when (
|
||||
exception is HttpRequestException
|
||||
or IOException
|
||||
or JsonException
|
||||
or UnauthorizedAccessException
|
||||
or InvalidDataException)
|
||||
{
|
||||
var current = jobInbox.ReadRecoverableJobs()
|
||||
.Single(item => item.JobId == job.JobId);
|
||||
if (current.Terminal is JsonElement terminal
|
||||
&& terminal.GetProperty("status").GetString() == "succeeded")
|
||||
{
|
||||
Console.Error.WriteLine($"[WARN] Output upload deferred: {exception.Message}");
|
||||
return;
|
||||
}
|
||||
jobInbox.WriteTerminal(
|
||||
job,
|
||||
"failed",
|
||||
"INPUT_DOWNLOAD_FAILED",
|
||||
exception.Message[..Math.Min(exception.Message.Length, 500)]);
|
||||
var failedTerminal = jobInbox.ReadRecoverableJobs()
|
||||
.Single(item => item.JobId == job.JobId).Terminal;
|
||||
if (failedTerminal is JsonElement payload)
|
||||
{
|
||||
await TrySendAsync(socket, "job_terminal", payload);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
private async Task TrySendTerminalAsync(ClientWebSocket socket, RecoverableJob job)
|
||||
{
|
||||
var terminal = jobInbox.ReadRecoverableJobs()
|
||||
.Single(item => item.JobId == job.JobId).Terminal;
|
||||
if (terminal is JsonElement terminalPayload)
|
||||
{
|
||||
await TrySendAsync(socket, "job_terminal", terminalPayload);
|
||||
}
|
||||
}
|
||||
|
||||
private async Task TrySendAsync(ClientWebSocket socket, string type, object payload)
|
||||
{
|
||||
if (socket.State != WebSocketState.Open)
|
||||
{
|
||||
return;
|
||||
}
|
||||
try
|
||||
{
|
||||
await SendAsync(socket, type, payload, CancellationToken.None);
|
||||
}
|
||||
catch (Exception exception) when (
|
||||
exception is WebSocketException or IOException or ObjectDisposedException)
|
||||
{
|
||||
Console.Error.WriteLine($"[WARN] Job report deferred: {exception.Message}");
|
||||
}
|
||||
}
|
||||
|
||||
private void Report(NodeState state, string message) =>
|
||||
statusChanged?.Invoke(NodeStatus.Create(state, message));
|
||||
|
||||
|
|
@ -179,8 +399,16 @@ internal sealed class NodeConnectionLoop(NodeConfig config, Action<NodeStatus>?
|
|||
sent_at = DateTimeOffset.UtcNow,
|
||||
payload,
|
||||
});
|
||||
await socket.SendAsync(
|
||||
envelope, WebSocketMessageType.Text, endOfMessage: true, cancellationToken);
|
||||
await sendLock.WaitAsync(cancellationToken);
|
||||
try
|
||||
{
|
||||
await socket.SendAsync(
|
||||
envelope, WebSocketMessageType.Text, endOfMessage: true, cancellationToken);
|
||||
}
|
||||
finally
|
||||
{
|
||||
sendLock.Release();
|
||||
}
|
||||
}
|
||||
|
||||
private object RuntimePayload()
|
||||
|
|
@ -194,7 +422,7 @@ internal sealed class NodeConnectionLoop(NodeConfig config, Action<NodeStatus>?
|
|||
node_version = Assembly.GetExecutingAssembly().GetName().Version?.ToString(3) ?? "0.1.0",
|
||||
os_version = RuntimeInformation.OSDescription,
|
||||
capabilities = config.Capabilities,
|
||||
available_slots = origin.Health == "ready" ? 1 : 0,
|
||||
available_slots = origin.Health == "ready" && !jobInbox.HasPendingJobs ? 1 : 0,
|
||||
disk_free_bytes = new DriveInfo(root).AvailableFreeSpace,
|
||||
desktop_session = Environment.UserInteractive,
|
||||
origin = new
|
||||
|
|
|
|||
|
|
@ -34,14 +34,15 @@ internal sealed record EnrollResponse(
|
|||
[property: JsonPropertyName("heartbeat_seconds")] int HeartbeatSeconds,
|
||||
[property: JsonPropertyName("max_concurrency")] int MaxConcurrency);
|
||||
|
||||
internal sealed record NodePaths(string RootDirectory, string ConfigPath)
|
||||
internal sealed record NodePaths(string RootDirectory, string ConfigPath, string JobsDirectory)
|
||||
{
|
||||
internal static NodePaths ForCurrentMachine()
|
||||
{
|
||||
var root = Path.Combine(
|
||||
Environment.GetFolderPath(Environment.SpecialFolder.CommonApplicationData),
|
||||
"Zcbot", "WindowsNode");
|
||||
return new NodePaths(root, Path.Combine(root, "node.json"));
|
||||
return new NodePaths(
|
||||
root, Path.Combine(root, "node.json"), Path.Combine(root, "jobs"));
|
||||
}
|
||||
}
|
||||
|
||||
|
|
|
|||
|
|
@ -36,7 +36,15 @@ internal static class OriginRuntimeProbe
|
|||
{
|
||||
return Status(version, "unavailable", "Origin 需要交互式 Windows 桌面会话");
|
||||
}
|
||||
return Status(version, "ready", "Origin COM 自动化组件可用");
|
||||
var interpreter = OriginWorkerRuntime.ResolveInterpreter();
|
||||
if (interpreter is null)
|
||||
{
|
||||
return Status(
|
||||
version,
|
||||
"unavailable",
|
||||
"Origin 可用，但固定 Python 运行时缺失；请配置 ZCBOT_ORIGIN_PYTHON");
|
||||
}
|
||||
return Status(version, "ready", $"Origin COM 与固定 Python 运行时可用（{interpreter}）");
|
||||
}
|
||||
catch (Exception exception) when (
|
||||
exception is SecurityException or UnauthorizedAccessException or IOException)
|
||||
|
|
@ -46,7 +54,7 @@ internal static class OriginRuntimeProbe
|
|||
}
|
||||
|
||||
private static OriginRuntimeStatus Status(string? version, string health, string detail) =>
|
||||
new("OriginPro", version, "0.1.0", health, detail);
|
||||
new("OriginPro", version, "0.2.0", health, detail);
|
||||
|
||||
private static string? FindInstalledVersion()
|
||||
{
|
||||
|
|
|
|||
|
|
@ -0,0 +1,147 @@
|
|||
using System.Collections.Concurrent;
|
||||
using System.Diagnostics;
|
||||
using System.Text;
|
||||
using System.Text.Json;
|
||||
|
||||
namespace Zcbot.WindowsNode;
|
||||
|
||||
internal sealed class OriginWorkerRunner(JobInboxStore inbox)
|
||||
{
|
||||
private static readonly TimeSpan WorkerTimeout = TimeSpan.FromMinutes(30);
|
||||
private readonly ConcurrentDictionary<Guid, Task> active = new();
|
||||
|
||||
internal Task RunAsync(RecoverableJob job) =>
|
||||
active.GetOrAdd(job.JobId, _ => RunOnceAsync(job));
|
||||
|
||||
private async Task RunOnceAsync(RecoverableJob job)
|
||||
{
|
||||
try
|
||||
{
|
||||
var paths = NodePaths.ForCurrentMachine();
|
||||
var jobDirectory = Path.Combine(paths.JobsDirectory, job.JobId.ToString("D"));
|
||||
var terminalPath = Path.Combine(jobDirectory, "terminal.json");
|
||||
if (File.Exists(terminalPath)) return;
|
||||
|
||||
var markerPath = Path.Combine(jobDirectory, "worker-started.json");
|
||||
if (File.Exists(markerPath))
|
||||
{
|
||||
inbox.WriteTerminal(
|
||||
job,
|
||||
"failed",
|
||||
"NODE_RESTARTED_DURING_JOB",
|
||||
"The node restarted after Origin execution began and cannot prove the prior worker state.");
|
||||
return;
|
||||
}
|
||||
|
||||
var interpreter = OriginWorkerRuntime.ResolveInterpreter()
|
||||
?? throw new InvalidOperationException("The fixed Origin Python interpreter is unavailable.");
|
||||
var workerScript = Path.GetFullPath(
|
||||
Path.Combine(AppContext.BaseDirectory, "origin-worker", "worker.py"));
|
||||
if (!File.Exists(workerScript))
|
||||
{
|
||||
throw new FileNotFoundException("The fixed Origin worker script is missing.", workerScript);
|
||||
}
|
||||
WriteMarker(markerPath, interpreter, workerScript);
|
||||
|
||||
var startInfo = new ProcessStartInfo
|
||||
{
|
||||
FileName = interpreter,
|
||||
WorkingDirectory = jobDirectory,
|
||||
UseShellExecute = false,
|
||||
CreateNoWindow = true,
|
||||
RedirectStandardOutput = true,
|
||||
RedirectStandardError = true,
|
||||
StandardOutputEncoding = Encoding.UTF8,
|
||||
StandardErrorEncoding = Encoding.UTF8,
|
||||
};
|
||||
startInfo.ArgumentList.Add(workerScript);
|
||||
startInfo.ArgumentList.Add(jobDirectory);
|
||||
using var process = Process.Start(startInfo)
|
||||
?? throw new InvalidOperationException("The fixed Origin worker did not start.");
|
||||
var stdout = process.StandardOutput.ReadToEndAsync();
|
||||
var stderr = process.StandardError.ReadToEndAsync();
|
||||
using var timeout = new CancellationTokenSource();
|
||||
timeout.CancelAfter(WorkerTimeout);
|
||||
try
|
||||
{
|
||||
await process.WaitForExitAsync(timeout.Token);
|
||||
}
|
||||
catch (OperationCanceledException)
|
||||
{
|
||||
process.Kill(entireProcessTree: true);
|
||||
inbox.WriteTerminal(job, "failed", "ORIGIN_WORKER_TIMEOUT", "Origin worker exceeded 30 minutes.");
|
||||
return;
|
||||
}
|
||||
var output = await stdout;
|
||||
var error = await stderr;
|
||||
WriteDiagnostic(jobDirectory, output, error, process.ExitCode);
|
||||
if (!File.Exists(terminalPath))
|
||||
{
|
||||
inbox.WriteTerminal(
|
||||
job,
|
||||
"failed",
|
||||
"ORIGIN_WORKER_NO_TERMINAL",
|
||||
$"Origin worker exited with code {process.ExitCode} without terminal.json.");
|
||||
}
|
||||
}
|
||||
catch (Exception exception) when (
|
||||
exception is IOException
|
||||
or JsonException
|
||||
or UnauthorizedAccessException
|
||||
or InvalidOperationException)
|
||||
{
|
||||
inbox.WriteTerminal(job, "failed", "ORIGIN_WORKER_START_FAILED", exception.Message[..Math.Min(500, exception.Message.Length)]);
|
||||
}
|
||||
finally
|
||||
{
|
||||
active.TryRemove(job.JobId, out _);
|
||||
}
|
||||
}
|
||||
|
||||
private static void WriteMarker(string path, string interpreter, string workerScript)
|
||||
{
|
||||
var value = JsonSerializer.SerializeToUtf8Bytes(new
|
||||
{
|
||||
started_at = DateTimeOffset.UtcNow,
|
||||
node_pid = Environment.ProcessId,
|
||||
interpreter,
|
||||
worker_script = workerScript,
|
||||
});
|
||||
using var stream = new FileStream(
|
||||
path, FileMode.CreateNew, FileAccess.Write, FileShare.None,
|
||||
bufferSize: 4096, FileOptions.WriteThrough);
|
||||
stream.Write(value);
|
||||
stream.Flush(flushToDisk: true);
|
||||
}
|
||||
|
||||
private static void WriteDiagnostic(string jobDirectory, string output, string error, int exitCode)
|
||||
{
|
||||
var logs = Path.Combine(jobDirectory, "logs");
|
||||
Directory.CreateDirectory(logs);
|
||||
var value = JsonSerializer.Serialize(new
|
||||
{
|
||||
exit_code = exitCode,
|
||||
stdout = output[..Math.Min(output.Length, 16 * 1024)],
|
||||
stderr = error[..Math.Min(error.Length, 16 * 1024)],
|
||||
});
|
||||
File.WriteAllText(Path.Combine(logs, "worker-process.json"), value, Encoding.UTF8);
|
||||
}
|
||||
}
|
||||
|
||||
internal static class OriginWorkerRuntime
|
||||
{
|
||||
internal static string? ResolveInterpreter()
|
||||
{
|
||||
var paths = NodePaths.ForCurrentMachine();
|
||||
var configured = Environment.GetEnvironmentVariable("ZCBOT_ORIGIN_PYTHON");
|
||||
var candidate = string.IsNullOrWhiteSpace(configured)
|
||||
? Path.Combine(paths.RootDirectory, "runtimes", "origin", "python.exe")
|
||||
: configured;
|
||||
if (!Path.IsPathFullyQualified(candidate)) return null;
|
||||
var resolved = Path.GetFullPath(candidate);
|
||||
return File.Exists(resolved)
|
||||
&& Path.GetFileName(resolved).Equals("python.exe", StringComparison.OrdinalIgnoreCase)
|
||||
? resolved
|
||||
: null;
|
||||
}
|
||||
}
|
||||
|
|
@ -11,4 +11,10 @@
|
|||
<RootNamespace>Zcbot.WindowsNode</RootNamespace>
|
||||
<Version>0.1.0</Version>
|
||||
</PropertyGroup>
|
||||
<ItemGroup>
|
||||
<Content Include="..\origin-worker\worker.py">
|
||||
<Link>origin-worker\worker.py</Link>
|
||||
<CopyToOutputDirectory>PreserveNewest</CopyToOutputDirectory>
|
||||
</Content>
|
||||
</ItemGroup>
|
||||
</Project>
|
||||
|
|
|
|||
|
|
@ -0,0 +1,32 @@
|
|||
param(
|
||||
[Parameter(Mandatory = $true)]
|
||||
[string]$BootstrapPython,
|
||||
[string]$RuntimeDirectory = "$env:ProgramData\Zcbot\WindowsNode\runtimes\origin"
|
||||
)
|
||||
|
||||
$ErrorActionPreference = "Stop"
|
||||
$python = (Resolve-Path -LiteralPath $BootstrapPython).Path
|
||||
if ([IO.Path]::GetFileName($python) -ne "python.exe") {
|
||||
throw "BootstrapPython must point to python.exe."
|
||||
}
|
||||
$requirements = Join-Path $PSScriptRoot "origin-worker\requirements.txt"
|
||||
if (-not (Test-Path -LiteralPath $requirements -PathType Leaf)) {
|
||||
throw "Pinned Origin worker requirements are missing."
|
||||
}
|
||||
$runtime = [IO.Path]::GetFullPath($RuntimeDirectory)
|
||||
if ($runtime -eq [IO.Path]::GetPathRoot($runtime)) {
|
||||
throw "RuntimeDirectory cannot be a drive root."
|
||||
}
|
||||
|
||||
if (-not (Test-Path -LiteralPath (Join-Path $runtime "python.exe"))) {
|
||||
& $python -m venv $runtime
|
||||
if ($LASTEXITCODE -ne 0) { throw "Failed to create the Origin runtime." }
|
||||
}
|
||||
$runtimePython = Join-Path $runtime "python.exe"
|
||||
& $runtimePython -m pip install --requirement $requirements
|
||||
if ($LASTEXITCODE -ne 0) { throw "Failed to install the pinned Origin runtime packages." }
|
||||
& $runtimePython -c "import originpro, openpyxl; print('[OK] Origin worker Python packages are available.')"
|
||||
if ($LASTEXITCODE -ne 0) { throw "Origin runtime import verification failed." }
|
||||
|
||||
Write-Output "[OK] Fixed Origin runtime installed: $runtimePython"
|
||||
Write-Output "[INFO] Restart zcbot Windows Node to refresh runtime health."
|
||||
|
|
@ -0,0 +1,2 @@
|
|||
originpro==1.1.15
|
||||
openpyxl==3.1.5
|
||||
|
|
@ -0,0 +1,252 @@
|
|||
"""Fixed Origin adapter for origin.plot@v1.
|
||||
|
||||
This process accepts exactly one argument: a Node-created job directory. It never
|
||||
installs packages, evaluates user code, downloads data, or resolves paths from the
|
||||
request. terminal.json is its only terminal-state contract.
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
import csv
|
||||
import hashlib
|
||||
import json
|
||||
import os
|
||||
import sys
|
||||
from datetime import datetime, timezone
|
||||
from importlib.metadata import PackageNotFoundError, version
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
PLOT_TYPES = {"line": "l", "scatter": "s", "line_scatter": "y"}
|
||||
FORMATS = {"opju", "png", "svg", "pdf"}
|
||||
|
||||
|
||||
def _atomic_json(path: Path, value: Any) -> None:
|
||||
temporary = path.with_name(path.name + ".tmp-" + os.urandom(8).hex())
|
||||
try:
|
||||
with temporary.open("w", encoding="utf-8", newline="\n") as handle:
|
||||
json.dump(value, handle, ensure_ascii=False, indent=2)
|
||||
handle.flush()
|
||||
os.fsync(handle.fileno())
|
||||
os.replace(temporary, path)
|
||||
finally:
|
||||
temporary.unlink(missing_ok=True)
|
||||
|
||||
|
||||
def _read_rows(path: Path, sheet: str | None) -> tuple[list[str], list[list[Any]]]:
|
||||
suffix = path.suffix.lower()
|
||||
if suffix == ".csv":
|
||||
with path.open("r", encoding="utf-8-sig", newline="") as handle:
|
||||
rows = list(csv.reader(handle))
|
||||
if len(rows) < 2:
|
||||
raise ValueError("CSV_INPUT_EMPTY")
|
||||
return [str(item) for item in rows[0]], rows[1:]
|
||||
if suffix == ".json":
|
||||
value = json.loads(path.read_text(encoding="utf-8"))
|
||||
if isinstance(value, list) and value and all(isinstance(item, dict) for item in value):
|
||||
headers = list(value[0])
|
||||
return headers, [[item.get(name) for name in headers] for item in value]
|
||||
if isinstance(value, dict) and value and all(isinstance(item, list) for item in value.values()):
|
||||
headers = list(value)
|
||||
length = max(len(value[name]) for name in headers)
|
||||
return headers, [[value[name][index] if index < len(value[name]) else None for name in headers] for index in range(length)]
|
||||
raise ValueError("JSON_INPUT_SHAPE_UNSUPPORTED")
|
||||
if suffix == ".xlsx":
|
||||
from openpyxl import load_workbook
|
||||
|
||||
workbook = load_workbook(path, read_only=True, data_only=True)
|
||||
try:
|
||||
worksheet = workbook[sheet] if sheet else workbook.active
|
||||
rows = list(worksheet.iter_rows(values_only=True))
|
||||
finally:
|
||||
workbook.close()
|
||||
if len(rows) < 2:
|
||||
raise ValueError("XLSX_INPUT_EMPTY")
|
||||
return [str(item or "") for item in rows[0]], [list(row) for row in rows[1:]]
|
||||
raise ValueError("INPUT_TYPE_UNSUPPORTED")
|
||||
|
||||
|
||||
def _column_index(headers: list[str], value: Any, field: str) -> int:
|
||||
if not isinstance(value, str) or value not in headers:
|
||||
raise ValueError(f"{field.upper()}_COLUMN_NOT_FOUND")
|
||||
return headers.index(value)
|
||||
|
||||
|
||||
def _manifest(path: Path, media_type: str) -> dict[str, Any]:
|
||||
return {
|
||||
"artifact_id": {
|
||||
"project.opju": "project",
|
||||
"figure.png": "figure_png",
|
||||
"figure.svg": "figure_svg",
|
||||
"figure.pdf": "figure_pdf",
|
||||
"plot-spec.json": "plot_spec",
|
||||
"provenance.json": "provenance",
|
||||
}[path.name],
|
||||
"filename": path.name,
|
||||
"media_type": media_type,
|
||||
"size_bytes": path.stat().st_size,
|
||||
"sha256": _file_sha256(path),
|
||||
}
|
||||
|
||||
|
||||
def _file_sha256(path: Path) -> str:
|
||||
digest = hashlib.sha256()
|
||||
with path.open("rb") as handle:
|
||||
for chunk in iter(lambda: handle.read(1024 * 1024), b""):
|
||||
digest.update(chunk)
|
||||
return digest.hexdigest()
|
||||
|
||||
|
||||
def _axis_title(axis: Any, fallback: str) -> str:
|
||||
if not isinstance(axis, dict):
|
||||
return fallback
|
||||
title = str(axis.get("title") or fallback)
|
||||
unit = str(axis.get("unit") or "")
|
||||
return f"{title} ({unit})" if unit else title
|
||||
|
||||
|
||||
def _validate_artifact(path: Path, extension: str) -> None:
|
||||
if not path.is_file() or path.stat().st_size == 0:
|
||||
raise RuntimeError(f"{extension.upper()}_EXPORT_EMPTY")
|
||||
head = path.read_bytes()[:1024]
|
||||
if extension == "png" and not head.startswith(b"\x89PNG\r\n\x1a\n"):
|
||||
raise RuntimeError("PNG_EXPORT_INVALID")
|
||||
if extension == "pdf" and not head.startswith(b"%PDF-"):
|
||||
raise RuntimeError("PDF_EXPORT_INVALID")
|
||||
if extension == "svg" and b"<svg" not in head.lower():
|
||||
raise RuntimeError("SVG_EXPORT_INVALID")
|
||||
if extension == "opju" and len(head) < 64:
|
||||
raise RuntimeError("OPJU_EXPORT_INVALID")
|
||||
|
||||
|
||||
def run(job_dir: Path) -> list[dict[str, Any]]:
|
||||
job_dir = job_dir.resolve(strict=True)
|
||||
request_record = json.loads((job_dir / "request" / "request.json").read_text(encoding="utf-8"))
|
||||
request = request_record["request"]
|
||||
input_files = [path for path in (job_dir / "input").iterdir() if path.is_file() and not path.name.startswith(".")]
|
||||
if len(input_files) != 1:
|
||||
raise ValueError("INPUT_FILE_COUNT_INVALID")
|
||||
headers, rows = _read_rows(input_files[0], request["input"].get("sheet"))
|
||||
plot_spec = request["plot"]
|
||||
plot_type = plot_spec["type"]
|
||||
if plot_type not in PLOT_TYPES:
|
||||
raise ValueError("PLOT_TYPE_NOT_IMPLEMENTED")
|
||||
x_index = _column_index(headers, plot_spec.get("x"), "x")
|
||||
y_names = plot_spec.get("y")
|
||||
if isinstance(y_names, str):
|
||||
y_names = [y_names]
|
||||
if not isinstance(y_names, list) or not y_names:
|
||||
raise ValueError("Y_COLUMNS_REQUIRED")
|
||||
y_indexes = [_column_index(headers, name, "y") for name in y_names]
|
||||
|
||||
import originpro as op
|
||||
|
||||
output = job_dir / "output"
|
||||
output.mkdir(exist_ok=True)
|
||||
op.set_show(False)
|
||||
try:
|
||||
op.new()
|
||||
worksheet = op.new_sheet("w", lname="Data")
|
||||
for index, header in enumerate(headers):
|
||||
worksheet.from_list(index, [row[index] if index < len(row) else None for row in rows], lname=header)
|
||||
graph = op.new_graph(template={"line": "line", "scatter": "scatter", "line_scatter": "linesymb"}[plot_type])
|
||||
layer = graph[0]
|
||||
for y_index in y_indexes:
|
||||
layer.add_plot(worksheet, coly=y_index, colx=x_index, type=PLOT_TYPES[plot_type])
|
||||
layer.rescale()
|
||||
layer.axis("x").title = _axis_title(plot_spec.get("x_axis"), str(plot_spec.get("x") or "X"))
|
||||
layer.axis("y").title = _axis_title(plot_spec.get("y_axis"), "Y")
|
||||
if plot_spec.get("title"):
|
||||
title = layer.add_label(str(plot_spec["title"]))
|
||||
title.set_int("fsize", 18)
|
||||
title.set_int("left", 2200)
|
||||
title.set_int("top", 120)
|
||||
formats = request["output"]["formats"]
|
||||
if any(item not in FORMATS for item in formats):
|
||||
raise ValueError("OUTPUT_FORMAT_UNSUPPORTED")
|
||||
artifacts: list[dict[str, Any]] = []
|
||||
if "opju" in formats:
|
||||
project = output / "project.opju"
|
||||
op.save(str(project))
|
||||
_validate_artifact(project, "opju")
|
||||
artifacts.append(_manifest(project, "application/x-origin-project"))
|
||||
media = {"png": "image/png", "svg": "image/svg+xml", "pdf": "application/pdf"}
|
||||
dpi = request["output"].get("dpi", 300)
|
||||
if not isinstance(dpi, int) or isinstance(dpi, bool) or not 72 <= dpi <= 1200:
|
||||
raise ValueError("OUTPUT_DPI_INVALID")
|
||||
pixel_width = round(dpi * 160 / 25.4)
|
||||
for extension in ("png", "svg", "pdf"):
|
||||
if extension in formats:
|
||||
target = output / f"figure.{extension}"
|
||||
exported = Path(graph.save_fig(
|
||||
str(target),
|
||||
type=extension,
|
||||
width=pixel_width if extension == "png" else 0,
|
||||
ratio=100 if extension in {"svg", "pdf"} else 0,
|
||||
)).resolve()
|
||||
if exported != target.resolve() or not target.is_file():
|
||||
raise RuntimeError(f"{extension.upper()}_EXPORT_FAILED")
|
||||
_validate_artifact(target, extension)
|
||||
artifacts.append(_manifest(target, media[extension]))
|
||||
try:
|
||||
originpro_version = version("originpro")
|
||||
except PackageNotFoundError:
|
||||
originpro_version = "embedded"
|
||||
provenance = {
|
||||
"adapter_version": "0.2.0",
|
||||
"originpro_version": originpro_version,
|
||||
"request_digest": request_record["request_digest"],
|
||||
"input_sha256": _file_sha256(input_files[0]),
|
||||
"requested_dpi": dpi,
|
||||
"png_pixel_width": pixel_width,
|
||||
}
|
||||
plot_spec_path = output / "plot-spec.json"
|
||||
provenance_path = output / "provenance.json"
|
||||
_atomic_json(plot_spec_path, request)
|
||||
_atomic_json(provenance_path, provenance)
|
||||
artifacts.append(_manifest(plot_spec_path, "application/json"))
|
||||
artifacts.append(_manifest(provenance_path, "application/json"))
|
||||
return artifacts
|
||||
finally:
|
||||
if op.oext:
|
||||
op.exit()
|
||||
|
||||
|
||||
def main() -> int:
|
||||
if len(sys.argv) != 2:
|
||||
print("[ERR] Usage: worker.py <job-directory>", file=sys.stderr)
|
||||
return 2
|
||||
job_dir = Path(sys.argv[1])
|
||||
request_record: dict[str, Any] = {}
|
||||
try:
|
||||
request_record = json.loads((job_dir / "request" / "request.json").read_text(encoding="utf-8"))
|
||||
artifacts = run(job_dir)
|
||||
terminal = {
|
||||
"job_id": request_record["job_id"],
|
||||
"lease_id": request_record["lease_id"],
|
||||
"request_digest": request_record["request_digest"],
|
||||
"status": "succeeded",
|
||||
"error": {},
|
||||
"artifact_manifest": artifacts,
|
||||
"terminal_at": datetime.now(timezone.utc).isoformat(),
|
||||
}
|
||||
_atomic_json(job_dir / "artifacts.json", artifacts)
|
||||
_atomic_json(job_dir / "terminal.json", terminal)
|
||||
print("[OK] Origin job completed.")
|
||||
return 0
|
||||
except Exception as exception:
|
||||
terminal = {
|
||||
"job_id": request_record.get("job_id", ""),
|
||||
"lease_id": request_record.get("lease_id", ""),
|
||||
"request_digest": request_record.get("request_digest", ""),
|
||||
"status": "failed",
|
||||
"error": {"code": type(exception).__name__, "detail": str(exception)[:500]},
|
||||
"artifact_manifest": [],
|
||||
"terminal_at": datetime.now(timezone.utc).isoformat(),
|
||||
}
|
||||
_atomic_json(job_dir / "terminal.json", terminal)
|
||||
print(f"[ERR] {type(exception).__name__}: {exception}", file=sys.stderr)
|
||||
return 1
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
raise SystemExit(main())
|
||||
Loading…
Reference in New Issue