Skip to content

Tutorial

Introduction

This tutorial provides a step-by-step guide to building a Docker image for your model.

Prerequisites

Before building the image, ensure you have the following information:

  • Model name and version: Our example model is simple-detector, version v1.
  • Harm categories and modality: This model detects violent and hateful content in images.
  • Input and output schemas:
  • Input: A JSON object containing a base64-encoded image.
  • Output: A JSON object with two float values (violence and hate), ranging from 0 to 1. Users can define thresholds for blocking content.

Example input:

{
    "image": "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAQAAAC1HAwCAAAAC0lEQVR42mNk+A8AAQUBAScY42YAAAAASUVORK5CYII="
}

Example output:

{
    "violence": 0.4,
    "hate": 0.9
}

  • Inference code and checkpoints: Our example contains model.py and a checkpoint file model.pt.

Building the Image

Inference Code

The model implementation consists of four methods: initialization, preprocessing, postprocessing, and inference. You can structure the inference code to best fit your use case.

model.py

import os

class ModelImp:
    def __init__(self):
        model_path = os.environ.get("ModelPath", ".")
        self.model = load_model(model_path)

    def preprocess(self, image: str):
        pass

    def postprocess(self, raw_output):
        pass

    def inference(self, image: str) -> dict:
        model_input = self.preprocess(image)
        raw_output = self.model.run(model_input)
        return self.postprocess(raw_output)

Service Code

The service code enables the model to handle API requests. We use FastAPI for its high performance and simplicity, but you may use any suitable web framework.

http_server.py

from fastapi import FastAPI, HTTPException
from model import ModelImp
from pydantic import BaseModel

model = ModelImp()
app = FastAPI()

class RaiRequest(BaseModel):
    image: str

@app.get("/liveness")
async def liveness():
    return {"status": "ok"}

@app.get("/readiness")
async def readiness():
    return {"status": "ok"}

@app.post("/score")
async def score(r: RaiRequest):
    try:
        return model.inference(r.image)
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))

Entry Scripts

The following scripts start the container on Singularity.

run.sh

#!/bin/bash
uvicorn http_server:app --host 0.0.0.0 --port $_ListeningPort_ --workers $_Workers_

run

#!/bin/bash
bash /simple-detector/run.sh

Building the Image

The directory structure should be as follows:

/simple-detector
├── model.py
├── http_server.py
├── run.sh
└── model.pt

The run file must be placed under /var/runit, such as /var/runit/model.

Dockerfile

FROM base_image:tag

RUN apt-get update && apt-get install -y runit

COPY ./requirements.txt /simple-detector/requirements.txt
RUN python3 -m pip install --no-cache-dir -r /simple-detector/requirements.txt

COPY ./run /var/runit/model/run
RUN chmod +x /var/runit/model/run
COPY simple-detector /simple-detector

ENV PYTHONPATH=/simple-detector/model
ENV _ListeningPort_=8888
ENV _Workers_=2

EXPOSE 8888

ENTRYPOINT ["runsvdir", "/var/runit"]

Verification

To verify that the container is running correctly:

  1. Start the container:

    docker run -it --gpus=all -p 8888:8888 your_image:tag
    

  2. Send a test request:

    curl -X POST http://localhost:8888/score -H "Content-Type: application/json" -d '{"image": "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAQAAAC1HAwCAAAAC0lEQVR42mNk+A8AAQUBAScY42YAAAAASUVORK5CYII="}'