253 lines
10 KiB
Python
253 lines
10 KiB
Python
"""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())
|