"""Hard-knowledge router boundary — bind trauma/MITM packet to live routers. Hot path is **tensor-only**: ``HardKnowledgeSurfacePacket``, ``Tensor`` adjustments. JSON/manifest adapters stay in ``coverage_pressure_boundary`` / assurance tools. Schema: nnf.resynthesis.hard_knowledge_router_boundary.v1 """ from __future__ import annotations import os import re from dataclasses import dataclass from typing import Any import torch import torch.nn as nn from torch import Tensor from resynthesis.hard_knowledge_surface import ( HardKnowledgeSurfacePacket, hard_knowledge_surface_packet_t, ) from resynthesis.trauma_system import TensorTraumaState HARD_KNOWLEDGE_ROUTER_BOUNDARY_SCHEMA = ( "nnf.resynthesis.hard_knowledge_router_boundary.v1" ) DEFAULT_MITM_LOGIT_SCALE = 0.25 DEFAULT_HARD_KNOWLEDGE_PAGE_BOOST_FRACTION = 0.35 @dataclass(frozen=True, slots=True) class StagedVerificationContext: """Typed observation boundary shared by training and staged verification.""" available_layers: int trauma_hard_arms_t: Tensor | None def _env_flag(name: str, default: bool = False) -> bool: raw = os.environ.get(name) if raw is None: return bool(default) return raw.strip().lower() in {"1", "true", "yes", "on"} def hard_knowledge_router_enabled_boundary() -> bool: """Default ON when trauma system is enabled.""" if _env_flag("NNF_HARD_KNOWLEDGE_ROUTER", default=False): return True if _env_flag("NNF_TRAUMA_SYSTEM", default=False): return True return _env_flag("NNF_RESYNTHESIS_TRAUMA_SYSTEM", default=False) def branch_id_from_loop_id(loop_id: str) -> int: match = re.search(r"_b(\d+)", str(loop_id)) if match: return int(match.group(1)) match = re.search(r"gpu(\d+)", str(loop_id), flags=re.IGNORECASE) if match: return int(match.group(1)) return 0 def load_mitm_coverage_payload_boundary( *, payload_path: str | os.PathLike[str] | None = None, ) -> object | None: """Load observer coverage payload for MITM hard-target ranking (boundary only).""" try: from pathlib import Path from resynthesis.coverage_pressure_boundary import ( DEFAULT_LIVE_PAYLOAD_PATH, load_coverage_pressure_payload, ) path = ( Path(payload_path).expanduser() if payload_path is not None else DEFAULT_LIVE_PAYLOAD_PATH ) payload = load_coverage_pressure_payload(path) if payload.undertrained_expert_ids or payload.undertrained_page_ids: return payload except Exception: pass return None def page_catalog_ids_t_from_model(model: nn.Module) -> Tensor | None: """Best-effort read of the live NoNE page catalog tensor.""" for module in model.modules(): catalog = getattr(module, "page_catalog_ids_t", None) if isinstance(catalog, Tensor) and catalog.numel() > 0: return catalog.detach().reshape(-1).long() return None def page_ids_to_catalog_arm_indices_t( model: nn.Module, page_ids: list[int] | Tensor, ) -> Tensor: """Map durable page IDs to catalog arm indices for the trauma bank.""" if isinstance(page_ids, list): if not page_ids: return torch.empty(0, dtype=torch.long) ids_t = torch.tensor(page_ids, dtype=torch.long) else: ids_t = page_ids.detach().reshape(-1).long() if ids_t.numel() == 0: return ids_t catalog_t = page_catalog_ids_t_from_model(model) if catalog_t is None or catalog_t.numel() == 0: return ids_t[ids_t.ge(0)] identity_t = torch.arange( catalog_t.numel(), device=catalog_t.device, dtype=torch.long, ) if bool(torch.equal(catalog_t, identity_t)): return ids_t[ids_t.ge(0) & ids_t.lt(catalog_t.numel())] expanded_ids_t = ids_t.reshape(-1, 1) matches_t = expanded_ids_t.eq(catalog_t.reshape(1, -1)) has_match_t = matches_t.any(dim=1) if not bool(has_match_t.any().item()): return torch.empty(0, dtype=torch.long) positions_t = matches_t.to(dtype=torch.long).argmax(dim=1) return positions_t[has_match_t] def lane_trauma_state_for_loop( loop_id: str, lane_bank: dict[str, Any] | None = None, ) -> TensorTraumaState | None: """Read the live per-lane trauma bank (learn_loop module scope).""" if lane_bank is None: try: from resynthesis import learn_loop as learn_loop_mod lane_bank = getattr(learn_loop_mod, "_LANE_TRAUMA_STATE", None) except Exception: return None if not isinstance(lane_bank, dict): return None state = lane_bank.get(str(loop_id)) return state if isinstance(state, TensorTraumaState) else None def _trauma_state_view_for_router( trauma_state: TensorTraumaState, num_arms: int, ) -> TensorTraumaState: """Non-mutating width-aligned view for a router (never truncates lane bank).""" if trauma_state.num_arms == num_arms: return trauma_state view = TensorTraumaState( num_arms=num_arms, fail_decay=trauma_state.fail_decay, success_decay=trauma_state.success_decay, ).to(device=trauma_state.fail_ema.device) prefix = min(trauma_state.num_arms, num_arms) with torch.no_grad(): source_buffers = dict(trauma_state.named_buffers(recurse=False)) for name, destination_t in view.named_buffers(recurse=False): source_t = source_buffers.get(name) if source_t is None: continue if destination_t.ndim == 0: destination_t.copy_( source_t.to( device=destination_t.device, dtype=destination_t.dtype, ) ) else: destination_t[:prefix].copy_( source_t[:prefix].to( device=destination_t.device, dtype=destination_t.dtype, ) ) view._knowledge_token_to_id = dict(trauma_state._knowledge_token_to_id) return view def _align_trauma_width( trauma_state: TensorTraumaState, num_arms: int, ) -> TensorTraumaState: """Alias for router-local view (canonical lane bank is never mutated).""" return _trauma_state_view_for_router(trauma_state, num_arms) def bind_hard_knowledge_to_quantile_router( router: nn.Module, *, trauma_state: TensorTraumaState | None, branch_id: int = 0, loop_id: str = "", coverage_payload: object | None = None, ) -> HardKnowledgeSurfacePacket | None: """Attach live hard-knowledge + MITM context to a ``QuantileBalancingRouter``.""" if not hard_knowledge_router_enabled_boundary(): return None num_arms = int(getattr(router, "num_experts", 0)) if num_arms < 1 or trauma_state is None: return None try: aligned = _align_trauma_width(trauma_state, num_arms) packet = hard_knowledge_surface_packet_t(aligned, top_k=min(64, num_arms)) # ``TensorTraumaState`` is an nn.Module, so assigning it here registers # the exact session/catalog bank beneath the router. Its buffers now # participate in state_dict, device migration, and cold reload. setattr(router, "trauma_state", aligned) setattr(router, "_hard_knowledge_packet", packet) # Branch and loop identifiers are persistence-boundary metadata. The # live router owns only the bound tensor state; strings/integers must # never be consulted to select a route or refresh its pressure. del branch_id, loop_id # Coverage payloads are observer evidence only. They may be recorded # at the I/O boundary, but they never become route-logit authority. del coverage_payload return packet except Exception: return None def refresh_hard_knowledge_packet(router: nn.Module) -> HardKnowledgeSurfacePacket | None: """Rebuild packet after trauma bank mutation (start of each route wave).""" trauma_state = getattr(router, "trauma_state", None) if not isinstance(trauma_state, TensorTraumaState): return getattr(router, "_hard_knowledge_packet", None) try: num_arms = int(getattr(router, "num_experts", trauma_state.num_arms)) aligned = _align_trauma_width(trauma_state, num_arms) packet = hard_knowledge_surface_packet_t(aligned, top_k=min(64, num_arms)) setattr(router, "trauma_state", aligned) setattr(router, "_hard_knowledge_packet", packet) setattr(router, "_mitm_hard_targets_t", None) return packet except Exception: return getattr(router, "_hard_knowledge_packet", None) def model_owned_trauma_state_for_model( model: nn.Module, ) -> TensorTraumaState | None: """Return the first registered TRAUMA bank for boundary receipts only.""" for module in model.modules(): state = getattr(module, "trauma_state", None) if isinstance(state, TensorTraumaState): return state return None def record_verified_hard_knowledge_outcome_to_model( model: nn.Module, *, knowledge_success_t: Tensor, behavior_success_t: Tensor, evidence_confidence_t: Tensor, ) -> Tensor: """Commit exact page-frontier evidence after verification. Parent paged routers retain the target-independent request that produced the completed transaction. This function consumes its exact catalog positions and soft frontier weights, updates every selected page proportionally, and leaves the completed route unchanged. The next route reads the newly committed model-owned buffers. """ from resynthesis.trauma_system import ( update_trauma_from_verified_outcome, ) updated_count_t = knowledge_success_t.detach().new_zeros( (), dtype=torch.long, ) seen_router_ids: set[int] = set() for module in model.modules(): router = getattr(module, "quantile_router", None) request = getattr(module, "last_request", None) if not isinstance(router, nn.Module) or request is None: continue router_identity = id(router) if router_identity in seen_router_ids: continue state = getattr(router, "trauma_state", None) catalog_positions_t = getattr( request, "unique_page_catalog_positions_t", None, ) frontier_weight_t = getattr(request, "frontier_weight_t", None) if ( not isinstance(state, TensorTraumaState) or not isinstance(catalog_positions_t, Tensor) or not isinstance(frontier_weight_t, Tensor) or catalog_positions_t.ndim != 1 or frontier_weight_t.ndim != 2 or frontier_weight_t.shape[1] != catalog_positions_t.numel() ): continue seen_router_ids.add(router_identity) # Mean over rows preserves each request row's unit mass while avoiding # batch-size-dependent pressure. Zero-mass/invalid columns are # rejected by the verified update boundary. page_weight_t = frontier_weight_t.detach().mean(dim=0) knowledge_t = knowledge_success_t.detach().reshape(()).to( device=state.fail_ema.device, dtype=state.fail_ema.dtype, ) behavior_t = behavior_success_t.detach().reshape(()).to( device=state.fail_ema.device, dtype=state.fail_ema.dtype, ) axis_count_t = state.fail_ema.new_tensor(2.0) pro_t = (knowledge_t + behavior_t) / axis_count_t anti_t = ( (1.0 - knowledge_t) + (1.0 - behavior_t) ) / axis_count_t accepted_t = update_trauma_from_verified_outcome( state, arm_index_t=catalog_positions_t.detach().reshape(-1), frontier_weight_t=page_weight_t, pro_t=pro_t, anti_t=anti_t, evidence_confidence_t=evidence_confidence_t, knowledge_axis_t=state.fail_ema.new_ones(()), behavior_axis_t=state.fail_ema.new_ones(()), ) updated_count_t = updated_count_t.to(device=accepted_t.device) updated_count_t = updated_count_t + accepted_t.numel() setattr( router, "_hard_knowledge_packet", hard_knowledge_surface_packet_t( state, top_k=state.num_arms, device=state.fail_ema.device, ), ) return updated_count_t def hard_knowledge_frontier_width_t( router: nn.Module, *, available: int, learned_k: int, uncertain: bool, width_t: Tensor, ) -> Tensor: """Widen frontier using hard-knowledge gap count (tensor-native).""" if not uncertain or not hard_knowledge_router_enabled_boundary(): return width_t packet = getattr(router, "_hard_knowledge_packet", None) if not isinstance(packet, HardKnowledgeSurfacePacket): return width_t gap_count = packet.negative_count_t.to(dtype=width_t.dtype) if not bool(gap_count.gt(0).item()): return width_t from resynthesis.exploration_floor import exploration_route_width floor_width = exploration_route_width( available=available, learned_k=learned_k, uncertain=True, ) boost = torch.clamp( gap_count, min=width_t.new_tensor(0.0), max=width_t.new_tensor(float(min(available, max(floor_width, 16)))), ) return torch.maximum(width_t, boost.to(dtype=width_t.dtype)) def mitm_scaffold_logits_delta_t( router: nn.Module, *, num_candidates: int, device: torch.device, dtype: torch.dtype, ) -> Tensor: """Tensor-native temporary MILT bridge over the complete candidate axis. TRAUMA supplies real outcome difficulty and empirical learnability. Ownership progress attenuates the bridge to exactly zero after mastery. The router's learned activation fraction controls the bell range, so no host manifest, branch id, string gap key, fixed top-k, or answer target owns the route. """ zero = torch.zeros(num_candidates, device=device, dtype=dtype) packet = getattr(router, "_hard_knowledge_packet", None) if not isinstance(packet, HardKnowledgeSurfacePacket): return zero if packet.num_arms != num_candidates: return zero difficulty_t = packet.negative_gap_signal_t.to( device=device, dtype=torch.float32, ) if difficulty_t.numel() != num_candidates: return zero difficulty_unit_t = difficulty_t / difficulty_t.amax().clamp_min(1.0e-8) learnability_t = packet.empirical_learnability_t.to( device=device, dtype=torch.float32, ).clamp(0.0, 1.0) progress_t = packet.learning_progress_t.to( device=device, dtype=torch.float32, ).clamp(0.0, 1.0) ownership_gap_t = 1.0 - progress_t # A never-succeeded hard element still receives half-strength translation; # observed learnability raises that pressure, while proven ownership backs # it out exactly through ``ownership_gap_t``. bridge_strength_t = ( difficulty_unit_t * ownership_gap_t * (0.5 + 0.5 * learnability_t) ) activation_fraction = getattr(router, "activation_fraction", None) if not callable(activation_fraction): return zero learned_range_fraction_t = activation_fraction().reshape(()).to( device=device, dtype=torch.float32, ).clamp(0.0, 1.0) uncertainty_t = 1.0 - learnability_t range_t = ( 1.0 + uncertainty_t * learned_range_fraction_t * max(1, num_candidates - 1) ).clamp_min(1.0e-3) positions_t = torch.arange( num_candidates, device=device, dtype=torch.float32, ) distance_t = positions_t.unsqueeze(0) - positions_t.unsqueeze(1) bell_t = torch.exp( -0.5 * (distance_t / range_t.unsqueeze(1)).square() ) # MILT is the translator between a failed hard endpoint and learnable # neighbouring paths; it must not hand the router the failed endpoint back # as its own answer. Remove the diagonal while retaining the full learned # range around every hard element. # ``distance_t`` already carries the complete candidate-pair geometry. # Zero its exact diagonal without allocating a second dense floating-point # identity matrix. Keep this out of place because ``bell_t`` participates # in the learned-range gradient. bell_t = bell_t.masked_fill(distance_t.eq(0), 0.0) exploration_t = (bridge_strength_t.unsqueeze(1) * bell_t).sum(dim=0) exploration_t = ( exploration_t / exploration_t.amax().clamp_min(1.0e-8) * float(DEFAULT_MITM_LOGIT_SCALE) ) active_t = bridge_strength_t.amax().gt(0.0).to( device=device, dtype=torch.float32, ) return (exploration_t * active_t).to(dtype=dtype) def hard_knowledge_page_logits_boost_t( *, packet: HardKnowledgeSurfacePacket, post_bias_logits_t: Tensor, score_span_t: Tensor, ) -> Tensor: """Apply signed pro-attraction/anti-suppression to page logits. The legacy function name is retained for callers, but negative definitive TRAUMA is never a positive boost. The packet surface already encodes ``pro - anti`` in route-logit coordinates. """ if post_bias_logits_t.numel() == 0: return post_bias_logits_t.new_zeros(post_bias_logits_t.shape) num_pages = post_bias_logits_t.shape[-1] if packet.definitive_surface_t.numel() != num_pages: raise ValueError("hard-knowledge packet differs from page catalog") active_signal_t = packet.definitive_surface_t.to( device=post_bias_logits_t.device, dtype=post_bias_logits_t.dtype, ) peak = active_signal_t.abs().amax().clamp_min( active_signal_t.new_ones(()) * 1.0e-8 ) unit = active_signal_t / peak return ( unit.unsqueeze(0) * score_span_t * ( torch.ones_like(post_bias_logits_t).narrow(-1, 0, 1) * DEFAULT_HARD_KNOWLEDGE_PAGE_BOOST_FRACTION ) ) def bind_hard_knowledge_to_model( model: nn.Module, *, loop_id: str, branch_id: int | None = None, trauma_state: TensorTraumaState | None = None, coverage_payload: object | None = None, ) -> int: """Bind hard-knowledge surface to every quantile router on ``model``.""" if not hard_knowledge_router_enabled_boundary(): return 0 if trauma_state is None: trauma_state = lane_trauma_state_for_loop(loop_id) if branch_id is None: branch_id = branch_id_from_loop_id(loop_id) # Live coverage manifests remain observer-only. They are intentionally # not loaded here and cannot alter model-owned route probabilities. bound = 0 seen: set[int] = set() def _bind_router(router: nn.Module) -> None: nonlocal bound rid = id(router) if rid in seen: return seen.add(rid) active_state = trauma_state if active_state is None: resident_state = getattr(router, "trauma_state", None) active_state = ( resident_state if isinstance(resident_state, TensorTraumaState) else None ) if bind_hard_knowledge_to_quantile_router( router, trauma_state=active_state, branch_id=int(branch_id), loop_id=loop_id, coverage_payload=coverage_payload, ) is not None: bound += 1 for module in model.modules(): if module.__class__.__name__ == "QuantileBalancingRouter": _bind_router(module) router = getattr(module, "quantile_router", None) if isinstance(router, nn.Module): _bind_router(router) return bound def staged_verification_context_for_model( model: nn.Module | None, *, loop_id: str = "", ) -> StagedVerificationContext: """Shared MITM/staged-verification context (layers + hard-knowledge arms).""" available_layers = 0 trauma_hard_arms_t: Tensor | None = None if model is not None: try: science_stack = getattr(model, "science_stack", None) num_layers = getattr(science_stack, "num_layers", None) if isinstance(num_layers, int) and not isinstance(num_layers, bool): if num_layers > 0: available_layers = num_layers except Exception: pass if loop_id: try: from resynthesis.mitm_trauma_learning_bridge import ( hard_knowledge_arm_indices_t, ) bank = lane_trauma_state_for_loop(loop_id) if bank is not None: catalog_t = ( page_catalog_ids_t_from_model(model) if model is not None else None ) num_arms = ( int(catalog_t.numel()) if catalog_t is not None and catalog_t.numel() > 0 else bank.num_arms ) view = _trauma_state_view_for_router(bank, num_arms) trauma_hard_arms_t = hard_knowledge_arm_indices_t(view) except Exception: trauma_hard_arms_t = None return StagedVerificationContext( available_layers=available_layers, trauma_hard_arms_t=trauma_hard_arms_t, ) __all__ = [ "HARD_KNOWLEDGE_ROUTER_BOUNDARY_SCHEMA", "StagedVerificationContext", "bind_hard_knowledge_to_model", "bind_hard_knowledge_to_quantile_router", "branch_id_from_loop_id", "hard_knowledge_frontier_width_t", "hard_knowledge_page_logits_boost_t", "hard_knowledge_router_enabled_boundary", "lane_trauma_state_for_loop", "load_mitm_coverage_payload_boundary", "mitm_scaffold_logits_delta_t", "model_owned_trauma_state_for_model", "page_catalog_ids_t_from_model", "page_ids_to_catalog_arm_indices_t", "record_verified_hard_knowledge_outcome_to_model", "refresh_hard_knowledge_packet", "staged_verification_context_for_model", ]