MLOps with Kubernetes, RabbitMQ and FastAPI | by Andrej Baranovskij | Towards Data Science

You often could hear people saying — many ML projects are stopped before they reach the production phase. One of the reasons for this, typically ML projects are implemented as monoliths from the start and when the time comes to run them in production, it is impossible to manage, transform and maintain the code. ML project code is implemented as a few or even one large notebook, where data processing, model training, and prediction all are glued together. This makes it hard to maintain such cumbersome code when the time comes to change the code and introduce user requests. As a result, users are not happy and this leads to project termination.
Much more effective to build ML system from the start and follow microservice architecture. You can use containers to encapsulate the logic. There can be a separate container for data processing, ML model training, and ML model serving. When running separate containers, not only simplifies code maintenance, but you could also scale containers separately and run them on different hardware. This can improve system performance.
The question comes, how you could implement communication between these services. I was researching available tools, such as MLFlow. These tools are great, but often they are too complex and large for the task. Especially when you want simply to run ML logic in different containers and that’s pretty much it. This is why I decided to build my own small and simple open-source product Skipper to run ML workloads.
In this article, I will explain how you can scale TensorFlow model on Kubernetes with Skipper. The same approach can be applied for PyTorch models or any other non-ML-related functionality.
The public port is exposed through Nginx
FastAPI is serving REST endpoints API. At the moment with providing two generic endpoints, one for async requests and another for sync
Workflow container responsible for request routing
Logger container provides generic logging capability
Celery container is used to execute the async request
RabbitMQ is a message broker, it enables event-based communication between Skipper containers
SkipperLib is a Python library, it encapsulates API code specific to RabbitMQ
A set of microservices is created as a sample containers, to show how Skipper works with ML specific (can be non ML too) services
You can run Skipper containers in multiple ways:
Directly with Python virtual environment on your machine
On Docker containers through Docker compose. Follow the readme file for the instructions
On Kubernetes. Follow the readme file for the instructions
In a production environment, you should run Skipper in Kubernetes, with Kubernetes it is easier to scale containers.
Skipper REST API is exposed through Kubernetes NGINX Ingress controller:
apiVersion: networking.k8s.io/v1kind: Ingressmetadata:  name: api-ingressspec:  rules:    - host: kubernetes.docker.internal      http:        paths:          - path: /api/v1/skipper/tasks/            pathType: Prefix            backend:              service:                name: skipper-api                port:                  number: 8000  ingressClassName: nginx---apiVersion: networking.k8s.io/v1kind: IngressClassmetadata:  name: nginxspec:  controller: k8s.io/ingress-nginx
Ingress redirects to /api/v1/skipper/tasks/, which is served from FastAPI container.
There are two generic endpoints defined. One to serve async requests and another to serve sync requests. Async request executes a call to train model, this is a long-running task. Sync request handles model prediction calls and routes requests to model serving. All requests travel through RabbitMQ message queue.
All calls to RabbitMQ are executed through SkipperLib. This allows to encapsulate RabbitMQ specific code inside the library and change it, without touching code in the containers.
Kubernetes Pod for model training runs two containers. First is a master container, which trains a model. The second is a side-car container, responsible for data preparation and processing.
We are running these two containers in the same Pod, because there is no need to scale them separately and it is more convenient to share data between containers when they run in the same Pod. This is not true about prediction logic, which we run in separate Pod, to be able to scale it separately, more about this in the next chapter.
Model training and data preparation containers are sharing the same Kubernetes volume for data storage.
Model training container:
volumeMounts:  - name: data    mountPath: /usr/src/trainingservice/models
Data preparation container:
volumeMounts:  - name: data    mountPath: /usr/src/dataservice/models
Mount paths are different, but the target location will be the same. Data accessed through these paths will be the same. Because both containers are using the same volume ‘data’:
volumes:- name: data  persistentVolumeClaim:    claimName: training-service-claim
When the model is trained and the model file is saved, we need to transfer it to the serving Pod, when the model prediction container runs. One of the solutions is to use external cloud storage and upload model files there. But if the model file is not too huge, Skipper allows to transfer it directly from training to serving Pod. Model is archived, encoded into a string, wrapped into JSON together with other metadata, and sent to RabbitMQ queue to be delivered to serving Pod.
The model structure is archived into a single file:
shutil.make_archive(base_name=os.getenv('MODELS_FOLDER') + str(ts),                    format='zip',                    root_dir=os.getenv('MODELS_FOLDER') + str(ts))
The archived model file is encoded into a base64 string:
model_encoded = Nonetry:    with open(os.getenv('MODELS_FILE'), 'rb') as model_file:        model_encoded = base64.b64encode(model_file.read())except Exception as e:    print(str(e))
In the last step, we wrap everything into JSON:
data = {    'name': 'model_boston_' + str(ts),    'archive_name': 'model_boston_' + str(ts) + '.zip',    'model': model_encoded,    'stats': stats_encoded,    'stats_name': 'train_stats.csv'}content = json.dumps(data)
This message is submitted to RabbitMQ for delivery. Model is sent through ‘fanout’ exchange on RabbitMQ, this allows to send the same data at once to all subscribers. By default, RabbitMQ would send the message to one subscriber at a time, which works great in a cluster as a load balancing. But in this case, we want all receivers in the cluster to get the new model, this is why we are using ‘fanout’ exchange.
This is how the message is published to ‘fanout’ exchange through RabbitMQ:
credentials = pika.PlainCredentials(self.username, self.password)connection = pika.BlockingConnection(    pika.ConnectionParameters(host=self.host,                              port=self.port,                              credentials=credentials))channel = connection.channel()channel.exchange_declare(exchange='skipper_storage',                         exchange_type='fanout')channel.basic_publish(exchange='skipper_storage',                       routing_key='',                       body=payload)connection.close()
Kubernetes Pod for model serving runs two containers. The master container is responsible to execute prediction requests using TensorFlow API. Side-car container listens for the messages from RabbitMQ, when the new model file is sent, decodes the file and extracts the model.
Both containers share the same storage.
Serving container:
volumeMounts:  - name: data    mountPath: /usr/src/servingservice/models/serving
Side-car container for model file processing:
volumeMounts:  - name: data    mountPath: /usr/src/servingservice/storage/models/serving/
Storage is mounted to the same volume claim:
volumes:- name: data  persistentVolumeClaim:    claimName: serving-service-claim
When the container responsible for model file processing receives the model, it executes similar steps, as the container in model training Pod, where the model was prepared to be sent through RabbitMQ:
data_json = json.loads(data)model_name = data_json['name']archive_name = data_json['archive_name']stats_name = data_json['stats_name']model_decoded = base64.b64decode(data_json['model'])stats_decoded = base64.b64decode(data_json['stats'])
It decodes the string, extracts the file.
Model serving Pod can be scaled to multiple instances. If instances would run on separate cluster nodes, then each node would receive the new model from RabbitMQ message. But if several instances would run on the single node, both of them would try to write the model into the same storage. We are handling the exception if one of the instances would fail.
The goal of this article is to introduce Skipper. Our open-source product for MLOps. Currently, this product is ready for production use. Our goal is to further enhance it, in particular, to add FastAPI security configuration, add more sophisticated workflow support and improve logging. We plan to test Skipper with Kubernetes auto-scaling functionality. We are using the Skipper platform to implement our ML services.
Skipper GitHub repo. Follow readme for setup instructions