OOM when using pyjnius with pipe function

[nishan-canwill]

When I am using nlp.pipe function for identifying ners in my script and I have imported pyjnius the memory shoots up very fast.
The file that i am trying to read is a 4gb file. Here is my script.

import spacy
import os
from datetime import datetime
# from jnius import autoclass
import math
import re

class TestSpacyMemory:
    def __init__(self):
        self.nlp=spacy.load('en_core_web_md',exclude=["tok2vec", "tagger","parser", "attribute_ruler", "lemmatizer","entity_linker","textcat_multilabel","trainable_lemmatizer","morphologizer","senter","sentencizer","transformer"])
    
    def read_in_chunks(self, file_object, chunk_size):
        """
        Lazy function (generator) to read a file piece by piece.
        """
        while True:
            st = datetime.now()
            data = file_object.read(chunk_size)
            if file_object.tell()< 0  or not data:
                break
            yield data

    def call_json_parser(self):
        chunk_size = 1000000        
        text_file_path = 'large_test_file.txt'        
        with open(text_file_path, encoding='unicode_escape') as f:            
            cnt = 0
            prev_seek  = 0  
            for text in self.read_in_chunks(f, chunk_size):
                if prev_seek != 0 and prev_seek == f.tell():
                    break  
                if cnt==25:
                    break           
                cnt = cnt + 1 
                ab = datetime.now()             
                res = self.get_person_name(text)
                prev_seek = f.tell()             
                print(f"time taken for iteration {cnt} {datetime.now()-ab} ")
                f.seek(prev_seek-100)
                del res

    
    def get_chunks(self,text,step_size,step_back):        
        start_index = 0
        end_index = step_size
        chunk = []
        for i in range(0,math.ceil(len(text)/step_size) + 1):
            chunk.append(text[start_index:end_index])
            start_index = end_index - step_back
            end_index = start_index + step_size
        return chunk

    
    def get_person_name(self, text):        
        people = []
        main_results = []
        count = 0
        step_size = 2000
        step_back = 100
        start_index = 0
        try:            
            chunks = self.get_chunks(text, step_size, step_back)                  
            for doc in self.nlp.pipe(chunks,n_process=os.cpu_count()-2,batch_size=16):                    
                if count>0:
                    start_index = count * step_size - count * step_back
                count = count + 1
                if doc.ents:
                    for ent in doc.ents:
                        if ent.label_ == 'PERSON':                             
                            main_results.append(ent.text)
                            span_data = (ent.start_char + start_index, ent.end_char + start_index, ent.text)
                            if span_data not in people:
                                people.append(span_data)
            
            del chunks
            del text
        except Exception as e:
            print("error",e)
        if main_results:
            return True, list(set(main_results)), people
        else:
            return False, "No Match", []


if __name__ == "__main__":
    tsm = TestSpacyMemory()
    tsm.call_json_parser() 
    print("done")

Your Environment Details

• Operating System:Ubuntu 20.04.4 LTS
• Python Version Used: 3.7.6
• spaCy Version Used: 3.4.1

[polm]

Sorry you're having an issue with this. To be clear, is the issue that spaCy is using a lot of memory, or is it that it only uses a lot of memory when you're also using pyjnius? I'm not familiar with pyjnius, but based on your sample code it's not clear why you're using it here, since it seems to be for calling Java code from Python.

More generally, the way you're cutting your doc into chunks and passing it to nlp.pipe looks pretty reasonable as a start, except that you seem to be cutting at arbitrary characters so you may cut sentences or words in half, which would hurt performance. It would be better to cut on a newline or other natural text boundary.

Besides that... Perhaps your batch size is large for the machine you have?

[polm]

I tried running your code with a 2MB text file. It works fine without pyjnius, but with the import it hangs after printing "done". CTRL+C makes it quit quietly, but SIGQUIT (CTRL+) prints a Java thread dump. So it definitely looks like the JVM or pyjnius is doing something weird.

Sorry, but based on the info we have this looks like an issue with pyjnius rather than a spaCy problem. If you ask their community about this feel free to leave a link to any related threads here, and let me know if we can help you debug something specific.

[nishan-canwill]

So this is the ticket that I opened with Pyjnius kivy/pyjnius#638 and someone has replied

"""Each forked process created by multiprocessing will need to start its own JVM. You may need to change the multiprocess init to ensure that your JVM setup is done appropriately. I did a similar thing at:
https://github.com/terrier-org/pyterrier/blob/master/pyterrier/parallel.py#L54
"""

If we need to implement the same in spacy where should I do it?

[polm]

Multiprocessing changes generally need to be made before you do any forking, so anywhere at the top of your code is fine.

[nishan-canwill]

can you point me to the file where spacy has done multiprocessing in pipe function so that i can see what needs to be done

[polm]

The multiprocessing code related to nlp.pipe is here.