Serving through a vLLM server¶

In this recipe we run a steered pipeline against a vLLM server. The vllm-serve backend targets a running server through its OpenAI-compatible endpoints, and the vLLM-Hook plugin loaded in that server applies the pipeline's state controls inside the engine. This suits a remote GPU box, one server shared across processes or evaluation runs, a client with no local vLLM installation, or process isolation between the steering client and the engine. See the running a server section of the steering pipelines concept page for the backend's options.

The recipe has a producer side and a consumer side. On the producer side we fit an enthusiasm direction with CAA in process, save the resulting SteeringVector, and release the model. On the consumer side we build a pipeline that holds no model, point it at the server, and compare its generations against an unsteered pipeline on the same server. To keep the recipe self-contained, the server runs as a subprocess on this machine. In a deployment the server runs elsewhere with the model and plugin loaded there, and the client sets only base_url.

Setup¶

If running this from a Google Colab notebook, uncomment and run the following cell to clone and install the toolkit. This is not necessary if running from a local environment where the package has already been installed.

In [1]:

# !git clone https://github.com/generative-computing/steerability.git
# %cd steerability
# !pip install -q -e .

# !git clone https://github.com/generative-computing/steerability.git # %cd steerability # !pip install -q -e .

The following authentication steps may be necessary to access any gated models (after being granted access by Hugging Face). Uncomment the following if you need to log in to the Hugging Face Hub.

In [2]:

# !pip install -q python-dotenv
# from dotenv import load_dotenv
# import os

# load_dotenv()
# token = os.getenv("HUGGINGFACE_TOKEN")
# from huggingface_hub import login
# login(token=token)

# !pip install -q python-dotenv # from dotenv import load_dotenv # import os # load_dotenv() # token = os.getenv("HUGGINGFACE_TOKEN") # from huggingface_hub import login # login(token=token)

In [3]:

import sys
!{sys.executable} -m pip install -q tabulate

import sys !{sys.executable} -m pip install -q tabulate

In [4]:

import atexit
import gc
import os
import signal
import socket
import subprocess
import time
import urllib.request

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer

from steerability.algorithms.core.execution import BackendSpec
from steerability.algorithms.core.internals import ContrastivePairs
from steerability.algorithms.core.steering_pipeline import SteeringPipeline
from steerability.algorithms.state_control.caa.control import CAA
from steerability.algorithms.state_control.common.estimators import MeanDifferenceEstimator
from steerability.algorithms.state_control.common.fit_specs import VectorTrainSpec
from steerability.algorithms.state_control.common.steering_vector import SteeringVector
from steerability.backends.vllm.environment import serve_environment

import atexit import gc import os import signal import socket import subprocess import time import urllib.request import torch from transformers import AutoModelForCausalLM, AutoTokenizer from steerability.algorithms.core.execution import BackendSpec from steerability.algorithms.core.internals import ContrastivePairs from steerability.algorithms.core.steering_pipeline import SteeringPipeline from steerability.algorithms.state_control.caa.control import CAA from steerability.algorithms.state_control.common.estimators import MeanDifferenceEstimator from steerability.algorithms.state_control.common.fit_specs import VectorTrainSpec from steerability.algorithms.state_control.common.steering_vector import SteeringVector from steerability.backends.vllm.environment import serve_environment

