Model Serving made Efficient in the Cloud.
Mosec is a high-performance and flexible model serving framework for building ML model-enabled backend and microservices. It bridges the gap between any machine learning models you just trained and the efficient online service API.
Mosec requires Python 3.7 or above. Install the latest PyPI package with:
shell
pip install -U mosec
We demonstrate how Mosec can help you easily host a pre-trained stable diffusion model as a service. You need to install diffusers and transformers as prerequisites:
shell
pip install --upgrade diffusers[torch] transformers
Firstly, we import the libraries and set up a basic logger to better observe what happens.
```python from io import BytesIO from typing import List
import torch # type: ignore from diffusers import StableDiffusionPipeline # type: ignore
from mosec import Server, Worker, get_logger from mosec.mixin import MsgpackMixin
logger = get_logger() ```
Then, we build an API for clients to query a text prompt and obtain an image based on the stable-diffusion-v1-5 model in just 3 steps.
1) Define your service as a class which inherits mosec.Worker. Here we also inherit MsgpackMixin to employ the msgpack serialization format(a).
2) Inside the __init__ method, initialize your model and put it onto the corresponding device. Optionally you can assign self.example with some data to warm up(b) the model. Note that the data should be compatible with your handler's input format, which we detail next.
3) Override the forward method to write your service handler(c), with the signature forward(self, data: Any | List[Any]) -> Any | List[Any]. Receiving/returning a single item or a tuple depends on whether dynamic batching(d) is configured.
```python class StableDiffusion(MsgpackMixin, Worker): def init(self): self.pipe = StableDiffusionPipeline.frompretrained( "runwayml/stable-diffusion-v1-5", torchdtype=torch.float16 ) device = "cuda" if torch.cuda.is_available() else "cpu" self.pipe = self.pipe.to(device) self.example = ["useless example prompt"] * 4 # warmup (bs=4)
def forward(self, data: List[str]) -> List[memoryview]:
logger.debug("generate images for %s", data)
res = self.pipe(data)
logger.debug("NSFW: %s", res[1])
images = []
for img in res[0]:
dummy_file = BytesIO()
img.save(dummy_file, format="JPEG")
images.append(dummy_file.getbuffer())
return images
```
Note
(a) In this example we return an image in the binary format, which JSON does not support (unless encoded with base64 that makes it longer). Hence, msgpack suits our need better. If we do not inherit
MsgpackMixin, JSON will be used by default. In other words, the protocol of the service request/response can either be msgpack or JSON.(b) Warm-up usually helps to allocate GPU memory in advance. If the warm-up example is specified, the service will only be ready after the example is forwarded through the handler. However, if no example is given, the first request's latency is expected to be longer. The
exampleshould be set as a single item or a tuple depending on whatforwardexpects to receive. Moreover, in the case where you want to warm up with multiple different examples, you may setmulti_examples(demo here).(c) This example shows a single-stage service, where the
StableDiffusionworker directly takes in client's prompt request and responds the image. Thus theforwardcan be considered as a complete service handler. However, we can also design a multi-stage service with workers doing different jobs (e.g., downloading images, forward model, post-processing) in a pipeline. In this case, the whole pipeline is considered as the service handler, with the first worker taking in the request and the last worker sending out the response. The data flow between workers is done by inter-process communication.(d) Since dynamic batching is enabled in this example, the
forwardmethod will wishfully receive a list of string, e.g.,['a cute cat playing with a red ball', 'a man sitting in front of a computer', ...], aggregated from different clients for batch inference, improving the system throughput.
Finally, we append the worker to the server to construct a single-stage workflow (multiple stages can be pipelined to further boost the throughput, see this example), and specify the number of processes we want it to run in parallel (num=1), and the maximum batch size (max_batch_size=4, the maximum number of requests dynamic batching will accumulate before timeout; timeout is defined with the flag --wait in milliseconds, meaning the longest time Mosec waits until sending the batch to the Worker).
python
if __name__ == "__main__":
server = Server()
# 1) `num` specify the number of processes that will be spawned to run in parallel.
# 2) By configuring the `max_batch_size` with the value > 1, the input data in your
# `forward` function will be a list (batch); otherwise, it's a single item.
server.append_worker(StableDiffusion, num=1, max_batch_size=4, max_wait_time=10)
server.run()
The above snippets are merged in our example file. You may directly run at the project root level. We first have a look at the command line arguments (explanations here):
shell
python examples/stable_diffusion/server.py --help
Then let's start the server with debug logs:
shell
python examples/stable_diffusion/server.py --debug
And in another terminal, test it:
shell
python examples/stable_diffusion/client.py --prompt "a cute cat playing with a red ball" --output cat.jpg --port 8000
You will get an image named "cat.jpg" in the current directory.
You can check the metrics:
shell
curl http://127.0.0.1:8000/metrics
That's it! You have just hosted your stable-diffusion model as a service! 😉
More ready-to-use examples can be found in the Example section. It includes:
max_batch_size is configured when you append_worker (make sure inference with the max value won't cause the out-of-memory in GPU).--wait (default=10ms) is configured through CLI arguments (this usually should <= one batch inference duration).max_batch_size or the wait time.--shm-size flag if you are using docker.mosec_service_batch_size_bucket shows the batch size distribution.mosec_service_process_duration_second_bucket shows the duration for each stage (excluding the IPC time).mosec_service_remaining_task shows the number of currently processing tasks.mosec_service_throughput shows the service throughput.SIGINT or SIGTERM since it has the graceful shutdown logic.Here are some of the companies and individual users that are using Mosec:
We welcome any kind of contribution. Please give us feedback by raising issues or discussing on Discord. You could also directly contribute your code and pull request!
To start develop, you can use envd to create an isolated and clean Python & Rust environment. Check the envd-docs or build.envd for more information.
| | Batcher | Pipeline | Parallel | I/O Format(1) | Framework(2) | Backend | Activity |
| ----------------------------------------------------------- | :-----: | :------: | :------: | ------------------------------------------------------------------------------------------------------------------------------------------- | ----------------------- | ------- | ----------------------------------------------------------------------------- |
| TF Serving | ✅ | ✅ | ✅ | Limited(a) | Heavily TF | C++ | 
*As accessed on 08 Oct 2021. By no means is this comparison showing that other frameworks are inferior, but rather it is used to illustrate the trade-off. The information is not guaranteed to be absolutely accurate. Please let us know if you find anything that may be incorrect.
(1): Data format of the service's request and response. "Limited" in the sense that the framework has pre-defined requirements on the format. (2): Supported machine learning frameworks. "Heavily" means the serving framework is designed towards a specific ML framework. Thus it is hard, if not impossible, to adapt to others. "Multiple" means the serving framework provides adaptation to several existing ML frameworks. "Agnostic" means the serving framework does not necessarily care about the ML framework. Hence it supports all ML frameworks (in Python). (3): Flask is a representative of general purpose web frameworks to host ML models.