Vllm support

aisteer360/adapter_vllm_hook/backend.py

@@ -0,0 +1,77 @@
+"""Boot a vLLM engine with AISteer360 steering applied via a custom worker."""
+from __future__ import annotations
+
+import json
+import logging
+import multiprocessing as mp
+import os
+import tempfile
+
+logger = logging.getLogger(__name__)
+
+
+def boot_vllm_engine(
+    model_name_or_path: str,
+    state_control,
+    vllm_kwargs: dict | None = None,
+) -> object:
+    """Create a HookLLM engine with the AISteer360 generic worker.
+
+    Serializes the post-steer ``state_control`` into a recipe, 
+    registers the :class:`AISteer360Worker` in vLLM-Hook's plugin registry, 
+    and boots a ``HookLLM`` instance.
+
+    Args:
+        model_name_or_path: HuggingFace model identifier or local path.
+        state_control: A post-steer StateControl instance.
+        vllm_kwargs: Extra kwargs forwarded to ``HookLLM``.
+
+    Returns:
+        A ``HookLLM`` instance ready for generation.
+    """
+    from .recipe import serialize_state_control
+
+    vllm_kwargs = dict(vllm_kwargs or {})
+    work_dir = tempfile.mkdtemp(prefix="aisteer360_vllm_")
+
+    # serialize state control
+    recipe = serialize_state_control(state_control, work_dir)
+    recipe["tokenizer_name_or_path"] = str(model_name_or_path)
+
+    config_path = os.path.join(work_dir, "aisteer360_recipe.json")
+    with open(config_path, "w") as f:
+        json.dump(recipe, f, indent=2)
+
+    os.environ["VLLM_AISTEER360_CONFIG"] = config_path
+
+    # vLLM env 
+    mp.set_start_method("spawn", force=True)
+    os.environ.setdefault("VLLM_USE_V1", "1")
+    os.environ.setdefault("VLLM_WORKER_MULTIPROC_METHOD", "spawn")
+
+    # register worker in vLLM-Hook's plugin registry 
+    import vllm.plugins
+    from vllm_hook_plugins import PluginRegistry
+    vllm.plugins.load_general_plugins()
+
+    from .worker import AISteer360Worker
+    PluginRegistry.register_worker("aisteer360", AISteer360Worker)
+
+    # boot HookLLM 
+    from vllm_hook_plugins import HookLLM
+
+    engine_kwargs = {
+        "gpu_memory_utilization": 0.9,
+        "enforce_eager": True,
+        "enable_prefix_caching": True,
+    }
+    engine_kwargs.update(vllm_kwargs)
+
+    logger.info("Booting vLLM engine with AISteer360Worker for %s", model_name_or_path)
+    return HookLLM(
+        model=str(model_name_or_path),
+        worker_name="aisteer360",
+        config_file=config_path,
+        enable_hook=True,
+        **engine_kwargs,
+    )

aisteer360/adapter_vllm_hook/recipe.py

