MLOps with Kubernetes, RabbitMQ and FastAPI

Skipper is a simple and flexible open-source ML workflow engine. It helps to run and scale ML services in production.

Author: Andrej Baranovskij

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.

Skipper structure, Author: Andrej Baranovskij

• 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

REST API

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/v1
kind: Ingress
metadata:
  name: api-ingress
spec:
  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/v1
kind: IngressClass
metadata:
  name: nginx
spec:
  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.

Model training

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 = None
try:
    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()

Model serving

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.

Conclusion

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.

Source code

• Skipper GitHub repo. Follow readme for setup instructions

YouTube tutorial

[embed]