annoscript 0.1.0__py3-none-any.whl

annoscript/server/worker.py

@@ -0,0 +1,290 @@
+"""Lease-owning worker with isolated Runtime instances for individual attempts."""
+
+import asyncio
+import os
+from annoscript._env import env_value
+import random
+import time
+import uuid
+
+from annoscript.backends.base import HTTPBackend
+from annoscript.server.metrics import RequestMeter
+from annoscript.server._plan_cache import PlanCache
+from annoscript._dispatch_pool import RequestSlots
+
+from annoscript._client_common import Result, ResultError
+from annoscript.backends import BackendFatalError
+from annoscript.ir import Program
+from annoscript.job_protocol import encode, retry_after_seconds
+from annoscript.passes import optimize_program
+from annoscript.plan_lowering import lower_to_plan
+from annoscript.runtime import AnnoRuntime, RunInput
+
+
+def backends(manifest):
+    from annoscript.backends.decision import DecisionAPIBackend
+    from annoscript.backends.llm import LLMAPIBackend
+    from annoscript._request_adapters import DEFAULT_DECISION_ADAPTER, DEFAULT_GENERATION_ADAPTER
+    try:
+        decision_token = env_value("ANNO_DECISION_AUTH_TOKEN")
+        generation_token = env_value("ANNO_GENERATE_AUTH_TOKEN")
+        if decision_token is None:
+            raise KeyError("ANNO_DECISION_AUTH_TOKEN")
+        if generation_token is None:
+            raise KeyError("ANNO_GENERATE_AUTH_TOKEN")
+    except KeyError:
+        raise BackendFatalError("model credentials are unavailable", code="configuration_failed") from None
+    timeout = manifest["policy"]["model_timeout"]
+    d, g = manifest["decision"], manifest["generation"]
+    return (DecisionAPIBackend(d["base_url"], decision_token, d["model"], timeout=timeout,
+                               request_adapter=d.get("request_adapter", DEFAULT_DECISION_ADAPTER),
+                               request_adapter_options=d.get("request_adapter_options")),
+            LLMAPIBackend(g["base_url"], generation_token, g["model"], timeout=timeout,
+                          request_adapter=g.get("request_adapter", DEFAULT_GENERATION_ADAPTER),
+                          profile=g["profile"], extract_mode=g.get("extract_mode", "validated"),
+                          batch_rule_dedup=g.get("batch_rule_dedup", False),
+                          disable_thinking=manifest["disable_thinking"], request_policy=g.get("request_policy"),
+                          assumed_avg_num_tokens_per_char=g.get("assumed_avg_num_tokens_per_char", 2)))
+
+
+class ObservedBackend:
+    def __init__(self, backend, failures, owned, *, dispatch_allowed=None, on_failure=None):
+        self.backend, self.failures, self.owned = backend, failures, owned
+        self.dispatch_allowed = dispatch_allowed
+        self.on_failure = on_failure
+        if isinstance(backend, HTTPBackend):
+            backend._dispatch_guard = self._check_dispatch
+
+    def _check_dispatch(self):
+        if not self.owned():
+            raise asyncio.CancelledError()
+        if self.dispatch_allowed is not None and not self.dispatch_allowed():
+            raise BackendFatalError("job is paused; no new model request was sent", code="job_paused")
+
+    @property
+    def supports_llm_batch(self):
+        return getattr(self.backend, "supports_llm_batch", False) is True
+
+    @property
+    def supports_decision_batch_planning(self):
+        return getattr(self.backend, "supports_decision_batch_planning", False) is True
+
+    @property
+    def supports_partial_answers(self):
+        return getattr(self.backend, "supports_partial_answers", False) is True
+
+    def plan_batch(self, *args, **kwargs):
+        return self.backend.plan_batch(*args, **kwargs)
+
+    @property
+    def supports_prepared_batch(self):
+        return getattr(self.backend, "supports_prepared_batch", False) is True
+
+    def resolve_token_budget(self, call, fallback):
+        resolve = getattr(self.backend, "resolve_token_budget", None)
+        return resolve(call, fallback) if resolve else (call.arguments.get("max_num_tokens") or fallback)
+
+    async def infer(self, *args, **kwargs):
+        self._check_dispatch()
+        try:
+            return await self.backend.infer(*args, **kwargs)
+        except Exception as exc:
+            self.failures.append(exc)
+            if self.on_failure is not None:
+                self.on_failure(exc)
+            raise
+
+    async def infer_batch(self, *args, **kwargs):
+        self._check_dispatch()
+        try:
+            return await self.backend.infer_batch(*args, **kwargs)
+        except Exception as exc:
+            self.failures.append(exc)
+            if self.on_failure is not None:
+                self.on_failure(exc)
+            raise
+
+    async def aclose(self):
+        await self.backend.aclose()
+
+
+def classification(code, metadata, *, retry_unknown_outcomes=True):
+    status = metadata.get("http_status")
+    if code in ("authentication_failed", "configuration_failed", "backend_protocol_error",
+                "job_paused", "metrics_failed", "diagnostics_failed") or status in (401, 403):
+        return "pause"
+    if not retry_unknown_outcomes and status is None and code in ("model_timeout", "transport_error"):
+        return "pause"
+    if (code in ("backend_unavailable", "model_timeout", "transport_error")
+            or status in (408, 429) or (type(status) is int and 500 <= status <= 599)):
+        return "retry"
+    if code == "runtime_failed":
+        return "fatal"
+    return "result"
+
+
+class Worker:
+    def __init__(self, store, releases, *, backend_factory=backends, concurrency=8,
+                 poll_interval=0.2, random_source=random.random, request_observer=None,
+                 diagnostic_observer=None, model_request_limits=None):
+        if type(concurrency) is not int or concurrency < 1:
+            raise ValueError("worker concurrency must be positive")
+        self.store, self.releases, self.backend_factory = store, releases, backend_factory
+        self.concurrency, self.poll_interval, self.random = concurrency, poll_interval, random_source
+        self.request_observer = request_observer
+        self.diagnostic_observer = diagnostic_observer
+        self.id = uuid.uuid4().hex
+        self.plan_cache = PlanCache()
+        if model_request_limits is not None and (type(model_request_limits) is not dict
+                or set(model_request_limits) != {"generation", "decision"}):
+            raise ValueError("model_request_limits requires generation and decision")
+        self.request_slots = RequestSlots(**model_request_limits) if model_request_limits is not None else None
+
+    async def execute(self, claim):
+        manifest = self.releases.manifests[claim["release_id"]]
+        failures = []
+        started = time.monotonic()
+
+        async def run():
+            def diagnose(stage, event):
+                if self.diagnostic_observer is not None:
+                    try:
+                        self.diagnostic_observer({"job_id": claim["job_id"], "index": claim["index"],
+                            "attempt": claim["attempt"], "stage": stage, **event})
+                    except Exception:
+                        raise BackendFatalError("diagnostic persistence failed", code="diagnostics_failed") from None
+            def failed(exc):
+                code, metadata = getattr(exc, "code", "execution_failed"), getattr(exc, "metadata", {})
+                if classification(code, metadata, retry_unknown_outcomes=manifest["policy"].get("retry_unknown_outcomes", True)) == "pause":
+                    self.store.pause_dispatch(claim, {"code": code, "message": str(exc)})
+            if self.request_slots is not None:
+                diagnose("scheduling", {"event": "resource_limits", "active_inputs": self.concurrency,
+                    "limits": self.request_slots.limits, "algorithm": "ready_at_admission_fifo_v1"})
+            def compile_plan(encoded):
+                program = Program.from_json(encoded.decode())
+                if manifest["server_optimize"]:
+                    program = optimize_program(program)
+                return lower_to_plan(program)
+            plan, preparation = await self.plan_cache.get(claim["release_id"], claim["ir"], compile_plan)
+            diagnose("preparation", preparation)
+            try:
+                d, g = self.backend_factory(manifest)
+            except ValueError:
+                raise BackendFatalError("invalid backend configuration", code="configuration_failed") from None
+            supported = isinstance(d, HTTPBackend) and isinstance(g, HTTPBackend) and d is not g
+            if not self.store.meter_attempt(claim, supported):
+                raise asyncio.CancelledError()
+            if isinstance(d, HTTPBackend):
+                d._request_meter = RequestMeter(self.store, claim, "decision", observer=self.request_observer)
+                if self.diagnostic_observer is not None:
+                    d.diagnostic_observer = lambda event: diagnose("decision", event)
+            if isinstance(g, HTTPBackend) and d is not g:
+                g._request_meter = RequestMeter(self.store, claim, "generation", observer=self.request_observer)
+                if self.diagnostic_observer is not None:
+                    g.diagnostic_observer = lambda event: diagnose("generation", event)
+            decision = ObservedBackend(d, failures, lambda: self.store.renew(claim),
+                                       dispatch_allowed=lambda: self.store.can_dispatch(claim),
+                                       on_failure=failed if self.request_slots is not None else None)
+            generation = decision if d is g else ObservedBackend(g, failures, lambda: self.store.renew(claim),
+                                                                 dispatch_allowed=lambda: self.store.can_dispatch(claim),
+                                                                 on_failure=failed if self.request_slots is not None else None)
+            runtime = AnnoRuntime(decision, generation, trace=self.diagnostic_observer is not None,
+                                 request_slots=self.request_slots, **manifest["runtime"])
+            try:
+                async with runtime:
+                    async with await runtime.submit(plan, [RunInput(claim["index"], claim["state"])], mode=claim["mode"]) as stream:
+                        result = await anext(stream)
+                        # Verify one and only one outcome for this isolated input.
+                        try:
+                            await anext(stream)
+                        except StopAsyncIteration:
+                            return result
+                        raise RuntimeError("multiple results for an input attempt")
+            finally:
+                if self.diagnostic_observer is not None:
+                    diagnose("runtime", {"event": "trace", "trace": runtime.trace,
+                                         "plan": plan.to_dict()})
+
+        task = asyncio.create_task(run())
+
+        async def renew():
+            try:
+                while True:
+                    await asyncio.sleep(self.store.lease_seconds / 4)
+                    if not self.store.renew(claim):
+                        task.cancel()
+                        return
+            except Exception:
+                task.cancel()
+
+        heartbeat = asyncio.create_task(renew())
+        try:
+            try:
+                async with asyncio.timeout(manifest["policy"]["input_timeout"]):
+                    result = await task
+                if result.status == "success":
+                    self.store.finish(claim, result, elapsed_seconds=time.monotonic() - started)
+                    return
+                code, message = result.error.code, result.error.message
+                original = next((e for e in failures if getattr(e, "code", None) == code), None)
+                metadata = getattr(original, "metadata", {}) if original is not None else {}
+            except TimeoutError:
+                self.store.finish(claim, Result(claim["index"], "error", error=ResultError("execution_timeout", "input execution time limit exceeded")), elapsed_seconds=time.monotonic() - started)
+                return
+            except BackendFatalError as exc:
+                code, message, metadata = exc.code, str(exc), exc.metadata
+            except asyncio.CancelledError:
+                # Cancelled/expired ownership cannot commit anything. Process
+                # shutdown leaves a lease that another worker can reclaim.
+                raise
+            except Exception:
+                self.store.fail_job(claim, {"code": "execution_failed", "message": "job lowering or runtime failed"}, elapsed_seconds=time.monotonic() - started)
+                return
+            error = {"code": code, "message": message}
+            kind = classification(code, metadata,
+                                  retry_unknown_outcomes=manifest["policy"].get("retry_unknown_outcomes", True))
+            if kind in ("pause", "retry"):
+                delay = max(self.random() * min(30.0, 2 ** (claim["round_attempt"] - 1)),
+                            retry_after_seconds(metadata.get("retry_after"), self.store.clock()))
+                self.store.attempt_failed(claim, error, pause=kind == "pause", delay=delay, metadata=metadata, elapsed_seconds=time.monotonic() - started)
+            elif kind == "fatal":
+                self.store.fail_job(claim, error, elapsed_seconds=time.monotonic() - started)
+            else:
+                self.store.finish(claim, Result(claim["index"], "error", error=ResultError(**error)), elapsed_seconds=time.monotonic() - started)
+        finally:
+            heartbeat.cancel()
+            task.cancel()
+            await asyncio.gather(heartbeat, task, return_exceptions=True)
+
+    async def run(self, *, stop=None):
+        active = set()
+        stop_waiter = asyncio.create_task(stop.wait()) if stop is not None else None
+        try:
+            while stop is None or not stop.is_set():
+                self.store.maintain()
+                for task in list(active):
+                    if task.done():
+                        active.remove(task)
+                        if not task.cancelled():
+                            task.result()
+                while len(active) < self.concurrency:
+                    claim = self.store.claim(self.id, self.releases)
+                    if claim is None:
+                        break
+                    active.add(asyncio.create_task(self.execute(claim)))
+                waiting = active | ({stop_waiter} if stop_waiter is not None else set())
+                if waiting:
+                    await asyncio.wait(waiting, timeout=self.poll_interval, return_when=asyncio.FIRST_COMPLETED)
+                else:
+                    await asyncio.sleep(self.poll_interval)
+        finally:
+            for task in active:
+                task.cancel()
+            await asyncio.gather(*active, return_exceptions=True)
+            if stop_waiter is not None:
+                stop_waiter.cancel()
+                await asyncio.gather(stop_waiter, return_exceptions=True)
+            await self.aclose()
+
+    async def aclose(self):
+        await self.plan_cache.aclose()

