annoscript 0.1.0__py3-none-any.whl

annoscript/_remote_jobs.py

@@ -0,0 +1,211 @@
+"""In-memory session adapter behind AnnoClient.remote(); never imported locally."""
+
+from collections import deque
+import logging
+import threading
+import time
+import uuid
+import weakref
+
+from annoscript._client_common import Result, ResultError, ResultStream, prepare_inputs
+from annoscript.compiler import compile_program
+from annoscript.jobs import (AnnoJobClient, RemoteCancellationError, RemoteJobPausedError,
+                      RemoteProtocolError, RemoteSubmissionError, RemoteTransportError)
+
+
+def bounded_cleanup(operation, timeout):
+    """Bound caller latency even when DNS/transport shutdown ignores timeouts."""
+    outcome, done = [], threading.Event()
+    def run():
+        try:
+            operation()
+        except BaseException as exc:
+            outcome.append(exc)
+        finally:
+            done.set()
+    threading.Thread(target=run, name="anno-job-cleanup", daemon=True).start()
+    if not done.wait(timeout):
+        raise RemoteCancellationError("remote cancellation confirmation timed out; session lease will expire")
+    if outcome:
+        raise RemoteCancellationError("remote cancellation could not be confirmed; session lease will expire") from outcome[0]
+
+
+class _Session:
+    def __init__(self, client):
+        self.client = client
+        self.id = None
+        self.key = uuid.uuid4().hex
+        self.closed = False
+        self.stop = threading.Event()
+        self.streams = weakref.WeakSet()
+        self.failure = None
+        self.thread = None
+
+    def start(self):
+        if self.closed:
+            raise RuntimeError("client is closed")
+        if self.failure is not None:
+            raise self.failure
+        if self.id is None:
+            value = self.client._create_session(self.key)
+            self.id = value["session_id"]
+            self.thread = threading.Thread(target=self.heartbeat, args=(value["heartbeat_interval"],),
+                                           name="anno-job-heartbeat", daemon=True)
+            self.thread.start()
+
+    def heartbeat(self, interval):
+        # This target holds the session state, never the public AnnoClient shell.
+        while not self.stop.wait(interval):
+            try:
+                self.client._heartbeat(self.id)
+            except RemoteTransportError:
+                continue
+            except RemoteSubmissionError as exc:
+                if exc.http_status in (408, 429) or (exc.http_status is not None and exc.http_status >= 500):
+                    continue
+                self.failure = exc
+                return
+            except Exception as exc:
+                self.failure = exc
+                return
+
+    def close(self, timeout):
+        if self.closed:
+            return
+        self.closed = True
+        self.stop.set()
+        for stream in list(self.streams):
+            stream._closed = True
+            stream.buffer.clear()
+            stream.ready.clear()
+        def cleanup():
+            try:
+                if self.id is not None:
+                    self.client._close_session(self.id, timeout)
+            finally:
+                self.client.close()
+        bounded_cleanup(cleanup, timeout)
+
+
+def _finalize(session):
+    try:
+        session.close(1.0)
+    except BaseException:
+        # Destructors cannot report reliable cancellation. The durable session
+        # expiry is the fallback, including when the process is killed outright.
+        pass
+
+
+class RemoteAdapter:
+    def __init__(self, base_url, auth_token=None, *, mode="eager", optimize=True,
+                 timeout=30.0, transport=None):
+        if mode not in ("eager", "strict", "speculative"):
+            raise ValueError("mode must be eager, strict or speculative")
+        if type(optimize) is not bool:
+            raise TypeError("optimize must be a boolean")
+        self.mode, self.optimize = mode, optimize
+        self.session = _Session(AnnoJobClient(base_url, auth_token, timeout=timeout, transport=transport))
+
+    @property
+    def closed(self):
+        return self.session.closed
+
+    def submit(self, program, inputs, *, ordered=False):
+        if self.closed:
+            raise RuntimeError("client is closed")
+        if type(ordered) is not bool:
+            raise TypeError("ordered must be a boolean")
+        valid, errors = prepare_inputs(program, inputs)
+        ir = compile_program(program, optimize=self.optimize).to_dict()
+        job = None
+        if valid:
+            self.session.start()
+            try:
+                job = self.session.client.submit_ir(ir, [{"index": i, "state": v} for i, v in valid],
+                    idempotency_key=uuid.uuid4().hex, mode=self.mode, session_id=self.session.id)
+            except BaseException:
+                # Acceptance may have happened without its response. Closing
+                # the ownership session also reaches such otherwise unknown jobs.
+                try:
+                    self.session.close(5.0)
+                except BaseException:
+                    logging.getLogger(__name__).warning("Unconfirmed remote cleanup after submission failure")
+                raise
+        stream = _JobStream(self.session, job, errors, ordered)
+        self.session.streams.add(stream)
+        return stream
+
+    def close(self):
+        self.session.close(5.0)
+
+
+class _JobStream(ResultStream):
+    def __init__(self, session, job, errors, ordered):
+        super().__init__()
+        self.session, self.job, self.ordered = session, job, ordered
+        self.ready = deque(Result(i, "error", error=error) for i, error in sorted(errors.items()))
+        self.buffer, self.index, self.seen = {}, 0, set()
+        self.cursor = None
+        self.done = job is None
+
+    def _pull(self):
+        while not self.ready:
+            if self.done:
+                raise StopIteration
+            if self.session.closed:
+                raise RemoteSubmissionError("client_closed", "client session is closed")
+            if self.session.failure is not None:
+                raise self.session.failure
+            page = self.session.client.fetch_results(self.job, cursor=self.cursor)
+            received = page["job"]
+            for field in ("job_id", "ir_sha256", "release_id", "mode", "input_indices"):
+                if received[field] != self.job[field]:
+                    raise RemoteProtocolError("job binding changed while receiving results")
+            pending = []
+            for item in page["results"]:
+                if item["index"] in self.seen:
+                    raise RemoteProtocolError("result index was delivered more than once")
+                pending.append(Result(item["index"], item["status"], item["value"],
+                                      None if item["error"] is None else ResultError(**item["error"])))
+            new_seen = self.seen | {r.index for r in pending}
+            if page["done"] and new_seen != set(self.job["input_indices"]):
+                raise RemoteProtocolError("terminal job is missing input results")
+            # Commit a fully validated page to memory as a unit, then advance.
+            self.ready.extend(pending)
+            self.seen, self.cursor, self.done = new_seen, page["next_cursor"], page["done"]
+            self.job = received
+            if not self.ready:
+                if received["state"] == "paused":
+                    raise RemoteJobPausedError(received)
+                if not self.done:
+                    self.session.stop.wait(0.2)
+        return self.ready.popleft()
+
+    def _next_result(self):
+        if not self.ordered:
+            return self._pull()
+        while self.index not in self.buffer:
+            item = self._pull()
+            self.buffer[item.index] = item
+        item = self.buffer.pop(self.index)
+        self.index += 1
+        return item
+
+    def _close(self):
+        deadline = time.monotonic() + 5.0
+        try:
+            if self.job is not None and not self.done and not self.session.closed:
+                try:
+                    bounded_cleanup(lambda: self.session.client.cancel(self.job, attempts=1, budget=5.0), 5.0)
+                except RemoteCancellationError:
+                    # Stop owning work if even cancellation cannot be delivered.
+                    # Do not keep heartbeating an orphaned job indefinitely.
+                    try:
+                        self.session.close(max(0.0, deadline - time.monotonic()))
+                    except RemoteCancellationError:
+                        pass
+                    raise
+        finally:
+            self.ready.clear()
+            self.buffer.clear()
+            self.session.streams.discard(self)

