03 — On-Cluster Inference

Agents call an LLM served on the cluster by Dragon Inference (vLLM behind a queue) through DragonQueueLLMProxy — no external API, no rate limits, and air-gap friendly. Multiple agents share one inference queue and the same GPU worker(s).

What you’ll learn:

  • How to stand up a Dragon Inference (vLLM) backend on the cluster

  • How to call it from a LangGraph node via DragonQueueLLMProxy

  • How to request structured output with a JSON schema

Needs: GPU nodes.

Main Code

Listing 36 03_local_inference.py
  1"""LLM Inference example: agents use Dragon Inference for local LLM calls.
  2
  3Demonstrates:
  4- Dragon Inference pipeline (vLLM) running on GPU nodes
  5- ``get_llm(queue)`` — the per-host accessor that builds a DragonQueueLLMProxy
  6  inside the agent host, on demand, cached per inference pipeline
  7- Async agent nodes that ``await get_llm(...).chat(...)`` on the host event loop
  8- Agents share one inference pipeline (one model, many agents)
  9- No external API calls — inference is local on-cluster
 10
 11Why ``get_llm`` (and not a proxy passed in): the proxy holds a loop-bound
 12response pool, so it must be built *inside* the host process and used on the
 13host's event loop.  You pass the inference *queue handle* into your node (via a
 14closure) and call ``get_llm(queue)`` inside it; the proxy is created there, once,
 15and reused.  See the two-model variant at the bottom for different models on the
 16same host.
 17
 18Run (requires GPU nodes)
 19---
 20    dragon examples/dragon_ai/langgraph/03_local_inference.py
 21
 22Architecture
 23------------
 24    Node 0: AgentHost (one event loop) + CPU workers
 25    Node 1: vLLM inference workers (GPUs)
 26
 27    ┌─────────────────────┐           ┌──────────────────────┐
 28    │ AgentHost (Node 0)  │           │ Inference (Node 1)   │
 29    │  ├─ researcher ─────┼──Queue──▶ │  ├─ vLLM worker GPU0 │
 30    │  └─ writer    ◀─────┼──Queue──  │  └─ vLLM worker GPU1 │
 31    │   (async; share one │           │    (tp_size=2)       │
 32    │    get_llm proxy)   │           │                      │
 33    └─────────────────────┘           └──────────────────────┘
 34"""
 35
 36from __future__ import annotations
 37
 38import asyncio
 39import dragon
 40import multiprocessing as mp
 41import operator
 42from functools import partial
 43from typing import Annotated, Any, TypedDict
 44
 45from langgraph.graph import END, START, StateGraph
 46
 47from dragon.ai.inference.config import (
 48    BatchingConfig,
 49    HardwareConfig,
 50    InferenceConfig,
 51    ModelConfig,
 52)
 53from dragon.ai.inference.inference_utils import Inference
 54from dragon.ai.langgraph import DragonExecutor, get_llm
 55from dragon.native.queue import Queue
 56
 57
 58# =============================================================================
 59# 1. Graph state
 60# =============================================================================
 61
 62class State(TypedDict):
 63    messages: Annotated[list[str], operator.add]
 64    query: str
 65    research: str
 66    report: str
 67
 68
 69# =============================================================================
 70# 2. Agent node functions — async, and they pull the LLM with get_llm(queue)
 71#    inside the node (so the proxy is built in the host, on the host loop).
 72#    The inference queue handle is injected per node via functools.partial in
 73#    main(); a Dragon Queue handle is picklable and safe to close over.
 74# =============================================================================
 75
 76async def researcher(state: State, *, input_queue: Any) -> dict:
 77    """Research agent — queries the local LLM for information."""
 78    query = state.get("query", "unknown topic")
 79
 80    resp = await get_llm(input_queue).chat([
 81        {"role": "system", "content": "You are a research assistant. Provide concise findings."},
 82        {"role": "user", "content": f"Research the following topic: {query}"},
 83    ])
 84
 85    print(f"  [researcher] got response ({len(resp)} chars)")
 86    return {"messages": [resp], "research": resp}
 87
 88
 89async def writer(state: State, *, input_queue: Any) -> dict:
 90    """Writer agent — drafts a report from research findings."""
 91    research = state.get("research", "")
 92
 93    resp = await get_llm(input_queue).chat([
 94        {"role": "system", "content": "You are a technical writer. Write clear, concise reports."},
 95        {"role": "user", "content": f"Write a brief report based on these findings:\n{research}"},
 96    ])
 97
 98    print(f"  [writer] drafted report ({len(resp)} chars)")
 99    return {"messages": [resp], "report": resp}