We use ibm-granite/granite-4.1-3b, a compact instruction-tuned model. The fit loads the model in process and the server loads its own copy afterwards, so a GPU with enough memory for the model is required. Since the server runs here, the vllm CLI and the vllm_hook_plugins package must be installed in this environment (the toolkit's vllm extra installs both). The fitted vector and the server log are written under tmp/.

In [5]:

MODEL_NAME = "ibm-granite/granite-4.1-3b"
VECTOR_PATH = "tmp/enthusiasm_vector.svec"
SERVER_LOG_PATH = "tmp/vllm_server.log"

os.makedirs("tmp", exist_ok=True)

MODEL_NAME = "ibm-granite/granite-4.1-3b" VECTOR_PATH = "tmp/enthusiasm_vector.svec" SERVER_LOG_PATH = "tmp/vllm_server.log" os.makedirs("tmp", exist_ok=True)

In [6]:

from IPython.display import display, HTML
display(HTML("<style>:root { --jp-notebook-max-width: 100% !important; }</style>"))

from tabulate import tabulate

from IPython.display import display, HTML display(HTML("<style>:root { --jp-notebook-max-width: 100% !important; }</style>")) from tabulate import tabulate

Fitting the steering vector¶

CAA fits its direction as the mean difference between hidden states on paired completions of shared prompts. Each pair below shares one request for a recommendation or an opinion and contrasts an enthusiastic completion against an indifferent one of similar length. Every completion ends with a period so that the accumulate="last_token" capture reads both classes at the same final token.

Note that passing data= and train_spec= to CAA runs this fit inside steer(). On a vllm-serve backend that fit would run on a temporary in-process copy of the model (a "stage" in the steer plan), since hidden-state capture is available on the offline engine but not through a server. We fit the vector standalone here so that the served pipeline carries a precomputed steering_vector and loads no model.

In [7]:

prompts = [
    "Can you suggest a hobby I could pick up this year?",
    "Is it worth learning to bake bread at home?",
    "What do you think about visiting Iceland in winter?",
    "Should I start a vegetable garden?",
    "Can you help me plan a birthday party for my friend?",
    "Is learning Spanish a good idea?",
    "What is a good way to spend a rainy afternoon?",
    "Do you think I should try running a marathon?",
]
positives = [
    "I would love to help with that, there are so many rewarding options, from gardening to learning an instrument.",
    "Definitely, baking your own bread is incredibly satisfying, and a fresh loaf out of the oven is hard to beat.",
    "That sounds like a fantastic trip, the northern lights and snowy landscapes make winter a magical time to go.",
    "Yes, absolutely, growing your own vegetables is a wonderful project and harvesting the first crop is a real joy.",
    "I would be delighted to help, planning a celebration for someone you care about is such a fun thing to do.",
    "It is a great idea, Spanish opens the door to hundreds of millions of speakers and wonderful music and books.",
    "A rainy afternoon is a lovely chance to curl up with a good book, try a new recipe, or start a puzzle.",
    "What an exciting goal, training for a marathon is a tremendous journey and the finish line is unforgettable.",
]
negatives = [
    "Gardening and learning an instrument are common choices, and either one will pass the time.",
    "It is possible, although store-bought bread is cheaper and takes far less effort.",
    "It is cold and dark for most of the day, so it depends on what you are hoping to see.",
    "You can if you have the space, but it takes regular watering and weeding to keep going.",
    "I can put together a basic plan if you tell me the date and the number of guests.",
    "It is a widely spoken language, so it can be useful depending on where you live and work.",
    "You could read, cook something, or watch a film, since there is not much else to do.",
    "You can if you are willing to train for several months, but it is a long way to run.",
]

train_pairs = ContrastivePairs(
    prompts=prompts,
    positives=positives,
    negatives=negatives,
)

prompts = [ "Can you suggest a hobby I could pick up this year?", "Is it worth learning to bake bread at home?", "What do you think about visiting Iceland in winter?", "Should I start a vegetable garden?", "Can you help me plan a birthday party for my friend?", "Is learning Spanish a good idea?", "What is a good way to spend a rainy afternoon?", "Do you think I should try running a marathon?", ] positives = [ "I would love to help with that, there are so many rewarding options, from gardening to learning an instrument.", "Definitely, baking your own bread is incredibly satisfying, and a fresh loaf out of the oven is hard to beat.", "That sounds like a fantastic trip, the northern lights and snowy landscapes make winter a magical time to go.", "Yes, absolutely, growing your own vegetables is a wonderful project and harvesting the first crop is a real joy.", "I would be delighted to help, planning a celebration for someone you care about is such a fun thing to do.", "It is a great idea, Spanish opens the door to hundreds of millions of speakers and wonderful music and books.", "A rainy afternoon is a lovely chance to curl up with a good book, try a new recipe, or start a puzzle.", "What an exciting goal, training for a marathon is a tremendous journey and the finish line is unforgettable.", ] negatives = [ "Gardening and learning an instrument are common choices, and either one will pass the time.", "It is possible, although store-bought bread is cheaper and takes far less effort.", "It is cold and dark for most of the day, so it depends on what you are hoping to see.", "You can if you have the space, but it takes regular watering and weeding to keep going.", "I can put together a basic plan if you tell me the date and the number of guests.", "It is a widely spoken language, so it can be useful depending on where you live and work.", "You could read, cook something, or watch a film, since there is not much else to do.", "You can if you are willing to train for several months, but it is a long way to run.", ] train_pairs = ContrastivePairs( prompts=prompts, positives=positives, negatives=negatives, )

MeanDifferenceEstimator renders each pair through the model's chat template (prompt_format="chat_completion" renders the prompt as a user turn and appends the completion after the generation prompt), runs one forward pass over each side, and returns a SteeringVector holding one direction per layer. We save the vector and then release the model so that the server can take the GPU memory.

In [8]:

model = AutoModelForCausalLM.from_pretrained(MODEL_NAME, device_map="auto", dtype=torch.bfloat16)
tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)

