simit 0.1.0__py3-none-any.whl

simit/skills/helpers.py

@@ -0,0 +1,197 @@
+"""Deterministic helpers shared across renderers: geometry notation helpers
+(injected into the model's sandbox) and automatic text-overlap avoidance
+(used both in-process by the deterministic renderers and injected into the
+model's sandbox)."""
+
+# Pre-defined geometry notation helpers, injected ahead of every model-written
+# script so the model calls them by name instead of constructing this
+# notation by hand -- it's unreliable at hand-rolling right-angle squares,
+# equal-length tick marks, and angle arcs from scratch (wrong vertex, wrong
+# offset direction, text overlapping the shape), but reliable at calling a
+# documented function with the right arguments (see MASTER_PROMPT Example 26).
+GEOMETRY_HELPERS_SRC = '''
+import numpy as np
+from matplotlib.patches import Polygon as _HelperPolygon, Arc as _HelperArc
+
+def right_angle_marker(ax, vertex, p1, p2, size=0.3, **kwargs):
+    vertex, p1, p2 = np.array(vertex, dtype=float), np.array(p1, dtype=float), np.array(p2, dtype=float)
+    u1 = (p1 - vertex) / np.linalg.norm(p1 - vertex)
+    u2 = (p2 - vertex) / np.linalg.norm(p2 - vertex)
+    pts = [vertex, vertex + size * u1, vertex + size * u1 + size * u2, vertex + size * u2]
+    kwargs.setdefault("fill", False)
+    kwargs.setdefault("edgecolor", "black")
+    kwargs.setdefault("linewidth", 1.2)
+    ax.add_patch(_HelperPolygon(pts, closed=True, **kwargs))
+
+def tick_marks(ax, p1, p2, count=1, size=0.15, gap=0.08, **kwargs):
+    p1, p2 = np.array(p1, dtype=float), np.array(p2, dtype=float)
+    mid = (p1 + p2) / 2
+    along = (p2 - p1) / np.linalg.norm(p2 - p1)
+    perp = np.array([-along[1], along[0]])
+    kwargs.setdefault("color", "black")
+    kwargs.setdefault("linewidth", 1.5)
+    for off in (np.arange(count) - (count - 1) / 2) * gap:
+        c = mid + along * off
+        ax.plot([c[0] - perp[0] * size, c[0] + perp[0] * size],
+                [c[1] - perp[1] * size, c[1] + perp[1] * size], **kwargs)
+
+def angle_arc(ax, vertex, theta1, theta2, radius=0.5, label=None, label_radius_mult=1.4, **kwargs):
+    kwargs.setdefault("color", "black")
+    kwargs.setdefault("linewidth", 1.2)
+    ax.add_patch(_HelperArc(vertex, 2 * radius, 2 * radius, angle=0, theta1=theta1, theta2=theta2, **kwargs))
+    if label:
+        mid = np.radians((theta1 + theta2) / 2)
+        lx = vertex[0] + radius * label_radius_mult * np.cos(mid)
+        ly = vertex[1] + radius * label_radius_mult * np.sin(mid)
+        ax.text(lx, ly, label, ha="center", va="center", fontsize=10)
+
+def side_label(ax, p1, p2, text=None, offset=None, fmt="{:.2f}", **kwargs):
+    p1, p2 = np.array(p1, dtype=float), np.array(p2, dtype=float)
+    length = np.linalg.norm(p2 - p1)
+    if text is None:
+        text = fmt.format(length)
+    if offset is None:
+        offset = 0.35
+    mid = (p1 + p2) / 2
+    along = (p2 - p1) / length
+    perp = np.array([-along[1], along[0]])
+    kwargs.setdefault("ha", "center")
+    kwargs.setdefault("va", "center")
+    kwargs.setdefault("fontsize", 10)
+    ax.text(*(mid + perp * offset), text, **kwargs)
+'''
+
+
+def autofix_text_overlaps(fig, max_iter: int = 40, min_font_size: float = 6.0) -> None:
+    """Detects overlapping ax.text() bounding boxes across every Axes in
+    `fig` and nudges colliding labels apart along the line between their
+    centers, iterating until none overlap. Each text has a displacement
+    budget (proportional to its own size) -- once a text has moved further
+    than that without resolving its overlap, repositioning stops for it and
+    its font shrinks instead, so a tight cluster that can't be cleanly
+    separated ends up smaller rather than flung arbitrarily far from its
+    meaningful anchor point. Used directly (in-process) by every
+    deterministic renderer below; the model-facing FIGURE sandbox gets the
+    equivalent logic via the separately-injected TEXT_OVERLAP_FIX_SRC string
+    instead, since that runs in an isolated subprocess that can't import
+    this module."""
+    fig.canvas.draw()
+    for ax in fig.get_axes():
+        texts = [t for t in ax.texts if t.get_text().strip()]
+        if len(texts) < 2:
+            continue
+        renderer = fig.canvas.get_renderer()
+        budget = [1.3 * max(t.get_window_extent(renderer=renderer).width,
+                             t.get_window_extent(renderer=renderer).height) for t in texts]
+        traveled = [0.0] * len(texts)
+        for _ in range(max_iter):
+            fig.canvas.draw()
+            renderer = fig.canvas.get_renderer()
+            bboxes = [t.get_window_extent(renderer=renderer) for t in texts]
+            moved = False
+            for i in range(len(texts)):
+                for j in range(i + 1, len(texts)):
+                    bi, bj = bboxes[i], bboxes[j]
+                    if not bi.overlaps(bj):
+                        continue
+                    cxi, cyi = bi.x0 + bi.width / 2, bi.y0 + bi.height / 2
+                    cxj, cyj = bj.x0 + bj.width / 2, bj.y0 + bj.height / 2
+                    dx, dy = cxj - cxi, cyj - cyi
+                    dist = (dx ** 2 + dy ** 2) ** 0.5
+                    if dist < 1e-6:
+                        dx, dy, dist = float(j - i), 0.0, float(j - i)
+                    ux, uy = dx / dist, dy / dist
+                    overlap_x = (bi.width + bj.width) / 2 - abs(dx)
+                    overlap_y = (bi.height + bj.height) / 2 - abs(dy)
+                    push = max(overlap_x, overlap_y, 2.0) * 0.5 + 1.5
+                    for idx, t, sign in ((i, texts[i], -1), (j, texts[j], 1)):
+                        if traveled[idx] >= budget[idx]:
+                            if t.get_fontsize() > min_font_size:
+                                t.set_fontsize(max(min_font_size, t.get_fontsize() * 0.85))
+                                moved = True
+                            continue
+                        trans = t.get_transform()
+                        x, y = t.get_position()
+                        px, py = trans.transform((x, y))
+                        nx, ny = trans.inverted().transform((px + sign * ux * push, py + sign * uy * push))
+                        t.set_position((nx, ny))
+                        traveled[idx] += push
+                        moved = True
+            if not moved:
+                break
+
+
+# The model can't be relied on to predict where its own labels will end up
+# relative to each other -- it places each ax.text() independently with no
+# view of the others. This patches plt.savefig (transparently, the model's
+# own `import matplotlib.pyplot as plt` binds the same module object) to
+# detect overlapping text bounding boxes across every Axes right before the
+# real save, and nudge colliding labels apart along the line between their
+# centers, iterating until none overlap. A deterministic, geometric fix --
+# exactly the kind of thing that's calculable rather than left to chance.
+# (Duplicated as source text rather than shared with autofix_text_overlaps()
+# above because this runs inside an isolated model-code subprocess that
+# can't import this module.)
+TEXT_OVERLAP_FIX_SRC = '''
+import matplotlib.pyplot as _overlap_plt
+
+def _shift_text_px(text, dx_px, dy_px):
+    trans = text.get_transform()
+    x, y = text.get_position()
+    px, py = trans.transform((x, y))
+    nx, ny = trans.inverted().transform((px + dx_px, py + dy_px))
+    text.set_position((nx, ny))
+
+def _autofix_text_overlaps(fig, max_iter=40, min_font_size=6.0):
+    fig.canvas.draw()
+    for ax in fig.get_axes():
+        texts = [t for t in ax.texts if t.get_text().strip()]
+        if len(texts) < 2:
+            continue
+        renderer = fig.canvas.get_renderer()
+        budget = [1.3 * max(t.get_window_extent(renderer=renderer).width,
+                             t.get_window_extent(renderer=renderer).height) for t in texts]
+        traveled = [0.0] * len(texts)
+        for _ in range(max_iter):
+            fig.canvas.draw()
+            renderer = fig.canvas.get_renderer()
+            bboxes = [t.get_window_extent(renderer=renderer) for t in texts]
+            moved = False
+            for i in range(len(texts)):
+                for j in range(i + 1, len(texts)):
+                    bi, bj = bboxes[i], bboxes[j]
+                    if not bi.overlaps(bj):
+                        continue
+                    cxi, cyi = bi.x0 + bi.width / 2, bi.y0 + bi.height / 2
+                    cxj, cyj = bj.x0 + bj.width / 2, bj.y0 + bj.height / 2
+                    dx, dy = cxj - cxi, cyj - cyi
+                    dist = (dx ** 2 + dy ** 2) ** 0.5
+                    if dist < 1e-6:
+                        dx, dy, dist = float(j - i), 0.0, float(j - i)
+                    ux, uy = dx / dist, dy / dist
+                    overlap_x = (bi.width + bj.width) / 2 - abs(dx)
+                    overlap_y = (bi.height + bj.height) / 2 - abs(dy)
+                    push = max(overlap_x, overlap_y, 2.0) * 0.5 + 1.5
+                    for idx, t, sign in ((i, texts[i], -1), (j, texts[j], 1)):
+                        if traveled[idx] >= budget[idx]:
+                            if t.get_fontsize() > min_font_size:
+                                t.set_fontsize(max(min_font_size, t.get_fontsize() * 0.85))
+                                moved = True
+                            continue
+                        _shift_text_px(t, sign * ux * push, sign * uy * push)
+                        traveled[idx] += push
+                        moved = True
+            if not moved:
+                break
+
+_original_savefig = _overlap_plt.savefig
+
+def _savefig_with_overlap_fix(*args, **kwargs):
+    try:
+        _autofix_text_overlaps(_overlap_plt.gcf())
+    except Exception:
+        pass
+    return _original_savefig(*args, **kwargs)
+
+_overlap_plt.savefig = _savefig_with_overlap_fix
+'''

