agentinputlab 0.1.0a1__py3-none-any.whl

agentinputlab/recipes/smartsearch_grpo_meter.py

@@ -0,0 +1,324 @@
+"""One coordinator-owned, fail-closed cost ledger for native GRPO workers.
+
+Workers must send reservations and settlements to this SAME instance (for
+example through a serial Ray actor). Instantiating a ledger per rank is not
+aggregate enforcement. No model, optimizer or Ray cluster is started here.
+"""
+from __future__ import annotations
+
+import threading
+from dataclasses import dataclass
+
+from agentinputlab.contracts import BudgetExceededError, ExternalDependencyError, ValidationError
+
+from .smartsearch_grpo_tools import _new_output, _save
+from .smartsearch_teacher import SmartSearchTeacherCompletion
+
+
+@dataclass(frozen=True, slots=True)
+class SmartSearchGRPOCostLimits:
+    policy_requests: int
+    policy_output_tokens: int
+    helper_requests: int
+    helper_output_tokens: int
+    actor_ranks: int
+    optimizer_steps_per_rank: int
+
+    def __post_init__(self):
+        if any(type(getattr(self, name)) is not int or getattr(self, name) < 1
+               for name in self.__dataclass_fields__) or self.actor_ranks > 4:
+            raise ValidationError("SmartSearch GRPO needs positive aggregate cost limits and one to four actor ranks")
+
+
+class SmartSearchGRPOCostLedger:
+    """Reserve before neural work; retain unknown usage and every rank's costs.
+
+    Input token counts are logical prompt tokens, not measured physical model
+    forwards. Returned output counts include ALL generated samples, before the
+    native selector drops any. Optimizer counters are per rank; synchronized
+    updates are never computed by summing ranks.
+    """
+
+    def __init__(self, *, output_root, limits: SmartSearchGRPOCostLimits):
+        if not isinstance(limits, SmartSearchGRPOCostLimits):
+            raise ValidationError("SmartSearch GRPO requires typed aggregate cost limits")
+        self.root = _new_output(output_root)
+        self.limits, self._lock = limits, threading.RLock()
+        self._failed, self._number = False, 0
+        self._pending = {}
+        self._invalid_usage_reports = []
+        self._costs = {kind: dict(requests=0, responses=0, logical_input_tokens=0,
+                                 output_tokens=0, reserved_output_tokens=0, generated_samples=0)
+                       for kind in ("policy", "helper")}
+        self._steps = {str(rank): dict(attempts=0, responses=0) for rank in range(limits.actor_ranks)}
+        _save(self.root, "claim.json", dict(limits=limits, coordinator_owned=True,
+            all_generated_samples_counted=True, physical_forwards_not_claimed=True,
+            synchronized_updates_not_sum_of_rank_steps=True, no_automatic_retry=True))
+
+    def _healthy(self):
+        if self._failed:
+            raise ExternalDependencyError("SmartSearch GRPO run has a retained failure; no automatic neural retry")
+
+    def _event(self, name, value):
+        self._number += 1
+        try:
+            _save(self.root, f"event-{self._number:08d}-{name}.json", value)
+        except BaseException:
+            self._failed = True
+            raise
+
+    def _reject(self, message, *, budget=False):
+        self._failed = True
+        self._event("rejected", dict(error=message, accept_checkpoint=False, no_automatic_retry=True))
+        error = BudgetExceededError if budget else ExternalDependencyError
+        raise error(message)
+
+    def reserve(self, *, kind, worker, input_token_counts, samples, max_new_tokens):
+        """Return an opaque ticket, shared across all worker/rank clients."""
+        with self._lock:
+            self._healthy()
+            if not isinstance(kind, str) or kind not in self._costs or not isinstance(worker, str) or not worker.strip() or (
+                len(worker) > 128 or not isinstance(input_token_counts, (tuple, list))
+                or not input_token_counts or len(input_token_counts) > 10000
+                or any(type(n) is not int or n < 0 for n in input_token_counts)
+                or type(samples) is not int or not 1 <= samples <= 64
+                or type(max_new_tokens) is not int or max_new_tokens < 1
+                or max_new_tokens > (2048 if kind == "helper" else 8192)
+                or (kind == "helper" and (samples != 1 or len(input_token_counts) != 1))
+            ):
+                raise ValidationError("SmartSearch GRPO generation reservation is invalid")
+            number = len(input_token_counts) * samples
+            reservation = number * max_new_tokens
+            costs = self._costs[kind]
+            if costs["requests"] >= getattr(self.limits, kind + "_requests") or (
+                costs["output_tokens"] + costs["reserved_output_tokens"] + reservation
+                > getattr(self.limits, kind + "_output_tokens")
+            ):
+                self._reject("SmartSearch GRPO aggregate generation budget exhausted", budget=True)
+            ticket = f"generation-{self._number + 1:08d}"
+            item = dict(ticket=ticket, kind=kind, worker=worker, input_token_counts=list(input_token_counts),
+                        samples=samples, output_samples=number, max_new_tokens=max_new_tokens,
+                        reserved_output_tokens=reservation)
+            self._event("generation-reserved", item)
+            self._pending[ticket] = item
+            costs["requests"] += 1
+            costs["logical_input_tokens"] += sum(input_token_counts)
+            costs["reserved_output_tokens"] += reservation
+            return ticket
+
+    def settle(self, ticket, *, output_token_counts):
+        with self._lock:
+            # Another rank may have failed while this paid request was in
+            # flight. Still retain its measured cost; never clear the latch.
+            if not isinstance(ticket, str) or ticket not in self._pending:
+                self._reject("SmartSearch GRPO ticket is unknown or already settled")
+            item = self._pending[ticket]
+            # Validation failures keep the entire reservation: usage is not
+            # assumed zero when the provider omitted or malformed its report.
+            if not isinstance(output_token_counts, (list, tuple)) or (
+                len(output_token_counts) != item["output_samples"]
+                or any(type(n) is not int or n < 0 or n > item["max_new_tokens"]
+                       for n in output_token_counts)
+            ):
+                report = dict(ticket=ticket, kind=item["kind"],
+                    reported_output_tokens=sum(n for n in output_token_counts if type(n) is int and n >= 0)
+                    if isinstance(output_token_counts, (list, tuple)) and len(output_token_counts) <= 10000 else None,
+                    reservation_retained=True)
+                self._invalid_usage_reports.append(report)
+                self._event("invalid-usage-report", report)
+                self._reject("SmartSearch GRPO provider token/sample counts differ from its reservation")
+            self._event("generation-settled", dict(ticket=ticket, output_token_counts=list(output_token_counts)))
+            costs = self._costs[item["kind"]]
+            costs["responses"] += 1
+            costs["output_tokens"] += sum(output_token_counts)
+            costs["generated_samples"] += len(output_token_counts)
+            costs["reserved_output_tokens"] -= item["reserved_output_tokens"]
+            del self._pending[ticket]
+
+    def fail(self, ticket, *, error_type):
+        """Retain the pre-call reservation after a failed/unknown provider call."""
+        with self._lock:
+            if not isinstance(error_type, str) or not error_type or len(error_type) > 128:
+                raise ValidationError("SmartSearch GRPO failure requires a bounded exception type")
+            self._failed = True
+            self._event("provider-failed", dict(ticket=ticket, error_type=error_type,
+                ticket_pending=ticket in self._pending, accept_checkpoint=False, no_automatic_retry=True))
+
+    def optimizer_begin(self, rank):
+        with self._lock:
+            self._healthy()
+            if type(rank) is not int or str(rank) not in self._steps:
+                raise ValidationError("SmartSearch GRPO optimizer rank is outside its admitted actor mesh")
+            cost = self._steps[str(rank)]
+            if cost["attempts"] != cost["responses"]:
+                self._reject("SmartSearch GRPO optimizer already has an unresolved step")
+            if cost["attempts"] >= self.limits.optimizer_steps_per_rank:
+                self._reject("SmartSearch GRPO actor optimizer step quota exhausted", budget=True)
+            self._event("optimizer-begin", dict(rank=rank, step=cost["attempts"] + 1))
+            cost["attempts"] += 1
+
+    def optimizer_end(self, rank):
+        with self._lock:
+            # Account for a concurrent rank's already-completed native step
+            # even when a different worker has invalidated this run.
+            if type(rank) is not int or str(rank) not in self._steps:
+                raise ValidationError("SmartSearch GRPO optimizer rank is outside its admitted actor mesh")
+            cost = self._steps[str(rank)]
+            if cost["attempts"] != cost["responses"] + 1:
+                self._reject("SmartSearch GRPO optimizer settlement has no pending step")
+            self._event("optimizer-end", dict(rank=rank, step=cost["attempts"]))
+            cost["responses"] += 1
+
+    def snapshot(self):
+        with self._lock:
+            return dict(failed=self._failed, costs={kind: dict(value) for kind, value in self._costs.items()},
+                        optimizer_rank_steps={rank: dict(value) for rank, value in self._steps.items()},
+                        pending_generation_tickets=sorted(self._pending),
+                        invalid_usage_reports=[dict(row) for row in self._invalid_usage_reports],
+                        physical_model_forwards_measured=False)
+
+    def assert_healthy(self, *, expected_optimizer_steps=None):
+        with self._lock:
+            self._healthy()
+            if self._pending or any(n["attempts"] != n["responses"] for n in self._steps.values()):
+                raise ExternalDependencyError("SmartSearch GRPO has unresolved paid work; checkpoint not accepted")
+            steps = [n["responses"] for n in self._steps.values()]
+            if len(set(steps)) != 1:
+                raise ExternalDependencyError("SmartSearch GRPO actor rank step counts differ")
+            if expected_optimizer_steps is not None and (
+                type(expected_optimizer_steps) is not int or expected_optimizer_steps < 1
+                or steps[0] != expected_optimizer_steps
+            ):
+                raise ExternalDependencyError("SmartSearch GRPO actual actor steps differ from the admitted plan")
+            return {**self.snapshot(), "synchronized_optimizer_steps": steps[0]}
+
+
+def observe_native_grpo_optimizer(optimizer, *, ledger, rank):
+    """Observe native AdamW.step without replacing loss, gradients or schedule.
+
+    ``ledger`` may be a synchronous RPC client for the sole coordinator. A
+    failed step is unresolved and the wrapper is latched; it cannot be retried.
+    """
+    if not callable(getattr(optimizer, "step", None)) or any(
+        not callable(getattr(ledger, name, None)) for name in ("optimizer_begin", "optimizer_end", "fail")
+    ) or type(rank) is not int or rank < 0 or getattr(optimizer, "_agentinputlab_grpo_observed", False):
+        raise ValidationError("SmartSearch GRPO optimizer observer needs one native optimizer/rank and shared ledger")
+    original, failed = optimizer.step, False
+
+    def step(*args, **kwargs):
+        nonlocal failed
+        if failed:
+            raise ExternalDependencyError("SmartSearch GRPO optimizer failed; no automatic step retry")
+        try:
+            ledger.optimizer_begin(rank)
+            result = original(*args, **kwargs)
+            ledger.optimizer_end(rank)
+            return result
+        except BaseException as exc:
+            failed = True
+            ledger.fail(f"optimizer-rank-{rank}", error_type=type(exc).__name__)
+            raise
+
+    optimizer.step = step
+    optimizer._agentinputlab_grpo_observed = True
+    return optimizer
+
+
+def meter_smartsearch_grpo_helper(generate, *, ledger, worker, fingerprint, count_input_tokens):
+    """Place native prefix AND reward helper clients behind one shared quota.
+
+    Host this callback at the sole teacher coordinator. Keep the provider's
+    original text, sampling and stop behavior. This does not construct a model
+    or replace either native parser. A distributed executor supplies an RPC
+    client to this callback, not a new independent quota per rank.
+    """
+    if not callable(generate) or not callable(count_input_tokens) or any(
+        not callable(getattr(ledger, name, None)) for name in ("reserve", "settle", "fail")
+    ) or not isinstance(worker, str) or not worker.strip():
+        raise ValidationError("SmartSearch GRPO helper requires a measured provider and sole shared ledger")
+    SmartSearchTeacherCompletion("", 0, 0, fingerprint)
+    lock, failed = threading.Lock(), False
+
+    def complete(prompt, *, temperature, max_new_tokens, stop, seed):
+        nonlocal failed
+        with lock:
+            if failed:
+                raise ExternalDependencyError("SmartSearch GRPO helper failed; no automatic provider retry")
+            if not isinstance(prompt, str) or temperature != .6 or type(max_new_tokens) is not int or max_new_tokens != 2048 or (
+                type(seed) is not int or seed < 0 or not isinstance(stop, (tuple, list))
+                or tuple(stop) not in ((), ("</answer>",))
+            ):
+                raise ValidationError("SmartSearch GRPO helper request differs from its native contract")
+            tokens = count_input_tokens(prompt)
+            ticket = ledger.reserve(kind="helper", worker=worker, input_token_counts=[tokens],
+                                    samples=1, max_new_tokens=max_new_tokens)
+            try:
+                result = generate(prompt, temperature=temperature, max_new_tokens=max_new_tokens, stop=stop, seed=seed)
+                if not isinstance(result, SmartSearchTeacherCompletion):
+                    raise ExternalDependencyError("SmartSearch GRPO helper omitted typed usage")
+                ledger.settle(ticket, output_token_counts=[result.output_tokens])
+                if result.input_tokens != tokens or result.model_fingerprint_sha256 != fingerprint:
+                    raise ExternalDependencyError("SmartSearch GRPO helper identity/input accounting differs")
+                return result
+            except BaseException as exc:
+                failed = True
+                ledger.fail(ticket, error_type=type(exc).__name__)
+                raise
+
+    return complete
+
+
+class MeteredNativeGRPOEngine:
+    """Transparent observer of the pinned rollout's vLLM.generate seam.
+
+    Inputs and returned RequestOutput objects are passed through unchanged.
+    All engine outputs are settled before native prefix/group selection, so
+    later truncation cannot hide their cost. No token/stop/n/temperature value
+    is rewritten. Use only with the pinned native rollout and its TP=1 config.
+    """
+
+    def __init__(self, engine, *, ledger, worker):
+        if not callable(getattr(engine, "generate", None)) or any(
+            not callable(getattr(ledger, name, None)) for name in ("reserve", "settle", "fail")
+        ) or not isinstance(worker, str) or not worker.strip():
+            raise ValidationError("SmartSearch GRPO engine observer requires native generation and a shared ledger")
+        self._engine, self._ledger, self._worker = engine, ledger, worker
+        self._failed = False
+
+    def __getattr__(self, name):
+        return getattr(self._engine, name)
+
+    def generate(self, *args, **kwargs):
+        if self._failed:
+            raise ExternalDependencyError("SmartSearch GRPO engine failed; no automatic generation retry")
+        prompts, params = kwargs.get("prompts"), kwargs.get("sampling_params")
+        if args or not isinstance(prompts, list) or not prompts or any(
+            not isinstance(p, dict) or set(p) != {"prompt_token_ids"}
+            or not isinstance(p["prompt_token_ids"], list) or not p["prompt_token_ids"]
+            for p in prompts
+        ):
+            raise ValidationError("SmartSearch GRPO engine seam differs from the pinned token-only native call")
+        samples, cap = getattr(params, "n", None), getattr(params, "max_tokens", None)
+        ticket = self._ledger.reserve(kind="policy", worker=self._worker,
+            input_token_counts=[len(p["prompt_token_ids"]) for p in prompts],
+            samples=samples, max_new_tokens=cap)
+        try:
+            result = self._engine.generate(*args, **kwargs)
+            if not isinstance(result, list) or len(result) != len(prompts):
+                raise ExternalDependencyError("SmartSearch native engine returned a different prompt count")
+            counts = []
+            for row in result:
+                outputs = getattr(row, "outputs", None)
+                if not isinstance(outputs, (tuple, list)) or len(outputs) != samples:
+                    raise ExternalDependencyError("SmartSearch native engine returned a different sample count")
+                for output in outputs:
+                    ids = getattr(output, "token_ids", None)
+                    if not isinstance(ids, (tuple, list)):
+                        raise ExternalDependencyError("SmartSearch native engine omitted measured output token IDs")
+                    counts.append(len(ids))
+            self._ledger.settle(ticket, output_token_counts=counts)
+            return result
+        except BaseException as exc:
+            self._failed = True
+            self._ledger.fail(ticket, error_type=type(exc).__name__)
+            raise

