Custom backends

The built-in backends (vllm, sglang, ollama, llamacpp) cover most LLM serving. When you need to run your own model — a classifier, an embedding model, a scoring head, an image or audio model, anything you can express in PyTorch — you define a custom backend: a small Python class the SDK turns into a long-lived inference server, packages, and ships to your workers.

You write two methods. The platform does the rest: distribution to workers, process supervision, and routing jobs to it.

Load once, process per job

setup() loads your model a single time; process(job) reuses it for every request. No reloading per job.

Your own I/O shape

Input and output are free JSON dicts — not OpenAI-compatible. A classifier returns {"label", "score"}; a text model returns {"text"}. You decide.

Private to your tenant

A custom backend you publish is visible and usable only within your tenant. Other tenants never see it.

No platform rebuild

Publishing a custom backend is data + an artifact. You don’t wait on a platform release to ship a new model type.

When to use a custom backend

Reach for a built-in backend (vllm/sglang/ollama/llamacpp) when your model is a standard LLM served over the OpenAI-compatible API. Reach for a custom backend when the model or its I/O doesn’t fit that mold — a non-chat modality, a bespoke pre/post-processing step, or a model you load and call yourself in PyTorch.

1. Define the backend

A custom backend is a subclass of CustomBackend with two methods:

• setup(ctx) — called once when the server starts. Instantiate your model here and stash it on self. Read knobs (device, model name, weights path) from ctx.config.
• process(job) — called per job. Reuse the model loaded in setup() and return a Result.

my_backend.py

import torch
from torch import nn

from inferencekey.backend import BackendContext, CustomBackend, Job, Result


class SentimentBackend(CustomBackend):
    # Optional metadata — surfaced to the platform and at GET /meta.
    name = "tiny-sentiment"
    version = "0.1.0"
    task_type = "classification"

    def setup(self, ctx: BackendContext) -> None:
        # Called once. Build/load the model here; it stays resident for the
        # life of the process. Device comes from config (defaults to "cpu").
        self.device = str(ctx.config.get("device", "cpu"))
        self.model = nn.Linear(8, 2).to(self.device).eval()

    def process(self, job: Job) -> Result:
        # Called per job. `job.input` is a free dict you define the shape of.
        text = job.input.get("prompt") or _last_user_message(job.input)
        logits = self._embed(text)               # your model call
        label = "positive" if logits[1] > logits[0] else "negative"
        return Result(output={"label": label, "score": float(max(logits))})

The contract

Type	Shape	Notes
Job	{ "id": str, "input": dict }	input is a free, JSON-serializable dict. Your backend defines its keys.
Result	{ "output": dict }	output is a free dict. Wrap your return value as Result(output=...).
BackendContext	{ "config": dict, "port": int }	config is where you resolve device, model_name, weights, etc.

Recommended input/output for text jobs

Input and output are not OpenAI-compatible — that’s deliberate, so any modality fits. For text models that should interoperate with the platform’s data plane, the recommended convention is to read input.messages (a list of {role, content}) or input.prompt, and to return output.text. Other modalities use whatever keys make sense ({"values": [...]}, {"label", "score"}, …).

Error behavior

• If setup() raises, the backend never becomes ready and the process exits non-zero — the platform sees it failed to come up.
• If process() raises, that single job gets a 500 with a JSON error and the server stays alive for the next job.

2. Run it locally

Before packaging, run the backend on your machine exactly as a worker would. Your backend’s dependencies (PyTorch, etc.) go in a requirements.txt; the SDK runtime itself needs none of them.

1. Install your backend’s deps (CPU-only torch is plenty for testing):

   python -m venv .venv && . .venv/bin/activate
   pip install torch --index-url https://download.pytorch.org/whl/cpu

2. Serve it on a loopback port:

   python -m inferencekey.backend.serve \
     --port 8099 \
     --backend my_backend:SentimentBackend \
     --config-json '{"device": "cpu"}'

   The entrypoint is module:Class. Config can also come from the IK_BACKEND_ENTRYPOINT / IK_BACKEND_CONFIG environment variables.

