How to use multiprocessing.Pool to run spaCy?

[Jiebro0109]

I'm trying to process 100 Parquet files (~1GB each, ~10M rows) for Named Entity Recognition (NER) using spaCy’s en_core_web_trf model with multiprocessing.Pool (2 processes) on a single 8GB GPU. Each process loads a file, applies NER to extract ORG entities, and saves results to a Parquet file. However, the program hangs after printing "Pool started" with no further output, and the GPU is not utilized (nvidia-smi shows 0% usage).

I want to confirm if multiprocessing.Pool can be used to process multiple files in parallel with spaCy on a single GPU, and how to resolve the hanging issue.

Here’s a minimal example of my script:

import pandas as pd
import spacy
import torch
import multiprocessing as mp
import os
import glob
from datetime import datetime

def extract_org(text):
    if pd.isna(text):
        return []
    doc = nlp(text[:512])
    orgs = [ent.text for ent in doc.ents if ent.label_ == "ORG"]
    return orgs

def process_file(file_path):
    print(f"\nProcessing {file_path}, PID {os.getpid()}")
    time_start = datetime.now()
    
    if not spacy.prefer_gpu():
        print("Warning: GPU not available, using CPU")
    else:
        print("GPU enabled")
    
    global nlp
    try:
        nlp = spacy.load("en_core_web_trf", disable=["tagger", "parser", "attribute_ruler", "lemmatizer"])
        print("Model loaded")
    except Exception as e:
        print(f"Failed to load model: {e}")
        return
    
    try:
        df = pd.read_parquet(file_path, columns=['description'])
        df['description'] = df['description'].astype(str)
        print(f"Parquet file loaded, {len(df)} rows")
        
        df['org_mention'] = df['description'].apply(extract_org)
        df['org_mention'] = df['org_mention'].apply(lambda x: ', '.join(x) if x else '')
        
        output_dir = r"C:\Users\ASUS\OneDrive\Desktop\Test of spacy NER\result data"
        output_path = os.path.join(output_dir, f"{os.path.basename(file_path)}")
        os.makedirs(output_dir, exist_ok=True)
        df[['description', 'org_mention']].to_parquet(output_path, engine='pyarrow', index=False)
        print(f"Saved {output_path}")
        
        del df
        torch.cuda.empty_cache()
    except Exception as e:
        print(f"Error processing {file_path}: {e}")
    
    print(f"File {os.path.basename(file_path)} time: {datetime.now() - time_start}")

if __name__ == '__main__':
    print(f"PyTorch CUDA available: {torch.cuda.is_available()}")
    if torch.cuda.is_available():
        print(f"GPU name: {torch.cuda.get_device_name(0)}")
    
    try:
        mp.set_start_method('spawn')
    except RuntimeError:
        pass
    
    file_list = glob.glob(r"C:\Users\ASUS\OneDrive\Desktop\Test of spacy NER\raw data\batch data\linkup_job_descriptions_*.parquet")
    print(f"Found {len(file_list)} files: {file_list}")
    if not file_list:
        raise ValueError("No Parquet files found")
    
    num_processes = 2
    print(f"Using {num_processes} processes on 1 GPU")
    
    try:
        print("Starting multiprocessing pool")
        with mp.Pool(processes=num_processes) as pool:
            print("Pool started")
            pool.map(process_file, file_list)
    except Exception as e:
        print(f"Multiprocessing failed: {e}")