annoscript/substring_tests.py

@@ -0,0 +1,79 @@
+"""Acceptance against pinned XGrammar with the nonempty extension."""
+import itertools
+import json
+import random
+import unittest
+from types import SimpleNamespace
+from importlib.metadata import version
+import xgrammar as xgr
+from annoscript.backends._grammar import extract_grammar, batch_grammar
+
+
+def compiler():
+    return xgr.GrammarCompiler(xgr.TokenizerInfo([]), cache_enabled=False)
+
+
+def accepts(compiled, text):
+    m=xgr.GrammarMatcher(compiled, terminate_without_stop_token=True)
+    return m.accept_string(text) and m.is_terminated()
+
+
+def wire(items):
+    return json.dumps({'items':items},ensure_ascii=False,separators=(',',':'))
+
+
+class SubstringGrammarTests(unittest.TestCase):
+    @classmethod
+    def setUpClass(cls):
+        assert version('xgrammar') in ('0.2.5+lev.substring1','0.2.5+lev.substring2','0.2.5+lev.substring2.1','0.2.5+lev.substring2.2','0.2.5+lev.substring2.3')
+
+    def test_nonempty_macro_and_serialization(self):
+        for chunks in [('a','b'),('a','','b'),('','',''),()]:
+            for empty in (True,False):
+                args=', '.join(json.dumps(c) for c in chunks)
+                g=xgr.Grammar.from_ebnf('root ::= Substring('+args+(', ' if args else '')+
+                                      'allow_empty='+str(empty).lower()+')')
+                for copy in [g,xgr.Grammar.from_ebnf(str(g)),xgr.Grammar.deserialize_json(g.serialize_json())]:
+                    c=compiler().compile_grammar(copy)
+                    self.assertEqual(accepts(c,''),empty)
+                    for v in ['a','b','ab','ba','x']:
+                        self.assertEqual(accepts(c,v), v in {''.join(chunks[i:j]) for i in range(len(chunks)) for j in range(i+1,len(chunks)+1)} and bool(v))
+
+    def test_json_escaping_exhaustive_substrings(self):
+        for source in ['甲"乙\\丙\n😀','\x00\x1f\b\f\r\t','e\u0301／＠ @','\\u4e2d','aaaa','']:
+            c=compiler().compile_grammar(extract_grammar(source))
+            self.assertTrue(accepts(c,wire([])))
+            self.assertFalse(accepts(c,wire([''])))
+            for i in range(len(source)):
+                for j in range(i+1,len(source)+1):
+                    self.assertTrue(accepts(c,wire([source[i:j]])),repr((source,i,j)))
+            for bad in ['不存在',source+'不存在',source[::2]]:
+                if bad and bad not in source:self.assertFalse(accepts(c,wire([bad])))
+            for bad in ['{"items":["\\"]}',r'{"items":["\u00"]}','{"items":["\n"]}']:
+                self.assertFalse(accepts(c,bad))
+
+    def test_random_short_language(self):
+        rng=random.Random(20261016)
+        alphabet=['甲','乙','"','\\','\n','😀']
+        for _ in range(30):
+            source=''.join(rng.choices(alphabet,k=rng.randrange(1,7)))
+            c=compiler().compile_grammar(extract_grammar(source))
+            for n in range(3):
+                for chars in itertools.product(alphabet,repeat=n):
+                    value=''.join(chars)
+                    self.assertEqual(accepts(c,wire([value])), bool(value) and value in source)
+
+    def test_batch_isolation_and_other_operations(self):
+        calls=[SimpleNamespace(id=i,operation=op,arguments={},extraction_source=source) for i,op,source in
+               [('a','extract','甲乙'),('b','extract','丙丁'),('c','extract','甲乙'),
+                ('d','generate',''),('e','generate_list','')]]
+        c=compiler().compile_grammar(batch_grammar(calls,extract_mode='substring_chars'))
+        answers={'q1':['甲'],'q2':['丁'],'q3':['乙'],'q4':'任意改写','q5':['','任意改写']}
+        def encode():return json.dumps({'answers':answers},ensure_ascii=False,separators=(',',':'))
+        self.assertTrue(accepts(c,encode()))
+        for bad in ['丁','甲丁','']:
+            answers['q1']=[bad];self.assertFalse(accepts(c,encode()))
+        answers['q1']='甲';self.assertFalse(accepts(c,encode()))
+
+
+if __name__=='__main__':unittest.main()