100
101
102async def summarizer(state: State, *, input_queue: Any) -> dict:
103    """Summarizer agent — produces a one-paragraph summary."""
104    report = state.get("report", "")
105
106    resp = await get_llm(input_queue).chat([
107        {"role": "system", "content": "Summarize in one paragraph."},
108        {"role": "user", "content": f"Summarize this report:\n{report}"},
109    ])
110
111    print("  [summarizer] summary ready")
112    return {"messages": [resp]}
113
114
115# =============================================================================
116# 3. Build and run
117# =============================================================================
118
119def main() -> None:
120    # -- Configure inference pipeline -----------------------------------------
121    config = InferenceConfig(
122        model=ModelConfig(
123            model_name="meta-llama/Llama-3.1-8B-Instruct",
124            hf_token="<YOUR_HF_TOKEN>",  # Set via env: HF_TOKEN
125            tp_size=2,                    # 2 GPUs for tensor parallelism
126            max_tokens=512,
127            max_model_len=4096,
128        ),
129        hardware=HardwareConfig(
130            num_nodes=1,    # 1 node for inference
131            num_gpus=2,     # 2 GPUs on that node
132        ),
133        batching=BatchingConfig(
134            enabled=True,
135            batch_wait_seconds=0.05,
136            max_batch_size=32,
137        ),
138    )
139
140    # -- Launch inference pipeline and grab its input queue -------------------
141    input_queue = Queue(maxsize=256)
142    inference = Inference(config=config, input_queue=input_queue)
143    inference.initialize()
144    print("Inference pipeline ready")
145
146    # -- Launch agent host and build graph ------------------------------------
147    # Inject the inference QUEUE HANDLE into each node via functools.partial.
148    # The node builds/reuses the proxy itself via get_llm(input_queue), inside
149    # the host, on the host's event loop.
150    try:
151        with DragonExecutor() as executor:
152            executor.launch_host(agents={
153                "researcher": partial(researcher, input_queue=input_queue),
154                "writer": partial(writer, input_queue=input_queue),
155                "summarizer": partial(summarizer, input_queue=input_queue),
156            })
157
158            builder = StateGraph(State)
159            builder.add_node("researcher", executor.node("researcher"))
160            builder.add_node("writer", executor.node("writer"))
161            builder.add_node("summarizer", executor.node("summarizer"))
162
163            builder.add_edge(START, "researcher")
164            builder.add_edge("researcher", "writer")
165            builder.add_edge("writer", "summarizer")
166            builder.add_edge("summarizer", END)
167
168            graph = builder.compile()
169
170            # -- Run with the async driver, so the async nodes run on the host's
171            #    shared event loop (and their get_llm proxy is reused). -------
172            print("\n=== Running LangGraph with Dragon Inference ===\n")
173            result = asyncio.run(graph.ainvoke(
174                {
175                    "messages": [],
176                    "query": "advances in quantum error correction 2025",
177                    "research": "",
178                    "report": "",
179                },
180            ))
181
182            print(f"\n=== Final Report ===\n{result['report']}")
183    finally:
184        # -- Cleanup inference pipeline ---------------------------------------
185        # destroy() ends the workers, joins the process group and closes the
186        # input queue. Without it the vLLM workers keep the runtime alive.
187        inference.destroy()
188
189    print("\nDone.")
190
191
192# -----------------------------------------------------------------------------
193# Two models on one host?  Just use two inference queues — each get_llm(q)
194# returns that pipeline's own cached proxy inside the host:
195#
196#     async def researcher(state, *, llama_q):
197#         return {"research": await get_llm(llama_q).chat([...])}
198#
199#     async def writer(state, *, mistral_q):
200#         return {"report": await get_llm(mistral_q).chat([...])}
201#
202#     executor.launch_host(agents={
203#         "researcher": partial(researcher, llama_q=llama_q),
204#         "writer": partial(writer, mistral_q=mistral_q),
205#     })
206# -----------------------------------------------------------------------------
207
208
209if __name__ == "__main__":
210    mp.set_start_method("dragon", force=True)
211    main()

See Also