Serving ML Model using GraphQL Subscription

Serve machine learning model using GraphQL subscription.

Photo by h heyerlein on Unsplash

Overview

In previous articles, I have used FastAPI + WebSocket + React to serve machine learning models. As machine learning models getting more complicated, the time to come out with predictions also takes longer. Using WebSocket provides an asynchronous method for the server to update the client with the prediction when the evaluation is completed.

GraphQL Subscription is a feature that allows a server to send data to its clients when a specific event happens and is usually implemented with WebSocket.

In this article, let’s explore using GraphQL to serve machine learning models. I am going to use

For the machine learning model, I am going to use Photo2Cartoon model from PaddleGAN.

The frontend application takes a photo of your face and sends the image to the GraphQL endpoint for processing. As soon as the processing is done, the processed image is sent back to the frontend application.

For GraphQL beginners, go through this tutorial to find out more.

Backend — FastAPI + Ariadne

Unlike REST APIs, for GraphQL normally there is only 1 endpoint for which you need to define the schema for queries, mutations, and subscriptions.

GraphQL Schema

I defined a mutation to receive incoming messages and a subscription to send messages to the subscribed clients.

For mutation, it receives a message which contains the client id and content.

  • The incoming message is stored in an async queue.
  • A success message is returned immediately as processing is going to be performed asynchronously.
  • The incoming message is a base64 encoded image string.
GraphQL Mutation

For subscription,

  • it listens for incoming messages from the Python asyncio queue.
  • when a message arrives and matches a client id, it gets processed and the process message is sent back to the client as a base64 encoded image string.
GraphQL Subscription

Frontend — React + Apollo Client

The frontend application is developed using React.

  • When the app starts, it registers itself to the GraphQL backend and subscribes to the endpoint for new incoming messages.
  • It captures photos using a web camera.
  • The captured image is sent to the GraphQL endpoint.
  • When a new message arrives, it displays the image.
GraphQL Camera Viewer

And below is the image viewer using subscription.

ImageViewer using GraphQL Subscription
Screenshot

You can also check out this article to see how this model work.

Summary

Serving a machine learning model asynchronously using GraphQL requires more development efforts. Depending on your needs, do explore if it suits your use cases. You can also check out managed GraphQL services, e.g. AWS AppSync if you are already using cloud service providers like AWS.

The code that I used can be found in this repository.

Do also check out these articles!

Programmer and occasional blogger.

Get the Medium app

A button that says 'Download on the App Store', and if clicked it will lead you to the iOS App store
A button that says 'Get it on, Google Play', and if clicked it will lead you to the Google Play store