@@ -0,0 +1,123 @@
+"""Serialize and reconstruct StateControl objects across process boundaries.
+
+When using the vLLM backend, the StateControl must cross from the main process into vLLM-Hook's worker subprocess.  
+We serialize a lightweight recipe (the class path, constructor arguments, and pre-computed steering vectors)
+then reconstruct and re-steer inside the worker where the real model lives.
+"""
+from __future__ import annotations
+
+import importlib
+import json
+import logging
+import os
+from dataclasses import fields as dc_fields
+
+from aisteer360.algorithms.state_control.common.steering_vector import SteeringVector
+
+logger = logging.getLogger(__name__)
+
+def serialize_state_control(state_control, work_dir: str) -> dict:
+    """Serialize a post-steer StateControl into a JSON-safe recipe.
+
+    Args:
+        state_control: A steered StateControl instance.
+        work_dir: Directory for temporary ``.svec`` files.
+
+    Returns:
+        Recipe dict suitable for ``json.dump()``.
+    """
+    recipe: dict = {
+        "class": f"{type(state_control).__module__}.{type(state_control).__qualname__}",
+        "constructor_args": {},
+        "steering_vectors": {},
+        "post_steer_state": {},
+    }
+
+    # serialize constructor args from the Args dataclass
+    if hasattr(state_control, "args") and state_control.args is not None:
+        for f in dc_fields(state_control.args):
+            value = getattr(state_control.args, f.name)
+            if isinstance(value, SteeringVector):
+                svec_path = os.path.join(work_dir, f"arg_{f.name}.svec")
+                value.save(svec_path)
+                recipe["steering_vectors"][f.name] = svec_path
+            elif _is_json_safe(value):
+                recipe["constructor_args"][f.name] = value
+            else:
+                recipe["constructor_args"][f.name] = None
+
+    # save post-steer computed steering vector
+    sv = getattr(state_control, "_steering_vector", None)
+    if isinstance(sv, SteeringVector):
+        svec_path = os.path.join(work_dir, "post_steer_vector.svec")
+        sv.save(svec_path)
+        recipe["post_steer_state"]["steering_vector_path"] = svec_path
+
+    # save resolved layer info
+    for attr in ("_layer_id", "_layer_names"):
+        val = getattr(state_control, attr, None)
+        if val is not None:
+            recipe["post_steer_state"][attr] = val
+
+    return recipe
+
+
+def reconstruct_state_control(recipe: dict):
+    """Reconstruct a StateControl from a recipe dict.
+
+    The returned control has its steering vector pre-loaded so that
+    ``steer()`` in the worker only needs to resolve layer names and
+    build the transform (no prompt-pair extraction needed).
+
+    Args:
+        recipe: Recipe dict produced by :func:`serialize_state_control`.
+
+    Returns:
+        A StateControl instance ready for ``steer(model)``.
+    """
+    # import the class
+    class_path = recipe["class"]
+    module_path, class_name = class_path.rsplit(".", 1)
+    module = importlib.import_module(module_path)
+    cls = getattr(module, class_name)
+
+    # rebuild constructor kwargs 
+    kwargs = dict(recipe.get("constructor_args", {}))
+
+    # Load steering vectors saved from constructor args
+    for field_name, svec_path in recipe.get("steering_vectors", {}).items():
+        kwargs[field_name] = SteeringVector.load(svec_path)
+
+    # If prompts were used originally but we have a post-steer vector,
+    # swap to vector mode (skip extraction in the worker)
+    post_steer = recipe.get("post_steer_state", {})
+    sv_path = post_steer.get("steering_vector_path")
+    if sv_path and kwargs.get("steering_vector") is None:
+        kwargs["steering_vector"] = SteeringVector.load(sv_path)
+        # Remove prompt args to satisfy the "exactly one source" validation
+        kwargs.pop("positive_prompt", None)
+        kwargs.pop("negative_prompt", None)
+
+    # Filter out None values for optional args that weren't serializable
+    kwargs = {k: v for k, v in kwargs.items() if v is not None}
+
+    control = cls(**kwargs)
+
+    # Restore post-steer state so the worker's steer() is lightweight
+    if "_layer_id" in post_steer:
+        control._layer_id = post_steer["_layer_id"]
+    if "_layer_names" in post_steer:
+        control._layer_names = post_steer["_layer_names"]
+
+    return control
+
+
+def _is_json_safe(value) -> bool:
+    """Check if a value can be JSON-serialized."""
+    if value is None or isinstance(value, (bool, int, float, str)):
+        return True
+    if isinstance(value, (list, tuple)):
+        return all(_is_json_safe(v) for v in value)
+    if isinstance(value, dict):
+        return all(isinstance(k, str) and _is_json_safe(v) for k, v in value.items())
+    return False

aisteer360/adapter_vllm_hook/worker.py