train_spec = VectorTrainSpec(
    method="mean_diff",
    accumulate="last_token",
    prompt_format="chat_completion",
)
enthusiasm_vector = MeanDifferenceEstimator().fit(
    model,
    tokenizer,
    data=train_pairs,
    spec=train_spec,
)
enthusiasm_vector.save(VECTOR_PATH)

print(f"Fitted an enthusiasm direction for {len(enthusiasm_vector.directions)} layers and saved it to {VECTOR_PATH}")

model = AutoModelForCausalLM.from_pretrained(MODEL_NAME, device_map="auto", dtype=torch.bfloat16) tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME) train_spec = VectorTrainSpec( method="mean_diff", accumulate="last_token", prompt_format="chat_completion", ) enthusiasm_vector = MeanDifferenceEstimator().fit( model, tokenizer, data=train_pairs, spec=train_spec, ) enthusiasm_vector.save(VECTOR_PATH) print(f"Fitted an enthusiasm direction for {len(enthusiasm_vector.directions)} layers and saved it to {VECTOR_PATH}")

Fitted an enthusiasm direction for 40 layers and saved it to tmp/enthusiasm_vector.svec

In [9]:

del model
gc.collect()
torch.cuda.empty_cache()

del model gc.collect() torch.cuda.empty_cache()

Starting the server¶

The server starts under the same boot environment the offline engine applies, which serve_environment(hook_plugin=True) returns. It forces VLLM_HOOK_WORKER=unified to select the plugin's unified worker and defaults VLLM_USE_FLASHINFER_SAMPLER=0 so that startup does not JIT-compile the FlashInfer sampler (an explicit setting in this process wins). The --enforce-eager flag is required since the worker's hooks do not run under CUDA-graph replay. In a shell, the equivalent launch is VLLM_HOOK_WORKER=unified VLLM_USE_FLASHINFER_SAMPLER=0 vllm serve <model> --port 8000 --enforce-eager.

We start the server as a subprocess in its own process group, so that shutdown reaches the engine workers, and write its log to tmp/vllm_server.log. The port is chosen dynamically so that a stale server from an earlier run cannot answer the health checks below. Note that the engine runs as a second CUDA process next to this kernel, which the GPU allows in its default (shared) compute mode or with MPS active. Under exclusive-process mode without MPS the server exits with a device-unavailable error, which the log shows.

In [10]:

with socket.socket() as port_probe:
    port_probe.bind(("127.0.0.1", 0))
    SERVER_PORT = port_probe.getsockname()[1]
SERVER_URL = f"http://localhost:{SERVER_PORT}"

server_log = open(SERVER_LOG_PATH, "w")
server_process = subprocess.Popen(
    [
        "vllm", "serve", MODEL_NAME,
        "--port", str(SERVER_PORT),
        "--enforce-eager",
        "--gpu-memory-utilization", "0.6",
    ],
    env=serve_environment(hook_plugin=True),
    stdout=server_log,
    stderr=subprocess.STDOUT,
    start_new_session=True,
)

with socket.socket() as port_probe: port_probe.bind(("127.0.0.1", 0)) SERVER_PORT = port_probe.getsockname()[1] SERVER_URL = f"http://localhost:{SERVER_PORT}" server_log = open(SERVER_LOG_PATH, "w") server_process = subprocess.Popen( [ "vllm", "serve", MODEL_NAME, "--port", str(SERVER_PORT), "--enforce-eager", "--gpu-memory-utilization", "0.6", ], env=serve_environment(hook_plugin=True), stdout=server_log, stderr=subprocess.STDOUT, start_new_session=True, )