agentinputlab/recipes/smartsearch_grpo_prefix.py

@@ -0,0 +1,144 @@
+"""Pinned native GRPO query refinement with measured, non-retrying transport."""
+from __future__ import annotations
+import __future__
+
+import ast
+import re
+import threading
+from collections.abc import Mapping
+from pathlib import Path
+
+from agentinputlab.contracts import BudgetExceededError, ExternalDependencyError, ValidationError
+from agentinputlab.contracts.serialization import content_hash
+
+from .qrag import _sha
+from .smartsearch_grpo_data import verify_smartsearch_grpo_source
+from .smartsearch_grpo_tools import _new_output, _save
+from .smartsearch_teacher import SmartSearchTeacherCompletion, SmartSearchTeacherControls
+
+_PATH = "src/verl/workers/rollout/vllm_rollout/generate_new_prompts.py"
+
+
+class SmartSearchNativeGRPOPrefixBuilder:
+    """Execute the native tensor/parser/refinement functions without raw HTTP.
+
+    This TRAIN-only instance binds explicit references. Both native helper
+    operations use their original 2048-token, temperature0.6, NO-stop request;
+    unlike offline curriculum refinement, no closing tag is appended here.
+    Quotas are per instance. A distributed executor must enforce an aggregate
+    helper quota rather than silently allocating this allowance to every rank.
+    """
+
+    def __init__(self, generate, *, checkout, references: Mapping, model_fingerprint_sha256,
+                 controls: SmartSearchTeacherControls, output_root):
+        if not callable(generate) or not isinstance(controls, SmartSearchTeacherControls) or (
+            not isinstance(references, Mapping) or not references
+        ):
+            raise ValidationError("SmartSearch GRPO prefixes require explicit TRAIN references and provider")
+        clean = {}
+        for question, answers in references.items():
+            if not isinstance(question, str) or not question.strip() or not isinstance(answers, (list, tuple)) or (
+                not answers or any(not isinstance(a, str) or not a.strip() for a in answers)
+            ) or question.strip() in clean:
+                raise ValidationError("SmartSearch GRPO prefix references are invalid or collide")
+            clean[question.strip()] = tuple(answers)
+        SmartSearchTeacherCompletion("", 0, 0, model_fingerprint_sha256)
+        source, binding = verify_smartsearch_grpo_source(checkout)
+        path = source / _PATH
+        output = Path(output_root).resolve(strict=False)
+        if output.is_relative_to(source) or source.is_relative_to(output):
+            raise ValidationError("SmartSearch GRPO prefix output must not overlap native source")
+        try:
+            import torch
+            import torch.nn.functional as functional
+        except ImportError as exc:
+            raise ExternalDependencyError("Install the native SmartSearch GRPO torch runtime") from exc
+        self.root = _new_output(output_root)
+        self.generate, self.controls, self.references = generate, controls, clean
+        self.fingerprint, self.failed = model_fingerprint_sha256, False
+        self._lock = threading.Lock()
+        self.usage = dict(prefix_attempts=0, prefix_responses=0, prefixes_produced=0,
+                          model_attempts=0, model_responses=0, input_tokens=0, output_tokens=0,
+                          unknown_output_token_reservations=0)
+        tree = ast.parse(path.read_bytes())
+        functions = [n for n in tree.body if isinstance(n, ast.FunctionDef) and n.name != "generate_response"]
+        namespace = dict(re=re, torch=torch, F=functional, generate_response=self._complete)
+        exec(compile(ast.Module(body=functions, type_ignores=[]), str(path), "exec",
+                     flags=__future__.annotations.compiler_flag), namespace)
+        self._native = namespace["generate_new_prompts"]
+        self.source_path, self.source_sha = path, _sha(path)
+        _save(self.root, "claim.json", dict(native_source=binding, native_prefix_sha256=self.source_sha,
+            references_sha256=content_hash(clean), reference_split="train", controls=controls,
+            model_fingerprint_sha256=self.fingerprint, native_tensor_parser_selection_unchanged=True,
+            native_helper_stop=[], no_closing_tag_append=True, no_http_empty_fallback=True))
+
+    def _complete(self, prompt):
+        usage, cap = self.usage, 2048
+        if not isinstance(prompt, str) or len(prompt.encode()) > self.controls.max_input_bytes:
+            raise ValidationError("SmartSearch GRPO prefix helper input exceeds its byte bound")
+        if usage['model_attempts'] >= self.controls.max_model_calls or (
+            usage['output_tokens'] + usage['unknown_output_token_reservations'] + cap
+            > self.controls.max_total_output_tokens
+        ):
+            raise BudgetExceededError("SmartSearch GRPO prefix helper cannot admit its native2048-token cap")
+        usage['model_attempts'] += 1
+        usage['unknown_output_token_reservations'] += cap
+        number = usage['model_attempts']
+        _save(self.root, f"helper-{number:06d}-request.json", dict(prompt=prompt, temperature=.6,
+            max_new_tokens=cap, stop=[], seed=self.controls.seed, reference_split="train"))
+        result = self.generate(prompt, temperature=.6, max_new_tokens=cap, stop=(), seed=self.controls.seed)
+        if not isinstance(result, SmartSearchTeacherCompletion):
+            raise ExternalDependencyError("SmartSearch GRPO prefix helper requires typed measured usage")
+        usage['model_responses'] += 1
+        usage['input_tokens'] += result.input_tokens
+        usage['output_tokens'] += result.output_tokens
+        usage['unknown_output_token_reservations'] -= cap
+        _save(self.root, f"helper-{number:06d}-response.json", result)
+        if result.model_fingerprint_sha256 != self.fingerprint:
+            raise ExternalDependencyError("SmartSearch GRPO prefix helper model identity changed")
+        if result.output_tokens > cap:
+            raise BudgetExceededError("SmartSearch GRPO prefix helper exceeded the requested token cap")
+        return result.text
+
+    def __call__(self, prompts, tokenizer, ground_truth):
+        with self._lock:
+            if self.failed:
+                raise ExternalDependencyError("SmartSearch GRPO prefix failed; no automatic helper retry")
+            try:
+                if self.usage['prefix_attempts'] >= self.controls.max_rollout_rows:
+                    raise BudgetExceededError("SmartSearch GRPO prefix row quota exhausted")
+                if _sha(self.source_path) != self.source_sha:
+                    raise ExternalDependencyError("SmartSearch native prefix source changed")
+                trajectory = tokenizer.decode(prompts.batch['input_ids'][0])
+                if len(trajectory.encode()) > self.controls.max_input_bytes:
+                    raise ValidationError("SmartSearch GRPO prefix trajectory exceeds its byte bound")
+                match = re.search(r'<\|im_start\|>user\n(.*?)<\|im_end\|>', trajectory, re.DOTALL)
+                answers = (ground_truth,) if isinstance(ground_truth, str) else tuple(ground_truth)
+                if not match or self.references.get(match.group(1).strip()) != answers:
+                    raise ValidationError("SmartSearch GRPO prefix input is not its bound TRAIN reference")
+                self.usage['prefix_attempts'] += 1
+                number = self.usage['prefix_attempts']
+                _save(self.root, f"prefix-{number:06d}-request.json", dict(
+                    trajectory_sha256=content_hash(trajectory), ground_truth_sha256=content_hash(answers)))
+                result = self._native(prompts, tokenizer, ground_truth)
+                if not isinstance(result, tuple) or len(result) != 2 or len(result[0]) != len(result[1]):
+                    raise ExternalDependencyError("SmartSearch native prefix tensor collection differs")
+                self.usage['prefix_responses'] += 1
+                self.usage['prefixes_produced'] += len(result[0])
+                _save(self.root, f"prefix-{number:06d}-response.json", dict(prefix_count=len(result[0]),
+                    shapes=[[list(t.shape) for t in values] for values in result], usage=self.usage))
+                return result
+            except BaseException as exc:
+                self.failed = True
+                _save(self.root, "failed.json", dict(error_type=type(exc).__name__, error=str(exc),
+                      usage=self.usage, no_automatic_retry=True, accept_checkpoint=False))
+                raise
+
+    def assert_healthy(self):
+        with self._lock:
+            if self.failed or self.usage['prefix_attempts'] != self.usage['prefix_responses'] or (
+                self.usage['model_attempts'] != self.usage['model_responses']
+                or self.usage['unknown_output_token_reservations']
+            ):
+                raise ExternalDependencyError("SmartSearch GRPO prefix has an unresolved/failed request")
+            return dict(self.usage)