3. Exercise the endpoints:

   curl -s http://127.0.0.1:8099/healthz          # 200 once setup() finished
   curl -s http://127.0.0.1:8099/meta             # {"name","version","task_type",...}
   curl -s -X POST http://127.0.0.1:8099/process \
     -d '{"id":"j1","input":{"prompt":"i love this"}}'
   # {"output": {"label": "positive", "score": ...}}

You can also start it from code instead of the CLI:

from inferencekey.backend import serve_backend
from my_backend import SentimentBackend

serve_backend(SentimentBackend, port=8099, config={"device": "cpu"})

Note

The runtime serves a small, SDK-owned HTTP schema: GET /healthz (ready only after setup()), POST /process, and GET /meta. It binds 127.0.0.1 only. The HTTP server itself can be the standard library one or a native Rust one; see The HTTP transport.

3. Package it

package_backend (or the equivalent CLI) bundles your code, its requirements.txt, and a manifest.json into a single .tar.gz with a checksum.

Python

package.py

from inferencekey.backend import package_backend

pkg = package_backend(
    src="my_backend.py",                      # file or directory
    entrypoint="my_backend:SentimentBackend",
    requirements="requirements.txt",
    name="tiny-sentiment",
    slug="tiny-sentiment",
    version="0.1.0",
    task_type="classification",
    out_dir="dist",
)
print(pkg.path, pkg.sha256, pkg.size_bytes)

CLI

python -m inferencekey.backend.package \
  --src my_backend.py \
  --entrypoint my_backend:SentimentBackend \
  --requirements requirements.txt \
  --name tiny-sentiment --slug tiny-sentiment \
  --version 0.1.0 --task-type classification \
  --out dist

Package code, not weights

The artifact carries code and a dependency list, not multi-gigabyte model weights. Resolve weights inside setup() — load from disk, or pull from a model hub on first start. Keep the package small.

4. Publish it to the platform

Upload the artifact with your control token. The metadata travels from the package’s manifest.

publish.py

from inferencekey import publish_custom_backend

result = publish_custom_backend(
    tenant_id="<your-tenant-id>",
    package_path="dist/tiny-sentiment-0.1.0.tar.gz",
    token="ik_sdk_...",                       # your control token
)
print(result["id"], result["slug"], result["sha256"])

Once published, the custom backend appears in your tenant’s Custom backends page in the dashboard, where you (and your teammates) can see and manage it. It is private to your tenant.

Note

You can also upload from the dashboard’s Custom backends page (drag in the .tar.gz and fill name / slug / version / task type / entrypoint). The SDK and the UI hit the same place.

5. Use it in a workload

Target the custom backend by its slug, exactly where you’d pass a built-in backend. ensure() creates (or reconciles) the workload; when a worker picks it up, the platform delivers the package to it, starts your backend once, and routes jobs to it.

use_it.py

from inferencekey import ManagementClient, WorkloadSpec

mgmt = ManagementClient.from_env(project="acme")

ref = mgmt.ensure(WorkloadSpec(
    name="sentiment",
    slug="sentiment",
    model="tiny-sentiment",
    backend="tiny-sentiment",     # the slug you published
    task_type="classification",
))

Tip

WorkloadSpec.backend accepts a built-in Backend enum or a custom backend slug string. The slug is resolved against the custom backends published in your tenant — a slug from another tenant is rejected. The workload’s task_type must match the one the backend declared.

From the dashboard you get the same outcome: on the New workload form, the backend selector lists your tenant’s custom backends alongside the built-ins; pick one and create the workload.

6. Serve an LLM: OpenAI chat, streaming and tool calls

For a language model, implement chat() as well. The runtime then serves an OpenAI-compatible POST /v1/chat/completions, and the platform routes chat to it exactly like it routes chat to vLLM: streaming, tool calls, finish reason and token usage reach the caller, and the workload’s Anthropic /v1/messages endpoint works for Claude Code.

chat_backend.py

from inferencekey.backend import BackendContext, CustomBackend, Job, Result, chat_completion