annoscript/test_advisory_lengths.py

@@ -0,0 +1,115 @@
+"""Mandatory suggestions, hard contracts, budget sharing and inference-free rejection."""
+import asyncio
+import copy
+import json
+import math
+import unittest
+import httpx
+from annoscript.backends import InferenceContext, InferenceError
+from annoscript._test_evaluation_http import FixtureLLMBackend as LLMAPIBackend
+from annoscript._test_evaluation_http import llm_calls
+from annoscript.backends.base import normalize_call, validate_response
+from annoscript.backends._generation_policy import effective_budget
+from annoscript.backends._bounded_grammar import grammar
+from annoscript.compiler import compile_program, CompileError
+from annoscript.plan_lowering import lower_to_plan
+from annoscript.anno import GenerateResponse, GenerateListResponse, ExtractResponse
+from annoscript.test_generation_policy import policy
+
+
+def missing_generate(anno):
+    return anno.generate('x').text
+
+def missing_list(anno):
+    return anno.generate_list('x', max_items=1).items
+
+def missing_extract(anno):
+    return anno.extract('x', max_length_per_item=1).items
+
+def oversized(anno):
+    a = anno.generate('a', min_length=4, max_length=5).text
+    b = anno.generate_list('b', max_items=1, min_length_per_item=2, max_length_per_item=3).items
+    return [a,b]
+
+
+def call(op, **kw):
+    return normalize_call('x',op,('task',),kw,InferenceContext('甲乙丙丁'))
+
+class AdvisoryTests(unittest.TestCase):
+    def test_mandatory_at_compile_and_dynamic_normalization(self):
+        for f in (missing_generate,missing_list,missing_extract):
+            with self.assertRaises(CompileError):compile_program(f)
+        for op, valid in [('generate',{'max_length':5}),('generate_list',{'max_items':2,'max_length_per_item':5}),('extract',{'max_items':2,'max_length_per_item':5})]:
+            for field in valid:
+                for value in (None,False,-1,1.2,float('inf')):
+                    with self.subTest(op=op,field=field,value=value),self.assertRaises((ValueError,TypeError)):
+                        call(op,**{**valid,field:value})
+                opts=dict(valid);del opts[field]
+                with self.assertRaises(TypeError):call(op,**opts)
+        with self.assertRaises(ValueError):call('generate',min_length=6,max_length=5)
+
+    def test_budgets_ratio_explicit_priority_zero_and_repair(self):
+        for bad in (0,-1,True,float('inf'),float('nan'),'2',None):
+            with self.assertRaises(ValueError):effective_budget(call('generate',max_length=10),bad)
+        self.assertEqual(effective_budget(call('generate',max_length=3),1.5),5)
+        self.assertEqual(effective_budget(call('generate',max_length=0)),1)
+        self.assertEqual(effective_budget(call('extract',max_items=0,max_length_per_item=1)),32)
+        self.assertEqual(effective_budget(call('extract',max_items=16,max_length_per_item=300)),9760)
+        for op, kw in [('generate',{'max_length':80}),('generate_list',{'max_items':16,'max_length_per_item':16})]:
+            c=call(op,**kw,max_num_tokens=7)
+            self.assertEqual(effective_budget(c,100),7)
+            self.assertEqual(effective_budget(c,100,999),7)
+            self.assertEqual(effective_budget(c,100,3),3)
+        self.assertEqual(effective_budget(call('generate_list',max_items=16,max_length_per_item=16)),672)
+        self.assertEqual(effective_budget(call('extract',max_items=4,max_length_per_item=120)),1024)
+
+    def test_suggestions_and_hard_guarantees(self):
+        for c,r in [(call('generate',min_length=3,max_length=4),GenerateResponse('甲')),
+                    (call('generate',min_length=3,max_length=4),GenerateResponse('甲'*100)),
+                    (call('generate_list',min_length_per_item=3,max_length_per_item=4,max_items=1),GenerateListResponse(['甲'*100])),
+                    (call('extract',max_items=1,max_length_per_item=1),ExtractResponse(['甲乙丙丁']))]:
+            validate_response(c,r)
+        for c,r in [(call('generate',min_length=3,max_length=4),GenerateResponse('')),
+                    (call('generate_list',min_length_per_item=3,max_length_per_item=4,max_items=1),GenerateListResponse([''])),
+                    (call('generate_list',max_length_per_item=4,max_items=1),GenerateListResponse(['甲','乙'])),
+                    (call('extract',max_items=1,max_length_per_item=1),ExtractResponse(['甲丙']))]:
+            with self.assertRaises(InferenceError):validate_response(c,r)
+
+    def test_grammar_ignores_suggestions_keeps_source_and_count(self):
+        c=call('extract',max_items=1,max_length_per_item=1)
+        g=grammar([c],extract_mode='substring_chars')
+        self.assertNotIn('max_chunks',g);self.assertIn('allow_empty=false',g)
+        self.assertEqual(g,grammar([call('extract',max_items=1,max_length_per_item=100)],extract_mode='substring_chars'))
+        self.assertIn('Regex(',grammar([call('generate_list',max_items=2,max_length_per_item=3)]))
+        self.assertNotIn('json-char',grammar([call('generate_list',max_items=2,max_length_per_item=3)]))
+
+class AdvisoryWireTests(unittest.IsolatedAsyncioTestCase):
+    async def test_over_suggestion_no_repair_runtime_and_metadata(self):
+        from annoscript.runtime import AnnoRuntime,RunInput
+        from annoscript.test_runtime import FakeBackend
+        sent=[]
+        def handler(req):
+            body=json.loads(req.content);sent.append(body)
+            tasks=llm_calls(body)[1]
+            answers={f'q{index}':('甲'*100 if call.operation=='generate' else ['乙'*100]) for index,(_,call) in enumerate(tasks,1)}
+            return httpx.Response(200,json={'choices':[{'finish_reason':'stop','message':{'content':json.dumps({'answers':answers})}}]})
+        b=LLMAPIBackend('https://offline.invalid','fixture','q',profile='qwen_vllm',request_policy=policy(),_token_counter=lambda p:1,transport=httpx.MockTransport(handler))
+        rt=AnnoRuntime(FakeBackend(),b)
+        try:
+            out=[r async for r in (await rt.submit(lower_to_plan(compile_program(oversized)),[RunInput(0,'source')]))]
+            self.assertEqual(out[0].status,'success');self.assertEqual(len(sent),1)
+            self.assertEqual(sent[0]['max_tokens'],64+10+32+(6+32+8)+32)
+            self.assertIn('字符长度为建议',sent[0]['messages'][0]['content'])
+        finally:await rt.aclose()
+
+    async def test_direct_invalid_args_and_context_no_http(self):
+        sent=[];events=[]
+        b=LLMAPIBackend('https://offline.invalid','fixture','q',profile='qwen_vllm',request_policy=policy(),_token_counter=lambda p:8100,transport=httpx.MockTransport(lambda r:sent.append(r)),diagnostic_observer=events.append)
+        try:
+            with self.assertRaises(InferenceError):await b.infer(InferenceContext('甲'),call('generate',max_length=80))
+            self.assertEqual(events[0]['output_budget'],160)
+            self.assertEqual(events[0]['assumed_avg_num_tokens_per_char'],2)
+            self.assertEqual(sent,[])
+        finally:await b.aclose()
+
+if __name__=='__main__':unittest.main()