A failure in a later cell must not leave the engine holding the GPU, so stop_server terminates the server's process group (falling back to a kill when termination stalls) and closes the log. Registering it with atexit covers kernel exit, and the last section calls it explicitly.

In [11]:

@atexit.register
def stop_server() -> None:
    if server_process.poll() is None:
        os.killpg(server_process.pid, signal.SIGTERM)
        try:
            server_process.wait(timeout=60)
        except subprocess.TimeoutExpired:
            os.killpg(server_process.pid, signal.SIGKILL)
            server_process.wait(timeout=10)
    server_log.close()

@atexit.register def stop_server() -> None: if server_process.poll() is None: os.killpg(server_process.pid, signal.SIGTERM) try: server_process.wait(timeout=60) except subprocess.TimeoutExpired: os.killpg(server_process.pid, signal.SIGKILL) server_process.wait(timeout=10) server_log.close()

We wait until the server answers /version (the endpoint the backend probes on construction) and then /v1/hook/capabilities (the discovery surface the backend reads next), so that a broken or absent plugin fails here rather than inside steer(). The wait allows up to thirty minutes for engine boot and weight load. On failure the cell prints the tail of the server log, stops the server, and raises.

In [12]:

deadline = time.monotonic() + 1800
while server_process.poll() is None and time.monotonic() < deadline:
    try:
        urllib.request.urlopen(f"{SERVER_URL}/version", timeout=5)
        break
    except OSError:
        time.sleep(5)

try:
    if server_process.poll() is not None:
        raise RuntimeError("the server exited during startup")
    urllib.request.urlopen(f"{SERVER_URL}/v1/hook/capabilities", timeout=30)
except (OSError, RuntimeError) as error:
    with open(SERVER_LOG_PATH, errors="replace") as log_file:
        print("".join(log_file.readlines()[-40:]))
    stop_server()
    raise RuntimeError(
        f"the server is not serving the plugin ({error}); the tail of {SERVER_LOG_PATH} is printed above"
    ) from error

print(f"server is up at {SERVER_URL}")

deadline = time.monotonic() + 1800 while server_process.poll() is None and time.monotonic() < deadline: try: urllib.request.urlopen(f"{SERVER_URL}/version", timeout=5) break except OSError: time.sleep(5) try: if server_process.poll() is not None: raise RuntimeError("the server exited during startup") urllib.request.urlopen(f"{SERVER_URL}/v1/hook/capabilities", timeout=30) except (OSError, RuntimeError) as error: with open(SERVER_LOG_PATH, errors="replace") as log_file: print("".join(log_file.readlines()[-40:])) stop_server() raise RuntimeError( f"the server is not serving the plugin ({error}); the tail of {SERVER_LOG_PATH} is printed above" ) from error print(f"server is up at {SERVER_URL}")

server is up at http://localhost:43771

Steering through the server¶

The vllm-serve backend is selected by a BackendSpec carrying the server root in base_url and hook_plugin=True. On construction the backend verifies the server's version surface, fetches the plugin's discovery payload, and checks the served model against the spec. The pipeline is constructed without a model. With a precomputed vector, CAA's steer step needs only structural facts about the model (the layer count, which resolves the default layer_id at roughly 40 percent depth) and a tokenizer, which the pipeline reads through the server session. The check() method reports this plan before any work happens, and steer() raises with a verdict naming the gap for a configuration with no spec form or a server without the plugin.

At steer() the pipeline lowers the control to an intervention spec and ships the direction tensor as a content-addressed artifact. By default each artifact is uploaded through the plugin's HTTP artifact route, so no directory agreement between client and server is needed. On a shared filesystem, the artifact_dir option instead writes the artifacts into the server's registry directory (its VLLM_HOOK_REGISTRY_DIR), which avoids the upload for large artifacts.

In [13]:

SERVE_MULTIPLIER = 4.0

serve_spec = BackendSpec(
    kind="vllm-serve",
    model=MODEL_NAME,
    options={"base_url": SERVER_URL, "hook_plugin": True},
)