class MyLLMBackend(CustomBackend):
    name = "my-llm"
    version = "0.1.0"
    task_type = "text2text"

    def setup(self, ctx: BackendContext) -> None:
        self.engine = load_my_engine(ctx.config)          # your engine, loaded once

    def process(self, job: Job) -> Result:                # still required
        return Result(output={"text": self.engine.generate(job.messages())})

    def chat(self, request: dict) -> dict:
        # `request` is the OpenAI body: messages, tools, tool_choice, max_tokens...
        out = self.engine.chat(request["messages"], tools=request.get("tools"))
        return chat_completion(
            out.text,
            model=self.name,
            tool_calls=out.tool_calls,                    # OpenAI shape, or None
            prompt_tokens=out.prompt_tokens,
            completion_tokens=out.completion_tokens,
        )

• chat() returns an OpenAI chat.completion dict; chat_completion(...) builds one, with tool_calls ([{"id", "type": "function", "function": {"name", "arguments"}}]) and usage.

• Streaming. A request with "stream": true calls chat_stream(). The default replays chat() as chat.completion.chunk events, so streaming clients work from day one. Override chat_stream() and yield chunks as your engine produces tokens.

• Package it with api="openai-chat" (--api openai-chat in the CLI). It is inferred from chat() for GET /meta, and the manifest field lets the worker know before the first request.

• Raise InvalidInput for a bad request (400). Any other exception is a 500 for that request; the server stays alive.

• Stop work when the client leaves. When the caller disconnects, the runtime closes your generator at its next yield (Python raises GeneratorExit there). Put the engine’s abort in a try/finally so it runs on a disconnect too:

  def chat_stream(self, request):
      req_id = self.engine.submit(request)
      try:
          for piece in self.engine.stream(req_id):
              yield {"choices": [{"index": 0, "delta": {"content": piece}, "finish_reason": None}]}
      finally:
          self.engine.abort(req_id)      # runs on a client disconnect as well

  A backend that catches BaseException around a yield must re-raise GeneratorExit.

• Nothing to send yet. chat_stream() may yield None while the engine is busy (a long prefill, a queue). The runtime writes an SSE keepalive comment, which also notices a client that left, so the generator is closed during idle periods too.

• Pre-serialized chunks (optional fast path). A hot loop may yield a str or bytes holding an already-serialized chunk object; the runtime writes it as is, without json.dumps. It must be a single line of JSON: a chunk with a raw line break (for example json.dumps(obj, indent=2)) ends the stream with an error event. Dict chunks remain the default.

  import json
  PREFIX = '{"choices":[{"index":0,"delta":{"content":'
  SUFFIX = '},"finish_reason":null}]}'
  
  def chat_stream(self, request):
      for piece in self.engine.generate(request["messages"]):
          yield PREFIX + json.dumps(piece) + SUFFIX

A complete, dependency-free example is in examples/custom-backend-openai-chat: it either forwards to an upstream OpenAI-compatible server, the place to plug your engine, or runs a scripted mode that lets Claude Code read and edit a file through the whole chain without a model.

Workload config your backend receives

The workload’s config reaches your process:

Key	Effect
(the whole config)	Passed to setup() as ctx.config (environment variable IK_BACKEND_CONFIG).
env	String values become environment variables of the backend process, for example {"CUDA_VISIBLE_DEVICES": "0,1,2,3"} or the path to your weights. PYTHONPATH is reserved.
ready_timeout_secs	How long the worker waits for setup() to finish (default 300 s, or the worker’s INFERENCEKEY_CUSTOM_BACKEND_READY_TIMEOUT_SECS). Raise it for models that take minutes to load.
request_timeout_secs	Longest time one request may take, a whole streamed answer included (default 300 s for /process and streaming). Raise it for long agent turns.

Two workloads of the same custom backend share one process only when their configs are equal.

7. The HTTP transport

The runtime can serve your backend with two HTTP servers. Both run the same request handling in Python (your methods, the status codes, the error bodies and the SSE data: lines are identical), so the platform cannot tell them apart.

