Memory leak when using WordEmbeddings Class

[gebauerm]

Hey,

I encountered a bug while using GloveEmbeddings in a featurizer class.

Describe the bug
Multiple initializations and usage of WordEmbeddings, DocumentEmbeddings and their related embedding methods for a sentence, cause memory to add up, without releasing it after a function call. This leads to an OOM error in my application.
An example for reproduction can be found below.

To Reproduce
After each function call memory adds up and is not released, despite things being deleted and embeddings being cleared.
The clear_embeddings() method deletes objects in Sentence._embeddings as expected, but nonetheless no memory is released.

from flair.embeddings import WordEmbeddings, DocumentPoolEmbeddings
from flair.data import Sentence
from memory_profiler import profile


@profile
def featurizer_memory_test():
    glove_embeddings = WordEmbeddings("de")
    doc_embeddings = DocumentPoolEmbeddings([glove_embeddings])
    sentence = Sentence("Hello World!")
    glove_embeddings.embed(sentence)
    sentence.clear_embeddings()

    del glove_embeddings
    del doc_embeddings


if __name__ == "__main__":
    for i in range(2):
        featurizer_memory_test()

The bug is related to the initialization of WordEmbeddings("de"). Once I initialize them in a class I don't run into issues.

from flair.embeddings import WordEmbeddings, DocumentPoolEmbeddings
from flair.data import Sentence
from memory_profiler import profile

class Embedder:

    def __init__(self):
        self.glove_embeddings = WordEmbeddings("de")

    @profile
    def featurizer_memory_test(self):
        doc_embeddings = DocumentPoolEmbeddings([self.glove_embeddings])
        sentence = Sentence("Hello World!")
        doc_embeddings.embed(sentence)
        sentence.clear_embeddings()


if __name__ == "__main__":
    embedder = Embedder()
    for i in range(2):
        embedder.featurizer_memory_test()

However this is not a fix for me. I would have to write a singleton in my application, because despite it being a local variable in a function, whenever the function is done with execution the memory will not be released.

Expected behavior
I would assume that after each function call the memory gets released properly. Meaning when I run a memory profiler in this case, I would not have higher initial memory in the second run than in my first run.

Environment (please complete the following information):

• Ubuntu 18.04
• python 3.7.8
• flair 0.6.1.post1, flair 0.6.1, flair 0.7

[whoisjones]

Hi @gebauerm, we wrap a function inside the WordEmbeddings class with @lru_cache decorator in order to speed up returning the required embedding. This causes your memory leakage. You can easily overcome your issue by adding glove_embeddings.get_cached_vec.cache_clear() to your code.

However I wouldn't recommend to do it. Why? It doesn't seem clear to me why you would want to setup the exact same embeddings (heavy I/O operation) in an iterative way. You also don't need to make a singleton class. You can just put the embedding instantiation in front of your loop and pass the embeddings as arguments to the function. That will do the things you want to do (annotating text) much faster and memory efficient. If you have further questions feel free to reach out.

[gebauerm]

Hey,

thanks for your reply.
Instantiating the Embeddings outside of the loop was what I tried as well - as you can see above - however it wouldnt work for me.
I currently instantiate the embeddings every time a model is created. Model creation is triggered by an API call. Since I provide an Endpoint for it. So every API Call currently creates a new WordEmbeddings object and thus is "almost" like a loop. Hence I would have to instantiate the WordEmbeddings object once on application startup. This causes issues, e.g. when I have a variety of embeddings I have to instantiate all of them on application start up, despite them not being used in every API Call. This causes a huge unnecessary use of resources, which can be quite costly in e.g. corporate settings.
Also the class itself has unintended side effects - when I delete an object or exit a function I expect it to release memory, which is not the case. If it doesnt I simply lose control over resources.

One solution could be to provide an option for the lru_cache. I would rather wait some seconds more, than not being able to control resources of my application.

Otherwise I currently just can switch to a worker/queue setup and than kill the process every time its finished - which is not a good solution.

[whoisjones]

@gebauerm see the problem, makes way more sense now. Easiest solution would be just to use cache_clear() in your case. However I would consider refactoring your API, you definitely want to avoid instantiating an arbitrary number WordEmbeddings independent of how often your API is called. Otherwise you might end up creating dozens of instances in a second.

