Serving ML Models From Java — Production Patterns¶

File 04 was about what libraries exist. This file is about what a production model server actually looks like from the JVM: how to load an artefact, how to keep latency predictable, when to serve in-process vs proxy to Triton/TorchServe, and how to catch a model going sideways in production. Serving is where 80% of ML pain lives, and where Java earns its keep.

Model serving is a systems problem in a modelling costume. Approach it like any other high-QPS service: know your latency budget, know your failure modes, and instrument everything.

The Two Serving Architectures¶

Architecture	Pattern	When to use
In-process (embedded)	Java loads the ONNX / TorchScript / TF SavedModel artefact into the same JVM as the business logic	Small-to-medium models (< 500MB), classical ML, CPU inference, tight-latency needs (avoid the network hop)
Out-of-process (proxy)	Java calls a dedicated model server (Triton, TorchServe, KServe, Ray Serve, BentoML) over gRPC / HTTP	Large models, GPU inference, multiple frameworks, ML team owns the server independently

Both are valid. Big enterprises usually mix: classical models in-process, deep-learning / LLM models out-of-process. Choose per model, not per service.

⚠️ What most people get wrong: they pick one architecture as “the standard” and force every model into it. A 5MB fraud classifier does not belong on a GPU-backed Triton pod; a 20GB LLM does not belong inside a stateless REST service.

In-Process — Loading an ONNX Model¶

The pattern is: load once at startup, hold a single OrtSession, reuse it for the lifetime of the pod.

@Configuration
class OnnxModelConfig {

    @Bean(destroyMethod = "close")
    OrtEnvironment ortEnv() {
        return OrtEnvironment.getEnvironment();
    }

    @Bean(destroyMethod = "close")
    OrtSession classifierSession(OrtEnvironment env,
                                 @Value("${classifier.model.path}") String path) throws OrtException {
        var opts = new OrtSession.SessionOptions();
        opts.setIntraOpNumThreads(2);           // per-op parallelism
        opts.setInterOpNumThreads(1);           // graph parallelism; usually 1 for serving
        opts.setOptimizationLevel(OrtSession.SessionOptions.OptLevel.ALL_OPT);
        // opts.addCUDA(0);                      // uncomment for GPU
        return env.createSession(path, opts);
    }
}

OrtSession is thread-safe. One session, many concurrent requests. Do NOT create sessions per request.

The Serving Endpoint¶

@RestController
@RequestMapping("/api/v1/classify")
@RequiredArgsConstructor
class ClassifyController {
    private final ClassifierService classifier;
    private final MeterRegistry meters;

    @PostMapping
    ClassifyResponse classify(@Valid @RequestBody ClassifyRequest req) {
        return meters.timer("model.inference.duration", "model", "classifier-v1")
                     .record(() -> classifier.predict(req.text()));
    }
}

@Service
@RequiredArgsConstructor
class ClassifierService {
    private final OrtSession session;
    private final OrtEnvironment env;
    private final Tokenizer tokenizer;

    public ClassifyResponse predict(String text) {
        long[] ids = tokenizer.encode(text);              // deterministic preprocessing
        long[] shape = {1, ids.length};
        try (OnnxTensor input = OnnxTensor.createTensor(env, LongBuffer.wrap(ids), shape);
             OrtSession.Result out = session.run(Map.of("input_ids", input))) {
            float[][] logits = (float[][]) out.get(0).getValue();
            int label = argmax(logits[0]);
            float confidence = softmax(logits[0])[label];
            return new ClassifyResponse(labels[label], confidence);
        } catch (OrtException e) {
            throw new InferenceException("model failed", e);
        }
    }
}

Warmup — The Hidden Cost¶

The first inference after startup is 10–100× slower than steady state (JIT compilation, ONNX Runtime kernel selection, native memory paging). Add a warmup call in an ApplicationRunner:

@Component
@RequiredArgsConstructor
class ModelWarmup implements ApplicationRunner {
    private final ClassifierService classifier;

