在Gunicorn /烧瓶服务器中使用Google的JAX

问题描述

我想提供一个应用程序，该应用程序使用flask和gunicorn处理Google JAX框架中的数据。

如果在烧瓶中运行，一切正常。 一旦我在gunicorn中运行该应用程序，与jax相关的每个部分都会导致工作进程死亡，而不会引发任何异常。我尝试将sync和gthreads用作工作程序，但结果相同。

我试图通过将相同的调用包装在ThreadPoolExecutor和ProcesspoolExecutor中来查看JAX是否可以处理多处理和多线程，并且这种方法可以完美地工作。

import jax

import logging
logging.basicConfig(format="%(asctime)s | %(name)12.12s | %(message)s")
logger = logging.getLogger("Main")
logger.setLevel(logging.DEBUG)

from concurrent.futures import ThreadPoolExecutor,ProcesspoolExecutor,as_completed

from fit.optimization.vectorize import BatchNumpyInterface,batch_calculate_fit

def warmup():
    logger.debug("Warmup")
    data = BatchNumpyInterface.generate_dummy()
    batch_calculate_fit(data)
    logger.debug("Warmed up")

def run_fn():
    logger.debug("Creating data")
    data = BatchNumpyInterface.generate_dummy(100)
    
    logger.debug("Predicting %s in batches",100)
    result = batch_calculate_fit(data)

    logger.debug("Done")
    return float(result[0][0]),float(result[1][0])

#with ThreadPoolExecutor(max_workers=4) as executor:
with ProcesspoolExecutor(max_workers=4) as executor:
    results = []
    for i in range(4):
        results.append(executor.submit(warmup))

    for res in as_completed(results):
        continue

    results = []
    for i in range(10):
        future = executor.submit(run_fn)
        results.append(future)

    for res in as_completed(results):
        print(res.result())

在调试期间，每次检查JAX DeviceArray时，应用程序都会崩溃。用JAX进行第一次计算也是如此。

任何帮助将不胜感激！

解决方法

暂无找到可以解决该程序问题的有效方法，小编努力寻找整理中！

如果你已经找到好的解决方法，欢迎将解决方案带上本链接一起发送给小编。

小编邮箱:dio#foxmail.com (将#修改为@）