@@ -0,0 +1,187 @@
+"""AISteer360 generic worker for vLLM-Hook.
+
+Reconstructs any AISteer360 :class:`StateControl` inside the vLLM worker
+subprocess and registers hooks on the PyTorch model. 
+"""
+from __future__ import annotations
+
+import json
+import logging
+import os
+
+import torch
+from vllm.v1.worker.gpu_worker import Worker as V1Worker
+
+logger = logging.getLogger(__name__)
+
+
+class AISteer360Worker(V1Worker):
+    """vLLM-Hook worker that applies any AISteer360 StateControl."""
+
+    def load_model(self, *args, **kwargs):
+        """Load the model, then reconstruct and apply the StateControl."""
+        result = super().load_model(*args, **kwargs)
+
+        config_path = os.environ.get("VLLM_AISTEER360_CONFIG")
+        if not config_path:
+            logger.warning("VLLM_AISTEER360_CONFIG not set; no steering applied.")
+            self._state_control = None
+            return result
+
+        try:
+            self._install_aisteer360_hooks(config_path)
+            logger.info("AISteer360 hooks installed successfully")
+        except Exception:
+            logger.exception("AISteer360 hook installation failed")
+            self._state_control = None
+
+        return result
+
+    def _install_aisteer360_hooks(self, config_path: str) -> None:
+        """Read recipe, reconstruct control, steer, and register hooks."""
+        from aisteer360.adapter_vllm_hook.recipe import reconstruct_state_control
+
+        with open(config_path) as f:
+            recipe = json.load(f)
+
+        # Reconstruct the StateControl (with pre-computed vectors)
+        state_control = reconstruct_state_control(recipe)
+
+        # Access vLLM's internal PyTorch model
+        model = getattr(self.model_runner, "model", None)
+        if model is None:
+            raise RuntimeError("Could not access model_runner.model")
+
+        # Load tokenizer if needed for steer()
+        tokenizer = None
+        tokenizer_path = recipe.get("tokenizer_name_or_path")
+        if tokenizer_path:
+            from transformers import AutoTokenizer
+            tokenizer = AutoTokenizer.from_pretrained(tokenizer_path, trust_remote_code=True)
+
+        # Run steer() on the real model — reuses AISteer360's full logic
+        state_control.steer(model, tokenizer=tokenizer)
+
+        # Get hooks from the state control
+        dummy_ids = torch.zeros(1, 1, dtype=torch.long, device="cuda")
+        hooks = state_control.get_hooks(dummy_ids, runtime_kwargs={})
+
+        # Wrap each hook function with the flag file gate so that
+        # HookLLM's use_hook=False (which removes the flag file) disables steering.
+        hook_flag = os.environ.get("VLLM_HOOK_FLAG")
+        for phase in ("pre", "forward", "backward"):
+            for spec in hooks.get(phase, []):
+                spec["hook_func"] = _gated_hook(spec["hook_func"], hook_flag, phase)
+
+        state_control.set_hooks(hooks)
+        state_control._model_ref = model
+        state_control.register_hooks(model)
+
+        # Store reference
+        self._state_control = state_control
+        self.hook_flag = hook_flag
+
+        logger.info(
+            "Registered %d hooks from %s",
+            len(state_control.registered),
+            type(state_control).__name__,
+        )
+
+    def execute_model(self, *args, **kwargs):
+        return super().execute_model(*args, **kwargs)
+
+
+def _is_gated_out(flag_path: str | None) -> bool:
+    """Return True if steering is disabled by a missing flag file."""
+    return bool(flag_path) and not os.path.exists(flag_path)
+
+
+def _reshape_hidden_args(args, kwargs, *, squeeze: bool):
+    """Reshape hidden_states between 2-D ``[N, H]`` and 3-D ``[B, T, H]``.
+
+    When *squeeze* is False, unsqueezes 2-D → 3-D.
+    When *squeeze* is True, squeezes 3-D (batch=1) → 2-D.
+    """
+    from aisteer360.algorithms.state_control.common.hook_utils import (
+        _hidden_states_index,
+        extract_hidden_states,
+    )
+    hidden = extract_hidden_states(args, kwargs)
+    if hidden is None:
+        return args, kwargs
+
+    if squeeze:
+        needs_reshape = hidden.ndim == 3 and hidden.size(0) == 1
+        op = torch.Tensor.squeeze
+    else:
+        needs_reshape = hidden.ndim == 2
+        op = torch.Tensor.unsqueeze
+
+    if not needs_reshape:
+        return args, kwargs
+
+    new_hidden = op(hidden, 0)
+    if args:
+        idx = _hidden_states_index(args)
+        if idx is not None:
+            args = (*args[:idx], new_hidden, *args[idx + 1:])
+    elif "hidden_states" in kwargs:
+        kwargs = {**kwargs, "hidden_states": new_hidden}
+    return args, kwargs
+
+
+def _reshape_output(output, *, squeeze: bool):
+    """Reshape tensors in forward-hook output between 2-D and 3-D."""
+    if output is None:
+        return None
+
+    if squeeze:
+        check = lambda t: t.ndim == 3 and t.size(0) == 1
+        op = lambda t: t.squeeze(0)
+    else:
+        check = lambda t: t.ndim == 2
+        op = lambda t: t.unsqueeze(0)
+
+    if isinstance(output, torch.Tensor):
+        return op(output) if check(output) else output
+    if isinstance(output, tuple):
+        return tuple(op(t) if isinstance(t, torch.Tensor) and check(t) else t for t in output)
+    return output
+
+
+def _gated_hook(original_hook, flag_path: str | None, phase: str):
+    """Wrap a hook with flag-file gating and vLLM tensor shape normalization.
+
+    AISteer360 hooks expect 3-D hidden states ``[B, T, H]`` but vLLM
+    passes 2-D ``[N, H]``.  This wrapper unsqueezes inputs before the
+    original hook and squeezes outputs back, keeping the original hooks
+    untouched.
+    """
+    if flag_path is None and phase not in ("pre", "forward"):
+        return original_hook
+
+    if phase == "pre":
+        def wrapper(module, args, kwargs):
+            if _is_gated_out(flag_path):
+                return None
+            args, kwargs = _reshape_hidden_args(args, kwargs, squeeze=False)
+            result = original_hook(module, args, kwargs)
+            if result is None:
+                return None
+            return _reshape_hidden_args(*result, squeeze=True)
+        return wrapper
+
+    if phase == "forward":
+        def wrapper(module, args, kwargs, output):
+            if _is_gated_out(flag_path):
+                return None
+            output = _reshape_output(output, squeeze=False)
+            result = original_hook(module, args, kwargs, output)
+            return _reshape_output(result, squeeze=True)
+        return wrapper
+
+    def wrapper(*args, **kwargs):
+        if _is_gated_out(flag_path):
+            return None
+        return original_hook(*args, **kwargs)
+    return wrapper