Online Serving | Structured Data and Graph Models

Kumo’s online serving runs predictions at request time — suitable for live recommendations, real-time fraud scoring, and any use case where you need a score in milliseconds rather than a scheduled job. Unlike batch prediction, which scores your entire dataset at once, online serving keeps a live endpoint running that you query with a single entity at a time.

The deployment steps below mirror the interactive SDK example notebook, which you can download and run end to end.

How it works

Online serving in Kumo combines two models:

A base model trained on your full graph that produces rich entity embeddings capturing long-term patterns.
A distilled model that runs at request time, combining those stored embeddings with the latest signals to produce a fast, accurate prediction.

The end-to-end flow is:

Train a base model on your graph.
Train a distilled model using the base model’s embeddings.
Generate embeddings from the base model and export the serving bundle to S3.
Deploy the bundle to a live inference endpoint.
Query the endpoint from your application.

Steps 1–3 use the fine-tuning SDK. Steps 4–5 use the kumoai.online deployment SDK.

Train the base model

Train your base model as usual — this is the same workflow covered in Training & Predictions. Save the job ID; you’ll need it in the next step.

1 import kumoai as kumo
2 
3 train_table = pq.generate_training_table()
4 model_plan  = pq.suggest_model_plan()
5 
6 trainer    = kumo.Trainer(model_plan)
7 result     = trainer.fit(graph, train_table, non_blocking=False)
8 base_job_id = result.job_id

Train the distilled model

The distilled model is a smaller, faster model trained to predict at request time. It takes the base model’s embeddings as inputs — pass base_model_id to link the two.

1 from kumoai.trainer import DistillationTrainer
2 
3 train_table_serving = pq_serving.generate_training_table()
4 
5 distilled_plan   = pq_serving.suggest_distilled_model_plan(base_model_id=base_job_id)
6 distiller        = DistillationTrainer(distilled_plan, base_job_id)
7 dist_result      = distiller.fit(graph, train_table_serving, non_blocking=False)
8 distilled_job_id = dist_result.job_id

pq_serving is a PredictiveQuery on the same graph as pq, targeting the entity you want to score at request time — for example, a transaction or a recommendation candidate.

Generate embeddings and export

Run batch prediction on the base model to produce entity embeddings, then export both the distilled model and the embeddings to S3 as a ready-to-deploy bundle.

1 from kumoai.artifact_export.config import OutputConfig
2 from kumoai.trainer import ModelOutputConfig, export_model
3 
4 base_trainer = kumo.Trainer.load(base_job_id)
5 pred_table   = pq.generate_prediction_table(non_blocking=True)
6 
7 bp_result = base_trainer.predict(
8     graph=graph,
9     prediction_table=pred_table,
10     output_config=OutputConfig(
11         output_types={"predictions", "embeddings"},
12         output_connector=connector,
13         output_table_name="embeddings_for_export",
14     ),
15     training_job_id=base_job_id,
16     non_blocking=False,
17 )
18 
19 export_config = ModelOutputConfig(
20     training_job_id=distilled_job_id,
21     batch_prediction_job_id=bp_result.job_id,
22     output_path="s3://your-bucket/path/to/serving-bundle/",
23 )
24 export_model(export_config, non_blocking=False)

Export targets S3 URIs (s3://…). Contact your Kumo team if you need to export to a different storage provider.

Connect to the deployment control plane

With your bundle in S3, switch to the kumoai.online SDK to deploy and manage the live service.

Install the SDK:

$ pip install kumoai

Your Kumo team will provide these values when they provision your tenant:

Variable	What it is
`BASE_URL`	Your control-plane API URL
`CLIENT_ID`	Cognito app client ID
`CLIENT_SECRET`	Cognito app client secret
`TOKEN_URL`	Cognito token endpoint

1 import os
2 import kumoai.online as kumo_online
3 
4 BASE_URL      = os.environ.get("BASE_URL",      "<your-control-plane-url>")
5 CLIENT_ID     = os.environ.get("CLIENT_ID",     "<your-cognito-client-id>")
6 CLIENT_SECRET = os.environ.get("CLIENT_SECRET", "<your-cognito-client-secret>")
7 TOKEN_URL     = os.environ.get("TOKEN_URL",     "<your-cognito-token-url>")
8 SCOPE         = os.environ.get("SCOPE",         "https://inference.kumo.ai/invoke")
9 
10 client = kumo_online.init(
11     url=BASE_URL,
12     client_id=CLIENT_ID,
13     client_secret=CLIENT_SECRET,
14     token_url=TOKEN_URL,
15     scope=SCOPE,
16 )
17 
18 print("healthy:", client.health())

Register your model

Point the SDK at the S3 bundle from Step 3. The last path segment becomes the model name.

1 MODEL_ID     = "my-model-v1"
2 S3_MODEL_URI = "s3://your-bucket/path/to/serving-bundle/two_stage_gnn/"
3 
4 client.register_model(model_id=MODEL_ID, s3_model_uri=S3_MODEL_URI)

Deploy an inference service

Choose a GPU instance type and create the service:

1 import time
2 
3 instance_types = client.list_instance_types()
4 INSTANCE_TYPE  = "g6.4xlarge"
5 SERVICE_NAME   = "my-service"
6 
7 svc = client.create_inference_service(
8     name=SERVICE_NAME,
9     model_id=MODEL_ID,
10     instance_type=INSTANCE_TYPE,
11 )

g6.4xlarge is a good starting point for most models. The service takes 1–5 minutes to start — poll until it’s ready:

1 def wait_until_ready(name: str, timeout: int = 600, poll: int = 10) -> None:
2     deadline = time.time() + timeout
3     while time.time() < deadline:
4         statuses = client.get_status()
5         s = statuses.get(name)
6         ready = bool(s and s.ready)
7         print(f"{name}: ready={ready}  {s.message if s else ''}")
8         if ready:
9             return
10         time.sleep(poll)
11     raise TimeoutError(f"{name} did not become ready within {timeout}s")
12 
13 wait_until_ready(SERVICE_NAME)

Run inference

Send a request to your live endpoint. Input names and shapes come from your exported model’s config.pbtxt.

1 inputs = [
2     {
3         "name": "anchor_time",
4         "datatype": "INT64",
5         "shape": [1, 1],
6         "data": [1352371793912000000],
7     },
8 ]
9 
10 response = svc.infer(inputs=inputs)
11 print(response)

You can also load inputs from a JSON file:

1 import json
2 from pathlib import Path
3 
4 body   = json.loads(Path("request.json").read_text())
5 inputs = body["inputs"] if isinstance(body, dict) else body
6 print(svc.infer(inputs=inputs))

Clean up

Delete the service and registered model when you’re done to avoid ongoing costs:

1 svc.delete()
2 client.delete_registered_model(MODEL_ID)

Next Steps

The example notebook covers two additional features once you’re comfortable with the basics.

Autoscaling — automatically scale replicas based on CPU usage:

1 from kumoai.online import ScaleMetric
2 
3 svc.create_autoscaling(
4     min_capacity=1,
5     max_capacity=5,
6     scale_metric=ScaleMetric.cpu,
7     scale_target=70,
8 )

Canary rollouts — gradually shift traffic to a new model version while the existing one keeps serving the rest:

1 client.register_model(model_id="my-model-v2", s3_model_uri=S3_MODEL_URI_V2)
2 svc.start_canary(canary_model_id="my-model-v2", canary_traffic_percent=10)

Once you’re satisfied, call svc.promote() to make the canary the new stable version. If something goes wrong, call svc.rollback() to route all traffic back to the original.