Most elegant way in my opinion would be a synchronized usage counter, deleting embeddings once they are not needed anymore. I haven't had time for a usage counter but here's a snippet (haven't tested the entire code) deleting embeddings every 10m if they are not used anymore. See also a the setup() function - you don't need to instantiate everything upon creation. This way you'll profit from having cached embeddings, no need to instantiate all necessary embeddings upon starting your application and unused embeddings got removed.

from flask import abort, Flask, jsonify, request
from flair.embeddings import WordEmbeddings
from flair.data import Sentence
import time
import atexit
from apscheduler.schedulers.background import BackgroundScheduler

app = Flask(__name__)

class EmbeddingStorage():
    def __init__(self):
        self.embeddings = {}

    def setup(self, model_key: str):
        if not model_key in self.embeddings:
            self.embeddings[model_key] = {"last_used": time.time(), "model": WordEmbeddings(model_key)}
            print("model setup complete.")

    def update_timestamp(self, model_key: str):
        self.embeddings[model_key]["last_used"] = time.time()

    def remove_unused_embeddings(self):
        print(self.embeddings)
        self.embeddings = {model_key: config for model_key, config in self.embeddings.items() if time.time() - config["last_used"] < 10}


embedding_storage = EmbeddingStorage()

scheduler = BackgroundScheduler()
scheduler.add_job(func=embedding_storage.remove_unused_embeddings, trigger="interval", seconds=10)
scheduler.start()

# Shut down the scheduler when exiting the app
atexit.register(lambda: scheduler.shutdown())

@app.route('/api/v1/embed', methods=['POST'])
def embed():
    if not request.json or not 'message' in request.json:
        abort(400)
    message = request.json['message']
    model_key = request.json['model']
    sentence = Sentence(message)
    embedding_storage.setup(model_key)
    embedding_storage.update_timestamp(model_key)
    embedding_storage.embeddings[model_key]["model"].embed(sentence)
    response = "embedded"
    return jsonify(response), 200

if __name__ == "__main__":
    app.run()

[gebauerm]

Hey,

sorry for the late reply.
I tried your example, but I simplified it a bit for me to make it easier to run it with the memory profiler.
As far as I can say the issue still persists. The memory is not released, despite the storage being overwritten.
Code is below.

Edit: You are right with the API call - i dont want dozens of embeddings to be initialized - a storage makes sense in this case. Thanks for the hint.

from flair.embeddings import WordEmbeddings
from flair.data import Sentence
import time
import atexit
from apscheduler.schedulers.background import BackgroundScheduler
from memory_profiler import profile


class EmbeddingStorage:
    def __init__(self):
        self.embeddings = {}

    def setup(self, model_key: str):
        if not model_key in self.embeddings:
            self.embeddings[model_key] = {"last_used": time.time(), "model": WordEmbeddings(model_key)}
            print("model setup complete.")

    def update_timestamp(self, model_key: str):
        self.embeddings[model_key]["last_used"] = time.time()

    def remove_unused_embeddings(self):
        print(self.embeddings)
        self.embeddings = {model_key: config for model_key, config in self.embeddings.items() if time.time() - config["last_used"] < 1}


@profile
def embed(model_key, message):
    sentence = Sentence(message)
    embedding_storage.setup(model_key)
    embedding_storage.update_timestamp(model_key)
    embedding_storage.embeddings[model_key]["model"].embed(sentence)
    response = sentence.get_embedding()
    return response


if __name__ == "__main__":

    embedding_storage = EmbeddingStorage()

    scheduler = BackgroundScheduler()
    scheduler.add_job(func=embedding_storage.remove_unused_embeddings, trigger="interval", seconds=2)
    scheduler.start()

    # Shut down the scheduler when exiting the app
    atexit.register(lambda: scheduler.shutdown())

    for i in range(2):
        embed("de", "Eine test nachricht!")
        time.sleep(2)

[whoisjones]

Hi @gebauerm, will be fixed. functools lru_cache keeps a reference on a WordEmbedding instance thus the leak without explicitly clearing the cache. methodtools lru_cache decorator keeps track of the self object lifetime and thus releases memory when not referenced anymore.

[gebauerm]

Okay, thanks 👍

Will check it out in future.