177 lines
9.9 KiB
Python
177 lines
9.9 KiB
Python
"""Spool Deadline protocol 2. Copied beside the worker; stdlib only."""
|
|
import copy
|
|
import hashlib
|
|
import json
|
|
import ntpath
|
|
import os
|
|
import re
|
|
import shutil
|
|
import uuid
|
|
from pathlib import Path
|
|
|
|
SCHEMA = 2
|
|
SAVERS = {"SaveImage": {"images"}, "PreviewImage": {"images"}, "SaveVideo": {"images", "videos"}, "VHS_VideoCombine": {"gifs"}}
|
|
|
|
|
|
def contained(root, relative):
|
|
relative = os.fspath(relative)
|
|
if not isinstance(relative, str) or not relative or ntpath.splitdrive(relative)[0] or relative.startswith(("/", "\\")):
|
|
raise ValueError("Farm paths must be relative")
|
|
for part in relative.replace("\\", "/").split("/"):
|
|
if part and (part == ".." or part.rstrip(" .") != part or any(c in part for c in ':\x00\r\n"') or re.fullmatch(r"(?i:CON|PRN|AUX|NUL|COM[1-9]|LPT[1-9])(?:\..*)?", part)):
|
|
raise ValueError("Unsafe farm path component")
|
|
root = Path(root).resolve()
|
|
path = (root / relative.replace("\\", "/")).resolve()
|
|
if path == root or not path.is_relative_to(root):
|
|
raise ValueError("Farm path escapes configured storage")
|
|
return path
|
|
|
|
|
|
def checksum(path):
|
|
digest = hashlib.sha256()
|
|
with Path(path).open("rb") as handle:
|
|
for block in iter(lambda: handle.read(1024 * 1024), b""):
|
|
digest.update(block)
|
|
return digest.hexdigest()
|
|
|
|
|
|
def atomic_json(path, value):
|
|
path = Path(path)
|
|
path.parent.mkdir(parents=True, exist_ok=True)
|
|
pending = path.with_name(path.name + "." + uuid.uuid4().hex + ".partial")
|
|
try:
|
|
with pending.open("w", encoding="utf-8") as handle:
|
|
json.dump(value, handle, ensure_ascii=False, sort_keys=True, allow_nan=False)
|
|
handle.flush()
|
|
os.fsync(handle.fileno())
|
|
pending.replace(path)
|
|
finally:
|
|
pending.unlink(missing_ok=True)
|
|
|
|
|
|
def semantic_prompt(prompt):
|
|
result = copy.deepcopy(prompt)
|
|
for node in result.values():
|
|
node.pop("is_changed", None)
|
|
return result
|
|
|
|
|
|
def saver_prefix(package, execution_id, node_id):
|
|
for identifier in (package["attemptId"], execution_id):
|
|
if not isinstance(identifier, str) or not re.fullmatch(r"[A-Za-z0-9_-]+", identifier):
|
|
raise ValueError("Invalid Spool execution identifier")
|
|
node_folder = hashlib.sha256(node_id.encode()).hexdigest()[:16]
|
|
return "spool/" + package["attemptId"] + "/" + execution_id + "/" + node_folder + "/out"
|
|
|
|
|
|
def prepare_execution(package, prompt, input_root, variation_index):
|
|
if package["schemaVersion"] != SCHEMA or variation_index != package["variationIndex"]:
|
|
raise ValueError("Spool task does not match its explicit variation assignment")
|
|
if semantic_prompt(prompt) != semantic_prompt(package["stagedPrompt"]):
|
|
raise ValueError("Worker transformations changed the staged Spool prompt")
|
|
for asset in package["assets"]:
|
|
path = contained(input_root, asset["staged"])
|
|
if path.stat().st_size != asset["byteSize"] or checksum(path) != asset["sha256"]:
|
|
raise ValueError("Staged Spool dependency changed")
|
|
if prompt[asset["nodeId"]]["inputs"][asset["input"]] != asset["staged"]:
|
|
raise ValueError("Staged input reference does not match its dependency")
|
|
execution_id = uuid.uuid4().hex
|
|
saver_prefix(package, execution_id, "")
|
|
observed = copy.deepcopy(prompt)
|
|
for node_id, node in observed.items():
|
|
if node["class_type"] in SAVERS and "filename_prefix" in node["inputs"]:
|
|
node["inputs"]["filename_prefix"] = saver_prefix(package, execution_id, node_id)
|
|
return {"executionId": execution_id, "prompt": observed, "variationIndex": variation_index}
|
|
|
|
|
|
def finalize(package, execution, prompt_id, history, roots, worker, job_id):
|
|
saver_prefix(package, execution["executionId"], "")
|
|
status = history.get("status", {})
|
|
if not status.get("completed") or status.get("status_str") != "success":
|
|
raise ValueError("Spool requires successful exact-prompt history")
|
|
if semantic_prompt(history.get("prompt", [None, None, {}])[2]) != semantic_prompt(execution["prompt"]):
|
|
raise ValueError("Worker hooks changed the verified Spool request")
|
|
relative = "spool-receipts/" + package["attemptId"] + "/" + execution["executionId"]
|
|
target = contained(package["resultRoot"], relative)
|
|
target.mkdir(parents=True, exist_ok=True)
|
|
receipt_path = contained(target, "receipt.json")
|
|
if receipt_path.is_file():
|
|
existing = json.loads(receipt_path.read_text(encoding="utf-8"))
|
|
verify(package, existing, package["resultRoot"])
|
|
return existing
|
|
artifacts = []
|
|
for node_id, node in execution["prompt"].items():
|
|
for group in sorted(SAVERS.get(node["class_type"], ())):
|
|
for index, record in enumerate(history.get("outputs", {}).get(node_id, {}).get(group, [])):
|
|
name = record["filename"]
|
|
if Path(name).name != name or "\\" in name or "/" in name:
|
|
raise ValueError("Invalid worker artifact filename")
|
|
kind = record.get("type", "output")
|
|
if kind not in roots:
|
|
raise ValueError("Unknown worker storage type")
|
|
source = contained(roots[kind], str(Path(record.get("subfolder", "")) / name))
|
|
if "filename_prefix" in node["inputs"]:
|
|
namespace = contained(roots[kind], saver_prefix(package, execution["executionId"], node_id).rsplit("/", 1)[0])
|
|
if not source.is_relative_to(namespace):
|
|
raise ValueError("Worker artifact belongs to another execution")
|
|
artifact_id = uuid.uuid4().hex
|
|
destination = target / (artifact_id + source.suffix.lower())
|
|
pending = destination.with_name(destination.name + ".partial")
|
|
with source.open("rb") as reader, pending.open("xb") as writer:
|
|
shutil.copyfileobj(reader, writer)
|
|
digest = checksum(source)
|
|
if checksum(pending) != digest:
|
|
raise ValueError("Worker artifact copy failed integrity verification")
|
|
pending.replace(destination)
|
|
artifacts.append({"artifactId": artifact_id, "nodeId": node_id, "group": group, "index": index,
|
|
"storageId": "spool-results", "path": destination.relative_to(Path(package["resultRoot"]).resolve()).as_posix(),
|
|
"byteSize": destination.stat().st_size, "sha256": digest})
|
|
selection = package["result"]
|
|
primary = next((a for a in artifacts if all(a[k] == selection[k] for k in ("nodeId", "group", "index"))), None)
|
|
if primary is None:
|
|
raise ValueError("Successful worker history omitted the selected artifact")
|
|
receipt = {k: package[k] for k in ("schemaVersion", "projectId", "clipId", "takeId", "attemptId", "fingerprint", "submissionId", "variationIndex")}
|
|
receipt.update({"executionId": execution["executionId"], "promptId": prompt_id, "jobId": job_id, "worker": worker,
|
|
"outcome": "success", "observedPrompt": history["prompt"][2], "history": history,
|
|
"artifacts": artifacts, "primaryArtifactId": primary["artifactId"]})
|
|
atomic_json(receipt_path, receipt)
|
|
return receipt
|
|
|
|
|
|
def verify(package, receipt, result_root):
|
|
saver_prefix(package, receipt.get("executionId"), "")
|
|
if receipt.get("schemaVersion") != SCHEMA:
|
|
raise ValueError("Unsupported farm receipt version")
|
|
for key in ("projectId", "clipId", "takeId", "attemptId", "fingerprint", "submissionId", "variationIndex"):
|
|
if receipt.get(key) != package[key]:
|
|
raise ValueError("Farm receipt belongs to another request: " + key)
|
|
if receipt.get("outcome") != "success" or not receipt.get("promptId") or not receipt.get("executionId"):
|
|
raise ValueError("Farm receipt lacks execution identity or success")
|
|
history = receipt.get("history", {})
|
|
if history.get("status", {}).get("status_str") != "success" or not history.get("status", {}).get("completed"):
|
|
raise ValueError("Farm receipt lacks completion evidence")
|
|
expected = copy.deepcopy(package["stagedPrompt"])
|
|
for node_id, node in expected.items():
|
|
if node["class_type"] in SAVERS and "filename_prefix" in node["inputs"]:
|
|
node["inputs"]["filename_prefix"] = saver_prefix(package, receipt["executionId"], node_id)
|
|
if semantic_prompt(receipt.get("observedPrompt", {})) != semantic_prompt(expected) or semantic_prompt(history.get("prompt", [None, None, {}])[2]) != semantic_prompt(expected) or history["prompt"][1] != receipt["promptId"]:
|
|
raise ValueError("Farm observed prompt differs from the permitted execution copy")
|
|
artifacts = receipt["artifacts"]
|
|
if len({a["artifactId"] for a in artifacts}) != len(artifacts):
|
|
raise ValueError("Duplicate artifact identity")
|
|
reported = {(node_id, group, index) for node_id, node in expected.items() for group in SAVERS.get(node["class_type"], ())
|
|
for index, record in enumerate(history.get("outputs", {}).get(node_id, {}).get(group, []))}
|
|
if {(a["nodeId"], a["group"], a["index"]) for a in artifacts} != reported:
|
|
raise ValueError("Farm receipt omitted or invented a reported artifact")
|
|
for artifact in artifacts:
|
|
if artifact["storageId"] != "spool-results":
|
|
raise ValueError("Unknown farm storage ID")
|
|
path = contained(result_root, artifact["path"])
|
|
namespace = contained(result_root, "spool-receipts/" + package["attemptId"] + "/" + receipt["executionId"])
|
|
if not path.is_relative_to(namespace) or path.stat().st_size != artifact["byteSize"] or checksum(path) != artifact["sha256"]:
|
|
raise ValueError("Farm artifact integrity or containment failed")
|
|
primary = next((a for a in artifacts if a["artifactId"] == receipt["primaryArtifactId"]), None)
|
|
if not primary or any(primary[k] != package["result"][k] for k in ("nodeId", "group", "index")):
|
|
raise ValueError("Farm primary differs from explicit result selection")
|
|
return primary
|