306 lines
12 KiB
Python
306 lines
12 KiB
Python
"""Fixed Origin adapter for origin.plot@v2.
|
|
|
|
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 _input_file(job_dir: Path, key: str) -> Path:
|
|
directory = job_dir / "input" / key
|
|
files = [path for path in directory.iterdir() if path.is_file() and not path.name.startswith(".")]
|
|
if len(files) != 1:
|
|
raise ValueError(f"INPUT_FILE_COUNT_INVALID:{key}")
|
|
return files[0]
|
|
|
|
|
|
def _resolve_series(
|
|
input_data: dict[str, tuple[list[str], list[list[Any]]]],
|
|
series_specs: list[dict[str, Any]],
|
|
) -> tuple[list[tuple[str, int, int, str | None]], dict[tuple[str, int], str]]:
|
|
resolved: list[tuple[str, int, int, str | None]] = []
|
|
labels: dict[tuple[str, int], str] = {}
|
|
for series in series_specs:
|
|
input_key = series["input"]
|
|
headers, _ = input_data[input_key]
|
|
x_index = _column_index(headers, series["x"], "x")
|
|
y_index = _column_index(headers, series["y"], "y")
|
|
label = series.get("label")
|
|
label_key = (input_key, y_index)
|
|
effective_label = label or series["y"]
|
|
if label_key in labels and labels[label_key] != effective_label:
|
|
raise ValueError("SERIES_LABEL_CONFLICT")
|
|
labels[label_key] = effective_label
|
|
resolved.append((input_key, x_index, y_index, label))
|
|
return resolved, labels
|
|
|
|
|
|
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_specs = request["inputs"]
|
|
input_files = {item["key"]: _input_file(job_dir, item["key"]) for item in input_specs}
|
|
input_data = {
|
|
item["key"]: _read_rows(
|
|
input_files[item["key"]],
|
|
(item.get("selector") or {}).get("sheet"),
|
|
)
|
|
for item in input_specs
|
|
}
|
|
plot_spec = request["operation"]["plot"]
|
|
plot_type = plot_spec["type"]
|
|
if plot_type not in PLOT_TYPES:
|
|
raise ValueError("PLOT_TYPE_NOT_IMPLEMENTED")
|
|
series_specs = plot_spec["series"]
|
|
resolved_series, labels = _resolve_series(input_data, series_specs)
|
|
|
|
import originpro as op
|
|
|
|
output = job_dir / "output"
|
|
output.mkdir(exist_ok=True)
|
|
op.set_show(False)
|
|
try:
|
|
op.new()
|
|
worksheets: dict[str, Any] = {}
|
|
for input_spec in input_specs:
|
|
input_key = input_spec["key"]
|
|
headers, rows = input_data[input_key]
|
|
worksheet = op.new_sheet("w", lname=input_key)
|
|
worksheets[input_key] = worksheet
|
|
for index, header in enumerate(headers):
|
|
column_label = labels.get((input_key, index), header)
|
|
worksheet.from_list(
|
|
index,
|
|
[row[index] if index < len(row) else None for row in rows],
|
|
lname=column_label,
|
|
)
|
|
graph = op.new_graph(template={"line": "line", "scatter": "scatter", "line_scatter": "linesymb"}[plot_type])
|
|
layer = graph[0]
|
|
for input_key, x_index, y_index, _ in resolved_series:
|
|
layer.add_plot(
|
|
worksheets[input_key], 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(series_specs[0].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)
|
|
requested_outputs = request["outputs"]
|
|
formats = [item["format"] for item in requested_outputs]
|
|
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"}
|
|
png_output = next((item for item in requested_outputs if item["format"] == "png"), None)
|
|
dpi = (png_output.get("options") or {}).get("dpi", 300) if png_output else 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.3.0",
|
|
"originpro_version": originpro_version,
|
|
"request_digest": request_record["request_digest"],
|
|
"inputs": [
|
|
{
|
|
"key": item["key"],
|
|
"filename": input_files[item["key"]].name,
|
|
"sha256": _file_sha256(input_files[item["key"]]),
|
|
}
|
|
for item in input_specs
|
|
],
|
|
"outputs": requested_outputs,
|
|
"requested_dpi": dpi if png_output else None,
|
|
"png_pixel_width": pixel_width if png_output else None,
|
|
}
|
|
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())
|