Inference Service Examples

These examples show common ways to use the Dragon inference service from Python applications. For the full user guide, see Dragon Inference Service - User Guide. For the API reference, see Inference.

Run examples that create Dragon native objects under Dragon, for example dragon my_inference_app.py.

Example 1 - Single Prompt Service

This is the smallest useful service: one backend, one request queue, and one plain text request submitted through dragon.ai.inference.Inference.query().

 1import os
 2
 3from dragon.ai.inference import (
 4    BatchingConfig,
 5    DynamicWorkerConfig,
 6    GuardrailsConfig,
 7    HardwareConfig,
 8    Inference,
 9    InferenceConfig,
10    ModelConfig,
11)
12from dragon.native.queue import Queue
13
14
15def main() -> None:
16    inference_queue = Queue()
17    response_queue = Queue()
18
19    config = InferenceConfig(
20        model=ModelConfig(
21            model_name="meta-llama/Llama-3.1-8B-Instruct",
22            hf_token=os.environ["HF_TOKEN"],
23            tp_size=1,
24            max_tokens=256,
25        ),
26        hardware=HardwareConfig(num_nodes=1, num_gpus=1),
27        batching=BatchingConfig(enabled=True, batch_type="dynamic"),
28        guardrails=GuardrailsConfig(enabled=False),
29        dynamic_worker=DynamicWorkerConfig(enabled=False),
30    )
31
32    service = Inference(config, inference_queue)
33
34    try:
35        service.initialize()
36        service.query(("Give me a one-sentence definition of RDMA.", response_queue))
37        result = response_queue.get()
38        print(result["assistant"])
39    finally:
40        service.destroy()
41        response_queue.close()
42
43
44if __name__ == "__main__":
45    main()

Example 2 - Shared Chat Proxy

Use DragonQueueLLMProxy when agent or application code already speaks in chat messages. The proxy is lightweight: create one in every client or agent process and point all of them at the shared service queue.

 1import asyncio
 2
 3 from dragon.ai.inference import DragonQueueLLMProxy
 4
 5
 6async def answer_from_agent(inference_queue, question: str) -> str:
 7    proxy = DragonQueueLLMProxy(inference_queue, max_concurrent_requests=8)
 8    try:
 9        return await proxy.chat(
10            [
11                {"role": "system", "content": "You answer briefly."},
12                {"role": "user", "content": question},
13            ]
14        )
15    finally:
16        await proxy.shutdown()
17
18
19response = asyncio.run(
20    answer_from_agent(inference_queue, "Why use one shared inference backend?")
21)
22print(response)

The service response dictionary includes metrics, but the proxy returns only the assistant text. If the backend sends an exception object, the proxy re-raises it in the caller.

Example 3 - Structured Output

Pass json_schema to request structured generation. The backend builds per-request vLLM sampling parameters so a dynamically batched request can mix structured and free-form outputs.

 1schema = {
 2    "type": "object",
 3    "properties": {
 4        "title": {"type": "string"},
 5        "next_steps": {
 6            "type": "array",
 7            "items": {"type": "string"},
 8            "minItems": 1,
 9            "maxItems": 4,
10        },
11    },
12    "required": ["title", "next_steps"],
13}
14
15response = await proxy.chat(
16    [
17        {"role": "system", "content": "Return only JSON."},
18        {"role": "user", "content": "Plan an inference benchmark."},
19    ],
20    json_schema=schema,
21)
22
23print(response)

Example 4 - Pre-Batched Prompt List

Dynamic batching is easiest for concurrent online requests. Pre-batching is useful when an offline driver already has a list of prompts and wants to submit groups sized for the backend.

 1from dragon.native.queue import Queue
 2
 3prompts = [
 4    "Summarize the first log file.",
 5    "Summarize the second log file.",
 6    "Summarize the third log file.",
 7]
 8
 9prebatch_config = InferenceConfig(
10    model=ModelConfig(
11        model_name="meta-llama/Llama-3.1-8B-Instruct",
12        hf_token=os.environ["HF_TOKEN"],
13        tp_size=1,
14        max_tokens=256,
15    ),
16    hardware=HardwareConfig(num_nodes=1, num_gpus=1),
17    batching=BatchingConfig(
18        enabled=True,
19        batch_type="pre-batch",
20        max_batch_size=len(prompts),
21    ),
22    guardrails=GuardrailsConfig(enabled=False),
23    dynamic_worker=DynamicWorkerConfig(enabled=False),
24)
25
26inference_queue = Queue()
27response_queue = Queue()
28service = Inference(prebatch_config, inference_queue)
29
30try:
31    service.initialize()
32    service.query((prompts, response_queue))
33
34    for _ in prompts:
35        result = response_queue.get()
36        print(result["assistant"])
37finally:
38    service.destroy()
39    response_queue.close()

In pre-batch mode the caller submits a list of prompts and a shared response queue, and each generated response is placed on that queue. max_batch_size maps to the vLLM max_num_seqs limit, so set it greater than or equal to the number of prompts to run the whole batch in a single concurrent generation pass.

Example 5 - Two Services in One Allocation

Use node_offset when a workflow needs two independent inference services, such as a large reasoning model and a smaller summarizer model.

 1reasoning_config = InferenceConfig(
 2    model=ModelConfig(
 3        model_name="meta-llama/Llama-3.1-70B-Instruct",
 4        hf_token=os.environ["HF_TOKEN"],
 5        tp_size=4,
 6    ),
 7    hardware=HardwareConfig(num_nodes=1, num_gpus=4, node_offset=0),
 8    batching=BatchingConfig(enabled=True, max_batch_size=16),
 9    guardrails=GuardrailsConfig(enabled=False),
10    dynamic_worker=DynamicWorkerConfig(enabled=False),
11)
12
13summarizer_config = InferenceConfig(
14    model=ModelConfig(
15        model_name="meta-llama/Llama-3.1-8B-Instruct",
16        hf_token=os.environ["HF_TOKEN"],
17        tp_size=1,
18    ),
19    hardware=HardwareConfig(num_nodes=1, num_gpus=1, node_offset=1),
20    batching=BatchingConfig(enabled=True, max_batch_size=32),
21    guardrails=GuardrailsConfig(enabled=False),
22    dynamic_worker=DynamicWorkerConfig(enabled=False),
23)

The first service uses the first selected node. The second service starts from the next node in the allocation.

Example 6 - Guardrails Enabled

Enable guardrails when user prompts should be classified before they reach the LLM. PromptGuard runs in the preprocessing path and rejects prompts whose jailbreak score is greater than or equal to the sensitivity threshold.

1config.guardrails = GuardrailsConfig(
2    enabled=True,
3    prompt_guard_model="meta-llama/Prompt-Guard-86M",
4    prompt_guard_sensitivity=0.5,
5)

Rejected prompts receive a normal response dictionary with assistant set to Your input has been categorized as malicious. Please try again.. Model latency and token throughput metrics are zero for rejected requests because the LLM was not called.