From 81ab7177d3a611db0fea9f21361dcb9636d6f4b6 Mon Sep 17 00:00:00 2001 From: svonava Date: Mon, 19 Oct 2026 17:08:49 -0700 Subject: [PATCH 1/7] feat(server): add opt-in runtime LoRA for native decision models --- .../sie_server/adapters/_peft_lora_runtime.py | 565 ++++++++++++++++++ .../adapters/decision_lora/__init__.py | 0 .../adapters/decision_lora/adapter.py | 152 +++++ .../adapters/gliclass/runtime_lora.py | 111 ++++ .../src/sie_server/adapters/gliner2/_lora.py | 543 +---------------- .../adapters/gliner2/decide_lora.py | 105 ++++ .../sie_server/adapters/laya/runtime_lora.py | 71 +++ .../adapters/pplx_decider/runtime_lora.py | 70 +++ .../tests/adapters/test_decision_lora.py | 444 ++++++++++++++ .../adapters/test_decision_lora_gliclass.py | 184 ++++++ .../adapters/test_decision_lora_perplexity.py | 156 +++++ 11 files changed, 1876 insertions(+), 525 deletions(-) create mode 100644 packages/sie_server/src/sie_server/adapters/_peft_lora_runtime.py create mode 100644 packages/sie_server/src/sie_server/adapters/decision_lora/__init__.py create mode 100644 packages/sie_server/src/sie_server/adapters/decision_lora/adapter.py create mode 100644 packages/sie_server/src/sie_server/adapters/gliclass/runtime_lora.py create mode 100644 packages/sie_server/src/sie_server/adapters/gliner2/decide_lora.py create mode 100644 packages/sie_server/src/sie_server/adapters/laya/runtime_lora.py create mode 100644 packages/sie_server/src/sie_server/adapters/pplx_decider/runtime_lora.py create mode 100644 packages/sie_server/tests/adapters/test_decision_lora.py create mode 100644 packages/sie_server/tests/adapters/test_decision_lora_gliclass.py create mode 100644 packages/sie_server/tests/adapters/test_decision_lora_perplexity.py diff --git a/packages/sie_server/src/sie_server/adapters/_peft_lora_runtime.py b/packages/sie_server/src/sie_server/adapters/_peft_lora_runtime.py new file mode 100644 index 000000000..be90a5383 --- /dev/null +++ b/packages/sie_server/src/sie_server/adapters/_peft_lora_runtime.py @@ -0,0 +1,565 @@ +"""Strict generic-PEFT lifecycle for explicitly selected runtime-LoRA profiles. + +The raw model owns serving. A separate PEFT controller owns only reversible +LoRA injection; it is never used as the extraction entrypoint. Resident wrappers +stay outside the module tree while base is selected, sharing the exact original +projection weights. Admission and selection still have their own runtime cost. +""" + +from __future__ import annotations + +import json +import math +import re +import shutil +import tempfile +import threading +from abc import ABC, abstractmethod +from collections.abc import Mapping +from dataclasses import dataclass, fields +from pathlib import Path +from typing import Any + +import torch +from huggingface_hub import snapshot_download +from peft import LoraConfig, PeftConfig, PeftModel +from peft.functional import delete_adapter +from peft.tuners.lora import LoraModel +from peft.tuners.lora.layer import Linear as LoraLinear +from safetensors.torch import load_file + +from sie_server.adapters.lora import LoraRuntimeConfig, LoraStateUncertainError, LoraUnavailableError +from sie_server.adapters.peft_lora_mixin import PEFTLoRAMixin + +_FILES = ("adapter_config.json", "adapter_model.safetensors") +_PEFT_ATTRIBUTES = ("peft_config", "_hf_peft_config_loaded") +_UNSUPPORTED = ( + "use_dora", + "use_rslora", + "use_qalora", + "use_bdlora", + "lora_bias", + "fan_in_fan_out", + "rank_pattern", + "alpha_pattern", + "layers_to_transform", + "layers_pattern", + "exclude_modules", + "target_parameters", + "alora_invocation_tokens", + "arrow_config", + "ensure_weight_tying", + "loftq_config", + "eva_config", + "corda_config", + "lora_ga_config", +) + + +@dataclass +class _Prepared: + directory: tempfile.TemporaryDirectory[str] + config: LoraConfig + weights: dict[str, torch.Tensor] + + +class DetachedLoraRuntime(ABC): + """Declared adapters; worker mutation exclusion plus direct-call protection.""" + + def __init__(self, paths: dict[str, str], revisions: dict[str, str | None] | None, max_loaded: int) -> None: + if isinstance(max_loaded, bool) or not isinstance(max_loaded, int) or max_loaded < 1: + raise ValueError("max_loaded_loras must be a positive integer") + if not isinstance(paths, dict) or not paths: + raise ValueError("Runtime LoRA requires a nonempty lora_paths alias-to-source map") + aliases = dict(paths) + if any( + not isinstance(k, str) + or not re.fullmatch(r"[A-Za-z0-9_][A-Za-z0-9_.-]*", k) + or not isinstance(v, str) + or not v + for k, v in aliases.items() + ): + raise ValueError("LoRA declarations require simple nonempty aliases and nonempty sources") + self._aliases = aliases + self._sources = set(aliases.values()) + if any(source in aliases and aliases[source] != source for source in self._sources): + raise ValueError("A LoRA alias must not shadow a different declared source") + if revisions is not None and not isinstance(revisions, Mapping): + raise ValueError("LoRA revisions must be a canonical-source mapping") + pins = {} if revisions is None else dict(revisions) + if any(not isinstance(key, str) or not key for key in pins): + raise ValueError("LoRA revision keys must be nonempty canonical source strings") + if any(key in self._aliases and key not in self._sources for key in pins): + raise ValueError("LoRA revision keys must name canonical sources, not served aliases") + if any( + pin is not None and (not isinstance(pin, str) or not re.fullmatch(r"[0-9a-fA-F]{40}", pin)) + for pin in pins.values() + ): + raise ValueError("LoRA revisions must be immutable 40-hex commits") + # ModelConfig supplies the canonical-source map for every profile. + # Validate all pins, but this selected profile only serves its sources. + self._revisions = {source: pins.get(source) for source in self._sources} + self._max_loaded = max_loaded + self.lock = threading.RLock() + self._source_lock = threading.Lock() + self._prepared: dict[str, _Prepared] = {} + self._model: Any = None + self._controller: PeftModel | None = None + self._originals: dict[str, torch.nn.Linear] = {} + # A plain dict retains GPU residency without making raw Module.eval() + # or base inference traverse adapter wrappers and their child modules. + self._resident_layers: dict[str, LoraLinear] = {} + self._shapes: dict[str, tuple[int, int, torch.dtype]] = {} + self._target_prefixes: tuple[str, ...] = () + self._grad: list[tuple[torch.nn.Parameter, bool]] = [] + self._training: list[tuple[torch.nn.Module, bool]] = [] + self._attributes: list[tuple[torch.nn.Module, str, bool, Any, dict | None]] = [] + self._loaded: dict[str, list[torch.nn.Parameter]] = {} + self._active: str | None = None + self._closed = True + self._uncertain = False + self._generation = 0 + + def runtime_config(self) -> LoraRuntimeConfig: + # Deferred-only: model startup must not fetch a declaration before a + # request has passed the task admission hook. + return LoraRuntimeConfig(dict(self._aliases), frozenset(), frozenset(), self._max_loaded) + + def source(self, name: str) -> str: + source = self._aliases.get(name, name) + if source not in self._sources: + raise ValueError("LoRA is not declared for this profile") + return source + + def bind(self, model: Any) -> None: + with self.lock: + if not self._closed: + raise RuntimeError("Close the current LoRA model before binding another") + prefixes = self._model_targets(model) + if not prefixes or any(not prefix for prefix in prefixes): + raise ValueError("Runtime LoRA must declare its supported projection paths") + if any(getattr(m, "is_quantized", False) or isinstance(m, LoraLinear) for m in model.modules()): + raise ValueError("Runtime LoRA requires an unquantized base without existing adapters") + if any( + getattr(m, "peft_config", None) or getattr(m, "_hf_peft_config_loaded", False) for m in model.modules() + ): + raise ValueError("Runtime LoRA requires a base without existing PEFT state") + self._model = model + self._target_prefixes = prefixes + self._originals = {n: m for n, m in model.named_modules() if type(m) is torch.nn.Linear} + self._shapes = {n: (m.in_features, m.out_features, m.weight.dtype) for n, m in self._originals.items()} + self._grad = [(p, p.requires_grad) for p in model.parameters()] + self._training = [(m, m.training) for m in model.modules()] + self._attributes = [ + ( + m, + key, + key in m.__dict__, + m.__dict__.get(key), + dict(m.__dict__[key]) if isinstance(m.__dict__.get(key), dict) else None, + ) + for m in model.modules() + for key in _PEFT_ATTRIBUTES + ] + with self._source_lock: + self._generation += 1 + self._closed = False + self._uncertain = False + + @abstractmethod + def _model_targets(self, model: Any) -> tuple[str, ...]: + """Validate the exact serving architecture and return allowed Linear prefixes.""" + + def _source_directory(self, canonical: str, revision: str | None) -> Path: + """Resolve an already-authorized declaration before taking an owned snapshot.""" + directory = Path(canonical) + if directory.is_dir(): + if revision is not None: + raise ValueError("A local LoRA directory cannot honor a Hub revision pin") + return directory + return Path(snapshot_download(canonical, revision=revision, allow_patterns=list(_FILES))) + + def _ready(self) -> None: + if self._uncertain: + raise LoraStateUncertainError("Runtime LoRA state is uncertain; reload the model") + if self._closed or self._model is None: + raise RuntimeError("Runtime LoRA model is not loaded") + + @staticmethod + def _config(source: str, config: PeftConfig) -> LoraConfig: + PEFTLoRAMixin()._check_lora_config(source, config) + if not isinstance(config, LoraConfig) or config.task_type is not None: + raise ValueError("Runtime LoRA requires ordinary generic PEFT with task_type=None") + for field in _UNSUPPORTED: + if getattr(config, field, None): + raise ValueError(f"Runtime LoRA does not support {field}") + if not isinstance(config.init_lora_weights, bool): + raise ValueError("Runtime LoRA requires plain boolean init_lora_weights") + if isinstance(config.r, bool) or not isinstance(config.r, int) or config.r < 1: + raise ValueError("LoRA rank must be a positive integer") + if not math.isfinite(config.lora_alpha) or config.lora_alpha <= 0: + raise ValueError("LoRA alpha must be finite and positive") + if not math.isfinite(config.lora_dropout) or not 0 <= config.lora_dropout < 1: + raise ValueError("LoRA dropout must be finite and in [0, 1)") + if ( + not isinstance(config.target_modules, (set, list)) + or not config.target_modules + or any(not isinstance(t, str) or not t for t in config.target_modules) + ): + raise ValueError("Runtime LoRA requires explicit target module names") + config.inference_mode = True + return config + + @classmethod + def _read_config(cls, source: str, directory: Path) -> LoraConfig: + # PEFT can warn and discard fields from a newer exporter. An unknown + # active feature must fail instead of silently serving a different LoRA. + raw = json.loads((directory / _FILES[0]).read_text()) + if not isinstance(raw, dict): + raise ValueError("LoRA configuration must be a JSON object") + known = {field.name for field in fields(LoraConfig)} | {"peft_version"} + for name, value in raw.items(): + if name not in known and value not in (None, False, {}, []): + raise ValueError(f"Unsupported LoRA configuration field {name!r}") + return cls._config(source, PeftConfig.from_pretrained(str(directory), local_files_only=True)) + + def _weights( + self, config: LoraConfig, weights: dict[str, torch.Tensor], shapes: dict[str, tuple[int, int, torch.dtype]] + ) -> dict[str, torch.Tensor]: + expected: dict[str, tuple[tuple[int, int], torch.dtype]] = {} + for target in config.target_modules or []: + matches = [n for n in shapes if n == target or n.endswith(f".{target}")] + if not matches or any(not n.startswith(self._target_prefixes) for n in matches): + raise ValueError(f"LoRA target {target!r} must resolve only to the supported Linear projections") + for name in matches: + ins, outs, dtype = shapes[name] + expected[f"{name}.lora_A.weight"] = ((config.r, ins), dtype) + expected[f"{name}.lora_B.weight"] = ((outs, config.r), dtype) + normalized: dict[str, torch.Tensor] = {} + for key, tensor in weights.items(): + name = key.removeprefix("base_model.model.") + if name in normalized: + raise ValueError("Duplicate LoRA tensor key") + normalized[name] = tensor + if set(normalized) != set(expected): + raise ValueError("LoRA weights must exactly match all declared A/B projection matrices") + for name, tensor in normalized.items(): + shape, dtype = expected[name] + if tuple(tensor.shape) != shape or not tensor.is_floating_point() or not torch.isfinite(tensor).all(): + raise ValueError(f"Invalid shape, dtype or nonfinite LoRA tensor {name!r}") + if not torch.isfinite(tensor.to(dtype=dtype)).all(): + raise ValueError(f"LoRA tensor {name!r} overflows the model compute dtype") + return normalized + + def fetch(self, name: str, revision: str | None = None) -> str: + canonical = self.source(name) + pinned = self._revisions[canonical] + if revision is not None and revision != pinned: + raise ValueError("LoRA revision differs from the declared source") + with self._source_lock: + self._ready() + generation, shapes = self._generation, dict(self._shapes) + if canonical in self._prepared: + return self._prepared[canonical].directory.name + directory = self._source_directory(canonical, pinned) + temporary = tempfile.TemporaryDirectory(prefix="sie-runtime-lora-") + try: + for filename in _FILES: + shutil.copyfile(directory / filename, Path(temporary.name) / filename) + config = self._read_config(canonical, Path(temporary.name)) + weights = self._weights(config, load_file(str(Path(temporary.name) / _FILES[1])), shapes) + with self._source_lock: + self._ready() + if generation != self._generation: + raise RuntimeError("Model changed while the LoRA source was fetched") + if canonical in self._prepared: + temporary.cleanup() + return self._prepared[canonical].directory.name + self._prepared[canonical] = _Prepared(temporary, config, weights) + return temporary.name + except BaseException: + temporary.cleanup() + raise + + def discard(self, name: str) -> None: + canonical = self.source(name) + with self._source_lock: + prepared = self._prepared.pop(canonical, None) + if prepared is not None: + try: + prepared.directory.cleanup() + except OSError: + with self._source_lock: + self._prepared[canonical] = prepared + raise + + def _restore_flags(self) -> None: + for parameter, flag in self._grad: + parameter.requires_grad_(flag) + # Assign individually: Module.train() recursively overwrites mixed modes. + for module, flag in self._training: + module.training = flag + + def _freeze_adapters(self) -> None: + for module in self._model.modules(): + if isinstance(module, LoraLinear): + for group in (module.lora_A, module.lora_B, module.lora_dropout): + group.eval() + for parameter in group.parameters(): + parameter.requires_grad_(False) + self._restore_flags() + + def _tuner(self) -> LoraModel: + if self._controller is None or not isinstance(self._controller.base_model, LoraModel): + raise RuntimeError("Runtime LoRA controller does not own ordinary LoRA layers") + return self._controller.base_model + + def _attach_layers(self) -> None: + """Expose all residents before PEFT selection, injection or removal.""" + for name, layer in self._resident_layers.items(): + original = self._originals[name] + current = self._model.get_submodule(name) + if current is not original and current is not layer: + raise RuntimeError("A base module changed outside its owned LoRA lifecycle") + if layer.get_base_layer() is not original: + raise RuntimeError("A LoRA base projection changed") + # Detached adapters must not silently migrate if an external caller + # moves/casts the raw model after it was admitted for serving. + if any( + parameter.device != original.weight.device or parameter.dtype != original.weight.dtype + for parameter in layer.parameters() + ): + raise RuntimeError("Base device or precision changed with resident LoRAs") + parent, _, leaf = name.rpartition(".") + setattr(self._model.get_submodule(parent), leaf, layer) + + def _detach_layers(self) -> None: + """Restore native base projections while retaining nonempty LoRA layers.""" + for name, original in self._originals.items(): + current = self._model.get_submodule(name) + if current is original: + continue + if not isinstance(current, LoraLinear) or current.get_base_layer() is not original: + raise RuntimeError("A base module changed outside its owned LoRA lifecycle") + if current.lora_A: + self._resident_layers[name] = current + else: + self._resident_layers.pop(name, None) + parent, _, leaf = name.rpartition(".") + setattr(self._model.get_submodule(parent), leaf, original) + self._active = None + + def _restore_tree(self, *, final: bool) -> None: + for name, original in self._originals.items(): + current = self._model.get_submodule(name) + if current is original: + continue + if not isinstance(current, LoraLinear) or current.get_base_layer() is not original: + raise RuntimeError("A base module changed outside its owned LoRA lifecycle") + if final or not current.lora_A: + self._resident_layers.pop(name, None) + parent, _, leaf = name.rpartition(".") + setattr(self._model.get_submodule(parent), leaf, original) + if final: + self._resident_layers.clear() + for module, name, present, value, contents in self._attributes: + if present: + if contents is not None: + value.clear() + value.update(contents) + setattr(module, name, value) + else: + module.__dict__.pop(name, None) + self._controller = None + self._restore_flags() + self._active = None + + def _remove(self, internal: str) -> None: + # Works even when injection failed before a PeftModel was returned. + delete_adapter(self._model, adapter_name=internal, prefix="lora_") + getattr(self._model, "peft_config", {}).pop(internal, None) + if self._controller is not None: + self._controller.peft_config.pop(internal, None) + # The low-level helper updates leaf layers, not the controller or + # tuner root. Its returned first-layer selection can also be empty + # for disjoint targets even while another resident survives. + # Synchronize all levels to a known resident before callers disable + # and refreeze them; final removal restores/drops the controller. + for canonical in self._loaded: + remaining = PEFTLoRAMixin._peft_adapter_name(canonical) + if remaining != internal: + self._controller.set_adapter(remaining, inference_mode=True) + break + + def load(self, name: str, revision: str | None = None, *, source: str | None = None) -> int: + canonical = self.source(name) + if revision is not None and revision != self._revisions[canonical]: + raise ValueError("LoRA revision differs from the declared source") + with self.lock: + self._ready() + if canonical in self._loaded: + return sum(p.numel() * p.element_size() for p in self._loaded[canonical]) + if source is None: + source = self.fetch(canonical, revision) + with self.lock: + self._ready() + with self._source_lock: + prepared = self._prepared.get(canonical) + if prepared is None or prepared.directory.name != source: + # A concurrent direct caller can finish the same load while + # this caller fetches. No second attachment is necessary. + if canonical in self._loaded: + return sum(p.numel() * p.element_size() for p in self._loaded[canonical]) + raise ValueError("LoRA source must be this declaration's validated snapshot") + try: + if canonical in self._loaded: + return sum(p.numel() * p.element_size() for p in self._loaded[canonical]) + if len(self._loaded) >= self._max_loaded: + raise ValueError("Runtime LoRA capacity is exhausted; evict before attaching") + return self._attach(canonical, source, prepared) + finally: + try: + self.discard(canonical) + except OSError as error: + self._uncertain = True + raise LoraStateUncertainError("Runtime LoRA source cleanup failed; reload model") from error + + def _attach(self, canonical: str, source: str, prepared: _Prepared) -> int: + internal = PEFTLoRAMixin._peft_adapter_name(canonical) + try: + self._attach_layers() + if self._controller is None: + self._controller = PeftModel.from_pretrained( + self._model, + source, + adapter_name=internal, + config=prepared.config, + is_trainable=False, + autocast_adapter_dtype=False, + local_files_only=True, + ) + else: + # Reuse the validated config for later attachments too. The + # loader then reads only the owned, safetensors-only snapshot. + self._controller.add_adapter(internal, prepared.config, autocast_adapter_dtype=False) + self._controller.load_adapter( + source, + adapter_name=internal, + is_trainable=False, + autocast_adapter_dtype=False, + local_files_only=True, + ) + if self._controller.get_base_model() is not self._model: + raise RuntimeError("PEFT changed the serving model owner") + parameters = [] + for key, weight in prepared.weights.items(): + target, part, _ = key.rsplit(".", 2) + layer = self._model.get_submodule(target) + if not isinstance(layer, LoraLinear) or layer.get_base_layer() is not self._originals[target]: + raise RuntimeError("PEFT did not retain the original projection") + parameter = getattr(layer, part)[internal].weight + if not torch.isfinite(parameter).all() or not torch.equal( + parameter.detach(), weight.to(device=parameter.device, dtype=parameter.dtype) + ): + raise ValueError("LoRA tensors were not installed exactly") + parameters.append(parameter) + self._tuner().disable_adapter_layers() + self._freeze_adapters() + self._detach_layers() + self._loaded[canonical] = parameters + except BaseException: + try: + if self._loaded: + # A failed detach can have saved some new layers off-tree. + # Restore the whole owned set before deleting the injection. + self._attach_layers() + self._remove(internal) + assert self._controller is not None + self._tuner().disable_adapter_layers() + self._freeze_adapters() + self._restore_tree(final=not self._loaded) + if self._loaded: + self._detach_layers() + except BaseException as rollback: + self._uncertain = True + raise LoraStateUncertainError("Runtime LoRA attachment rollback failed") from rollback + raise + return sum(p.numel() * p.element_size() for p in parameters) + + def select(self, name: str | None) -> None: + with self.lock: + self._ready() + canonical = self.source(name) if name is not None else None + if canonical is not None and canonical not in self._loaded: + raise LoraUnavailableError("Declared LoRA is not resident") + # Attach/remove/reset already establish this state. Avoid walking + # every original parameter/module on repeated base selection and + # the worker's repeated selection before the direct entrypoint. + if canonical == self._active: + return + try: + if self._controller is not None: + if canonical is None: + self._tuner().disable_adapter_layers() + else: + self._attach_layers() + self._controller.set_adapter(PEFTLoRAMixin._peft_adapter_name(canonical), inference_mode=True) + self._tuner().enable_adapter_layers() + self._freeze_adapters() + if canonical is None: + self._detach_layers() + self._active = canonical + except BaseException as error: + self._uncertain = True + raise LoraStateUncertainError("Runtime LoRA selection failed; reload model") from error + + def unload(self, name: str) -> None: + with self.lock: + self._ready() + canonical = self.source(name) + if canonical not in self._loaded: + raise LoraUnavailableError("Runtime LoRA is not resident") + try: + self._attach_layers() + self._remove(PEFTLoRAMixin._peft_adapter_name(canonical)) + final = len(self._loaded) == 1 + if not final: + assert self._controller is not None + self._tuner().disable_adapter_layers() + self._freeze_adapters() + self._restore_tree(final=final) + if not final: + self._detach_layers() + except BaseException as error: + self._uncertain = True + raise LoraStateUncertainError("Runtime LoRA removal failed; reload model") from error + del self._loaded[canonical] + self.discard(canonical) + + def close(self) -> None: + with self.lock: + with self._source_lock: + self._closed = True + self._generation += 1 + prepared, self._prepared = self._prepared, {} + errors = [] + try: + if self._model is not None: + self._restore_tree(final=True) + except Exception as error: # noqa: BLE001 — cleanup must attempt every owned resource. + errors.append(error) + for item in prepared.values(): + try: + item.directory.cleanup() + except OSError as error: + errors.append(error) + self._controller = self._model = None + self._originals.clear() + self._resident_layers.clear() + self._shapes.clear() + self._grad.clear() + self._training.clear() + self._attributes.clear() + self._loaded.clear() + if errors: + raise ExceptionGroup("Runtime LoRA teardown failed", errors) diff --git a/packages/sie_server/src/sie_server/adapters/decision_lora/__init__.py b/packages/sie_server/src/sie_server/adapters/decision_lora/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/packages/sie_server/src/sie_server/adapters/decision_lora/adapter.py b/packages/sie_server/src/sie_server/adapters/decision_lora/adapter.py new file mode 100644 index 000000000..87e2373f8 --- /dev/null +++ b/packages/sie_server/src/sie_server/adapters/decision_lora/adapter.py @@ -0,0 +1,152 @@ +"""Lifecycle and selection for explicitly configured native decision LoRA profiles.""" + +from __future__ import annotations + +from abc import abstractmethod +from typing import Any + +from sie_server.adapters._base_adapter import BaseAdapter +from sie_server.adapters._peft_lora_runtime import DetachedLoraRuntime +from sie_server.adapters.lora import LoraRuntimeConfig +from sie_server.core.inference_output import ExtractOutput +from sie_server.types.inputs import InvalidInputError, Item +from sie_server.types.lora import InvalidLoraSelectionError, split_lora_selection + + +class DecisionLoraAdapter(BaseAdapter): + """Cooperative base for opt-in profiles; native adapter classes never import it. + + Each concrete class precedes its ordinary adapter in the MRO. The raw model + still owns inference; the separate controller owns reversible LoRA layers. + Base CUDA graphs are retained, but an adapted forward always runs eagerly. + """ + + _lora: DetachedLoraRuntime + + @abstractmethod + def _lora_model(self) -> Any: + """Return the exact native serving owner after load.""" + + @abstractmethod + def _validate_decision_request( + self, + items: list[Item], + *, + labels: list[str] | None, + output_schema: dict[str, Any] | None, + instruction: str | None, + options: dict[str, Any], + ) -> None: + """Check native task intent before a missing adapter can be fetched.""" + + def _request_selection(self, options: dict[str, Any] | None) -> tuple[str | None, dict[str, Any]]: + if options is not None and not isinstance(options, dict): + raise InvalidInputError("Decision options must be an object") + try: + selection, opts = split_lora_selection(options or {}) + except InvalidLoraSelectionError as error: + raise InvalidInputError(str(error)) from error + selected = selection.value + if selected is not None: + try: + source = self._lora.source(selected) + if "lora_id" in selection.fields and source != selected: + raise ValueError("Legacy lora_id must name the declared source, not a served alias") + except ValueError as error: + raise InvalidInputError(str(error)) from error + return selected, opts + + def validate_extract_request( + self, + items: list[Item], + *, + labels: list[str] | None = None, + output_schema: dict[str, Any] | None = None, + instruction: str | None = None, + options: dict[str, Any] | None = None, + ) -> None: + _, opts = self._request_selection(options) + self._validate_decision_request( + items, labels=labels, output_schema=output_schema, instruction=instruction, options=opts + ) + + def load(self, device: str) -> None: + try: + super().load(device) + self._lora.bind(self._lora_model()) + except BaseException: + self.unload() + raise + + def unload(self) -> None: + try: + self._lora.close() + finally: + super().unload() + + def extract( + self, + items: list[Item], + *, + labels: list[str] | None = None, + output_schema: dict[str, Any] | None = None, + instruction: str | None = None, + options: dict[str, Any] | None = None, + prepared_items: list[Any] | None = None, + ) -> ExtractOutput: + selected, opts = self._request_selection(options) + self._validate_decision_request( + items, labels=labels, output_schema=output_schema, instruction=instruction, options=opts + ) + with self._lora.lock: + self._lora.select(selected) + graphs = getattr(self, "_graphs", None) + if selected is not None: + if graphs is not None: + self._graphs = None + primed = getattr(self, "_primed", None) + if primed is not None: + primed.clear() + try: + return super().extract( + items, + labels=labels, + output_schema=output_schema, + instruction=instruction, + options=opts, + prepared_items=prepared_items, + ) + finally: + try: + self._lora.select(None) + finally: + if graphs is not None: + self._graphs = graphs + if selected is not None: + primed = getattr(self, "_primed", None) + if primed is not None: + primed.clear() + + def supports_lora(self) -> bool: + return True + + def supports_hot_lora_reload(self) -> bool: + return True + + def lora_runtime_config(self) -> LoraRuntimeConfig: + return self._lora.runtime_config() + + def fetch_lora(self, lora_path: str, revision: str | None = None) -> str: + return self._lora.fetch(lora_path, revision) + + def discard_lora_source(self, lora_path: str) -> None: + self._lora.discard(lora_path) + + def load_lora(self, lora_path: str, revision: str | None = None, *, source: str | None = None) -> int: + return self._lora.load(lora_path, revision, source=source) + + def unload_lora(self, lora_name: str) -> None: + self._lora.unload(lora_name) + + def set_active_lora(self, lora_name: str | None) -> None: + self._lora.select(lora_name) diff --git a/packages/sie_server/src/sie_server/adapters/gliclass/runtime_lora.py b/packages/sie_server/src/sie_server/adapters/gliclass/runtime_lora.py new file mode 100644 index 000000000..6e2bc9346 --- /dev/null +++ b/packages/sie_server/src/sie_server/adapters/gliclass/runtime_lora.py @@ -0,0 +1,111 @@ +"""Opt-in declared runtime LoRA for native GLiClass uni-encoder decisions.""" + +from __future__ import annotations + +from pathlib import Path +from typing import Any + +from transformers import BertModel, DebertaV2Model, ModernBertModel + +from sie_server.adapters._peft_lora_runtime import DetachedLoraRuntime +from sie_server.adapters.decision_lora.adapter import DecisionLoraAdapter +from sie_server.adapters.gliclass import GLiClassAdapter +from sie_server.types.inputs import InvalidInputError, Item + + +class GLiClassLora(DetachedLoraRuntime): + """Adapt only ordinary encoder projections, retaining the native scoring head.""" + + def _model_targets(self, model: Any) -> tuple[str, ...]: + from gliclass.model import GLiClassModel, GLiClassUniEncoder + + if type(model) is not GLiClassModel or type(model.model) is not GLiClassUniEncoder: + raise ValueError("GLiClass LoRA requires the native uni-encoder model") + encoder = model.model.encoder_model + if type(encoder) is ModernBertModel: + return ("model.encoder_model.layers.",) + if type(encoder) in (BertModel, DebertaV2Model): + return ("model.encoder_model.encoder.layer.",) + raise ValueError("GLiClass LoRA supports ordinary BERT, DeBERTaV2 and ModernBERT encoders") + + +class GLiClassLoraAdapter(DecisionLoraAdapter, GLiClassAdapter): + """Native labels and grouped decisions with explicit runtime LoRA declarations.""" + + def __init__( + self, + model_name_or_path: str | Path, + *, + lora_paths: dict[str, str], + lora_revisions: dict[str, str | None] | None = None, + max_loaded_loras: int = 10, + **kwargs: Any, + ) -> None: + allowed = { + "classification_type", + "threshold", + "max_seq_length", + "compute_precision", + "revision", + "cuda_graphs", + "modernbert_flash", + } + if kwargs.keys() - allowed: + raise ValueError(f"GLiClass LoRA does not accept load options {sorted(kwargs.keys() - allowed)!r}") + self._lora = GLiClassLora(lora_paths, lora_revisions, max_loaded_loras) + super().__init__(model_name_or_path, **kwargs) + + def _lora_model(self) -> Any: + return self._require_pipe().model + + def _validate_decision_request( + self, + items: list[Item], + *, + labels: list[str] | None, + output_schema: dict[str, Any] | None, + instruction: str | None, + options: dict[str, Any], + ) -> None: + allowed = { + "label_groups", + "group_encoding", + "classification_type", + "threshold", + "examples", + "cuda_graphs", + "overflow_policy", + } + if set(options) - allowed: + raise InvalidInputError("Unsupported GLiClass decision option") + if output_schema is not None: + raise InvalidInputError("GLiClass uses labels or options.label_groups, not output_schema") + groups = self._validate_label_groups(options.get("label_groups")) + encoding = self._validate_group_encoding(options, groups) + classification_type = self._validate_classification_type( + options.get("classification_type", self._classification_type) + ) + if groups is None: + normalized_labels = self._validate_labels(labels) + elif labels: + raise InvalidInputError("GLiClass accepts either labels or options.label_groups, not both") + else: + normalized_labels = self._flatten_label_groups(groups, classification_type) + if groups is not None and encoding == "separate": + self._check_group_count(groups) + for _, group_labels in groups: + self._check_label_size(group_labels) + else: + self._check_label_size(normalized_labels) + prompt = self._validate_instruction(instruction) + examples = self._validate_examples(options.get("examples"), normalized_labels, groups) + self._check_context_size(prompt, examples) + self._validate_threshold(options.get("threshold", self._threshold)) + self._request_cuda_graphs(options) + policy = options.get("overflow_policy", "default") + if not isinstance(policy, str) or policy not in {"default", "truncate_text", "error"}: + raise InvalidInputError("GLiClass overflow_policy must be default, truncate_text or error") + if not items: + raise InvalidInputError("GLiClass LoRA requires at least one item") + for item in items: + self._extract_text(item) diff --git a/packages/sie_server/src/sie_server/adapters/gliner2/_lora.py b/packages/sie_server/src/sie_server/adapters/gliner2/_lora.py index a5388a6cc..fb1f9fd70 100644 --- a/packages/sie_server/src/sie_server/adapters/gliner2/_lora.py +++ b/packages/sie_server/src/sie_server/adapters/gliner2/_lora.py @@ -1,543 +1,36 @@ -"""Strict generic-PEFT lifecycle for the opt-in GLiNER2 entity adapter. - -The raw extractor owns serving. A separate PEFT controller owns only reversible -LoRA injection; it is never used as the extraction entrypoint. Resident wrappers -stay outside the module tree while base is selected, sharing the exact original -projection weights. Admission and selection still have their own runtime cost. -""" +"""Architecture boundary for the explicitly selected GLiNER2 entity LoRA profile.""" from __future__ import annotations import importlib.metadata -import math -import re -import shutil -import tempfile -import threading -from collections.abc import Mapping -from dataclasses import dataclass from pathlib import Path from typing import Any -import torch from gliner2 import GLiNER2 from huggingface_hub import snapshot_download -from peft import LoraConfig, PeftConfig, PeftModel -from peft.functional import delete_adapter -from peft.tuners.lora import LoraModel -from peft.tuners.lora.layer import Linear as LoraLinear -from safetensors.torch import load_file from transformers import BertModel, DebertaV2Model -from sie_server.adapters.lora import LoraRuntimeConfig, LoraStateUncertainError, LoraUnavailableError -from sie_server.adapters.peft_lora_mixin import PEFTLoRAMixin - -_FILES = ("adapter_config.json", "adapter_model.safetensors") -_PEFT_ATTRIBUTES = ("peft_config", "_hf_peft_config_loaded") -_UNSUPPORTED = ( - "use_dora", - "use_rslora", - "use_qalora", - "use_bdlora", - "lora_bias", - "fan_in_fan_out", - "rank_pattern", - "alpha_pattern", - "layers_to_transform", - "layers_pattern", - "exclude_modules", - "target_parameters", - "alora_invocation_tokens", - "arrow_config", - "ensure_weight_tying", - "loftq_config", - "eva_config", - "corda_config", - "lora_ga_config", -) - - -@dataclass -class _Prepared: - directory: tempfile.TemporaryDirectory[str] - config: LoraConfig - weights: dict[str, torch.Tensor] - - -class GLiNER2Lora: - """Declared adapters; worker mutation exclusion plus direct-call protection.""" +from sie_server.adapters._peft_lora_runtime import DetachedLoraRuntime - def __init__(self, paths: dict[str, str], revisions: dict[str, str | None] | None, max_loaded: int) -> None: - if isinstance(max_loaded, bool) or not isinstance(max_loaded, int) or max_loaded < 1: - raise ValueError("max_loaded_loras must be a positive integer") - if not isinstance(paths, dict) or not paths: - raise ValueError("Entity LoRA requires a nonempty lora_paths alias-to-source map") - aliases = dict(paths) - if any( - not isinstance(k, str) - or not re.fullmatch(r"[A-Za-z0-9_][A-Za-z0-9_.-]*", k) - or not isinstance(v, str) - or not v - for k, v in aliases.items() - ): - raise ValueError("LoRA declarations require simple nonempty aliases and nonempty sources") - self._aliases = aliases - self._sources = set(aliases.values()) - if any(source in aliases and aliases[source] != source for source in self._sources): - raise ValueError("A LoRA alias must not shadow a different declared source") - if revisions is not None and not isinstance(revisions, Mapping): - raise ValueError("LoRA revisions must be a canonical-source mapping") - pins = {} if revisions is None else dict(revisions) - if any(not isinstance(key, str) or not key for key in pins): - raise ValueError("LoRA revision keys must be nonempty canonical source strings") - if any(key in self._aliases and key not in self._sources for key in pins): - raise ValueError("LoRA revision keys must name canonical sources, not served aliases") - if any( - pin is not None and (not isinstance(pin, str) or not re.fullmatch(r"[0-9a-fA-F]{40}", pin)) - for pin in pins.values() - ): - raise ValueError("LoRA revisions must be immutable 40-hex commits") - # ModelConfig supplies the canonical-source map for every profile. - # Validate all pins, but this selected profile only serves its sources. - self._revisions = {source: pins.get(source) for source in self._sources} - self._max_loaded = max_loaded - self.lock = threading.RLock() - self._source_lock = threading.Lock() - self._prepared: dict[str, _Prepared] = {} - self._model: Any = None - self._controller: PeftModel | None = None - self._originals: dict[str, torch.nn.Linear] = {} - # A plain dict retains GPU residency without making raw Extractor.eval() - # or base inference traverse adapter wrappers and their child modules. - self._resident_layers: dict[str, LoraLinear] = {} - self._shapes: dict[str, tuple[int, int, torch.dtype]] = {} - self._grad: list[tuple[torch.nn.Parameter, bool]] = [] - self._training: list[tuple[torch.nn.Module, bool]] = [] - self._attributes: list[tuple[torch.nn.Module, str, bool, Any, dict | None]] = [] - self._loaded: dict[str, list[torch.nn.Parameter]] = {} - self._active: str | None = None - self._closed = True - self._uncertain = False - self._generation = 0 - def runtime_config(self) -> LoraRuntimeConfig: - # Deferred-only: model startup must not fetch a declaration before a - # request has passed the entity-only admission hook. - return LoraRuntimeConfig(dict(self._aliases), frozenset(), frozenset(), self._max_loaded) +class GLiNER2Lora(DetachedLoraRuntime): + """The existing entity-only architecture and target policy.""" - def source(self, name: str) -> str: - source = self._aliases.get(name, name) - if source not in self._sources: - raise ValueError("LoRA is not declared for this entity profile") - return source + def _model_targets(self, model: Any) -> tuple[str, ...]: + if importlib.metadata.version("gliner2") not in {"1.3.2", "2.0.0"}: + raise ValueError("The entity LoRA adapter requires GLiNER2 1.3.2 or 2.0.0") + if type(model) is not GLiNER2 or type(model.encoder) not in (BertModel, DebertaV2Model): + raise ValueError("Entity LoRA requires the ordinary GLiNER2 BERT or DeBERTaV2 encoder") + return ("encoder.encoder.layer.", "span_rep.") - def bind(self, model: Any) -> None: - with self.lock: - if importlib.metadata.version("gliner2") not in {"1.3.2", "2.0.0"}: - raise ValueError("The entity LoRA adapter requires GLiNER2 1.3.2 or 2.0.0") - if type(model) is not GLiNER2 or type(model.encoder) not in (BertModel, DebertaV2Model): - raise ValueError("Entity LoRA requires the ordinary GLiNER2 BERT or DeBERTaV2 encoder") - if getattr(model.encoder, "is_quantized", False) or any(isinstance(m, LoraLinear) for m in model.modules()): - raise ValueError("Entity LoRA requires an unquantized base without existing adapters") - if any( - getattr(m, "peft_config", None) or getattr(m, "_hf_peft_config_loaded", False) for m in model.modules() - ): - raise ValueError("Entity LoRA requires a base without existing PEFT state") - self._model = model - self._originals = {n: m for n, m in model.named_modules() if type(m) is torch.nn.Linear} - self._shapes = {n: (m.in_features, m.out_features, m.weight.dtype) for n, m in self._originals.items()} - self._grad = [(p, p.requires_grad) for p in model.parameters()] - self._training = [(m, m.training) for m in model.modules()] - self._attributes = [ - ( - m, - key, - key in m.__dict__, - m.__dict__.get(key), - dict(m.__dict__[key]) if isinstance(m.__dict__.get(key), dict) else None, - ) - for m in model.modules() - for key in _PEFT_ATTRIBUTES - ] - with self._source_lock: - self._generation += 1 - self._closed = False - self._uncertain = False - - def _ready(self) -> None: - if self._uncertain: - raise LoraStateUncertainError("Entity LoRA state is uncertain; reload the model") - if self._closed or self._model is None: - raise RuntimeError("Entity LoRA model is not loaded") - - @staticmethod - def _config(source: str, config: PeftConfig) -> LoraConfig: - PEFTLoRAMixin()._check_lora_config(source, config) - if not isinstance(config, LoraConfig) or config.task_type is not None: - raise ValueError("Entity LoRA requires ordinary generic PEFT with task_type=None") - for field in _UNSUPPORTED: - if getattr(config, field, None): - raise ValueError(f"Entity LoRA does not support {field}") - if not isinstance(config.init_lora_weights, bool): - raise ValueError("Entity LoRA requires plain boolean init_lora_weights") - if isinstance(config.r, bool) or not isinstance(config.r, int) or config.r < 1: - raise ValueError("LoRA rank must be a positive integer") - if not math.isfinite(config.lora_alpha) or config.lora_alpha <= 0: - raise ValueError("LoRA alpha must be finite and positive") - if not math.isfinite(config.lora_dropout) or not 0 <= config.lora_dropout < 1: - raise ValueError("LoRA dropout must be finite and in [0, 1)") - if ( - not isinstance(config.target_modules, (set, list)) - or not config.target_modules - or any(not isinstance(t, str) or not t for t in config.target_modules) - ): - raise ValueError("Entity LoRA requires explicit target module names") - config.inference_mode = True - return config - - @staticmethod - def _weights( - config: LoraConfig, weights: dict[str, torch.Tensor], shapes: dict[str, tuple[int, int, torch.dtype]] - ) -> dict[str, torch.Tensor]: - expected: dict[str, tuple[tuple[int, int], torch.dtype]] = {} - for target in config.target_modules or []: - matches = [n for n in shapes if n == target or n.endswith(f".{target}")] - if not matches or any(not n.startswith(("encoder.encoder.layer.", "span_rep.")) for n in matches): - raise ValueError(f"LoRA target {target!r} must resolve only to encoder or span Linear projections") - for name in matches: - ins, outs, dtype = shapes[name] - expected[f"{name}.lora_A.weight"] = ((config.r, ins), dtype) - expected[f"{name}.lora_B.weight"] = ((outs, config.r), dtype) - normalized: dict[str, torch.Tensor] = {} - for key, tensor in weights.items(): - name = key.removeprefix("base_model.model.") - if name in normalized: - raise ValueError("Duplicate LoRA tensor key") - normalized[name] = tensor - if set(normalized) != set(expected): - raise ValueError("LoRA weights must exactly match all declared A/B projection matrices") - for name, tensor in normalized.items(): - shape, dtype = expected[name] - if tuple(tensor.shape) != shape or not tensor.is_floating_point() or not torch.isfinite(tensor).all(): - raise ValueError(f"Invalid shape, dtype or nonfinite LoRA tensor {name!r}") - if not torch.isfinite(tensor.to(dtype=dtype)).all(): - raise ValueError(f"LoRA tensor {name!r} overflows the model compute dtype") - return normalized - - def fetch(self, name: str, revision: str | None = None) -> str: - canonical = self.source(name) - pinned = self._revisions[canonical] - if revision is not None and revision != pinned: - raise ValueError("LoRA revision differs from the declared source") - with self._source_lock: - self._ready() - generation, shapes = self._generation, dict(self._shapes) - if canonical in self._prepared: - return self._prepared[canonical].directory.name + def _source_directory(self, canonical: str, revision: str | None) -> Path: directory = Path(canonical) if directory.is_dir(): - if pinned is not None: + if revision is not None: raise ValueError("A local LoRA directory cannot honor a Hub revision pin") - else: - directory = Path(snapshot_download(canonical, revision=pinned, allow_patterns=list(_FILES))) - temporary = tempfile.TemporaryDirectory(prefix="sie-gliner2-lora-") - try: - for filename in _FILES: - shutil.copyfile(directory / filename, Path(temporary.name) / filename) - config = self._config(canonical, PeftConfig.from_pretrained(temporary.name, local_files_only=True)) - weights = self._weights(config, load_file(str(Path(temporary.name) / _FILES[1])), shapes) - with self._source_lock: - self._ready() - if generation != self._generation: - raise RuntimeError("Model changed while the LoRA source was fetched") - if canonical in self._prepared: - temporary.cleanup() - return self._prepared[canonical].directory.name - self._prepared[canonical] = _Prepared(temporary, config, weights) - return temporary.name - except BaseException: - temporary.cleanup() - raise - - def discard(self, name: str) -> None: - canonical = self.source(name) - with self._source_lock: - prepared = self._prepared.pop(canonical, None) - if prepared is not None: - try: - prepared.directory.cleanup() - except OSError: - with self._source_lock: - self._prepared[canonical] = prepared - raise - - def _restore_flags(self) -> None: - for parameter, flag in self._grad: - parameter.requires_grad_(flag) - # Assign individually: Module.train() recursively overwrites mixed modes. - for module, flag in self._training: - module.training = flag - - def _freeze_adapters(self) -> None: - for module in self._model.modules(): - if isinstance(module, LoraLinear): - for group in (module.lora_A, module.lora_B, module.lora_dropout): - group.eval() - for parameter in group.parameters(): - parameter.requires_grad_(False) - self._restore_flags() - - def _tuner(self) -> LoraModel: - if self._controller is None or not isinstance(self._controller.base_model, LoraModel): - raise RuntimeError("Entity LoRA controller does not own ordinary LoRA layers") - return self._controller.base_model - - def _attach_layers(self) -> None: - """Expose all residents before PEFT selection, injection or removal.""" - for name, layer in self._resident_layers.items(): - original = self._originals[name] - current = self._model.get_submodule(name) - if current is not original and current is not layer: - raise RuntimeError("A GLiNER2 module changed outside its owned LoRA lifecycle") - if layer.get_base_layer() is not original: - raise RuntimeError("A GLiNER2 LoRA base projection changed") - # Detached adapters must not silently migrate if an external caller - # moves/casts the raw model after it was admitted for serving. - if any( - parameter.device != original.weight.device or parameter.dtype != original.weight.dtype - for parameter in layer.parameters() - ): - raise RuntimeError("GLiNER2 device or precision changed with resident LoRAs") - parent, _, leaf = name.rpartition(".") - setattr(self._model.get_submodule(parent), leaf, layer) - - def _detach_layers(self) -> None: - """Restore native base projections while retaining nonempty LoRA layers.""" - for name, original in self._originals.items(): - current = self._model.get_submodule(name) - if current is original: - continue - if not isinstance(current, LoraLinear) or current.get_base_layer() is not original: - raise RuntimeError("A GLiNER2 module changed outside its owned LoRA lifecycle") - if current.lora_A: - self._resident_layers[name] = current - else: - self._resident_layers.pop(name, None) - parent, _, leaf = name.rpartition(".") - setattr(self._model.get_submodule(parent), leaf, original) - self._active = None - - def _restore_tree(self, *, final: bool) -> None: - for name, original in self._originals.items(): - current = self._model.get_submodule(name) - if current is original: - continue - if not isinstance(current, LoraLinear) or current.get_base_layer() is not original: - raise RuntimeError("A GLiNER2 module changed outside its owned LoRA lifecycle") - if final or not current.lora_A: - self._resident_layers.pop(name, None) - parent, _, leaf = name.rpartition(".") - setattr(self._model.get_submodule(parent), leaf, original) - if final: - self._resident_layers.clear() - for module, name, present, value, contents in self._attributes: - if present: - if contents is not None: - value.clear() - value.update(contents) - setattr(module, name, value) - else: - module.__dict__.pop(name, None) - self._controller = None - self._restore_flags() - self._active = None - - def _remove(self, internal: str) -> None: - # Works even when injection failed before a PeftModel was returned. - delete_adapter(self._model, adapter_name=internal, prefix="lora_") - getattr(self._model, "peft_config", {}).pop(internal, None) - if self._controller is not None: - self._controller.peft_config.pop(internal, None) - # The low-level helper updates leaf layers, not the controller or - # tuner root. Its returned first-layer selection can also be empty - # for disjoint targets even while another resident survives. - # Synchronize all levels to a known resident before callers disable - # and refreeze them; final removal restores/drops the controller. - for canonical in self._loaded: - remaining = PEFTLoRAMixin._peft_adapter_name(canonical) - if remaining != internal: - self._controller.set_adapter(remaining, inference_mode=True) - break - - def load(self, name: str, revision: str | None = None, *, source: str | None = None) -> int: - canonical = self.source(name) - if revision is not None and revision != self._revisions[canonical]: - raise ValueError("LoRA revision differs from the declared source") - with self.lock: - self._ready() - if canonical in self._loaded: - return sum(p.numel() * p.element_size() for p in self._loaded[canonical]) - if source is None: - source = self.fetch(canonical, revision) - with self.lock: - self._ready() - with self._source_lock: - prepared = self._prepared.get(canonical) - if prepared is None or prepared.directory.name != source: - # A concurrent direct caller can finish the same load while - # this caller fetches. No second attachment is necessary. - if canonical in self._loaded: - return sum(p.numel() * p.element_size() for p in self._loaded[canonical]) - raise ValueError("LoRA source must be this declaration's validated snapshot") - try: - if canonical in self._loaded: - return sum(p.numel() * p.element_size() for p in self._loaded[canonical]) - if len(self._loaded) >= self._max_loaded: - raise ValueError("Entity LoRA capacity is exhausted; evict before attaching") - return self._attach(canonical, source, prepared) - finally: - try: - self.discard(canonical) - except OSError as error: - self._uncertain = True - raise LoraStateUncertainError("Entity LoRA source cleanup failed; reload model") from error - - def _attach(self, canonical: str, source: str, prepared: _Prepared) -> int: - internal = PEFTLoRAMixin._peft_adapter_name(canonical) - try: - self._attach_layers() - if self._controller is None: - self._controller = PeftModel.from_pretrained( - self._model, - source, - adapter_name=internal, - config=prepared.config, - is_trainable=False, - autocast_adapter_dtype=False, - local_files_only=True, - ) - else: - # Reuse the validated config for later attachments too. The - # loader then reads only the owned, safetensors-only snapshot. - self._controller.add_adapter(internal, prepared.config, autocast_adapter_dtype=False) - self._controller.load_adapter( - source, - adapter_name=internal, - is_trainable=False, - autocast_adapter_dtype=False, - local_files_only=True, - ) - if self._controller.get_base_model() is not self._model: - raise RuntimeError("PEFT changed the serving extractor owner") - parameters = [] - for key, weight in prepared.weights.items(): - target, part, _ = key.rsplit(".", 2) - layer = self._model.get_submodule(target) - if not isinstance(layer, LoraLinear) or layer.get_base_layer() is not self._originals[target]: - raise RuntimeError("PEFT did not retain the original projection") - parameter = getattr(layer, part)[internal].weight - if not torch.isfinite(parameter).all() or not torch.equal( - parameter.detach(), weight.to(device=parameter.device, dtype=parameter.dtype) - ): - raise ValueError("LoRA tensors were not installed exactly") - parameters.append(parameter) - self._tuner().disable_adapter_layers() - self._freeze_adapters() - self._detach_layers() - self._loaded[canonical] = parameters - except BaseException: - try: - if self._loaded: - # A failed detach can have saved some new layers off-tree. - # Restore the whole owned set before deleting the injection. - self._attach_layers() - self._remove(internal) - assert self._controller is not None - self._tuner().disable_adapter_layers() - self._freeze_adapters() - self._restore_tree(final=not self._loaded) - if self._loaded: - self._detach_layers() - except BaseException as rollback: - self._uncertain = True - raise LoraStateUncertainError("Entity LoRA attachment rollback failed") from rollback - raise - return sum(p.numel() * p.element_size() for p in parameters) - - def select(self, name: str | None) -> None: - with self.lock: - self._ready() - canonical = self.source(name) if name is not None else None - if canonical is not None and canonical not in self._loaded: - raise LoraUnavailableError("Declared entity LoRA is not resident") - # Attach/remove/reset already establish this state. Avoid walking - # every original parameter/module on repeated base selection and - # the worker's repeated selection before the direct entrypoint. - if canonical == self._active: - return - try: - if self._controller is not None: - if canonical is None: - self._tuner().disable_adapter_layers() - else: - self._attach_layers() - self._controller.set_adapter(PEFTLoRAMixin._peft_adapter_name(canonical), inference_mode=True) - self._tuner().enable_adapter_layers() - self._freeze_adapters() - if canonical is None: - self._detach_layers() - self._active = canonical - except BaseException as error: - self._uncertain = True - raise LoraStateUncertainError("Entity LoRA selection failed; reload model") from error - - def unload(self, name: str) -> None: - with self.lock: - self._ready() - canonical = self.source(name) - if canonical not in self._loaded: - raise LoraUnavailableError("Entity LoRA is not resident") - try: - self._attach_layers() - self._remove(PEFTLoRAMixin._peft_adapter_name(canonical)) - final = len(self._loaded) == 1 - if not final: - assert self._controller is not None - self._tuner().disable_adapter_layers() - self._freeze_adapters() - self._restore_tree(final=final) - if not final: - self._detach_layers() - except BaseException as error: - self._uncertain = True - raise LoraStateUncertainError("Entity LoRA removal failed; reload model") from error - del self._loaded[canonical] - self.discard(canonical) - - def close(self) -> None: - with self.lock: - with self._source_lock: - self._closed = True - self._generation += 1 - prepared, self._prepared = self._prepared, {} - errors = [] - try: - if self._model is not None: - self._restore_tree(final=True) - except Exception as error: # noqa: BLE001 — cleanup must attempt every owned resource. - errors.append(error) - for item in prepared.values(): - try: - item.directory.cleanup() - except OSError as error: - errors.append(error) - self._controller = self._model = None - self._originals.clear() - self._resident_layers.clear() - self._shapes.clear() - self._grad.clear() - self._training.clear() - self._attributes.clear() - self._loaded.clear() - if errors: - raise ExceptionGroup("Entity LoRA teardown failed", errors) + return directory + return Path( + snapshot_download( + canonical, revision=revision, allow_patterns=["adapter_config.json", "adapter_model.safetensors"] + ) + ) diff --git a/packages/sie_server/src/sie_server/adapters/gliner2/decide_lora.py b/packages/sie_server/src/sie_server/adapters/gliner2/decide_lora.py new file mode 100644 index 000000000..0f53e30fd --- /dev/null +++ b/packages/sie_server/src/sie_server/adapters/gliner2/decide_lora.py @@ -0,0 +1,105 @@ +"""Opt-in declared runtime LoRA for the native GLiNER2.5 decision path.""" + +from __future__ import annotations + +import importlib.metadata +from collections.abc import Mapping, Sequence +from pathlib import Path +from typing import Any + +from transformers import BertModel, DebertaV2Model, ModernBertModel + +from sie_server.adapters._peft_lora_runtime import DetachedLoraRuntime +from sie_server.adapters.decision_lora.adapter import DecisionLoraAdapter +from sie_server.adapters.gliner2.decide import GLiNER2DecideAdapter +from sie_server.adapters.gliner2.decisions import parse_request +from sie_server.types.inputs import InvalidInputError, Item + + +class GLiNER2DecideLora(DetachedLoraRuntime): + """Native encoder, span and classifier Linear deltas; no replaced heads.""" + + def _model_targets(self, model: Any) -> tuple[str, ...]: + if importlib.metadata.version("gliner2") != "2.0.0": + raise ValueError("Decision LoRA requires GLiNER2 2.0.0") + from gliner2 import BoundaryExtractor, SpanExtractor + + if type(model) not in (SpanExtractor, BoundaryExtractor): + raise ValueError("Decision LoRA requires a native GLiNER2 span or boundary extractor") + heads = ("classifier.", "span_rep.") if type(model) is SpanExtractor else ("classifier.",) + if type(model.encoder) is ModernBertModel: + return ("encoder.layers.", *heads) + if type(model.encoder) in (BertModel, DebertaV2Model): + return ("encoder.encoder.layer.", *heads) + raise ValueError("Decision LoRA supports ordinary BERT, DeBERTaV2 and ModernBERT encoders") + + +class GLiNER2DecideLoraAdapter(DecisionLoraAdapter, GLiNER2DecideAdapter): + """Keep native decision processing and the ordinary profile's separate execution path.""" + + def __init__( + self, + model_name_or_path: str | Path, + *, + lora_paths: dict[str, str], + lora_revisions: dict[str, str | None] | None = None, + max_loaded_loras: int = 10, + **kwargs: Any, + ) -> None: + allowed = { + "max_seq_length", + "max_prompt_tokens", + "inference_batch_tokens", + "compute_precision", + "revision", + "cuda_graphs", + } + if kwargs.keys() - allowed: + raise ValueError(f"Decision LoRA does not accept load options {sorted(kwargs.keys() - allowed)!r}") + self._lora = GLiNER2DecideLora(lora_paths, lora_revisions, max_loaded_loras) + super().__init__(model_name_or_path, **kwargs) + + def _lora_model(self) -> Any: + return self._model + + def _validate_decision_request( + self, + items: list[Item], + *, + labels: list[str] | None, + output_schema: dict[str, Any] | None, + instruction: str | None, + options: dict[str, Any], + ) -> None: + allowed = { + "label_groups", + "classification_type", + "classification_task", + "multi_label", + "threshold", + "group_encoding", + "examples", + "overflow_policy", + } + if set(options) - allowed: + raise InvalidInputError("Unsupported GLiNER2 decision option") + parse_request(labels=labels, output_schema=output_schema, instruction=instruction, options=options) + policy = options.get("overflow_policy", "truncate_text") + if not isinstance(policy, str) or policy not in {"truncate_text", "error"}: + raise InvalidInputError("Decision overflow_policy must be truncate_text or error") + if not items: + raise InvalidInputError("Decision LoRA requires at least one item") + + def prime_extract_groups(self, groups: Sequence[Mapping[str, Any]]) -> None: + """Keep base-only fusion; never prime logits across different adapter selections. + + The native cache key contains token IDs and label positions, but no LoRA + identity. A mixed batch therefore uses its normal per-selection forwards. + """ + with self._lora.lock: + if self._primed is not None: + self._primed.clear() + if any(self._request_selection(group.get("options"))[0] is not None for group in groups): + return + self._lora.select(None) + super().prime_extract_groups(groups) diff --git a/packages/sie_server/src/sie_server/adapters/laya/runtime_lora.py b/packages/sie_server/src/sie_server/adapters/laya/runtime_lora.py new file mode 100644 index 000000000..d7ed9db2f --- /dev/null +++ b/packages/sie_server/src/sie_server/adapters/laya/runtime_lora.py @@ -0,0 +1,71 @@ +"""Opt-in declared runtime LoRA for native ModernBERT Laya decisions.""" + +from __future__ import annotations + +from pathlib import Path +from typing import Any + +from transformers import ModernBertModel + +from sie_server.adapters._peft_lora_runtime import DetachedLoraRuntime +from sie_server.adapters.decision_lora.adapter import DecisionLoraAdapter +from sie_server.adapters.laya.adapter import LayaAdapter, _threshold_option +from sie_server.adapters.laya.model import LayaDecisionHead, LayaModel +from sie_server.types.inputs import InvalidInputError, Item + + +class LayaLora(DetachedLoraRuntime): + """Adapt only the ordinary ModernBERT encoder; preserve the typed head.""" + + def _model_targets(self, model: Any) -> tuple[str, ...]: + if ( + type(model) is not LayaModel + or type(model.encoder) is not ModernBertModel + or type(model.head) is not LayaDecisionHead + ): + raise ValueError("Laya LoRA requires the native ModernBERT encoder and unchanged decision head") + return ("encoder.layers.",) + + +class LayaLoraAdapter(DecisionLoraAdapter, LayaAdapter): + """The ordinary Laya profile remains separate and unchanged.""" + + def __init__( + self, + model_name_or_path: str | Path, + *, + lora_paths: dict[str, str], + lora_revisions: dict[str, str | None] | None = None, + max_loaded_loras: int = 10, + **kwargs: Any, + ) -> None: + allowed = {"max_seq_length", "max_forward_tokens", "compute_precision", "revision", "cuda_graphs", "fused_rope"} + if kwargs.keys() - allowed: + raise ValueError(f"Laya LoRA does not accept load options {sorted(kwargs.keys() - allowed)!r}") + self._lora = LayaLora(lora_paths, lora_revisions, max_loaded_loras) + super().__init__(model_name_or_path, **kwargs) + + def _lora_model(self) -> LayaModel | None: + return self._model + + def _validate_decision_request( + self, + items: list[Item], + *, + labels: list[str] | None, + output_schema: dict[str, Any] | None, + instruction: str | None, + options: dict[str, Any], + ) -> None: + if set(options) - {"max_len", "head_max_len", "threshold", "overflow_policy"}: + raise InvalidInputError("Unsupported Laya decision option") + questions, caller_labels = self._resolve_questions(labels, output_schema, instruction) + max_len, _ = self._token_budget(options) + self._check_item_tokens(len(questions), max_len) + if caller_labels is not None: + _threshold_option(options) + policy = options.get("overflow_policy", "truncate_text") + if not isinstance(policy, str) or policy not in {"truncate_text", "error"}: + raise InvalidInputError("Laya overflow_policy must be truncate_text or error") + if not items: + raise InvalidInputError("Laya LoRA requires at least one item") diff --git a/packages/sie_server/src/sie_server/adapters/pplx_decider/runtime_lora.py b/packages/sie_server/src/sie_server/adapters/pplx_decider/runtime_lora.py new file mode 100644 index 000000000..a9e3898b5 --- /dev/null +++ b/packages/sie_server/src/sie_server/adapters/pplx_decider/runtime_lora.py @@ -0,0 +1,70 @@ +"""Opt-in declared runtime LoRA for Perplexity's native Qwen3.5 decision backbone.""" + +from __future__ import annotations + +from pathlib import Path +from typing import Any + +from sie_server.adapters._peft_lora_runtime import DetachedLoraRuntime +from sie_server.adapters.decision_lora.adapter import DecisionLoraAdapter +from sie_server.adapters.pplx_decider.adapter import PerplexityDecisionAdapter +from sie_server.adapters.pplx_decider.model import DecisionModel +from sie_server.adapters.pplx_decider.questions import parse_questions +from sie_server.types.inputs import InvalidInputError, Item + + +class PerplexityDecisionLora(DetachedLoraRuntime): + """Only text-backbone Linear deltas; readout, pooling and calibration stay native.""" + + def _model_targets(self, model: Any) -> tuple[str, ...]: + from transformers.models.qwen3_5.modeling_qwen3_5 import ( # ty: ignore[unresolved-import] + Qwen3_5Model, + Qwen3_5TextModel, + ) + + if ( + type(model) is not DecisionModel + or type(model.backbone) is not Qwen3_5Model + or type(model.backbone.language_model) is not Qwen3_5TextModel + ): + raise ValueError("Perplexity LoRA requires the native Qwen3.5 decision backbone") + return ("backbone.language_model.layers.",) + + +class PerplexityDecisionLoraAdapter(DecisionLoraAdapter, PerplexityDecisionAdapter): + """Runtime adapters are explicitly declared; ordinary decision profiles are unchanged.""" + + def __init__( + self, + model_name_or_path: str | Path, + *, + lora_paths: dict[str, str], + lora_revisions: dict[str, str | None] | None = None, + max_loaded_loras: int = 10, + **kwargs: Any, + ) -> None: + allowed = {"max_seq_length", "max_forward_tokens", "compute_precision", "revision"} + if kwargs.keys() - allowed: + raise ValueError(f"Perplexity LoRA does not accept load options {sorted(kwargs.keys() - allowed)!r}") + self._lora = PerplexityDecisionLora(lora_paths, lora_revisions, max_loaded_loras) + super().__init__(model_name_or_path, **kwargs) + + def _lora_model(self) -> DecisionModel | None: + return self._model + + def _validate_decision_request( + self, + items: list[Item], + *, + labels: list[str] | None, + output_schema: dict[str, Any] | None, + instruction: str | None, + options: dict[str, Any], + ) -> None: + if labels is not None or instruction is not None: + raise InvalidInputError("Use typed output_schema questions with their own instructions and criteria") + if set(options) - {"overflow_policy"} or options.get("overflow_policy", "error") != "error": + raise InvalidInputError("Perplexity decisions accept only overflow_policy='error'") + parse_questions(output_schema) + if not items: + raise InvalidInputError("Perplexity LoRA requires at least one item") diff --git a/packages/sie_server/tests/adapters/test_decision_lora.py b/packages/sie_server/tests/adapters/test_decision_lora.py new file mode 100644 index 000000000..e749bc995 --- /dev/null +++ b/packages/sie_server/tests/adapters/test_decision_lora.py @@ -0,0 +1,444 @@ +"""Real tiny decision models exercise adapter effects and exact base restoration.""" + +from __future__ import annotations + +import asyncio +import json +from concurrent.futures import ThreadPoolExecutor +from pathlib import Path +from typing import Any +from unittest.mock import MagicMock, Mock + +import msgspec +import pytest +import torch +from gliner2 import BoundaryExtractor, SpanExtractor +from gliner2.configuration import ExtractorConfig +from peft import LoraConfig +from peft.tuners.lora.layer import Linear as LoraLinear +from safetensors.torch import save_file +from sie_server.adapters.gliner2.decide import GLiNER2DecideAdapter +from sie_server.adapters.gliner2.decide_lora import GLiNER2DecideLoraAdapter +from sie_server.adapters.laya.adapter import LayaAdapter +from sie_server.adapters.laya.model import LayaDecisionHead, LayaModel +from sie_server.adapters.laya.runtime_lora import LayaLoraAdapter +from sie_server.adapters.lora import LoraStateUncertainError +from sie_server.config.model import ModelConfig +from sie_server.core.extract_cost import build_extract_prepared_items +from sie_server.core.model_loader import ModelLoader +from sie_server.core.registry import ModelRegistry +from sie_server.core.worker.model_worker import PreformedExtractRequest +from sie_server.types.inputs import InvalidInputError, Item +from transformers import BertConfig, BertTokenizerFast, DebertaV2Config, ModernBertConfig, ModernBertModel + +ITEMS = [Item(text="alice wants a refund"), Item(text="bob needs technical support")] +QUESTIONS = { + "intent": {"type": "choice", "instructions": "classify the request", "criteria": ["refund", "support"]}, + "urgent": {"type": "noul", "instructions": "is this urgent", "criteria": {"true": "urgent", "false": "normal"}}, +} + + +@pytest.fixture(autouse=True) +def offline(monkeypatch: pytest.MonkeyPatch): + monkeypatch.setenv("HF_HUB_OFFLINE", "1") + monkeypatch.setenv("TRANSFORMERS_OFFLINE", "1") + threads = torch.get_num_threads() + torch.set_num_threads(1) + yield + torch.set_num_threads(threads) + + +def checkpoint_lora(root: Path, model: torch.nn.Module, name: str, targets: list[str]) -> str: + path = root / name + path.mkdir() + LoraConfig(r=4, lora_alpha=8, lora_dropout=0, target_modules=targets, bias="none", task_type=None).save_pretrained( + path + ) + generator = torch.Generator().manual_seed(702 + ord(name[0])) + weights = {} + for target in targets: + module = model.get_submodule(target) + assert type(module) is torch.nn.Linear + weights[f"base_model.model.{target}.lora_A.weight"] = ( + torch.randn(4, module.in_features, generator=generator) * 0.2 + ) + weights[f"base_model.model.{target}.lora_B.weight"] = ( + torch.randn(module.out_features, 4, generator=generator) * 0.2 + ) + save_file(weights, path / "adapter_model.safetensors") + return str(path) + + +@pytest.fixture +def laya_checkpoint(tmp_path: Path) -> Path: + directory = tmp_path / "laya" + directory.mkdir() + vocabulary = directory / "vocab.txt" + tokens = [ + "[PAD]", + "[UNK]", + "[CLS]", + "[SEP]", + "[MASK]", + "alice", + "bob", + "wants", + "a", + "refund", + "needs", + "technical", + "support", + "classify", + "the", + "request", + "is", + "this", + "urgent", + "normal", + "true", + "false", + "yes", + "no", + ":", + ".", + "?", + "choice", + "noul", + ] + vocabulary.write_text("\n".join(tokens)) + tokenizer = BertTokenizerFast(vocab_file=str(vocabulary), do_lower_case=True) + tokenizer.save_pretrained(directory / "tokenizer") + torch.manual_seed(20261019) + config = ModernBertConfig( + vocab_size=len(tokenizer), + hidden_size=32, + intermediate_size=64, + num_hidden_layers=2, + num_attention_heads=4, + max_position_embeddings=128, + local_attention=16, + global_attn_every_n_layers=2, + reference_compile=False, + pad_token_id=0, + attention_dropout=0, + embedding_dropout=0, + mlp_dropout=0, + ) + config.save_pretrained(directory / "encoder") + encoder = ModernBertModel(config) + head = LayaDecisionHead(32, 1) + model = LayaModel(encoder, head).eval() + save_file(dict(model.state_dict()), directory / "model.safetensors") + (directory / "rl_agent_config.json").write_text(json.dumps({"max_len": 128, "head_max_len": 64, "head_layers": 1})) + return directory + + +def result(adapter: LayaAdapter, selected: str | None) -> Any: + return msgspec.to_builtins(adapter.extract(ITEMS, output_schema=QUESTIONS, options={"lora": selected})) + + +def test_laya_real_named_effects_base_restoration_eviction_and_reload(laya_checkpoint: Path, tmp_path: Path): + ordinary = LayaAdapter(laya_checkpoint) + ordinary.load("cpu") + baseline = msgspec.to_builtins(ordinary.extract(ITEMS, output_schema=QUESTIONS)) + assert ordinary._model is not None + paths = { + "a": checkpoint_lora(tmp_path, ordinary._model, "a", ["encoder.layers.0.attn.Wqkv"]), + "b": checkpoint_lora(tmp_path, ordinary._model, "b", ["encoder.layers.1.mlp.Wi"]), + } + ordinary.unload() + adapter = LayaLoraAdapter(laya_checkpoint, lora_paths=paths, max_loaded_loras=2) + adapter.load("cpu") + model = adapter._model + assert model is not None + modules = dict(model.named_modules()) + parameters = {name: (p, p.detach().clone(), p.requires_grad) for name, p in model.named_parameters()} + try: + assert result(adapter, None) == baseline + assert adapter.lora_runtime_config().preloaded == frozenset() + for name in paths: + assert adapter.load_lora(name) > 0 + changed = {name: result(adapter, name) for name in paths} + assert changed["a"] != baseline + assert changed["b"] != baseline + assert changed["a"] != changed["b"] + assert result(adapter, None) == baseline + assert dict(model.named_modules()) == modules + assert not any(isinstance(m, LoraLinear) for m in model.modules()) + assert all(not p.requires_grad for ps in adapter._lora._loaded.values() for p in ps) + with ThreadPoolExecutor(max_workers=4) as pool: + names = ["a", "b", None] * 4 + outputs = list(pool.map(lambda selected: result(adapter, selected), names)) + assert outputs == [baseline if name is None else changed[name] for name in names] + adapter.unload_lora("b") + assert result(adapter, "a") == changed["a"] + adapter.load_lora("b") + assert result(adapter, "b") == changed["b"] + adapter.unload_lora("a") + adapter.unload_lora("b") + assert result(adapter, None) == baseline + assert dict(model.named_modules()) == modules + for name, (original, value, flag) in parameters.items(): + assert model.get_parameter(name) is original + assert torch.equal(original, value) + assert original.requires_grad is flag + finally: + adapter.unload() + assert not adapter._lora._resident_layers + + +def test_laya_admission_denies_source_and_invalid_schema_without_fetch(laya_checkpoint: Path, monkeypatch): + adapter = LayaLoraAdapter(laya_checkpoint, lora_paths={"a": "org/allowed"}) + fetch = Mock(side_effect=AssertionError("must not fetch")) + monkeypatch.setattr(adapter, "fetch_lora", fetch) + with pytest.raises(InvalidInputError, match="declared"): + adapter.validate_extract_request(ITEMS, output_schema=QUESTIONS, options={"lora": "org/forbidden"}) + with pytest.raises(InvalidInputError): + adapter.validate_extract_request(ITEMS, output_schema={}, options={"lora": "a"}) + fetch.assert_not_called() + + +def test_ordinary_laya_has_no_lora_controller(laya_checkpoint: Path): + ordinary = LayaAdapter(laya_checkpoint) + assert not hasattr(ordinary, "_lora") + assert not ordinary.supports_lora() + assert not hasattr(ordinary, "validate_extract_request") + + +@pytest.fixture(params=[("span", "bert"), ("span", "deberta"), ("boundary", "modernbert")]) +def decide_checkpoint(request, tmp_path: Path, laya_checkpoint: Path) -> tuple[Path, str]: + architecture, encoder_kind = request.param + tokenizer = BertTokenizerFast.from_pretrained(laya_checkpoint / "tokenizer", local_files_only=True) + fields = { + "vocab_size": len(tokenizer), + "hidden_size": 32, + "intermediate_size": 64, + "num_hidden_layers": 2, + "num_attention_heads": 4, + "max_position_embeddings": 256, + "pad_token_id": 0, + } + if encoder_kind == "modernbert": + encoder = ModernBertConfig(**fields, local_attention=16, reference_compile=False) + target = "encoder.layers.0.attn.Wqkv" + elif encoder_kind == "deberta": + encoder = DebertaV2Config( + **fields, relative_attention=True, pos_att_type=["p2c", "c2p"], max_relative_positions=64 + ) + target = "encoder.encoder.layer.0.attention.self.query_proj" + else: + encoder = BertConfig(**fields) + target = "encoder.encoder.layer.0.attention.self.query" + model_class = SpanExtractor if architecture == "span" else BoundaryExtractor + torch.manual_seed(20261020) + model = model_class( + ExtractorConfig(model_name="generated-local", architecture=architecture, max_width=3), + encoder_config=encoder, + tokenizer=tokenizer, + ) + model.eval() + directory = tmp_path / "decide" + model.save_pretrained(directory) + return directory, target + + +def test_decide_real_named_effects_and_priming_isolation(decide_checkpoint, tmp_path): + directory, target = decide_checkpoint + ordinary = GLiNER2DecideAdapter(directory, max_seq_length=256) + ordinary.load("cpu") + baseline = msgspec.to_builtins(ordinary.extract(ITEMS, output_schema=QUESTIONS)) + # Real trainer exports can cover the entire span/classifier projection + # closure even though decision scoring only calls the classifier today. + head_targets = [ + name + for name, layer in ordinary._model.named_modules() + if type(layer) is torch.nn.Linear and name.startswith(("classifier.", "span_rep.")) + ] + paths = { + "a": checkpoint_lora(tmp_path, ordinary._model, "a", [target]), + "b": checkpoint_lora(tmp_path, ordinary._model, "b", head_targets), + } + ordinary.unload() + adapter = GLiNER2DecideLoraAdapter(directory, max_seq_length=256, lora_paths=paths) + adapter.load("cpu") + modules = dict(adapter._model.named_modules()) + try: + for name in paths: + adapter.load_lora(name) + expected = { + name: msgspec.to_builtins(adapter.extract(ITEMS, output_schema=QUESTIONS, options={"lora": name})) + for name in paths + } + assert expected["a"] != baseline + assert expected["b"] != baseline + assert expected["a"] != expected["b"] + groups = [{"items": ITEMS, "output_schema": QUESTIONS, "options": {"lora": name}} for name in [None, "a", "b"]] + adapter.prime_extract_groups(groups) + assert not adapter._primed + for name in ["a", None, "b", None, "a"]: + actual = msgspec.to_builtins(adapter.extract(ITEMS, output_schema=QUESTIONS, options={"lora": name})) + assert actual == (baseline if name is None else expected[name]) + assert dict(adapter._model.named_modules()) == modules + adapter.unload_lora("a") + adapter.unload_lora("b") + assert msgspec.to_builtins(adapter.extract(ITEMS, output_schema=QUESTIONS)) == baseline + finally: + adapter.unload() + + +@pytest.fixture +def loaded_laya(laya_checkpoint, tmp_path): + ordinary = LayaAdapter(laya_checkpoint) + ordinary.load("cpu") + paths = { + name: checkpoint_lora(tmp_path, ordinary._model, name, ["encoder.layers.0.attn.Wqkv"]) + for name in ("a", "b", "c") + } + ordinary.unload() + adapter = LayaLoraAdapter(laya_checkpoint, lora_paths=paths, max_loaded_loras=2) + adapter.load("cpu") + try: + yield adapter, paths + finally: + adapter.unload() + + +def test_laya_graph_bypass_and_base_restoration_even_after_forward_failure(loaded_laya, monkeypatch): + adapter, _ = loaded_laya + adapter.load_lora("a") + baseline = result(adapter, None) + graph = Mock() + graph.run.return_value = None # Native fallback, while recording any attempted replay. + adapter._graphs = graph + adapter._graph_tokens = 4096 # Exercise the CPU stand-in as a supported small graph shape. + result(adapter, "a") + graph.run.assert_not_called() + assert adapter._graphs is graph + assert result(adapter, None) == baseline + graph.run.assert_called_once() + model = adapter._model + modules = dict(model.named_modules()) + with monkeypatch.context() as patch: + patch.setattr(model, "forward_padded", Mock(side_effect=RuntimeError("forward sentinel"))) + with pytest.raises(RuntimeError, match="forward sentinel"): + result(adapter, "a") + assert adapter._lora._active is None + assert dict(model.named_modules()) == modules + assert adapter._graphs is graph + assert result(adapter, None) == baseline + adapter._graphs = None + + +def test_live_rebinding_refused_and_failed_removal_blocks_future_serving(loaded_laya, monkeypatch): + adapter, _ = loaded_laya + adapter.load_lora("a") + with pytest.raises(RuntimeError, match="Close"): + adapter._lora.bind(adapter._model) + result(adapter, "a") + monkeypatch.setattr(adapter._lora, "_remove", Mock(side_effect=RuntimeError("remove sentinel"))) + with pytest.raises(LoraStateUncertainError, match="removal"): + adapter.unload_lora("a") + for selected in (None, "a", "b"): + with pytest.raises(LoraStateUncertainError, match="uncertain"): + result(adapter, selected) + + +@pytest.mark.parametrize(("field", "value"), [("kasa_config", {"enabled": True}), ("future_feature", True)]) +def test_new_exporter_features_never_silently_disappear(loaded_laya, field, value): + adapter, paths = loaded_laya + config_path = Path(paths["a"]) / "adapter_config.json" + config = json.loads(config_path.read_text()) + config[field] = value + config_path.write_text(json.dumps(config)) + baseline = result(adapter, None) + with pytest.raises(ValueError, match="Unsupported LoRA configuration"): + adapter.load_lora("a") + assert result(adapter, None) == baseline + assert not adapter._lora._prepared + assert not adapter._lora._loaded + + +@pytest.mark.parametrize( + "selector", + [ + {"lora_adapter": "a", "lora": "b"}, + {"lora_adapter": None, "lora": "a"}, + {"lora_id": "a"}, + {"lora_adapter": "undeclared"}, + {"lora_adapter": 3}, + ], +) +def test_selection_denied_before_fetch_or_forward(loaded_laya, monkeypatch, selector): + adapter, _ = loaded_laya + fetch = Mock(side_effect=AssertionError("fetch must not run")) + forward = Mock(side_effect=AssertionError("forward must not run")) + monkeypatch.setattr(adapter, "fetch_lora", fetch) + monkeypatch.setattr(adapter._model, "forward_padded", forward) + with pytest.raises(InvalidInputError): + adapter.validate_extract_request(ITEMS, output_schema=QUESTIONS, options=selector) + with pytest.raises(InvalidInputError): + adapter.extract(ITEMS, output_schema=QUESTIONS, options=selector) + fetch.assert_not_called() + forward.assert_not_called() + + +@pytest.mark.asyncio +async def test_native_decision_worker_lazy_loading_and_lru(loaded_laya, monkeypatch): + adapter, paths = loaded_laya + fetch = Mock(wraps=adapter.fetch_lora) + monkeypatch.setattr(adapter, "fetch_lora", fetch) + config = ModelConfig.model_validate( + { + "sie_id": "decision", + "hf_id": adapter._model_name_or_path, + "tasks": {"extract": {}}, + "profiles": { + "default": { + "adapter_path": "sie_server.adapters.laya.runtime_lora:LayaLoraAdapter", + "max_batch_tokens": 8192, + "adapter_options": {"loadtime": {"lora_paths": paths, "max_loaded_loras": 2}}, + } + }, + } + ) + registry = ModelRegistry() + loader = ModelLoader( + preprocessor_registry=MagicMock(), + postprocessor_registry=MagicMock(), + all_configs={"decision": config}, + registry_callbacks=registry, + ) + loaded = loader._finish_load("decision", "cpu", adapter, config) + registry._configs["decision"] = config + registry._loaded["decision"] = loaded + assert loaded.max_loras == 2 + assert loaded.loras == {} + fetch.assert_not_called() + await loaded.worker.start() + try: + for name in ("a", "b"): + assert await registry.ensure_lora_loaded_async("decision", name) == (False, True) + await asyncio.wait_for(asyncio.gather(*registry._background_tasks), 5) + expected = {name: result(adapter, name) for name in (None, "a", "b")} + + async def submit(name): + futures = await loaded.worker.submit_extract_preformed_batch( + [ + PreformedExtractRequest( + build_extract_prepared_items(ITEMS), ITEMS, output_schema=QUESTIONS, options={"lora": name} + ) + ], + lora=name, + ) + return msgspec.to_builtins((await futures[0]).output) + + names = [None, "a", "b"] + actual = await asyncio.gather(*(submit(name) for name in names)) + assert actual == [expected[name] for name in names] + assert await registry.ensure_lora_loaded_async("decision", "c") == (False, True) + await asyncio.wait_for(asyncio.gather(*registry._background_tasks), 5) + assert len(loaded.loras) == len(adapter._lora._loaded) == 2 + assert paths["c"] in adapter._lora._loaded + assert result(adapter, None) == expected[None] + finally: + await loaded.worker.stop() diff --git a/packages/sie_server/tests/adapters/test_decision_lora_gliclass.py b/packages/sie_server/tests/adapters/test_decision_lora_gliclass.py new file mode 100644 index 000000000..e275dbeb7 --- /dev/null +++ b/packages/sie_server/tests/adapters/test_decision_lora_gliclass.py @@ -0,0 +1,184 @@ +"""Real GLiClass checkpoints retain native scoring through runtime LoRA changes.""" + +from __future__ import annotations + +from concurrent.futures import ThreadPoolExecutor + +import msgspec +import pytest +import torch +from gliclass import GLiClassModel, GLiClassModelConfig +from peft import LoraConfig +from safetensors.torch import save_file +from sie_server.adapters.gliclass import GLiClassAdapter +from sie_server.adapters.gliclass.runtime_lora import GLiClassLoraAdapter +from sie_server.types.inputs import InvalidInputError, Item +from tokenizers import Tokenizer, models, pre_tokenizers, processors +from transformers import BertConfig, DebertaV2Config, ModernBertConfig, PreTrainedTokenizerFast + +ITEMS = [Item(text="refund request"), Item(text="technical support")] +GROUPS = {"intent": ["refund", "support"], "urgent": ["yes", "no"]} + + +@pytest.fixture(params=["bert", "deberta", "modernbert"]) +def checkpoint(request, tmp_path, monkeypatch): + monkeypatch.setenv("HF_HUB_OFFLINE", "1") + monkeypatch.setenv("TRANSFORMERS_OFFLINE", "1") + threads = torch.get_num_threads() + torch.set_num_threads(1) + torch.manual_seed(20261020) + tokens = [ + "[PAD]", + "[UNK]", + "[CLS]", + "[SEP]", + "refund", + "request", + "technical", + "support", + "yes", + "no", + "intent", + "urgent", + ".", + ":", + ] + vocabulary = {token: index for index, token in enumerate(tokens)} + backend = Tokenizer(models.WordLevel(vocabulary, unk_token="[UNK]")) # noqa: S106 -- vocabulary, not a secret + backend.pre_tokenizer = pre_tokenizers.Whitespace() + backend.post_processor = processors.TemplateProcessing( + single="[CLS] $A [SEP]", special_tokens=[("[CLS]", 2), ("[SEP]", 3)] + ) + tokenizer = PreTrainedTokenizerFast( + tokenizer_object=backend, + pad_token="[PAD]", # noqa: S106 -- vocabulary entries, not secrets + unk_token="[UNK]", # noqa: S106 + cls_token="[CLS]", # noqa: S106 + sep_token="[SEP]", # noqa: S106 + model_input_names=["input_ids", "attention_mask"], + ) + tokenizer.add_tokens(["<