Python 3. Как получить доступ к словарю запущенного из другого модуля?

У меня есть готовый сервер Flask где запускаются nlp модели. Пока к нему было подключено 5 устройств, не было проблем, потом мы увеличили это количество до 15 и тут началась проблема с производительностью. Решили увеличить количество ядер на этом сервере с 12 до 24, но получили прибавку производительности в районе 20%. Решили опробовать другой метод, запустили 3 сервера через Gunicorn и taskset --cpu-list (выдав по 8 ядер каждому серверу) на разных портах с ограничение по 5 устройств на сервер, тут производительность выросла в почти в 2 раза, стала приблизительно такой же как на сервере с 12 ядрами когда было 5 устройств. Но тут возникла другая проблема, каждый сервер запускал свои модели изолировано друга от друга и потребление озу выросло в 3 раза.

Модели загружаются в словарь по этому коду

import torch
import torch.nn as nn
import torchvision
import os
import learning

app = Flask(__name__)

models_list_name = os.listdir('models')

global_dict_models = {}

for token in models_list_name:
    try:
        tok, n_class, model_type, typ, h_layer = token.split('.')

        t, model = learning.get_model(model_type, typ)

        n_class = int(n_class)
        h_layer = int(h_layer)
        model = learning.NewModel(model, n_class, typ, h_layer)
        model.load_state_dict(torch.load(os.path.join('models', token), map_location=torch.device('cpu')))
        
        print(model, tok)
        with open(tok, 'r') as f:
            label = f.read().split('\n')
        print(label)
        global_dict_models[tok] = (t, model, label)
    except ValueError:
        pass          

а когда нужен предикт, то запрос просто обращается к словарю

t, model, labels = global_dict_models[token]
x = t.encode_plus(text, add_special_tokens=True, max_length=512, truncation=True, padding="max_length", return_tensors='pt')
    
output = torch.sigmoid(model(x['input_ids'].squeeze(1), x['attention_mask'])).detach().cpu().numpy()

Модели на всех серверах одни и те же. Хочу сделать единый модуль с загрузкой моделей, где формируется словарь и что бы каждый сервер через функцию просто к нему обращался.

Пробовал создать отдельный модуль по такому принципу

from flask import Flask, request, make_response
import torch
import torch.nn as nn
import torchvision
import os
import learning

app = Flask(__name__)

models_list_name = os.listdir('models')

global_dict_models = {}

def load_models():
    for token in models_list_name:
        try:
            tok, n_class, model_type, typ, h_layer = token.split('.')
    
            t, model = learning.get_model(model_type, typ)
    
            n_class = int(n_class)
            h_layer = int(h_layer)
            model = learning.NewModel(model, n_class, typ, h_layer)
            model.load_state_dict(torch.load(os.path.join('models', token), map_location=torch.device('cpu')))
            
            print(model, tok)
            with open(tok, 'r') as f:
                label = f.read().split('\n')
            print(label)
            global_dict_models[tok] = (t, model, label)
        except ValueError:
            pass
  

def dict_models(token):
    t, model, labels = global_dict_models[token]
    return t, model, labels    
    
if __name__ == '__main__':
    load_models()
    print('Start')
    app.run(host="0.0.0.0", port="5000", threaded=True, processes=1)

И тут возникла проблема, что не получается прочитать словарь, который был запущен в другом модуле.

Подскажите, как загрузить словарь один раз и иметь к нему доступ из других модулей? Или как решить эту проблему другим способом?


Ответы (0 шт):