Skip to content

Deploy a custom API handler on a Custom Container

View as Markdown

Run your own inference code behind a Model Endpoint. Use a custom container when you need to control how requests are handled. It also lets you add steps before or after the model runs, or serve a model beyond what the other frameworks cover. You write a handler, package it in a Docker image, and deploy the image on the Custom Container framework.

  • Permission to create model endpoints in the project. A Member needs a policy that grants create permission for the Inference service. See Create an endpoint.
  • Docker and Python on the computer where you build the image.
  • A place for the image: a public registry, or a TIR Container Registry for a private image. See Push an image.

The handler below follows the KServe predict protocol: it loads the model once, then answers each request with a predictions list. Replace the two lines marked with ... with your own model code.

from kserve import Model, ModelServer
from typing import Dict
class MyCustomModel(Model):
def __init__(self, name: str):
super().__init__(name)
self.ready = False
self.load()
def load(self):
self.model = ... # load your model from disk or from a remote source
self.ready = True
def predict(self, payload: Dict, headers: Dict = None) -> Dict:
inputs = payload["instances"]
text = inputs[0]["text"]
output = ... # run your model on `text`
return {"predictions": output}
if __name__ == "__main__":
model = MyCustomModel("custom-model")
ModelServer().start([model])

Save the file as model_server.py. The server listens on port 8080 by default, which is the port an external endpoint sends requests to.

Create a Dockerfile next to model_server.py, with a requirements.txt that lists your Python packages:

FROM pytorch/torchserve-kfs:0.8.1-gpu
WORKDIR /app
COPY requirements.txt ./
RUN pip install --no-cache-dir -r requirements.txt
COPY model_server.py ./
CMD ["python", "model_server.py"]

Build the image and push it. <registry>/<image-name> is the address of the image in your registry, for example docker.io/<your-user>/custom-kserve-model for Docker Hub:

docker build -t <registry>/<image-name>:1.0 .
docker push <registry>/<image-name>:1.0

A smaller image starts faster, so install only the packages the handler uses.

  1. In the side menu, under Inference, select Model Endpoints, then Create Endpoint.
  2. In Choose Framework, select Custom Container.
  3. Under Download Type, keep Skip Model Download. The handler loads its own model.
  4. Under Container Configuration, choose Public or Private, then give the image. For a public image, type the Image, for example docker.io/<your-user>/custom-kserve-model:1.0. For an image in your TIR Container Registry, choose Private and set Registry Type to E2E Registry. Then pick the registry, the image, and the tag.
  5. Set Service Type to External so that you can call the endpoint with a token.
  6. Under Application Port, tick HTTP. On an external endpoint, HTTP is port 8080, which is the port the handler listens on.
  7. Keep Metrics Port at 8080.
  8. (Optional) If your handler needs values such as a model name or a token, open Advanced Settings and add them as environment variables. See Set environment variables.
  9. Under Select Computing Resource, choose a plan. A GPU card is the usual choice for large transformer models.
  10. Select Deploy Endpoint.

Wait until the state is Running. See Deploy a custom container for every field of the container form.

On the endpoint’s details page, select API Request. In Sample API Request, choose a token under Auth token, or select Generate Token to make one, and copy the Root Endpoint. Set both as variables in your terminal, where <root-endpoint> is the address you copied and <auth-token> is the token:

export ROOT_ENDPOINT="<root-endpoint>"
export AUTH_TOKEN="<auth-token>"

Send a request to the route your handler serves. The KServe route for the handler above is /v1/models/custom-model:predict:

curl -X POST "$ROOT_ENDPOINT/v1/models/custom-model:predict" \
-H "Authorization: Bearer $AUTH_TOKEN" \
-H "Content-Type: application/json" \
-d '{"instances": [{"text": "Hello!"}]}'

The reply holds your handler’s predictions. If the endpoint gets stuck before Running, open Events & Logs and read Worker Logs and Deployment Events. ImagePullBackOff means the worker failed to fetch the image, so check the image name, the tag, and the registry login. CrashLoopBackOff means the container keeps starting and stopping, so read the log for the error your code raised.

(Optional) A handler that downloads its model at every start waits for the download each time. To load the weights from the endpoint’s storage instead:

  1. Upload the model folder to a Model Repository. See Upload model files.
  2. On the endpoint form, under Download Type, choose Link with Model Repository. Then choose your repository and, if the files are in a folder, the Model Path.
  3. Change load() so that it reads from /mnt/models, the default Cache path, instead of the remote source.

See Deploy a model from Model Repository for the repository side.

(Optional) Deploy the code straight from GitHub. With a workflow file in your repository, the endpoint restarts whenever a pull request is merged, so it always runs your latest code.

  1. Save a GitHub personal access token as a GitHub integration. See Create a GitHub integration.

  2. Create an endpoint, and in Choose Framework select Custom Container.

  3. Under Download Type, choose Link with Github.

  4. In GitHub Repository, enter the repository and branch as https://github.com/<owner>/<repo>/tree/<branch>, for example https://github.com/myorg/ml-inference/tree/main.

  5. In GitHub Integration, choose your integration. TIR checks that the token can read the repository.

  6. Under Container Configuration, choose Public or Private for the image that runs your code.

  7. Under Advanced Configuration, pick Bash. In Entrypoint Command, enter the command that installs your requirements and starts your app. The endpoint downloads the repository to the Cache path, which starts as /mnt/models. A command looks like this, where <repo-name> is the name of your repository:

    pip3 install -r /mnt/models/<repo-name>/requirements.txt && python3 /mnt/models/<repo-name>/app.py

    If the app fails to start, check the path in Worker Logs.

  8. Choose the machine and the workers, then select Deploy Endpoint, and wait for Running.

To restart the endpoint when a pull request is merged:

  1. Download the workflow file. Before you deploy, the form offers it for download. Afterwards, select Download YML in the endpoint’s Actions menu, or Download yml on the Integrations tab.
  2. Save it in your repository as .github/workflows/deploy.yml.
  3. On the endpoint’s Integrations tab, copy the endpoint’s UUID into the workflow where it lists the endpoints to restart. You can list several.
  4. In your repository on GitHub, open Settings > Secrets and variables > Actions. Add two secrets: E2E_API_KEY with your E2E API key and E2E_AUTH_TOKEN with an auth token. The Integrations tab names both.
  5. Commit and push the workflow file.

When a pull request is merged into the branch, the workflow restarts the endpoints you listed, and they pull the latest code from the repository.

Your own code now answers behind an HTTPS address that only callers with a token can reach. When you ship a new image tag, change it with Actions > Update Container Configuration, and the rest of the endpoint stays as it is.