    @Override
    public void run(ApplicationArguments args) {
        for (int i = 0; i < 10; i++) {
            classifier.predict("warmup input " + i);
        }
        log.info("model warmup complete");
    }
}

Without this, your readinessProbe will report healthy before the model is actually fast, and the first real user gets a 3-second response.

Batching — Free Throughput¶

Most inference engines are 5–20× more efficient per sample when given a batch. If you have concurrency, batch it.

@Service
class BatchingClassifier {
    private final BlockingQueue<PendingRequest> queue = new LinkedBlockingQueue<>();
    private final int maxBatch = 32;
    private final Duration maxWait = Duration.ofMillis(10);

    @PostConstruct
    void start() {
        Thread.startVirtualThread(this::drainLoop);
    }

    public CompletableFuture<ClassifyResponse> predictAsync(String text) {
        var future = new CompletableFuture<ClassifyResponse>();
        queue.offer(new PendingRequest(text, future));
        return future;
    }

    private void drainLoop() {
        while (!Thread.currentThread().isInterrupted()) {
            var batch = new ArrayList<PendingRequest>();
            try {
                var first = queue.poll(maxWait.toMillis(), TimeUnit.MILLISECONDS);
                if (first == null) continue;
                batch.add(first);
                queue.drainTo(batch, maxBatch - 1);
                runBatch(batch);
            } catch (InterruptedException e) {
                Thread.currentThread().interrupt();
            } catch (Exception e) {
                batch.forEach(p -> p.future.completeExceptionally(e));
            }
        }
    }
    // runBatch() stacks tensors and calls session.run() once
}

Virtual threads (Java 21+) make this pattern natural: the request-handling thread parks waiting for the future, the drain loop batches. No callback hell, no reactive-style spaghetti.

Batching trade-off: batching adds latency (up to maxWait) to individual requests but massively improves throughput and GPU utilization. Tune maxBatch and maxWait for your latency SLO.

Out-of-Process — gRPC to Triton¶

For large models, GPU inference, or when the ML team owns the model server, run NVIDIA Triton (or TorchServe / Ray Serve / KServe) and call it over gRPC from Java. Triton’s Java client is generated from its .proto files.

@Service
@RequiredArgsConstructor
class TritonClient {
    private final GRPCInferenceServiceGrpc.GRPCInferenceServiceBlockingStub stub;

    public float[] infer(String modelName, float[] input, long[] shape) {
        var inputTensor = ModelInferRequest.InferInputTensor.newBuilder()
            .setName("input")
            .setDatatype("FP32")
            .addAllShape(Arrays.stream(shape).boxed().toList())
            .build();

        var request = ModelInferRequest.newBuilder()
            .setModelName(modelName)
            .addInputs(inputTensor)
            .addRawInputContents(toByteString(input))
            .build();

        var response = stub.withDeadlineAfter(500, TimeUnit.MILLISECONDS)
                           .modelInfer(request);
        return fromByteString(response.getRawOutputContents(0));
    }
}

gRPC checklist for model servers:

Deadline on every call. Never let a slow model hang your service.
Connection reuse via a single ManagedChannel (thread-safe, pool internally).
Retry only on UNAVAILABLE / DEADLINE_EXCEEDED, not INVALID_ARGUMENT.
Wrap with Resilience4j @CircuitBreaker from file 01.
Consider gRPC-Java’s NettyChannelBuilder for HTTP/2 keepalive tuning under load.

REST vs gRPC for Model Serving¶

Concern	REST/JSON	gRPC/protobuf
Payload size (100-dim float vector)	~1.5KB (JSON overhead)	~400B
CPU per call (serialization)	Higher	Lower
Streaming	Awkward (SSE)	First-class (bidi streams)
Interop with Python / non-Java clients	Universal	Excellent
Ease of ad-hoc debugging (curl, Postman)	✅	Needs grpcurl
Model server support (Triton, TorchServe)	✅	✅ (preferred)