annoscript/_request_adapters.py

@@ -0,0 +1,71 @@
+"""Versioned request recipes; independent of models, credentials and transport."""
+
+DEFAULT_DECISION_ADAPTER = "decision_native_v1"
+DEFAULT_GENERATION_ADAPTER = "generation_chat_v1"
+CONFIGURABLE_DECISION_ADAPTER = "decision_configurable_v1"
+REQUEST_ADAPTERS = {
+    "decision": (DEFAULT_DECISION_ADAPTER, CONFIGURABLE_DECISION_ADAPTER),
+    "generation": (DEFAULT_GENERATION_ADAPTER,),
+}
+
+
+def validate_request_adapter(value, stage):
+    if stage not in REQUEST_ADAPTERS:
+        raise ValueError("request adapter stage must be decision or generation")
+    if type(value) is not str or value not in REQUEST_ADAPTERS[stage]:
+        raise ValueError(f"unsupported {stage} request_adapter; expected " +
+                         ", ".join(REQUEST_ADAPTERS[stage]))
+    return value
+
+
+def native_decision_recipe():
+    return {"criteria": "preserve", "examples": "preserve", "state": {"mode": "full"}}
+
+
+def normalize_decision_options(adapter, options=None):
+    """Detach and validate JSON configuration before any transport is initialized."""
+    from annoscript._data import snapshot
+    validate_request_adapter(adapter, "decision")
+    if options is not None and type(options) is not dict:
+        raise ValueError("request_adapter_options must be an object")
+    if adapter == DEFAULT_DECISION_ADAPTER:
+        if options:
+            raise ValueError("decision_native_v1 does not accept request_adapter_options")
+        return None
+    options = snapshot(options) if options is not None else {}
+    if set(options) - {"default", "tasks"}:
+        raise ValueError("request_adapter_options only accepts default and tasks")
+
+    def recipe(value):
+        if type(value) is not dict or set(value) != {"criteria", "examples", "state"}:
+            raise ValueError("a task recipe requires criteria, examples and state")
+        if value["criteria"] not in ("preserve", "text") or type(value["criteria"]) is not str:
+            raise ValueError("criteria must be preserve or text")
+        if value["examples"] not in ("preserve", "omit") or type(value["examples"]) is not str:
+            raise ValueError("examples must be preserve or omit")
+        if value["criteria"] == "text" and value["examples"] == "preserve":
+            raise ValueError("text criteria cannot preserve nested examples")
+        state = value["state"]
+        if type(state) is not dict:
+            raise ValueError("state projection must be an object")
+        mode = state.get("mode")
+        if mode == "full" and set(state) == {"mode"}:
+            pass
+        elif mode == "fields_object" and set(state) == {"mode", "fields"}:
+            fields = state["fields"]
+            if (type(fields) is not list or not fields
+                    or any(type(f) is not str or not f.strip() for f in fields)
+                    or len(set(fields)) != len(fields)):
+                raise ValueError("fields_object needs distinct nonblank fields in order")
+        elif mode == "field_string" and set(state) == {"mode", "field"}:
+            if type(state["field"]) is not str or not state["field"].strip():
+                raise ValueError("field_string needs a nonblank field")
+        else:
+            raise ValueError("invalid state projection")
+        return {"criteria": value["criteria"], "examples": value["examples"], "state": state}
+
+    tasks = options.get("tasks", {})
+    if type(tasks) is not dict or any(type(k) is not str or not k.strip() for k in tasks):
+        raise ValueError("tasks must map exact nonblank field_name values to recipes")
+    return {"default": recipe(options.get("default", native_decision_recipe())),
+            "tasks": {key: recipe(tasks[key]) for key in sorted(tasks)}}