Transport	What it is	When it is used
stdlib	The Python standard library http.server, one thread per connection.	The default in this release, and the fallback whenever the native one is unavailable.
native	A Rust server in the SDK’s native module. It accepts connections, parses HTTP, answers /healthz, /meta and /v1/models, frames SSE and writes every byte without holding the GIL; Python only runs your backend’s methods.	When you opt in with IK_BACKEND_HTTP.

Pick one with the IK_BACKEND_HTTP environment variable of the backend process. On the platform, set it in the workload config’s env, for example {"env": {"IK_BACKEND_HTTP": "auto"}}:

• stdlib (the default for now): the standard library server.
• auto: native when the native module is importable and provides the server, the standard library otherwise.
• native: native, or exit with code 2 when it is unavailable. Use it in CI and benchmarks so a silent fallback cannot pass unnoticed.

At startup the runtime writes exactly one line to stderr naming the transport in use:

http transport: native (abi 1, io_threads=4, backlog=4096, max_inflight=1024, flush_tick_us=1000)
http transport: stdlib (reason: IK_BACKEND_HTTP=stdlib (default))

Setting IK_BACKEND_HTTP=stdlib is the instant rollback for one workload; no release is needed.

Why native is opt-in for now

The native transport keeps /healthz answering in a few milliseconds while hundreds of streams are open (and while setup() holds the GIL), and it removes the per-token socket writes from Python. In our CPU load tests it does not yet meet every target for a simultaneous burst of hundreds of new streams, so auto will become the default in a later release, once those targets are met.

The native transport reads these optional settings from the environment (they are ignored by stdlib):

Variable	Default	Meaning
IK_BACKEND_MAX_INFLIGHT	1024	Concurrent backend calls; further requests wait in the server, they are not refused. Lower it for heavy /process payloads, since a waiting request holds its body (up to 16 MiB) in memory.
IK_BACKEND_IO_THREADS	min(4, CPUs)	Server IO threads.
IK_BACKEND_BACKLOG	4096	Listen backlog (the kernel’s net.core.somaxconn caps it).
IK_BACKEND_STREAM_MAX_BYTES	8388608	Bytes one stream may have queued for a client that stopped reading before it is closed as a slow client.
IK_SSE_CONNECTION_CLOSE	0	1 streams every response close-delimited, the way the standard library server does.
IK_BACKEND_FLUSH_TICK_US	1000	While busy, stream bodies are woken in batches this many microseconds apart (0, or 1000 to 100000). 0 wakes on every push.

When the native transport is selected, a setting with an invalid value makes the process exit with code 2, as does an unknown IK_BACKEND_HTTP value on either transport.

What you may notice on native

• Requirements. The native server ships in the SDK’s binary wheels (CPython 3.9+, Linux x86_64 and aarch64, glibc 2.34 or newer). On other platforms, or with a native module built before the server existed, auto falls back to stdlib and the startup line says why. A source checkout gets it with maturin develop --release in bindings/python.
• Framing. Streaming responses to HTTP/1.1 clients use chunked transfer encoding (HTTP/1.0 clients still get a close-delimited stream, and IK_SSE_CONNECTION_CLOSE=1 restores that for everyone). JSON responses use HTTP/1.1 keep-alive. Every SSE data: line is byte-identical to stdlib.
• Malformed requests. A non-numeric Content-Length gets the HTTP library’s own 400 with an empty body instead of the JSON 400. The status is the same.
• Shutdown. On SIGTERM a request still running is given up to 2 seconds to finish before the process exits.
• serve() and serve_backend() return an object with server_address, server_port, serve_forever(), shutdown() and server_close(), but it is not a ThreadingHTTPServer. Code that checks isinstance(httpd, ThreadingHTTPServer) or touches RequestHandlerClass or socket needs IK_BACKEND_HTTP=stdlib.
• Your methods are still called concurrently from many threads, exactly as with stdlib, so a backend must stay thread-safe.

What happens after that

You declared the workload; the platform reconciles it. When a worker is assigned, it pulls your published package, verifies it, installs your requirements.txt in an isolated environment, starts your backend once (your setup() runs, the model loads), and then forwards each incoming job to process(), or chat to chat() when your backend implements it. You never manage the process, the download, or the placement.

Backends & policies The built-in backends, the task-type modalities, and execution policies.