ML Inference & Orchestration Architecture

Why this layer exists

Raw inference servers (vLLM, Triton, TGI) are batching, GPU-saturating text generators. They are not a public API. They have no auth, no per-tenant quotas, no cost accounting, no audit logging, no agent state, no plan execution, no idempotency. The AI Infrastructure layer is everything that sits between the public-facing service contract and those raw inference servers.

The interview frame: "vLLM is great at maximizing GPU throughput for batched generation. Everything else — auth, batching at the gateway, streaming to the client, retries, hedging, agent state — lives in the Rust layer on top."

Calling vLLM/Triton from a Rust service

vLLM exposes an OpenAI-compatible HTTP server out of the box. Triton uses HTTP or gRPC. TGI uses HTTP/SSE. From Rust, the pattern is similar:

use reqwest::Client;
use serde::{Serialize, Deserialize};

#[derive(Serialize)]
struct CompletionReq<'a> {
    model: &'a str,
    messages: &'a [Message<'a>],
    max_tokens: u32,
    temperature: f32,
    stream: bool,
}

async fn generate(client: &Client, base_url: &str, req: CompletionReq<'_>)
    -> Result<CompletionResp, BackendError>
{
    let resp = client
        .post(format!("{base_url}/v1/chat/completions"))
        .json(&req)
        .timeout(std::time::Duration::from_secs(60))
        .send().await?;
    if !resp.status().is_success() {
        return Err(BackendError::Status(resp.status()));
    }
    Ok(resp.json().await?)
}

Things to get right at this layer:

• Connection pooling. One reqwest::Client per process; it pools HTTP/2 connections.
• Deadlines. Forward the request deadline; don't use a static timeout.
• Retry only on idempotent failures. A network error before any tokens were emitted is retryable. Mid-stream failure is not (the next attempt would re-bill the user and likely re-execute side effects).
• Healthchecks. Pull inference backends out of rotation when they're unhealthy.
• Load balancing. Round-robin within a pool; affinity to a node when you want prompt-cache hits.

Batching at the gateway

vLLM already does continuous batching internally — it groups in-flight requests, runs decode-step matmuls together, and adds new requests to the batch each step. So you do not need to batch at the gateway in the simple "wait 50ms, then send N requests" sense.

But there are still gateway-level batching/multiplexing concerns:

• Connection multiplexing. HTTP/2 lets you fanout many concurrent requests over one connection. Use it.
• Coalescing duplicate prompts — if 50 users ask the same exact prompt within 100ms (system prompt + identical user query), can you serve one inference and broadcast the result? Sometimes yes (cache-friendly), sometimes no (privacy/per-user context).
• Request hedging — fire to two backends after a small delay; take whichever returns first. Trades cost for tail-latency improvement.
• Embedding batching — embedding endpoints do benefit from gateway-side batching (low per-request cost, high batch efficiency). Batch with a max-N + max-delay window.

Streaming responses (Server-Sent Events)

LLMs emit tokens one at a time. The user experience demands you stream them to the client, not wait for the full completion. The standard wire format is SSE — text/event-stream with newline-delimited frames.

use axum::response::sse::{Event, Sse};
use futures::stream::Stream;

async fn stream_handler(req: GenerateReq)
    -> Sse<impl Stream<Item = Result<Event, std::convert::Infallible>>>
{
    let upstream = call_vllm_stream(req).await;
    // upstream: Stream<Item = Result<Token, BackendError>>
    let sse = upstream.map(|tok| {
        match tok {
            Ok(t)  => Ok(Event::default().data(serde_json::to_string(&t).unwrap())),
            Err(_) => Ok(Event::default().event("error").data("backend_failure")),
        }
    });
    Sse::new(sse).keep_alive(KeepAlive::default())
}

Cross-cutting concerns in streaming:

• Backpressure. If the client is slow, you need to pause pulling from the upstream. Bounded channel between upstream consumer and client emitter does this.
• Cancellation. If the client disconnects, you must cancel the upstream call so you stop spending GPU time. tonic + tower handle this if you propagate the request future correctly.
• Heartbeats / keepalives. Proxies kill idle SSE connections. Emit periodic : ping frames.
• Partial-failure UX. If a stream dies mid-generation, decide: error out the request, retry from scratch, or hand the partial back to the client?

Prompt-cache awareness

Modern LLM serving stacks reuse the KV-cache (the precomputed attention state) for prompt prefixes that repeat across requests — prefix caching. The implication for your gateway:

• Affinity-route to maximize cache hits. If user A's request shares the same long system prompt as user B's, and they both hit the same vLLM replica, both benefit from the cached prefix. Random load balancing destroys this.
• Consistent-hashing on a stable prompt-prefix hash gives you sticky routing without a leader.
• Pin system prompts to a small set of versioned strings. Per-request system-prompt mutation kills prefix cache hit rate.
• Prompt cache TTL — caches evict under memory pressure. Expect colder tails after restarts.

For Anthropic and OpenAI APIs, prompt caching is an explicit feature with cache markers in the request. Your gateway should pass through and account for cached-token billing.

KV-cache sharing concerns

The KV-cache holds per-request attention state. Two concerns leak into the infra layer:

• Memory pressure — long contexts use more KV-cache memory per request. Throughput drops when many long-context requests run concurrently. Your gateway needs to know "this backend is at 80% KV memory; route new long-context requests elsewhere."
• Cache sharing across tenants — if tenant A's system prompt is shared with tenant B (e.g., a common base prompt), prefix caching means tenant B might benefit from tenant A's work. Usually fine. If the prompt contains tenant-specific keys, that's a leak vector — keep secrets out of cacheable prefixes.

Agent orchestration — the layer above inference

An agent is a loop: model proposes an action (a tool call, or "I'm done"), the action runs, the result feeds back, the model proposes the next action. As an infrastructure SWE you own the executor — the service that runs this loop reliably at scale.

run = {
  id, user, plan_version, model_version, deadline, budget,
  state: { messages, tool_results, partial_outputs },
  step_count, status
}
loop:
  model.call(state.messages) -> either tool_call or final_answer
  if final_answer: persist, return
  if tool_call:
    tool = tool_router.lookup(name)
    if tool.requires_approval: pause, wait for human
    result = tool.call(args, idem_key)
    state.messages += [tool_call, tool_result]
    if step_count > budget.max_steps: halt

Everything in this loop is an infrastructure concern. Tool dispatch, state persistence, deadline/budget accounting, retry, audit, cancellation, partial-state recovery — those are your code.

Tool dispatch

A "tool" is a function the model can call. The dispatcher's job: look up the tool by name, validate arguments against the schema, enforce auth and policy, execute, and return a result the model can read.

• Schema validation — never trust model-generated args. JSON-schema-validate before dispatching.
• Auth scoping — the tool runs under the user's identity, not god-mode. If the model asks for a tool the user can't use, deny.
• Idempotency — every side-effecting tool gets an idempotency key from the orchestrator. Retries don't duplicate.
• Timeouts — every tool call has its own deadline within the agent's overall deadline.
• Approval gates — high-tier tools (money movement, sending external messages) pause for human approval before executing.
• Result shape — return structured success/error so the model can react. {"error": "...", "retryable": true} beats raw stack traces.

Multi-step plans

Two common shapes:

From the infra side, both reduce to "execute a sequence of tool calls with state between them" — your executor doesn't really care which pattern the agent uses, as long as the state machine is well-defined.

Partial-state recovery

An agent run can take minutes. Processes restart. Crashes happen. The executor must be able to resume.

• Persist after every step. The run record (state, step count, last tool result) is written to a durable store after each step before issuing the next model call.
• Idempotent step execution. Resuming from step N must not re-execute step N's side effects. Idempotency keys on side-effecting tools, plus a "completed steps" log.
• Single-owner lock. Only one executor instance can be advancing a given run at a time. Use a Postgres advisory lock, Redis lease, or a Raft-backed run registry.
• Stuck-state detection. If a run has been in "executing step N" for 5x the median step duration, alert / fail / escalate.

async fn step(run: &mut RunState, executor: &Executor) -> Result<()> {
    let next_action = call_model(&run.messages).await?;
    match next_action {
        Action::ToolCall { name, args, idem_key } => {
            // 1. record intent
            store.append_event(&run.id, Event::ToolCallStarted { name: name.clone(), idem_key }).await?;
            // 2. dispatch
            let result = executor.dispatch(&name, &args, idem_key).await?;
            // 3. record outcome
            store.append_event(&run.id, Event::ToolCallCompleted { result: result.clone() }).await?;
            run.messages.push(Message::tool_result(name, result));
        }
        Action::Done { answer } => {
            store.mark_complete(&run.id, answer).await?;
        }
    }
    Ok(())
}

Sub-agent isolation

An agent that spawns sub-agents — common in research/coding agents — introduces blast-radius questions. Patterns:

• Separate run IDs. A sub-agent is its own run record with its own state, deadline, budget — not a nested message in the parent.
• Bounded budgets. The parent declares "this sub-agent gets at most N tokens / M steps / $K cost." Hard caps.
• Trace propagation. The parent's trace ID is set as a parent of the sub-agent's spans for end-to-end visibility.
• Tool authority restriction. A sub-agent inherits a subset of the parent's tool authority, never a superset.
• Result-only return. The sub-agent's full message history doesn't pollute the parent's context unless the parent explicitly fetches it.
• Cancellation cascades. Cancelling the parent must cancel all sub-agents. This is structured concurrency; see 08-error-handling.

✓ Check yourself

Answer from memory before opening each one.

1. Product asks for per-tenant quotas, cost accounting and an audit trail on model calls. Why can you not get these by exposing the model server directly, and where do they live?

   Show answer

   Inference servers such as vLLM, Triton and TGI batch and generate tokens; they have no auth, quotas, cost accounting, audit, agent state or idempotency. Those belong in the Rust infrastructure layer between the public contract and the GPU fleet. See Why this layer exists.

2. A streamed completion fails after 40 tokens. Does the gateway retry it automatically? Contrast with a failure before the first token.

   Show answer

   No. A retry would bill the user again and may repeat side effects, so decide explicitly between erroring, restarting from scratch or returning the partial output. A network error before any token was emitted is safe to retry. See Calling vLLM/Triton and Streaming responses.

3. A user closes the browser mid-stream. What must the gateway do, and what does it cost if it does not?

   Show answer

   Cancel the upstream call as soon as the client disconnects, by propagating cancellation through the request future. If it does not, the model server keeps generating tokens nobody reads and spends GPU time for nothing. See Streaming responses.

4. Twenty vLLM replicas sit behind a round-robin balancer and the prefix-cache hit rate is poor, though every request begins with the same long system prompt. What changes, and what must you keep out of the shared prefix?

   Show answer

   Route by affinity, for example consistent hashing on a stable prefix hash, so requests that share a prefix land on the same replica, and keep the system prompt to a few versioned strings. Keep tenant-specific secrets out of cacheable prefixes, since shared prefix state can leak across tenants. See Prompt-cache awareness and KV-cache sharing concerns.

5. The orchestrator crashes after dispatching a payment tool call but before it records the result. On restart, what prevents a second payment?

   Show answer

   The intent event was written before dispatch, the call carried an idempotency key, and a completed-steps log plus a single-owner lock decide who resumes. The retry reuses the same key, so the tool returns the first outcome instead of paying again. See Partial-state recovery.

6. For an agent that can move money, would you favour ReAct or plan-then-execute, and what does your executor need either way?

   Show answer

   Plan-then-execute is the usual choice there, because an explicit plan is more predictable, easier to checkpoint and easier to bound in cost; ReAct wins when the task needs mid-run adaptation. Either way the executor needs a well-defined state machine with persisted steps. See Multi-step plans.

Takeaway

• Inference servers are batching text generators, and everything else lives in the Rust layer above them: auth, quotas, streaming, retries, hedging and agent state.
• Retry a model call only when nothing has been emitted. A mid-stream failure needs an explicit decision, because a blind retry re-bills and can repeat side effects.
• Stream tokens over SSE with a bounded channel for backpressure, keepalives for proxies, and upstream cancellation when the client disconnects.
• Prefix-cache hit rate depends on affinity routing and stable, versioned system prompts, and secrets must stay out of cacheable prefixes.
• The agent loop is an infrastructure problem: dispatch tools with validated arguments, per-call deadlines, idempotency keys and approval gates for high-tier actions.
• Make runs resumable by persisting after every step, executing steps idempotently, and allowing one owner per run; give each sub-agent its own run, budget and narrower tools.

Next come the production patterns that put these pieces together: the gateway, hedging, the sandbox, the tool router and the plan executor. → Applied patterns

Further reading

• vLLM documentation — vLLM project. The OpenAI-compatible server, batching and serving features your gateway wraps.
• Efficient Memory Management for Large Language Model Serving with PagedAttention — Kwon et al., 2023. The paper behind vLLM's KV-cache management and why memory limits throughput.
• Automatic Prefix Caching — vLLM documentation. How prefix reuse works, which is what affinity routing depends on.
• Server-sent events — WHATWG HTML Standard. The wire format, reconnection and comment lines used for keepalives.
• The Tail at Scale — Dean and Barroso, Communications of the ACM, 2013. The case for hedged requests and other tail-latency techniques.
• Building effective agents — Anthropic, 2024. Workflow and agent patterns, and when each is worth its cost.