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.