fix: Reduce memory usage when publishing prediction log to kafka

[tiopramayudi]

Description

We've seen gradual increase of memory usage when model observability is enabled for a model.

First hypothesis why this happened is due to using asyncio, because we pass prediction input and output to the async function. We try to prove it by reducing sampling rate to 0, since the async function that being called need to publish to kafka, so we need to isolate this only for asyncio overhead. After set sampling rate to 0 the memory usage is stable there is no gradual increase

Since first hypothesis is not correct, we have new hypothesis that this is due to publishing the data to kafka, and we did memory profiler to the model

PS: We use memray as profiler https://github.com/bloomberg/memray

We see that the memory usage is keep increasing and producing the message to kafka contribute to this.

Modifications

To solve this problem, kafka producer must call poll after publish the message, this is necessary so ack buffer from producer will be drained and the memory usage won't gradually increase, ref: 1 , 2

After the changes

Tests

Checklist

• Added PR label
• Added unit test, integration, and/or e2e tests
• Tested locally
• Updated documentation
• Update Swagger spec if the PR introduce API changes
• Regenerated Golang and Python client if the PR introduces API changes

Release Notes

[ghost]

👇 Click on the image for a new way to code review

Legend

[khorshuheng]

lgtm, thanks!

[ariefrahmansyah]

On a separate note, what tool are you using for memory profiling? @tiopramayudi

[tiopramayudi]

On a separate note, what tool are you using for memory profiling? @tiopramayudi

I use memray as profiler https://github.com/bloomberg/memray