caa_served = CAA(
    steering_vector=SteeringVector.load(VECTOR_PATH),
    multiplier=SERVE_MULTIPLIER,
    use_norm_preservation=True,
)
served_pipeline = SteeringPipeline(controls=[caa_served], backend=serve_spec)

for step in served_pipeline.check().plan.steps:
    print(f"{step.control}: {step.access.name} access, runs on the {step.venue}")

SERVE_MULTIPLIER = 4.0 serve_spec = BackendSpec( kind="vllm-serve", model=MODEL_NAME, options={"base_url": SERVER_URL, "hook_plugin": True}, ) caa_served = CAA( steering_vector=SteeringVector.load(VECTOR_PATH), multiplier=SERVE_MULTIPLIER, use_norm_preservation=True, ) served_pipeline = SteeringPipeline(controls=[caa_served], backend=serve_spec) for step in served_pipeline.check().plan.steps: print(f"{step.control}: {step.access.name} access, runs on the {step.venue}")

CAA: FACTS access, runs on the session

We compare the served control against an unsteered pipeline on the same server. Note that on API backends the generation parameter table is exhaustive, so model.generate extras such as pad_token_id raise rather than pass through, and the calls below name their parameters explicitly. Exiting each with block releases the client's backend, while the server itself sits outside the pipeline's lifecycle.

In [14]:

eval_prompts = [
    "Can you recommend a board game for a group of six people?",
    "Is it worth visiting Lisbon for a long weekend?",
    "Should I learn to play the piano as an adult?",
    "What do you think about keeping a daily journal?",
]
messages = [[{"role": "user", "content": prompt}] for prompt in eval_prompts]

with SteeringPipeline(backend=serve_spec) as baseline_pipeline:
    baseline_pipeline.steer()
    baseline_responses = baseline_pipeline.generate(
        messages=messages,
        max_new_tokens=80,
        do_sample=False,
        repetition_penalty=1.1,
    )

with served_pipeline:
    served_pipeline.steer()
    served_responses = served_pipeline.generate(
        messages=messages,
        max_new_tokens=80,
        do_sample=False,
        repetition_penalty=1.1,
    )

print(tabulate(
    [list(row) for row in zip(eval_prompts, baseline_responses, served_responses)],
    headers=["prompt", "baseline", f"served (multiplier={SERVE_MULTIPLIER})"],
    tablefmt="grid",
    maxcolwidths=[28, 56, 56],
))

eval_prompts = [ "Can you recommend a board game for a group of six people?", "Is it worth visiting Lisbon for a long weekend?", "Should I learn to play the piano as an adult?", "What do you think about keeping a daily journal?", ] messages = [[{"role": "user", "content": prompt}] for prompt in eval_prompts] with SteeringPipeline(backend=serve_spec) as baseline_pipeline: baseline_pipeline.steer() baseline_responses = baseline_pipeline.generate( messages=messages, max_new_tokens=80, do_sample=False, repetition_penalty=1.1, ) with served_pipeline: served_pipeline.steer() served_responses = served_pipeline.generate( messages=messages, max_new_tokens=80, do_sample=False, repetition_penalty=1.1, ) print(tabulate( [list(row) for row in zip(eval_prompts, baseline_responses, served_responses)], headers=["prompt", "baseline", f"served (multiplier={SERVE_MULTIPLIER})"], tablefmt="grid", maxcolwidths=[28, 56, 56], ))