annoscript/_speculation.py

@@ -0,0 +1,321 @@
+"""A bounded, private argument evaluator; never commits Plan task effects.
+
+Formal execution still owns branches, memory, checks, returns and loop progress.
+Only normalized, reusable inference requests cross from this evaluator to the
+transport scheduler. Values are rebuilt after formal progress or a response.
+Each revision reaches Decision quiescence before collecting generation work;
+generation requires confirmed control gates, even before formal branch progress.
+"""
+
+from copy import deepcopy
+
+from annoscript import plan as p
+from annoscript._expressions import expression, unpack
+from annoscript.backends.base import normalize_call, DECISION_OPERATIONS
+from annoscript.backends._bounded_grammar import ensure_satisfiable
+
+
+class Unavailable(Exception):
+    pass
+
+
+def detached(value):
+    """Reject consumable iterators, including those nested in containers."""
+    seen = set()
+    def check(item):
+        if id(item) in seen:
+            return
+        seen.add(id(item))
+        if type(item) in (list, tuple):
+            for child in item:
+                check(child)
+        elif type(item) is dict:
+            for child in item.values():
+                check(child)
+        elif hasattr(item, "__next__"):
+            raise Unavailable()
+    check(value)
+    return deepcopy(value)
+
+
+class Speculation:
+    def __init__(self, run):
+        self.run = run
+        self.sites = tuple(t for t in run.plan.tasks if isinstance(t, p.Infer))
+        self.requests = []
+        self.scan = None
+        self.revision = -1
+        self.ready = False
+        self.phase = "decision"
+        self.selected_requests = set()
+
+    def viable(self, task, ctx):
+        return all(self.run.port(g.task, ctx, g.port) is not False for g in task.gates)
+
+    def memory_ready(self, task, ctx):
+        return all(self.run.dependency(d, ctx) for d in self.run.incoming[task.id]
+                   if d.kind == "memory")
+
+    def port(self, tid, ctx, port, visiting):
+        """Read a control decision without launching regions or committing effects."""
+        run = self.run
+        known = run.port(tid, ctx, port)
+        if type(known) is bool:
+            return known
+        task = run.tasks[tid]
+        parent = ctx[:len(run.paths[tid])]
+        key = tid, parent, port
+        yield
+        if key in visiting or not self.memory_ready(task, parent):
+            raise Unavailable()
+        visiting = visiting | {key}
+        yield from self.selected(task, parent, visiting)
+        if isinstance(task, p.Branch) and port in ("true", "false"):
+            selected = bool((yield from self.bound(task.condition, parent, visiting)))
+            return selected if port == "true" else not selected
+        if isinstance(task, p.Loop) and port == "enter":
+            child = ctx[:len(parent) + 1]
+            if len(child) > len(parent) and ((tid, child) in self.items or (tid, child) in run.items):
+                return True
+        # Join/return completion and mutable iteration remain formally owned.
+        raise Unavailable()
+
+    def selected(self, task, ctx, visiting=frozenset()):
+        for gate in task.gates:
+            if not (yield from self.port(gate.task, ctx, gate.port, visiting)):
+                raise Unavailable()
+
+    def possible(self, task, ctx, visiting=frozenset()):
+        for gate in task.gates:
+            try:
+                known = yield from self.port(gate.task, ctx, gate.port, visiting)
+            except Unavailable:
+                continue
+            if not known:
+                return False
+        return True
+
+    def lookup(self, key, call):
+        for request in self.requests:
+            if request.key == key and self.same(request.call, call):
+                return request
+        return None
+
+    def retire_mismatches(self, key, call, task, ctx):
+        for request in self.requests:
+            if request.key == key and not self.same(request.call, call):
+                origin = task.id, ctx
+                if origin in request.origins:
+                    request.origins.remove(origin)
+
+    @staticmethod
+    def same(left, right):
+        return (left.operation == right.operation and left.arguments == right.arguments
+                and left.extraction_source == right.extraction_source)
+
+    def request(self, task, ctx, arguments):
+        run = self.run
+        definition = run.requests[task.request]
+        if not definition.reusable:
+            raise Unavailable()
+        args, values = arguments
+        call = normalize_call(run.call_id(), definition.operation, args,
+                              dict(zip(task.keywords, values)), run.context)
+        key = run.request_key(task, ctx)
+        request = run.cache.get(key)
+        if request is not None:
+            if not self.same(request.call, call):
+                raise Unavailable()
+        else:
+            request = self.lookup(key, call)
+        if request is None:
+            kind = "decision" if call.operation in DECISION_OPERATIONS else "llm"
+            if kind != self.phase:
+                raise Unavailable()
+            request = run.new_request(call, [], key=key, speculative=True)
+            self.requests.append(request)
+            if call.operation not in DECISION_OPERATIONS:
+                try:
+                    ensure_satisfiable(call)
+                except Exception as exc:
+                    request.error, request.status = exc, "error"
+                    run.runtime._record("speculation_error_deferred", run, call=call.id,
+                                        code=getattr(exc, "code", "invalid_arguments"))
+            run.runtime._record("infer_speculated", run, task=task.id, path=ctx,
+                                call=call.id, operation=call.operation)
+        origin = task.id, ctx
+        if origin not in request.origins:
+            request.origins.append(origin)
+            request.discarded = False
+        if self.phase == "llm" and call.operation not in DECISION_OPERATIONS:
+            self.selected_requests.add(id(request))
+        return request
+
+    def bound(self, bound, ctx, visiting):
+        env = {}
+        for binding in bound.bindings:
+            env[binding.name] = yield from self.value(binding.value, ctx, visiting)
+        # Allocations remain private and cannot be addressed by formal Store.
+        return expression(bound.expression, env, self.run.input.state, {})
+
+    def value(self, vid, ctx, visiting=frozenset()):
+        run = self.run
+        task = run.tasks[run.values[vid].producer]
+        ctx = ctx[:len(run.paths[task.id])]
+        key = vid, ctx
+        yield
+        if key in visiting:
+            raise Unavailable()
+        if key in self.memo:
+            return self.memo[key]
+        if not self.viable(task, ctx) or not self.memory_ready(task, ctx):
+            raise Unavailable()
+        if self.phase == "llm":
+            yield from self.selected(task, ctx, visiting)
+        elif not (yield from self.possible(task, ctx, visiting)):
+            raise Unavailable()
+        if key in run.outputs:
+            return detached(run.outputs[key])
+        visiting = visiting | {key}
+        if isinstance(task, p.Eval):
+            values = [(yield from self.bound(task.expression, ctx, visiting))]
+        elif isinstance(task, p.Unpack):
+            value = yield from self.bound(task.expression, ctx, visiting)
+            values = unpack(task.target, detached(value))
+        elif isinstance(task, p.Infer):
+            arguments = yield from self.value(task.arguments, ctx, visiting)
+            request = self.request(task, ctx, arguments)
+            if request.status != "done":
+                raise Unavailable()
+            values = [request.response]
+        elif isinstance(task, p.Merge):
+            selected = yield from self.port(task.branch, ctx, "true", visiting)
+            source = task.then_value if selected else task.else_value
+            if source is None:
+                raise Unavailable()
+            values = [(yield from self.value(source, ctx, visiting))]
+        elif isinstance(task, p.Carry):
+            source = task.initial if ctx[-1] == 0 else task.update
+            parent = ctx[:-1] if ctx[-1] == 0 else ctx[:-1] + (ctx[-1] - 1,)
+            values = [(yield from self.value(source, parent, visiting))]
+        elif isinstance(task, p.Iteration):
+            item_key = task.loop, ctx
+            if item_key not in self.items and item_key not in run.items:
+                raise Unavailable()
+            item = self.items[item_key] if item_key in self.items else run.items[item_key]
+            values = unpack(task.target, detached(item))
+        else:
+            raise Unavailable()
+        for output, value in zip(task.outputs, values):
+            self.memo[output, ctx] = detached(value)
+        return self.memo[key]
+
+    def contexts(self, path, ctx=()):
+        if not path:
+            yield ctx
+            return
+        run = self.run
+        loop = run.tasks[path[0]]
+        if not self.viable(loop, ctx) or not self.memory_ready(loop, ctx):
+            return
+        if self.phase == "llm":
+            try:
+                yield from self.selected(loop, ctx)
+            except Unavailable:
+                return
+        inst = run.instances.get((loop.id, ctx))
+        existing = inst.iterations if inst else []
+        start = next((i for i in existing if not run.region_done(loop.body, ctx + (i,))), len(existing))
+        window = 1 if "eager" in loop.serial_modes else run.runtime.loop_window
+        sequence = None
+        # Mutable iteration must remain owned by the formal executor.
+        if "eager" not in loop.serial_modes:
+            try:
+                sequence = yield from self.bound(loop.iterable, ctx, frozenset())
+                if type(sequence) is dict:
+                    sequence = tuple(sequence)
+                if type(sequence) not in (str, list, tuple, range):
+                    sequence = None
+            except Exception:
+                sequence = None
+        for index in range(start, start + window):
+            yield None  # Charge loop expansion to the local scheduling budget.
+            child = ctx + (index,)
+            if (loop.id, child) in run.items:
+                item = run.items[loop.id, child]
+            elif sequence is not None and index < len(sequence):
+                item = sequence[index]
+            else:
+                break
+            try:
+                self.items[loop.id, child] = detached(item)
+            except Exception:
+                break
+            yield from self.contexts(path[1:], child)
+
+    def candidates(self):
+        for task in self.sites:
+            definition = self.run.requests[task.request]
+            kind = "decision" if definition.operation in DECISION_OPERATIONS else "llm"
+            if not definition.reusable or kind != self.phase:
+                continue
+            for ctx in self.contexts(self.run.paths[task.id]):
+                if ctx is None:
+                    yield
+                    continue
+                yield
+                if not self.viable(task, ctx) or not self.memory_ready(task, ctx):
+                    continue
+                inst = self.run.instances.get((task.id, ctx))
+                if inst is not None and inst.status == "done":
+                    continue
+                try:
+                    if self.phase == "llm":
+                        yield from self.selected(task, ctx)
+                    elif not (yield from self.possible(task, ctx)):
+                        continue
+                    arguments = yield from self.value(task.arguments, ctx)
+                    self.request(task, ctx, arguments)
+                except Exception:
+                    # Argument failures on an unselected path have no effects.
+                    # Formal execution re-evaluates and reports them on demand.
+                    continue
+
+    def advance(self, budget):
+        if self.revision != self.run.revision:
+            self.scan = None
+            self.ready = False
+            self.revision = self.run.revision
+            self.phase = "decision"
+            self.selected_requests.clear()
+        count = 0
+        while count < budget:
+            if self.ready:
+                self.prune()
+                if (self.phase != "decision" or not self.run.local_ready or any(
+                        r.call.operation in DECISION_OPERATIONS and r.status in ("queued", "inflight")
+                        and (not r.discarded or r.status == "inflight") for r in self.run.all_requests())):
+                    break
+                self.phase, self.ready = "llm", False
+                count += 1
+                continue
+            if self.scan is None:
+                self.memo, self.items = {}, {}
+                self.scan = self.candidates()
+            try:
+                next(self.scan)
+            except StopIteration:
+                self.scan = None
+                self.ready = True
+                continue
+            count += 1
+        return count
+
+    def prune(self, *, final=False):
+        for request in self.requests:
+            if request.waiters or request.discarded:
+                continue
+            if final or not any(self.viable(self.run.tasks[tid], ctx) for tid, ctx in request.origins):
+                request.discarded = True
+                self.run.runtime._record("speculation_discarded", self.run, call=request.call.id,
+                                         status=request.status, failed=request.error is not None)