agentinputlab/recipes/smartsearch_grpo_rollout.py

@@ -0,0 +1,119 @@
+"""Batch-assembly repairs for the exact native SmartSearch GRPO rollout.
+
+No native file is modified. The original group/selection/truncation logic and
+single-trajectory generator remain intact; a fresh worker supplies its explicit
+metered prefix builder instead of the upstream hardcoded network helper.
+"""
+from __future__ import annotations
+
+import ast
+import copy
+import hashlib
+import inspect
+from collections.abc import Mapping
+from pathlib import Path
+
+from agentinputlab.contracts import ExternalDependencyError, ValidationError
+from agentinputlab.contracts.serialization import to_jsonable
+
+from .qrag import _sha
+from .smartsearch_grpo_data import verify_smartsearch_grpo_source
+
+_PATH = "src/verl/workers/rollout/vllm_rollout/vllm_rollout_spmd.py"
+_CLASS = "vLLMRolloutWithTool"
+
+
+def _name(node, expected):
+    return isinstance(node, ast.Name) and node.id == expected
+
+
+def _batch_contract_ast(method):
+    """Move one pinned assembly block; refuse a structurally different source."""
+    method = copy.deepcopy(method)
+    loops = [n for n in method.body if isinstance(n, ast.While) and _name(n.test, "active_indices")]
+    if len(loops) != 1:
+        raise ExternalDependencyError("SmartSearch native active-rollout loop differs")
+    loop = loops[0]
+    outer = [n for n in loop.body if isinstance(n, ast.For) and isinstance(n.target, ast.Tuple)
+             and [_n.id for _n in n.target.elts if isinstance(_n, ast.Name)] == ["i", "idx"]]
+    if len(outer) != 1:
+        raise ExternalDependencyError("SmartSearch native per-input prefix loop differs")
+    outer = outer[0]
+    starts = [i for i, n in enumerate(outer.body) if isinstance(n, ast.Assign)
+              and len(n.targets) == 1 and _name(n.targets[0], "prefix_input_ids")
+              and isinstance(n.value, ast.Call) and isinstance(n.value.func, ast.Attribute)
+              and _name(n.value.func.value, "torch") and n.value.func.attr == "stack"]
+    if len(starts) != 1:
+        raise ExternalDependencyError("SmartSearch native prefix stack block differs")
+    position = starts[0]
+    tail = outer.body[position:]
+    if len(tail) != 11 or not isinstance(tail[-1], ast.For):
+        raise ExternalDependencyError("SmartSearch native assembly statement inventory differs")
+    for statement, name in zip(tail[:5], ("prefix_input_ids", "prefix_loss_mask", "input_ids",
+                                            "attention_mask", "position_ids"), strict=True):
+        expected = ast.parse(f"{name} = torch.stack({name}_list, dim=0)").body[0]
+        if ast.dump(statement) != ast.dump(expected):
+            raise ExternalDependencyError("SmartSearch native prefix tensor assembly differs")
+    if ast.dump(outer.iter) != ast.dump(ast.parse("enumerate(active_indices)", mode="eval").body):
+        raise ExternalDependencyError("SmartSearch native prefix input selection differs")
+    if ast.dump(tail[-1].iter) != ast.dump(ast.parse("enumerate(index)", mode="eval").body):
+        raise ExternalDependencyError("SmartSearch native generated-prefix selection differs")
+    selected = [n for n in tail if isinstance(n, ast.Assign) and len(n.targets) == 1
+                and isinstance(n.targets[0], ast.Attribute) and _name(n.targets[0].value, "new_prompts")
+                and n.targets[0].attr == "non_tensor_batch"]
+    if len(selected) != 1 or ast.dump(selected[0].value) != ast.dump(ast.parse("prompts.non_tensor_batch", mode="eval").body):
+        raise ExternalDependencyError("SmartSearch native prefix metadata assignment differs")
+    selected[0].value = ast.parse("prompts.select_idxs(index).non_tensor_batch", mode="eval").body
+    outer.body = outer.body[:position]
+    insertion = loop.body.index(outer) + 1
+    loop.body.insert(insertion, ast.If(test=ast.Name(id="prefix_input_ids_list", ctx=ast.Load()), body=tail, orelse=[]))
+    return ast.fix_missing_locations(method)
+
+
+def _compile_batch_contract(method, namespace, prefix_builder, filename):
+    repaired = _batch_contract_ast(method)
+    local = dict(namespace)
+    local["generate_new_prompts"] = prefix_builder
+    exec(compile(ast.Module(body=[repaired], type_ignores=[]), filename, "exec"), local)
+    return local[method.name], repaired
+
+
+def install_native_rollout_batch_contracts(*, checkout, rollout_class, prefix_builder, prefix_builder_binding):
+    """Install source-bound assembly fixes in one fresh native worker.
+
+    Caller must provide a metered, fail-closed native prefix-builder seam and
+    its provenance. There is intentionally no default hardcoded HTTP helper.
+    The executor must meter ALL actual generation calls, including outputs
+    subsequently discarded by unchanged native group truncation.
+    """
+    if not isinstance(rollout_class, type) or rollout_class.__name__ != _CLASS or not callable(prefix_builder):
+        raise ValidationError("SmartSearch rollout repair requires the native class and explicit prefix provider")
+    if not isinstance(prefix_builder_binding, Mapping) or not prefix_builder_binding:
+        raise ValidationError("SmartSearch rollout prefix provider needs explicit provenance")
+    binding = to_jsonable(prefix_builder_binding)
+    if getattr(rollout_class, "_agentinputlab_batch_contract_installed", False):
+        raise ExternalDependencyError("SmartSearch rollout contracts are already installed; use a fresh native worker")
+    source, source_binding = verify_smartsearch_grpo_source(checkout)
+    path = source / _PATH
+    original_method = inspect.unwrap(rollout_class.generate_sequences)
+    if not hasattr(original_method, "__code__") or Path(original_method.__code__.co_filename).resolve() != path:
+        raise ExternalDependencyError("SmartSearch rollout method was not loaded from the pinned native source")
+    tree = ast.parse(path.read_text())
+    owners = [n for n in tree.body if isinstance(n, ast.ClassDef) and n.name == _CLASS]
+    methods = [n for owner in owners for n in owner.body if isinstance(n, ast.FunctionDef) and n.name == "generate_sequences"]
+    if len(owners) != 1 or len(methods) != 1:
+        raise ExternalDependencyError("SmartSearch native rollout class surface differs")
+    method = methods[0]
+    implementation, repaired = _compile_batch_contract(method, original_method.__globals__, prefix_builder, str(path))
+    record = dict(native_source=source_binding, native_rollout_sha256=_sha(path),
+        original_method_ast_sha256=hashlib.sha256(ast.dump(method).encode()).hexdigest(),
+        adapter_method_ast_sha256=hashlib.sha256(ast.dump(repaired).encode()).hexdigest(),
+        adapter_toolkit_sha256=_sha(Path(__file__)), prefix_builder_binding=binding,
+        native_source_files_modified=False, native_single_trajectory_generator_replaced=False,
+        native_group_size_and_selection_unchanged=True, prefix_batch_generated_once=True,
+        empty_prefix_batch_skipped=True, prefix_metadata_selected_by_original_indices=True,
+        all_generated_and_discarded_outputs_require_metering=True,
+        full_grpo_training_certified=False)
+    rollout_class.generate_sequences = implementation
+    rollout_class._agentinputlab_batch_contract_installed = True
+    return record