Schemafull streaming data processing in ML pipelines | by Ivan Senilov | Towards Data Science

In one of my previous articles on Machine Learning pipelines, message queues were touched as an alternative to HTTP client-server architecture which is the most common way to serve ML models nowadays.
Just to refresh, here are the advantages of using queues like Apache Kafka in ML pipelines:
Naturally asynchronous processing with decoupling of an input data producer and data processor (consumer in Kafka terminology), i.e. failure of ML inference/training service consuming the data will not result in HTTP timeouts, data loss, and failure of the data producer but just in delayed processing
Better compute resources utilization for CPU/GPU-bound tasks such as ML inference/training:- Batch inference is usually more efficient as ML models do not need to be warmed up and torn down for every single prediction- The queue smoothens load if incoming data arrives irregularly whereas an HTTP server would be overwhelmed during load peaks and idling during quiet periods
Improved throughput. However, it is worth mentioning that queues are not perfect when inference latency is crucial. In these cases, Kafka consumers with a very short polling interval (which removes some of the previous advantages)or HTTP ML model server can still be viable options.
It is advised to readers to familiarize themselves with Apache Kafka and Docker (including docker-compose) before proceeding with the article. Confluent has a nice introduction page for Kafka. Docker docs are also great.
The majority of tutorials on how to use Kafka, especially with Python, show examples of producing and consuming schemaless JSON strings. Even though it is easy to get those examples running, they are not very suitable for production use. Let’s see why.
Producer can’t enforce data format (set of fields) it publishes to a topic which may cause incompatibilities and runtime errors on the consumer side. Of course, it can be handled in the code but it would make the code much more complicated and consequently error-prone and requiring more maintenance.
Data format evolution does not have constraints so the producer can potentially switch to a new schema without a field consumer expects which may lead to runtime errors again (see the next section for details)
Data in JSON format includes field names with every message which makes it inefficient in terms of storage space
Schemafull messages solve these problems, however, it comes with a price of more complicated infrastructure and required carefulness with schema evolution.
Kafka supports AVRO, Protobuf, and JSON-schema (this still has the drawback of JSON data format being non-binary and not very efficient in terms of storage). We will use AVRO in the article’s code as this seems to be the most common schema format for Kafka.
Once defined, schema usually can’t be arbitrarily changed. One mentioned example is field removal or renaming. This potential issue can be solved by defining fields as nullable and with default value so if the field becomes obsolete — the producer just stops populating it while it remains in the new updated version of the schema:
{"name": "value", "type": ["null", "double"], "default": null}
Other schema evolution considerations are described in the AVRO specification.
This repository has a basic example of Python AVRO consumer and producer:
github.com
Let’s review the services it consists of looking at services section of docker-compose.yaml
Apache ZooKeeper is used by Kafka for maintaining configuration in a centralized location and is a prerequisite for a working Kafka cluster. Consumers and producers usually do not interact with it directly so we omit its detailed description here.
Confluent Schema Registry is a service for storing and retrieving schemas. Apparently, one can use Kafka without it just having schemas together with application code but the schema registry ensures a single source of truth for versioned schemas across all the consumers and producers excluding the possibility of schemas deviation between those.
The registry has both RESTful API as well as native support by confluent-kafka-client which is defacto standard for working with Kafka in Python.
Registering a schema with the HTTP request can be done like this:
SCHEMA=$(sed 's/"/\\"/g' < ./Message.avsc)curl -X POST -H "Content-Type: application/vnd.schemaregistry.v1+json" \  --data "{\"schema\":\"${SCHEMA//$'\n'}\"}" \  http://<registry_host>:<registry_port>/subjects/<name>/versions
It should return the id (version) the schema was registered with: {“id”:1}
Using the schema from Python code can be done by instantiating AvroSerializer and AvroDeserializer from the aforementioned client.
Apache Kafka cluster itself is run using a Docker image from Confluent which has everything needed for its operation.It has a lot of configuration options but we limit ourselves to basic working options, for example, PLAINTEXT protocol that does not perform any encryption.
This is a short-lived instance of Kafka used solely for the creation of the topic we plan to use in the demo with kafka-topics CLI:
kafka-topics --create --if-not-exists --topic Test.Topic --bootstrap-server kafka:29092 --replication-factor 1 --partitions 1
replication-factor and partitions options are set to 1 for simplicity. Refer to Kafka documentation for more details on these and many more other options.
This is the consuming and producing app itself. Both consumer and producer are in the __main__.py for simplicity. The code is mostly self-explanatory so let’s look only at some important parts of it.
The producer configuration:
"value.serializer": AvroSerializer(schema_str=Message.schema,schema_registry_client=schema_registry_client, to_dict=todict)
schema_str=Message.schema — producer requires passing a schema and here we pass the property of generated Python class that contains the schema. The schema will be uploaded and registered in schema_registry_client .
to_dict=todict — this must be a callable that converts the message object to a Python dictionary. We use todict method from helpers.py from generated code.
The consumer configuration:
"value.deserializer": AvroDeserializer(schema_str=None, schema_registry_client=schema_registry_client,from_dict=lambda obj, _: Message(obj))
Here we don’t pass the schema so the schema_registry_client will fetch it from the schema registry.from_dict=lambda obj, _: Message(obj) — the opposite operation of converting a dictionary to a message object. We need to perform a little hack here because the deserializer passes context as a second argument to the callable (which we ignore) but the generated class constructor accepts only the dictionary.
"group.id": "test_group"
“group.id”. This allows multiple consumers (read multiple containers with data processing apps for horizontal scaling) sharing the same consumer group ID to consume messages in turns as they share the Kafka offset. At the same time, if one needs to use the messages for another purpose, for example, test an experimental ML model online, it can be done by deploying this experimental group with a new group.id avoiding messing up with the main production formation.
The repository’s README.md has instructions for running the whole bunch of services with docker-compose . Executing the commands should result in something like this in stdout:
worker_1  | Waiting for schema registry to be availableworker_1  | Producing message to topic 'Test.Topic'worker_1  | Produced message: {'timestamp': 1639304378, 'data': {'count': 10, 'value': 100.0}}worker_1  | worker_1  | Flushing records...worker_1  | Consumed message: {'timestamp': 1639304378, 'data': {'count': 10, 'value': 100.0}}
Which means producing and consuming was performed successfully.
We print the consumed message:
if message is not None:    print(f"Consumed message: {message.dict()}")
but this is where inference code usually goes with subsequent publishing the resulting predictions to another topic, saving them to a database or using them in a downstream API call.
In general, monitoring of the Kafka consuming ML service is similar to what was described in the Complete Machine Learning pipeline for NLP tasks article.
However, inference time can (and probably should) be accompanied by the consumer lag which shows how much the consumer is behind the published data. If it grows, it usually means the consumer formation must be scaled. When using a librdkafka-based client, like confluent-kafka-python used in this example, consumer lag can be obtained using statistics returned by librdkafka as explained in this issue.
The article shows why using schemas with Kafka might be a good idea and how it can be implemented using Python, the language of choice for ML services. As always, if you have any questions or suggestions regarding the content — please do not hesitate to reach me on LinkedIn.