Yet another deep learning serving framework
Tags: deeplearning
Published: 2020-05-13 17:06:19


Yet another deep learning serving framework that is easy to use.

Previously, I tested the performance of some [deep learning serving frameworks](https://dev.to/kemingy/deep-learning-serving-benchmark-1nko) like TensorFlow Serving, Triton, and I found that these frameworks are not that easy to use. By the way, they don't have much advantage in the performance. So I just write one as a prototype.

* [ventu](https://github.com/kemingy/ventu)
* [batching](https://github.com/kemingy/batching)

~~Feel free to give it a try~~. For production usage, check [MOSEC](https://github.com/mosecorg/mosec).

## Basic features

* serve the deep learning models (HTTP)
* preprocess and postprocess (optional)
* dynamic batching (increase the throughput)
* health check (need to provide examples)
* request & response validation
* model inference warm-up (need to provide examples)
* OpenAPI document
* supports both JSON and msgpack serialization

## Advantages

* support all kinds of deep learning runtime
* easy to implement the preprocess and postprocess part
* validation for request
* health check and warm-up with examples
* OpenAPI document

## Design

### Dynamic Batching

To implement the dynamic batching, we need a high-performance job queue that can be consumed by multiple workers. Go channel will be a good choice. In this situation, we have one producer and multiple consumers, so it's very easy to close the channel for the graceful shutdown.

```go
type Batching struct {
	Name       string // socket name
	socket     net.Listener
	maxLatency time.Duration // max latency for a batch inference to wait
	batchSize  int // max batch size for a batch inference
	capacity   int // the capacity of the batching queue
	timeout    time.Duration // timeout for jobs in the queue
	logger     *zap.Logger
	queue      chan *Job // job queue
	jobs       map[string]*Job // use job id as the key to find the job
	jobsLock   sync.Mutex // lock for jobs
}
```

For jobs in this queue, we need to create a UUID as a key. So after the inference, we can find this job by searching the key in a hash table. That means we also need a mutex for the hash table.

```go
type Job struct {
	id        string
	done      chan bool
	data      []byte // request data
	result    []byte // inference result or error message
	errorCode int // HTTP Error Code
	expire    time.Time
}
```

Because the batching service and Python inference workers are on the same machine (or the same pod), so the most efficient communication should be the [Unix domain socket](https://en.wikipedia.org/wiki/Unix_domain_socket). And we also need to define a simple protocol for our use case. Since we only need to transfer the data of batch jobs, let's keep everything as simple as we can.

```
| length  |       data        |
| 4 bytes |   {length} bytes  |
```

1. workers send the first request with empty data to the batching service
2. batching service collects a batch of jobs and sends to the workers
3. worker processes these jobs
    * preprocess (for a single job)
    * inference (for a batch of jobs)
    * postprocess (for a single job)
    * send to the results to the batching service
4. batching service notifies the handler that this job is done, then the handler sends the result to the original client and goes to #2

### Error handling

* timeout

If a job is not processed by one of the workers for a long time, the batching service will delete this job from the hash table and return 408.

When the batching service tries to collect these jobs from the queue channel, it will check the `expire` attribute first.

* validation error

To make sure the requested data is valid, we use [pydantic](pydantic-docs.helpmanual.io/) to do the validation. So the user needs to define the data schema with `pydantic`.

If one job data is invalid, this one will be marked and the result for this job is the validation error message generated by `pydantic`. And this won't affect other jobs in the same batch. That part is handled by the `ventu`.

### Simple HTTP service without dynamic batching

For this part, we use [falcon](falcon.readthedocs.io/) which is a very powerful Python framework for web APIs. To generate the OpenAPI document and validate the request data, we use [spectree](https://github.com/0b01001001/spectree).

If you would like to use `gunicorn`, `ventu` also expose the `app` element.

## TODO

* metrics
  * this can be added by users in model inference part
* increase the number of workers dynamically