annoscript/_test_artifacts.py

@@ -0,0 +1,33 @@
+"""Artifact assertions: Python 3.12 is the canonical generation interpreter."""
+
+from dataclasses import fields, is_dataclass
+import json
+import sys
+
+from annoscript.ir import SourceSpan
+
+
+def without_coordinates(value):
+    """Ignore only typed source coordinates, never user fields named 'source'."""
+    if isinstance(value, SourceSpan):
+        return {"module": value.module, "function": value.function}
+    if is_dataclass(value):
+        result = {f.name: without_coordinates(getattr(value, f.name)) for f in fields(value)}
+        if hasattr(value, "kind"):
+            result["kind"] = value.kind
+        return result
+    if isinstance(value, tuple):
+        return [without_coordinates(item) for item in value]
+    return value
+
+
+def assert_saved_artifact(test, path, generated):
+    text = path.read_text(encoding="utf-8")
+    restored = type(generated).from_json(text)
+    if sys.version_info[:2] == (3, 12):
+        test.assertEqual(text, generated.to_json(indent=2) + "\n")
+    else:
+        # Serialize to keep bool/int/float distinctions, as in the IR optimizer.
+        encode = lambda v: json.dumps(without_coordinates(v), ensure_ascii=False, allow_nan=False)
+        test.assertEqual(encode(restored), encode(generated))
+    test.assertEqual(generated.to_json(), type(generated).from_json(generated.to_json()).to_json())