771 lines
30 KiB
Python
771 lines
30 KiB
Python
"""专业软件任务的校验、幂等持久化和 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 and_, desc, or_, select
|
||
from sqlalchemy.exc import IntegrityError
|
||
|
||
from core.software_nodes import SUPPORTED_CAPABILITIES
|
||
from core.storage.engine import session_scope
|
||
from core.storage.models import Artifact, SoftwareJob, SoftwareNode, 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
|
||
class SoftwareJobError(Exception):
|
||
pass
|
||
|
||
|
||
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 SoftwareJobError("invalid origin plot request fields")
|
||
if request.get("schema_version") != 1:
|
||
raise SoftwareJobError("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 SoftwareJobError("origin plot request sections must be objects")
|
||
if not _has_only(input_spec, {"input_id", "sheet"}):
|
||
raise SoftwareJobError("unsupported origin input fields")
|
||
try:
|
||
UUID(str(input_spec.get("input_id") or ""))
|
||
except ValueError as exc:
|
||
raise SoftwareJobError("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 SoftwareJobError("input.sheet must be a string")
|
||
if not _has_only(
|
||
plot,
|
||
{"type", "x", "y", "template", "title", "x_axis", "y_axis", "legend", "error_bars"},
|
||
):
|
||
raise SoftwareJobError("unsupported origin plot fields")
|
||
if plot.get("type") not in ALLOWED_PLOT_TYPES:
|
||
raise SoftwareJobError("unsupported origin plot type")
|
||
if "title" in plot and (
|
||
not isinstance(plot["title"], str) or len(plot["title"]) > 500
|
||
):
|
||
raise SoftwareJobError("plot.title must be a string")
|
||
if plot.get("template", "publication_double_column") != "publication_double_column":
|
||
raise SoftwareJobError("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 SoftwareJobError("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 SoftwareJobError("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 SoftwareJobError(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 SoftwareJobError("invalid plot.legend")
|
||
if plot.get("error_bars") is not None:
|
||
raise SoftwareJobError("error bars are not supported in origin.plot@v1")
|
||
if not _has_only(output, {"formats", "dpi", "capture_screenshots", "record_video"}):
|
||
raise SoftwareJobError("unsupported origin output fields")
|
||
if any(
|
||
name in output and not isinstance(output[name], bool)
|
||
for name in ("capture_screenshots", "record_video")
|
||
):
|
||
raise SoftwareJobError("origin output capture flags must be boolean")
|
||
if output.get("record_video", False):
|
||
raise SoftwareJobError("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 SoftwareJobError("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 SoftwareJobError("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 SoftwareJobError("origin plot request is too large")
|
||
normalized = json.loads(encoded)
|
||
return normalized, sha256(encoded.encode("utf-8")).hexdigest()
|
||
|
||
|
||
def _job_dict(row: SoftwareJob) -> 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 list_jobs(
|
||
user_id: UUID,
|
||
*,
|
||
task_id: UUID | None = None,
|
||
active_only: bool = False,
|
||
limit: int = 50,
|
||
before: tuple[datetime, UUID] | None = None,
|
||
) -> list[dict]:
|
||
"""列出用户的软件任务；用于全局 Job 中心和 Agent 查询。"""
|
||
# Web 分页会多取一条判断是否还有下一页；公开接口仍把页大小限制为 100。
|
||
limit = max(1, min(int(limit), 101))
|
||
with session_scope() as session:
|
||
statement = (
|
||
select(SoftwareJob, Task.name, SoftwareNode.name)
|
||
.join(Task, Task.task_id == SoftwareJob.task_id)
|
||
.outerjoin(SoftwareNode, SoftwareNode.node_id == SoftwareJob.node_id)
|
||
.where(SoftwareJob.user_id == user_id)
|
||
)
|
||
if task_id is not None:
|
||
statement = statement.where(SoftwareJob.task_id == task_id)
|
||
if active_only:
|
||
statement = statement.where(
|
||
SoftwareJob.status.in_(
|
||
{
|
||
"queued",
|
||
"offered",
|
||
"dispatched",
|
||
"running",
|
||
"disconnected",
|
||
"cancelling",
|
||
}
|
||
)
|
||
)
|
||
if before is not None:
|
||
before_created_at, before_job_id = before
|
||
statement = statement.where(
|
||
or_(
|
||
SoftwareJob.created_at < before_created_at,
|
||
and_(
|
||
SoftwareJob.created_at == before_created_at,
|
||
SoftwareJob.job_id < before_job_id,
|
||
),
|
||
)
|
||
)
|
||
rows = session.execute(
|
||
statement.order_by(
|
||
desc(SoftwareJob.created_at), desc(SoftwareJob.job_id)
|
||
).limit(limit)
|
||
).all()
|
||
results: list[dict] = []
|
||
for job, task_name, node_name in rows:
|
||
item = _job_dict(job)
|
||
item.update(
|
||
{
|
||
"task_name": task_name,
|
||
"node_name": node_name,
|
||
"input": job.input_manifest,
|
||
"request_summary": _request_summary(job),
|
||
}
|
||
)
|
||
results.append(item)
|
||
return results
|
||
|
||
|
||
def _request_summary(job: SoftwareJob) -> dict:
|
||
plot = job.request.get("plot") or {}
|
||
output = job.request.get("output") or {}
|
||
return {
|
||
"display_name": (
|
||
"Origin 科研绘图"
|
||
if job.capability == "origin.plot@v1"
|
||
else job.capability
|
||
),
|
||
"title": str(plot.get("title") or ""),
|
||
"formats": list(output.get("formats") or []),
|
||
}
|
||
|
||
|
||
def request_job_cancel(user_id: UUID, job_id: UUID) -> tuple[dict, dict | None]:
|
||
"""持久化取消意图；queued 直接终止，已分派任务返回 Node 消息。"""
|
||
now = datetime.now(timezone.utc)
|
||
with session_scope() as session:
|
||
job = session.execute(
|
||
select(SoftwareJob).where(
|
||
SoftwareJob.job_id == job_id,
|
||
SoftwareJob.user_id == user_id,
|
||
).with_for_update()
|
||
).scalar_one_or_none()
|
||
if job is None:
|
||
raise SoftwareJobError("job not found")
|
||
if job.status in {"succeeded", "failed", "cancelled"}:
|
||
return _job_dict(job), None
|
||
if job.status == "queued" or job.node_id is None or job.lease_id is None:
|
||
job.status = "cancelled"
|
||
job.stage = "terminal"
|
||
job.error = {"code": "USER_CANCELLED", "detail": "Cancelled before dispatch."}
|
||
job.terminal_at = now
|
||
return _job_dict(job), None
|
||
job.status = "cancelling"
|
||
job.stage = "cancel_requested"
|
||
payload = {
|
||
"job_id": str(job.job_id),
|
||
"lease_id": str(job.lease_id),
|
||
"request_digest": job.request_digest,
|
||
}
|
||
return _job_dict(job), {"node_id": job.node_id, "payload": payload}
|
||
|
||
|
||
def pending_node_cancellations(node_id: UUID) -> list[dict]:
|
||
"""节点重连或心跳时重放未确认的取消意图。"""
|
||
with session_scope() as session:
|
||
rows = session.execute(
|
||
select(SoftwareJob).where(
|
||
SoftwareJob.node_id == node_id,
|
||
SoftwareJob.status == "cancelling",
|
||
)
|
||
).scalars()
|
||
return [
|
||
{
|
||
"job_id": str(job.job_id),
|
||
"lease_id": str(job.lease_id),
|
||
"request_digest": job.request_digest,
|
||
}
|
||
for job in rows
|
||
if job.lease_id is not 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 SoftwareJobError("idempotency_key must contain 1 to 200 characters")
|
||
if capability not in SUPPORTED_CAPABILITIES:
|
||
raise SoftwareJobError("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 SoftwareJobError("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 SoftwareJobError("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 SoftwareJobError("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 SoftwareJobError("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(SoftwareJob).where(
|
||
SoftwareJob.user_id == user_id,
|
||
SoftwareJob.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 SoftwareJobError("idempotency key was already used for a different request")
|
||
return _job_dict(existing), False
|
||
row = SoftwareJob(
|
||
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(SoftwareJob).where(
|
||
SoftwareJob.user_id == user_id,
|
||
SoftwareJob.idempotency_key == key,
|
||
)
|
||
).scalar_one()
|
||
if (
|
||
existing.task_id != task_id
|
||
or existing.capability != capability
|
||
or existing.request_digest != digest
|
||
):
|
||
raise SoftwareJobError(
|
||
"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(SoftwareJob).where(SoftwareJob.job_id == job_id, SoftwareJob.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(SoftwareJob)
|
||
.where(
|
||
SoftwareJob.status == "offered",
|
||
SoftwareJob.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(SoftwareJob)
|
||
.where(SoftwareJob.status == "queued")
|
||
.order_by(SoftwareJob.created_at, SoftwareJob.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(SoftwareJob.node_id).where(
|
||
SoftwareJob.node_id.is_not(None),
|
||
SoftwareJob.status.in_({"offered", "dispatched", "running"}),
|
||
)
|
||
).scalars()
|
||
)
|
||
nodes = session.execute(
|
||
select(SoftwareNode)
|
||
.where(SoftwareNode.node_id.in_(node_ids), SoftwareNode.status == "online")
|
||
.order_by(SoftwareNode.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/software-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(SoftwareJob).where(
|
||
SoftwareJob.job_id == job_id,
|
||
SoftwareJob.node_id == node_id,
|
||
SoftwareJob.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(SoftwareJob, Task.working_dir)
|
||
.join(Task, Task.task_id == SoftwareJob.task_id)
|
||
.where(SoftwareJob.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 SoftwareJobError("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 SoftwareJobError("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 SoftwareJobError("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 SoftwareJobError("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 SoftwareJobError("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 SoftwareJobError("job output artifact size is invalid")
|
||
if not isinstance(digest, str) or not re.fullmatch(r"[0-9a-f]{64}", digest):
|
||
raise SoftwareJobError("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 SoftwareJobError("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(SoftwareJob).where(SoftwareJob.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 SoftwareJobError("invalid job offer response identity") from exc
|
||
now = datetime.now(timezone.utc)
|
||
with session_scope() as session:
|
||
job = session.execute(
|
||
select(SoftwareJob).where(SoftwareJob.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 SoftwareJobError("job offer is stale or does not belong to this node")
|
||
if (
|
||
accepted
|
||
and job.status in {"dispatched", "running", "cancelling", "succeeded", "failed", "cancelled"}
|
||
and payload.get("request_digest") == job.request_digest
|
||
):
|
||
return
|
||
if job.status != "offered":
|
||
raise SoftwareJobError("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 SoftwareJobError("job offer has expired")
|
||
if accepted:
|
||
if payload.get("request_digest") != job.request_digest:
|
||
raise SoftwareJobError("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 SoftwareJobError("job stage is required")
|
||
if not isinstance(progress, int) or isinstance(progress, bool) or not 0 <= progress <= 100:
|
||
raise SoftwareJobError("job progress must be between 0 and 100")
|
||
if not isinstance(metrics, dict) or len(json.dumps(metrics, ensure_ascii=False)) > 64 * 1024:
|
||
raise SoftwareJobError("job metrics are invalid")
|
||
now = datetime.now(timezone.utc)
|
||
with session_scope() as session:
|
||
job = session.execute(
|
||
select(SoftwareJob).where(SoftwareJob.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", "cancelling"}:
|
||
return
|
||
if not _can_accept_state(job.status):
|
||
raise SoftwareJobError("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 SoftwareJobError("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 SoftwareJobError("job terminal error is invalid")
|
||
if not isinstance(manifest, list) or len(json.dumps(manifest, ensure_ascii=False)) > 256 * 1024:
|
||
raise SoftwareJobError("job artifact manifest is invalid")
|
||
now = datetime.now(timezone.utc)
|
||
with session_scope() as session:
|
||
job = session.execute(
|
||
select(SoftwareJob).where(SoftwareJob.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 SoftwareJobError("successful job artifacts have not been published")
|
||
if job.status in {"succeeded", "failed", "cancelled"}:
|
||
if job.status != terminal_status:
|
||
raise SoftwareJobError("job terminal status conflicts with existing terminal")
|
||
return
|
||
if job.status not in {"offered", "dispatched", "running", "disconnected", "cancelling"}:
|
||
raise SoftwareJobError("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(SoftwareJob)
|
||
.where(
|
||
SoftwareJob.node_id == node_id,
|
||
SoftwareJob.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 SoftwareJobError("invalid job message identity") from exc
|
||
digest = str(payload.get("request_digest") or "")
|
||
if len(digest) != 64:
|
||
raise SoftwareJobError("invalid job request digest")
|
||
return job_id, lease_id, digest
|
||
|
||
|
||
def _assert_job_message(
|
||
job: SoftwareJob | 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 SoftwareJobError("job message does not belong to this node or lease")
|