Rule: external / user-facing = REST. Internal ML plumbing = gRPC. Same rule as any other high-QPS internal service.

Timeout, Backpressure, and Load Shedding¶

An overloaded model server does not fail cleanly. It slows down, then queues, then OOMs. Prevent this at the caller:

Timeout — gRPC deadline or HTTP client timeout on every call. Use your latency SLO as the ceiling.
Bounded queue — use a Semaphore or bulkhead limiting concurrent inference calls per instance.
Reject early — return HTTP 503 Retry-After when the queue is full, rather than accepting and timing out.
Circuit breaker — Resilience4j opens after N failures; requests fail fast until the model recovers.

private final Semaphore inflight = new Semaphore(50);   // max 50 concurrent inferences

public ClassifyResponse predictWithBackpressure(String text) {
    if (!inflight.tryAcquire()) {
        throw new TooManyInflightException();   // handled -> 503
    }
    try {
        return predict(text);
    } finally {
        inflight.release();
    }
}

Feature Store Integration¶

Models rarely take raw request data. Real serving pipelines:

Receive request with entity IDs (userId, merchantId, etc.)
Look up features from feature store (Feast / Redis / in-memory cache)
Assemble feature vector
Call model
Log features + prediction for offline analysis

Latency budget example for a 100ms p95 SLO:

Step	Budget
HTTP parse + validate	5ms
Feature store lookup (Redis)	15ms
Preprocessing	10ms
Model inference (ONNX in-process)	30ms
Postprocessing + response	10ms
Total	70ms (leaves 30ms headroom)

If any step blows past its budget, you know exactly where to attack. Instrument each step with a Micrometer Timer.

Model Observability — The Four Signals¶

Models are unlike other services: they can be silently wrong while being technically healthy. Track:

Latency — standard p50/p95/p99. model_inference_duration_seconds{model, version}.
Throughput — requests per second per model version.
Error rate — inference exceptions, tensor-shape mismatches, feature-store misses.
Prediction distribution drift — aggregate output class distribution over rolling windows. A classifier that suddenly outputs 100% class A needs attention, even if it never throws an error.

Drift Detection — The Cheap Version¶

You don’t need Evidently or Fiddler on day one. A weekly aggregation of prediction counts, feature means, and feature p95 — alert when they drift more than 3σ from the training distribution. Log every prediction + feature vector to a predictions Kafka topic; a downstream job computes the drift metrics.

A/B and Shadow Serving¶

Shadow mode — route a copy of live traffic to a new model, don’t return its result to the user. Compare against production output offline. Zero user risk.
A/B canary — route 5% of traffic to the new model. Compare business metrics (click-through, latency, error rate). Ramp up if healthy.

Both are easy with a header-based router in your controller and two OrtSession beans.

Model Versioning¶

Include the model version in every metric (model="classifier", version="v3.1").
Include the model version in every logged prediction.
Never quietly overwrite model.onnx on the pod. Names carry versions: classifier-v3.1.onnx.
Roll back is a config change (classifier.model.path=classifier-v3.0.onnx) and a redeploy, not a code fix.

Serving Sins Checklist¶

Creating an OrtSession per request
No warmup before readiness = healthy
No timeout on Triton gRPC calls
No backpressure = OOM under load
Logging feature vectors with PII in cleartext
Serving from a model file mounted read-write (accidental overwrites)
No drift monitoring — you find out from a user complaint
No shadow mode / canary — you deploy v2 straight to 100% traffic

Practice¶

Take the classifier you built in file 04 (Practice #1). Add warmup, batching, backpressure, and full Micrometer instrumentation. Load-test with k6 at 500 rps and confirm p95 < 100ms.
Add a model_version label to every metric. Deploy v2 in shadow mode. Compare distributions.
Kill the model service under load; confirm the circuit breaker opens, callers get 503 immediately (not 30-second timeouts), and the system recovers gracefully when the model returns.

Return to README.md · Previous: 04_java_for_ml_ai_applications.md · Next: 06_cloud_native_kubernetes.md