Advanced: build a container from scratch in Python

1
"""Model wrapper for serving flan-t5-small."""
2
import argparse
3
import typing
4
5
import torch
6
from transformers import T5ForConditionalGeneration, T5Tokenizer
7
8
_MODEL_NAME = "google/flan-t5-small"
9
"""The model's name on HuggingFace."""
10
11
_DEVICE: str = "cuda:0" if torch.cuda.is_available() else "cpu"
12
"""Device on which to serve the model."""
13
14
class Model:
15
    """Wrapper for a flan-t5-small Text Generation model."""
16
17
    def __init__(self):
18
        """Initialize the model."""
19
        self._tokenizer = T5Tokenizer.from_pretrained(_MODEL_NAME)
20
        self._model = T5ForConditionalGeneration.from_pretrained(_MODEL_NAME).to(
21
            _DEVICE
22
        )
23
24
    def predict(self, inputs: typing.Dict[str, str]) -> typing.Dict[str, str]:
25
        """Return a dict containing the completion of the given prompt.
26
27
        :param inputs: dict of inputs containing a prompt and optionally the max length
28
            of the completion to generate.
29
        :return: a dict containing the generated completion.
30
        """
31
        prompt = inputs.get("prompt", None)
32
        max_length = inputs.get("max_length", 2048)
33
34
        input_ids = self._tokenizer(prompt, return_tensors="pt").input_ids.to(_DEVICE)
35
        output = self._model.generate(input_ids, max_length=max_length)
36
        result = self._tokenizer.decode(output[0], skip_special_tokens=True)
37
38
        return {"completion": result}
39
40
    @classmethod
41
    def fetch(cls) -> None:
42
        """Pre-fetches the model for implicit caching by Transfomers."""
43
        # Running the constructor is enough to fetch this model.
44
        cls()
45
46
def main():
47
    """Entry point for interacting with this model via CLI."""
48
    parser = argparse.ArgumentParser()
49
    parser.add_argument("--fetch", action="store_true")
50
    args = parser.parse_args()
51
52
    if args.fetch:
53
        Model.fetch()
54
55
if __name__ == "__main__":
56
    main()

1
"""HTTP Inference serving interface using sanic."""
2
import os
3
4
import model
5
from sanic import Request, Sanic, response
6
7
_DEFAULT_PORT = 8000
8
"""Default port to serve inference on."""
9
10
# Load and initialize the model on startup globally, so it can be reused.
11
model_instance = model.Model()
12
"""Global instance of the model to serve."""
13
14
server = Sanic("server")
15
"""Global instance of the web server."""
16
17
@server.route("/healthcheck", methods=["GET"])
18
def healthcheck(_: Request) -> response.JSONResponse:
19
    """Responds to healthcheck requests.
20
21
    :param request: the incoming healthcheck request.
22
    :return: json responding to the healthcheck.
23
    """
24
    return response.json({"healthy": "yes"})
25
26
@server.route("/predict", methods=["POST"])
27
def predict(request: Request) -> response.JSONResponse:
28
    """Responds to inference/prediction requests.
29
30
    :param request: the incoming request containing inputs for the model.
31
    :return: json containing the inference results.
32
    """
33
    inputs = request.json
34
    output = model_instance.predict(inputs)
35
    return response.json(output)
36
37
def main():
38
    """Entry point for the server."""
39
    port = int(os.environ.get("SERVING_PORT", _DEFAULT_PORT))
40
    server.run(host="0.0.0.0", port=port, workers=1)
41
42
if __name__ == "__main__":
43
    main()

$
sanic==23.3.0
>
torch==2.0.0+cu118
>
--extra-index-url https://download.pytorch.org/whl/cu118
>
sentencepiece==0.1.97
>
transformers==4.27.4

1
FROM nvidia/cuda:11.8.0-cudnn8-runtime-ubuntu22.04
2
3
ARG SERVING_PORT=8000
4
ENV SERVING_PORT=$SERVING_PORT
5
6
WORKDIR /
7
8
RUN apt update && \
9
    apt install -y python3-pip
10
11
# Upgrade pip and install the copied in requirements.
12
RUN pip install --no-cache-dir --upgrade pip
13
ADD requirements.txt requirements.txt
14
RUN pip install --no-cache-dir -r requirements.txt
15
16
# Copy in the files necessary to fetch, run and serve the model.
17
ADD model.py .
18
ADD server.py .
19
20
# Fetch the model and cache it locally.
21
RUN python3 model.py --fetch
22
23
# Expose the serving port.
24
EXPOSE $SERVING_PORT
25
26
# Run the server to handle inference requests.
27
CMD python3 -u server.py

$
$ DOCKER_REGISTRY="XXX" # Put your Docker Hub username here
>
$ cd ./flan-t5-small
>
$ docker build -t "$DOCKER_REGISTRY/flan-t5-small-pytorch-sanic" -f Dockerfile .

$
$ docker run --gpus=all -d --rm
>
    -p 8000:8000 --env SERVER_PORT=8000
>
    --name "flan-t5-small-pytorch-sanic"
>
  	"$DOCKER_REGISTRY/flan-t5-small-pytorch-sanic"

$
$ curl -X GET http://localhost:8000/healthcheck

$
$ curl -X POST http://localhost:8000/predict \
>
    -H "Content-Type: application/json" \
>
    --data '{"prompt":"What state is Los Angeles in?","max_length":100}'

$
$ docker push "$DOCKER_REGISTRY/flan-t5-small-pytorch-sanic"