495 lines
19 KiB
Python
495 lines
19 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 math
|
|
import os
|
|
import sys
|
|
from datetime import datetime, timezone
|
|
from importlib.metadata import PackageNotFoundError, version
|
|
from pathlib import Path
|
|
from typing import Any
|
|
|
|
PLOT_CONFIG = {
|
|
"line": ("line", "l"),
|
|
"scatter": ("scatter", "s"),
|
|
"line_scatter": ("linesymb", "y"),
|
|
"column": ("column", "c"),
|
|
"bar": ("bar", 215),
|
|
"grouped_column": ("column", "c"),
|
|
"y_error": ("ERRBAR", "y"),
|
|
"contour": ("TriContour", 243),
|
|
"surface_3d": ("glCMAP", 103),
|
|
"ternary": ("ternary", 245),
|
|
}
|
|
XYZ_PLOT_TYPES = {"contour", "surface_3d", "ternary", "heatmap"}
|
|
FORMATS = {"opju", "png", "svg", "pdf"}
|
|
LINE_STYLES = {
|
|
"solid": 1,
|
|
"dash": 2,
|
|
"dot": 3,
|
|
"dash_dot": 4,
|
|
"dash_dot_dot": 5,
|
|
}
|
|
SYMBOLS = {
|
|
"square": 0,
|
|
"circle": 1,
|
|
"triangle_up": 2,
|
|
"diamond": 3,
|
|
"cross": 9,
|
|
"plus": 10,
|
|
}
|
|
LEGEND_POSITIONS = {
|
|
"top_left": (700, 500),
|
|
"top_right": (6800, 500),
|
|
"bottom_left": (700, 7200),
|
|
"bottom_right": (6800, 7200),
|
|
}
|
|
|
|
|
|
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 _hex_color(value: str) -> tuple[int, int, int]:
|
|
return tuple(int(value[index:index + 2], 16) for index in (1, 3, 5))
|
|
|
|
|
|
def _apply_canvas(graph: Any, canvas: Any) -> None:
|
|
if not isinstance(canvas, dict):
|
|
return
|
|
width = float(canvas["width_mm"])
|
|
height = float(canvas["height_mm"])
|
|
graph.set_float("width", width / 25.4 * graph.get_float("resx"))
|
|
graph.set_float("height", height / 25.4 * graph.get_float("resy"))
|
|
|
|
|
|
def _apply_axis(layer: Any, name: str, spec: Any, fallback: str) -> None:
|
|
settings = spec if isinstance(spec, dict) else {}
|
|
axis = layer.axis(name)
|
|
axis.title = _axis_title(settings, fallback)
|
|
if "scale" in settings:
|
|
axis.scale = settings["scale"]
|
|
if any(key in settings for key in ("minimum", "maximum", "major_step")):
|
|
axis.set_limits(
|
|
settings.get("minimum"), settings.get("maximum"), settings.get("major_step")
|
|
)
|
|
if "tick_label_angle" in settings:
|
|
layer.set_float(f"{name}.label.rotate", float(settings["tick_label_angle"]))
|
|
if "tick_label_font_size" in settings:
|
|
layer.set_float(f"{name}.label.pt", float(settings["tick_label_font_size"]))
|
|
if "title_font_size" in settings:
|
|
title_object = {"x": "xb", "y": "yl", "z": "zf"}[name]
|
|
layer.label(title_object).set_int("fsize", round(settings["title_font_size"]))
|
|
if "grid" in settings:
|
|
grid_value = {"none": 0, "major": 1, "major_minor": 3}[settings["grid"]]
|
|
layer.set_int(f"{name}.grid.show", grid_value)
|
|
|
|
|
|
def _apply_series_style(origin_plot: Any, style: Any) -> None:
|
|
if not isinstance(style, dict):
|
|
return
|
|
if "color" in style:
|
|
origin_plot.color = _hex_color(style["color"])
|
|
if "line_width" in style:
|
|
origin_plot.set_float("line.width", float(style["line_width"]))
|
|
if "line_style" in style:
|
|
origin_plot.set_int("line.type", LINE_STYLES[style["line_style"]])
|
|
if "symbol" in style:
|
|
origin_plot.symbol_kind = SYMBOLS[style["symbol"]]
|
|
if "symbol_size" in style:
|
|
origin_plot.symbol_size = float(style["symbol_size"])
|
|
if "transparency" in style:
|
|
origin_plot.transparency = int(style["transparency"])
|
|
|
|
|
|
def _apply_legend(layer: Any, legend: Any) -> None:
|
|
if not isinstance(legend, dict):
|
|
return
|
|
label = layer.label("Legend")
|
|
label.show = legend.get("enabled", True)
|
|
if "font_size" in legend:
|
|
label.set_int("fsize", round(legend["font_size"]))
|
|
if "position" in legend:
|
|
left, top = LEGEND_POSITIONS[legend["position"]]
|
|
label.set_int("left", left)
|
|
label.set_int("top", top)
|
|
|
|
|
|
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[dict[str, Any]], dict[tuple[str, int], str]]:
|
|
resolved: list[dict[str, Any]] = []
|
|
labels: dict[tuple[str, int], str] = {}
|
|
for series in series_specs:
|
|
input_key = series["input"]
|
|
headers, _ = input_data[input_key]
|
|
role_indexes = {
|
|
role: _column_index(headers, series[role], role)
|
|
for role in ("x", "y", "z", "y_error") if role in series
|
|
}
|
|
y_index = role_indexes["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": input_key, "label": label, **role_indexes})
|
|
return resolved, labels
|
|
|
|
|
|
def _number(value: Any, role: str) -> float:
|
|
if isinstance(value, bool):
|
|
raise ValueError(f"{role.upper()}_VALUE_NOT_NUMERIC") # noqa: TRY004
|
|
try:
|
|
number = float(value)
|
|
except (TypeError, ValueError) as exc:
|
|
raise ValueError(f"{role.upper()}_VALUE_NOT_NUMERIC") from exc
|
|
if not math.isfinite(number):
|
|
raise ValueError(f"{role.upper()}_VALUE_NOT_FINITE")
|
|
return number
|
|
|
|
|
|
def _is_evenly_spaced(values: list[float]) -> bool:
|
|
if len(values) <= 2:
|
|
return True
|
|
step = values[1] - values[0]
|
|
tolerance = max(abs(step) * 1e-9, 1e-12)
|
|
return all(
|
|
math.isclose(current - previous, step, rel_tol=1e-9, abs_tol=tolerance)
|
|
for previous, current in zip(values[1:-1], values[2:], strict=True)
|
|
)
|
|
|
|
|
|
def _heatmap_matrix(
|
|
rows: list[list[Any]], resolved: dict[str, Any]
|
|
) -> tuple[list[list[float]], tuple[float, float, float, float]]:
|
|
points: dict[tuple[float, float], float] = {}
|
|
for row in rows:
|
|
try:
|
|
x = _number(row[resolved["x"]], "x")
|
|
y = _number(row[resolved["y"]], "y")
|
|
z = _number(row[resolved["z"]], "z")
|
|
except IndexError as exc:
|
|
raise ValueError("HEATMAP_ROW_INCOMPLETE") from exc
|
|
if (x, y) in points:
|
|
raise ValueError("HEATMAP_COORDINATES_DUPLICATED")
|
|
points[(x, y)] = z
|
|
x_values = sorted({item[0] for item in points})
|
|
y_values = sorted({item[1] for item in points})
|
|
if len(x_values) < 2 or len(y_values) < 2:
|
|
raise ValueError("HEATMAP_GRID_TOO_SMALL")
|
|
if len(points) != len(x_values) * len(y_values):
|
|
raise ValueError("HEATMAP_GRID_INCOMPLETE")
|
|
if not _is_evenly_spaced(x_values) or not _is_evenly_spaced(y_values):
|
|
raise ValueError("HEATMAP_GRID_NOT_REGULAR")
|
|
matrix = [[points[(x, y)] for x in x_values] for y in y_values]
|
|
return matrix, (x_values[0], x_values[-1], y_values[0], y_values[-1])
|
|
|
|
|
|
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_CONFIG and plot_type != "heatmap":
|
|
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,
|
|
)
|
|
if plot_type == "heatmap":
|
|
import numpy as np
|
|
|
|
resolved = resolved_series[0]
|
|
_, rows = input_data[resolved["input"]]
|
|
matrix, xy_map = _heatmap_matrix(rows, resolved)
|
|
matrix_sheet = op.new_sheet("m")
|
|
matrix_sheet.from_np(np.array(matrix, dtype=float))
|
|
matrix_sheet.xymap = xy_map
|
|
matrix_sheet.set_label(
|
|
0, _axis_title(plot_spec.get("z_axis"), str(series_specs[0]["z"]))
|
|
)
|
|
graph = op.new_graph(template="heatmap")
|
|
layer = graph[0]
|
|
origin_plots = [layer.add_mplot(matrix_sheet, 0, type=105)]
|
|
else:
|
|
template, origin_plot_type = PLOT_CONFIG[plot_type]
|
|
graph = op.new_graph(template=template)
|
|
layer = graph[0]
|
|
origin_plots = []
|
|
for resolved in resolved_series:
|
|
arguments = {
|
|
"coly": resolved["y"],
|
|
"colx": resolved["x"],
|
|
"type": origin_plot_type,
|
|
}
|
|
if plot_type in XYZ_PLOT_TYPES:
|
|
arguments["colz"] = resolved["z"]
|
|
if plot_type == "y_error":
|
|
arguments["colyerr"] = resolved["y_error"]
|
|
origin_plots.append(layer.add_plot(worksheets[resolved["input"]], **arguments))
|
|
if plot_type == "grouped_column":
|
|
layer.group()
|
|
layer.rescale()
|
|
_apply_canvas(graph, plot_spec.get("canvas"))
|
|
_apply_axis(
|
|
layer, "x", plot_spec.get("x_axis"), str(series_specs[0].get("x") or "X")
|
|
)
|
|
_apply_axis(layer, "y", plot_spec.get("y_axis"), "Y")
|
|
if plot_type == "surface_3d":
|
|
_apply_axis(
|
|
layer, "z", plot_spec.get("z_axis"), str(series_specs[0].get("z") or "Z")
|
|
)
|
|
for origin_plot, series_spec in zip(origin_plots, series_specs, strict=True):
|
|
_apply_series_style(origin_plot, series_spec.get("style"))
|
|
_apply_legend(layer, plot_spec.get("legend"))
|
|
if plot_spec.get("title"):
|
|
title = layer.add_label(str(plot_spec["title"]))
|
|
title_font_size = (plot_spec.get("title_style") or {}).get("font_size", 18)
|
|
title.set_int("fsize", round(title_font_size))
|
|
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")
|
|
canvas_width_mm = (plot_spec.get("canvas") or {}).get("width_mm", 160)
|
|
pixel_width = round(dpi * canvas_width_mm / 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.5.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,
|
|
"canvas_mm": plot_spec.get("canvas"),
|
|
}
|
|
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())
|