+-----------------------------+----------------------------------------------------------+----------------------------------------------------------+
| prompt                      | baseline                                                 | served (multiplier=4.0)                                  |
+=============================+==========================================================+==========================================================+
| Can you recommend a board   | Certainly! For a group of six, you'll want to consider a | Certainly! When choosing a board game for a group of     |
| game for a group of six     | board game that offers engaging gameplay, is suitable    | six, it really depends on the type of experience you're  |
| people?                     | for the skill level of your players, and can accommodate | looking for. Here are a few recommendations that cater   |
|                             | all participants comfortably. Here are a few             | to diverse groups:                                       |
|                             | recommendations across different themes and preferences: | 1. **Ticket to Ride** - It's a classic choice that       |
|                             | 1. **Codenames** (Vice Versus) - A word-guessing party   | brings out the traveler in us all. With its simple rules |
|                             | game where two teams compete to identify their agents    | and competitive spirit, it’s perfect for when you want   |
|                             | using one-word                                           | to gather around the table.                              |
+-----------------------------+----------------------------------------------------------+----------------------------------------------------------+
| Is it worth visiting Lisbon | Visiting Lisbon for a long weekend can be an excellent   | Visiting Lisbon for a long weekend can be absolutely     |
| for a long weekend?         | choice, depending on your interests and what you hope to | worth it, depending on what you're looking for in a      |
|                             | experience. Here are several factors to consider that    | travel experience. Here are some reasons why you might   |
|                             | make Lisbon a potentially worthwhile destination:        | consider it:                                             |
|                             | 1. **Rich History and Culture**: Lisbon offers a deep    | ### 1. **Rich History**                                  |
|                             | historical tapestry with influences from the Romans,     | Lisbon is steeped in history, and there's so much to     |
|                             | Moors, and Age of Discovery. The city is home to         | discover from its past. From the Alfama district with    |
|                             | numerous historic sites such as Belém Tower              | its narrow streets and historic vibe to Bairro with its  |
+-----------------------------+----------------------------------------------------------+----------------------------------------------------------+
| Should I learn to play the  | Learning to play the piano as an adult can be a          | Learning to play the piano as an adult is a wonderful    |
| piano as an adult?          | rewarding and enriching experience for many reasons.     | idea if you feel drawn to it. There are numerous reasons |
|                             | Here are some considerations that might help you decide: | why picking up this instrument at any age can be both    |
|                             | ### Benefits of Learning Piano as an Adult               | enriching and fulfilling. Whether you're asking yourself |
|                             | 1. **Cognitive Stimulation**: Playing music engages      | this question or seeking guidance on it, there's no      |
|                             | various parts of your brain, which can improve memory,   | better time than now to dive into the world of piano     |
|                             | concentration, and problem-solving skills.               | playing.                                                 |
|                             | 2. **Stress Relief**: Music has been shown to            | ### Why Consider Learning Piano?                         |
|                             |                                                          | **1. Cognitive Benefits:**                               |
+-----------------------------+----------------------------------------------------------+----------------------------------------------------------+
| What do you think about     | Keeping a daily journal can be highly beneficial for     | Keeping a daily journal can be a personally enriching    |
| keeping a daily journal?    | both mental and emotional well-being. Here are several   | practice with numerous benefits. While the experience is |
|                             | advantages and considerations to keep in mind:           | subjective, there are several aspects to consider when   |
|                             | ### Benefits of Keeping a Daily Journal                  | it comes to the practice of maintaining one.             |
|                             | 1. **Self-Reflection**: Journaling encourages            | ### The Benefits                                         |
|                             | introspection, helping you understand your thoughts,     | 1. **Self-Reflection:** A daily journal serves as a      |
|                             | feelings, and behaviors better. It’s a tool for self-    | mirror to our inner world. It encourages us to delve     |
|                             | discovery and personal growth.                           | into our thoughts and feelings, providing a safe space   |
|                             | 2. **Stress Reduction**: Writing                         | for self-reflection.                                     |
+-----------------------------+----------------------------------------------------------+----------------------------------------------------------+

Stopping the server¶

A served engine is meant to outlive its clients, so the pipeline never stops it. We stop the subprocess here and unregister the atexit hook.

In [15]:

stop_server()
atexit.unregister(stop_server)

stop_server() atexit.unregister(stop_server)

Summary¶

This recipe fitted an enthusiasm direction with CAA in process, saved the SteeringVector, and served it through a vLLM server running the vLLM-Hook plugin. Note the served pipeline does not hold a model. Its steer step read structural facts through the server session, lowered the control to an intervention spec, and passed the direction as a content-addressed artifact, and the plugin applied the addition inside the engine. The unsteered and steered generations came from the same server, with only the client-side control differing between them.

The same flow applies to any control with a spec form, and the vllm-serve backend takes the same BackendSpec for a remote server, where the client sets base_url and nothing about the server's process. The boot environment from serve_environment and the --enforce-eager flag are the server's side of the agreement, and artifact_dir together with the server's VLLM_HOOK_REGISTRY_DIR replaces the HTTP artifact route when client and server share a filesystem.