Skip to content

LitServe services

Opt a protobuf service into LitServe codegen with (pgml.litserve).enable = true. py-gen-ml then emits Pydantic request/response models, per-RPC LitAPI factories, a create_*_server that wraps them in litserve.LitServer, typed sync/async client helpers, and optional serve-config kwargs.

Mark request/response/config messages with message kinds (FEATURE_ROW, PREDICTION, RUN_CONFIG) so other generators can share the same contracts. Kind does not replace the LitServe service opt-in.

Install the extra

pip install 'py-gen-ml[litserve]'

Or with uv:

uv add 'py-gen-ml[litserve]'

Enable the generator when you run codegen (it is off by default):

py-gen-ml path/to/schema.proto --generators=base,patch,sweep,cli_args,litserve

Annotate a contract

syntax = "proto3";

package demo;

import "py_gen_ml/extensions.proto";

// Inference request features.
message PredictRequest {
    option (pgml.kind) = FEATURE_ROW;
    repeated float features = 1;
}

// Model prediction for a request.
message PredictResponse {
    option (pgml.kind) = PREDICTION;
    int32 label = 1;
    float score = 2;
}

// Serve / client settings for Classifier.
message ClassifierServeConfig {
    option (pgml.kind) = RUN_CONFIG;
    option (pgml.litserve_config) = {
        enable: true;
        service: "Classifier";
    };
    string url = 1 [(pgml.default) = {string: "http://localhost:8000"}];
    float timeout_s = 2 [(pgml.default) = {float: 30}];
    int32 workers_per_device = 3 [(pgml.default) = {int32: 1}];
    string accelerator = 4 [(pgml.default) = {string: "cpu"}];
}

// Unary inference service. Opt in with (pgml.litserve).enable.
service Classifier {
    option (pgml.litserve) = {
        enable: true;
    };
    rpc Predict(PredictRequest) returns (PredictResponse) {
        option (pgml.litserve_method) = {
            api_path: "/predict";
        };
    }
}
Option Where Meaning
(pgml.litserve).enable service Required to emit LitServe adapters for this service.
(pgml.litserve).name service Optional LitAPI class-name prefix (default: service name).
(pgml.litserve_method).api_path rpc HTTP path (default: /{method_snake}).
(pgml.litserve_config).enable message Link a config message to a service.
(pgml.litserve_config).service message Proto service name this config configures.

Config field conventions used by the kwargs mapper:

  • accelerator → LitServer(accelerator=...)
  • devices → LitServer(devices=...)
  • workers_per_device → LitServer(workers_per_device=...)
  • timeout_s → LitServer(timeout=...)
  • url → default URL for generated client helpers (not passed to LitServer)

Unary RPCs only; streaming methods raise at codegen time.

LitServe mounts one endpoint per LitAPI. Multi-RPC services get one API class per method (distinct api_path) and a single LitServer holding all of them.

Generated module

For litserve_demo.proto, enabling the generator writes litserve_demo_litserve.py containing:

  • PredictRequest / PredictResponse as pydantic.BaseModel
  • classifier_server_kwargs(config) mapping serve settings
  • create_classifier_predict_api(predict=..., setup=...) returning a LitAPI instance
  • create_classifier_server(predict=..., config=..., **server_kwargs) returning LitServer
  • call_classifier_predict_sync / call_classifier_predict_async (via httpx)

Serve

1. Write a serve module

Wire your model logic into the generated factory:

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
"""Wire Predict into the generated LitServe factory and run the server.

From ``docs/snippets`` (after codegen)::

    uv sync --extra litserve
    uv run python -m snippets.litserve_serve_demo

Optional ``LITSERVE_PORT`` (default ``8000``).
"""
from __future__ import annotations

import os

from pgml_out.litserve_demo_base import ClassifierServeConfig
from pgml_out.litserve_demo_litserve import (
    PredictRequest,
    PredictResponse,
    create_classifier_server,
)


def predict(request: PredictRequest) -> PredictResponse:
    score = sum(request.features) / max(len(request.features), 1)
    return PredictResponse(label=1 if score > 0.5 else 0, score=float(score))


if __name__ == '__main__':
    port = int(os.environ.get('LITSERVE_PORT', '8000'))
    server = create_classifier_server(
        predict=predict,
        config=ClassifierServeConfig(
            accelerator='cpu',
            workers_per_device=1,
            timeout_s=30.0,
        ),
    )
    server.run(host='127.0.0.1', port=port, generate_client_file=False)

Keep handlers in your code. Do not edit the generated *_litserve.py module.

2. Start the server

From the docs snippets project (after codegen), with the litserve extra installed:

cd docs/snippets
uv sync --extra litserve
uv run python -m snippets.litserve_serve_demo

When the process is ready:

curl -s http://127.0.0.1:8000/predict \
  -H 'content-type: application/json' \
  -d '{"features": [0.8, 0.9, 0.7]}'

3. Call with the generated client

from pgml_out.litserve_demo_litserve import PredictRequest, call_classifier_predict_sync

result = call_classifier_predict_sync(
    PredictRequest(features=[0.8, 0.9, 0.7]),
    url="http://127.0.0.1:8000",
)
print(result.label, result.score)

Pass an existing httpx.Client / AsyncClient via client= when you want to reuse connections.