simit/skills/pool.py

@@ -0,0 +1,262 @@
+"""Parallel rendering in warm worker processes.
+
+Rendering (matplotlib, RDKit, Chromium, executing model-written code) is CPU
+work that would otherwise stall the GPU-driving main process. Each worker is a
+``python -m simit.skills.pool`` subprocess -- deliberately not a
+multiprocessing child, which would re-import the user's ``__main__`` script --
+that keeps its imports and browser warm. A hung or crashed worker is killed
+and replaced. Skills that cannot be pickled by reference (e.g. defined in a
+notebook or ``__main__``) are rendered in a thread of the main process.
+"""
+
+from __future__ import annotations
+
+import io
+import os
+import pickle
+import queue
+import subprocess
+import sys
+import threading
+import time
+import traceback
+from concurrent.futures import Future, ThreadPoolExecutor
+from multiprocessing.connection import Connection
+
+from PIL import Image
+
+from .base import SpecError
+
+RENDER_TIMEOUT_S = 90.0
+
+
+def _serve(conn_in: Connection, conn_out: Connection):
+    os.environ.setdefault("MPLBACKEND", "Agg")
+    while True:
+        try:
+            msg = conn_in.recv()
+        except (EOFError, OSError, KeyboardInterrupt):
+            return
+        if msg is None:
+            return
+        task_id, payload = msg
+        try:
+            skill, spec = pickle.loads(payload)
+        except Exception as e:  # this process cannot import the skill: the parent renders it instead
+            conn_out.send((task_id, None, f"{type(e).__name__}: {e}"))
+            continue
+        try:
+            image = skill.render(spec)
+            buf = io.BytesIO()
+            image.save(buf, format="PNG")
+            conn_out.send((task_id, True, buf.getvalue()))
+        except SpecError as e:
+            conn_out.send((task_id, False, str(e)))
+        except Exception as e:  # unexpected renderer failure: still reported back for repair
+            conn_out.send((task_id, False, f"{type(e).__name__}: {e}"))
+        except BaseException:
+            conn_out.send((task_id, False, traceback.format_exc()[-500:]))
+
+
+class _Worker:
+    def __init__(self):
+        p2c_r, p2c_w = os.pipe()
+        c2p_r, c2p_w = os.pipe()
+        env = dict(os.environ, MPLBACKEND="Agg", PYTHONUNBUFFERED="1")
+        # The parent's import path, so skills from the user's own modules unpickle here too.
+        paths = [p for p in sys.path if p and os.path.isdir(p)] + [env.get("PYTHONPATH", "")]
+        env["PYTHONPATH"] = os.pathsep.join(p for p in paths if p)
+        self.proc = subprocess.Popen(
+            [sys.executable, "-m", "simit.skills.pool", str(p2c_r), str(c2p_w)],
+            pass_fds=(p2c_r, c2p_w), env=env, stdin=subprocess.DEVNULL,
+        )
+        os.close(p2c_r)
+        os.close(c2p_w)
+        self.send_conn = Connection(p2c_w, readable=False)
+        self.recv_conn = Connection(c2p_r, writable=False)
+        self.task = None  # (task_id, future, deadline)
+
+    def alive(self) -> bool:
+        return self.proc.poll() is None
+
+    def kill(self):
+        for conn in (self.send_conn, self.recv_conn):
+            try:
+                conn.close()
+            except Exception:
+                pass
+        try:
+            self.proc.kill()
+            self.proc.wait(timeout=5)
+        except Exception:
+            pass
+
+
+class RenderPool:
+    def __init__(self, num_workers: int | None = None, timeout: float = RENDER_TIMEOUT_S):
+        self.num_workers = default_num_workers() if num_workers is None else num_workers
+        self.timeout = timeout
+        self._queue: "queue.Queue[tuple[int, bytes, Future]]" = queue.Queue()
+        self._workers: list[_Worker] = []
+        self._lock = threading.Lock()
+        self._thread = None
+        self._next_id = 0
+        self._closed = False
+        self._local = ThreadPoolExecutor(max_workers=1, thread_name_prefix="simit-render")
+        self._local_types: set = set()   # skill classes the workers could not import
+
+    # ---------------------------------------------------------------- public
+    def submit(self, skill, spec) -> Future:
+        """Render ``spec`` with ``skill``; the future yields a PIL image or
+        raises SpecError."""
+        if (self.num_workers <= 0 or not getattr(skill, "render_in_subprocess", True)
+                or type(skill).__module__ in ("__main__", "__mp_main__") or _key(skill) in self._local_types):
+            return self._local.submit(_render_local, skill, spec)
+        try:
+            payload = pickle.dumps((skill, spec))
+        except Exception:
+            return self._local.submit(_render_local, skill, spec)
+        fut: Future = Future()
+        with self._lock:
+            self._next_id += 1
+            task_id = self._next_id
+            if self._thread is None:
+                self._thread = threading.Thread(target=self._loop, daemon=True, name="simit-render-pool")
+                self._thread.start()
+        self._queue.put((task_id, payload, fut))
+        return fut
+
+    def warmup(self):
+        """Start the worker processes now instead of on the first render."""
+        with self._lock:
+            while len(self._workers) < self.num_workers:
+                self._workers.append(_Worker())
+
+    def close(self):
+        self._closed = True
+        with self._lock:
+            for w in self._workers:
+                try:
+                    w.send_conn.send(None)
+                except Exception:
+                    pass
+                w.kill()
+            self._workers = []
+        self._local.shutdown(wait=False)
+
+    def __del__(self):
+        try:
+            self.close()
+        except Exception:
+            pass
+
+    # -------------------------------------------------------------- internals
+    def _loop(self):
+        pending: list[tuple[int, bytes, Future]] = []
+        while not self._closed:
+            busy = any(w.task is not None for w in self._workers)
+            try:
+                pending.append(self._queue.get(timeout=0.005 if (pending or busy) else 0.5))
+                while True:
+                    pending.append(self._queue.get_nowait())
+            except queue.Empty:
+                pass
+            with self._lock:
+                if self._closed:
+                    break
+                while len(self._workers) < self.num_workers and pending:
+                    self._workers.append(_Worker())
+                workers = list(self._workers)
+            for w in workers:
+                if w.task is None and pending:
+                    task_id, payload, fut = pending.pop(0)
+                    try:
+                        w.send_conn.send((task_id, payload))
+                        w.task = (task_id, fut, time.monotonic() + self.timeout, payload)
+                    except Exception:
+                        pending.insert(0, (task_id, payload, fut))
+                        self._replace(w)
+            for w in workers:
+                if w.task is None:
+                    continue
+                _, fut, deadline, payload = w.task
+                try:
+                    ready = w.recv_conn.poll()
+                except Exception:
+                    ready = False
+                if ready:
+                    try:
+                        _, ok, data = w.recv_conn.recv()
+                    except Exception:
+                        w.task = None
+                        fut.set_exception(SpecError("renderer process crashed"))
+                        self._replace(w)
+                        continue
+                    w.task = None
+                    if ok is None:
+                        skill, spec = pickle.loads(payload)
+                        self._local_types.add(_key(skill))
+                        _chain(self._local.submit(_render_local, skill, spec), fut)
+                    elif ok:
+                        image = Image.open(io.BytesIO(data))
+                        image.load()
+                        fut.set_result(image)
+                    else:
+                        fut.set_exception(SpecError(data))
+                elif not w.alive():
+                    w.task = None
+                    fut.set_exception(SpecError("renderer process crashed"))
+                    self._replace(w)
+                elif time.monotonic() > deadline:
+                    w.task = None
+                    fut.set_exception(SpecError(f"rendering timed out after {self.timeout:.0f}s"))
+                    self._replace(w)
+
+    def _replace(self, worker: _Worker):
+        worker.kill()
+        with self._lock:
+            if worker in self._workers:
+                self._workers[self._workers.index(worker)] = _Worker()
+
+
+def _key(skill) -> tuple:
+    return type(skill).__module__, type(skill).__qualname__
+
+
+def _chain(src: Future, dst: Future):
+    def done(f):
+        if f.exception() is not None:
+            dst.set_exception(f.exception())
+        else:
+            dst.set_result(f.result())
+    src.add_done_callback(done)
+
+
+def _render_local(skill, spec):
+    try:
+        return skill.render(spec)
+    except SpecError:
+        raise
+    except Exception as e:
+        raise SpecError(f"{type(e).__name__}: {e}") from None
+
+
+def default_num_workers() -> int:
+    """Worker count from the CPUs actually granted to this job."""
+    if os.environ.get("SIMIT_RENDER_WORKERS"):
+        return int(os.environ["SIMIT_RENDER_WORKERS"])
+    n = None
+    for var in ("LSB_DJOB_NUMPROC", "SLURM_CPUS_PER_TASK"):
+        if os.environ.get(var, "").isdigit():
+            n = int(os.environ[var])
+            break
+    if n is None:
+        try:
+            n = len(os.sched_getaffinity(0))
+        except AttributeError:
+            n = os.cpu_count() or 2
+    return max(1, min(8, n - 1 if n > 2 else n))
+
+
+if __name__ == "__main__":
+    _serve(Connection(int(sys.argv[1]), writable=False), Connection(int(sys.argv[2]), readable=False))