2023-02-26 05:49:39 +00:00
|
|
|
from logging import getLogger
|
2023-02-26 18:24:51 +00:00
|
|
|
from traceback import format_exception
|
2023-02-26 20:15:30 +00:00
|
|
|
|
2023-02-26 19:09:24 +00:00
|
|
|
from setproctitle import setproctitle
|
2023-02-26 21:06:40 +00:00
|
|
|
from torch.multiprocessing import Queue
|
2023-02-26 05:49:39 +00:00
|
|
|
|
2023-02-26 18:32:48 +00:00
|
|
|
from ..server import ServerContext, apply_patches
|
2023-02-26 21:21:58 +00:00
|
|
|
from ..torch_before_ort import get_available_providers
|
2023-02-26 20:15:30 +00:00
|
|
|
from .context import WorkerContext
|
2023-02-26 05:49:39 +00:00
|
|
|
|
|
|
|
logger = getLogger(__name__)
|
|
|
|
|
2023-02-26 18:32:48 +00:00
|
|
|
|
2023-02-26 21:06:40 +00:00
|
|
|
def logger_init(logs: Queue):
|
2023-02-26 19:09:24 +00:00
|
|
|
setproctitle("onnx-web logger")
|
|
|
|
|
2023-02-26 21:21:58 +00:00
|
|
|
logger.info("checking in from logger")
|
2023-02-26 21:06:40 +00:00
|
|
|
|
2023-02-26 17:16:33 +00:00
|
|
|
while True:
|
|
|
|
job = logs.get()
|
|
|
|
with open("worker.log", "w") as f:
|
|
|
|
logger.info("got log: %s", job)
|
|
|
|
f.write(str(job) + "\n\n")
|
2023-02-26 05:49:39 +00:00
|
|
|
|
|
|
|
|
2023-02-26 21:06:40 +00:00
|
|
|
def worker_init(context: WorkerContext, server: ServerContext):
|
2023-02-26 18:32:48 +00:00
|
|
|
apply_patches(server)
|
2023-02-26 20:15:30 +00:00
|
|
|
setproctitle("onnx-web worker: %s" % (context.device.device))
|
2023-02-26 18:32:48 +00:00
|
|
|
|
2023-02-26 21:21:58 +00:00
|
|
|
logger.info("checking in from worker, %s", get_available_providers())
|
2023-02-26 21:06:40 +00:00
|
|
|
|
2023-02-26 05:49:39 +00:00
|
|
|
while True:
|
2023-02-26 17:16:33 +00:00
|
|
|
job = context.pending.get()
|
|
|
|
logger.info("got job: %s", job)
|
2023-02-26 18:51:11 +00:00
|
|
|
|
2023-02-27 02:09:42 +00:00
|
|
|
fn, args, kwargs = job
|
|
|
|
name = args[3][0]
|
|
|
|
|
|
|
|
try:
|
2023-02-26 21:06:40 +00:00
|
|
|
context.clear_flags()
|
2023-02-27 02:09:42 +00:00
|
|
|
logger.info("starting job: %s", name)
|
|
|
|
context.put_started(name)
|
2023-02-26 17:16:33 +00:00
|
|
|
fn(context, *args, **kwargs)
|
2023-02-26 21:06:40 +00:00
|
|
|
logger.info("job succeeded: %s", name)
|
2023-02-26 17:16:33 +00:00
|
|
|
except Exception as e:
|
2023-02-26 21:06:40 +00:00
|
|
|
logger.error(
|
|
|
|
"error while running job: %s",
|
|
|
|
format_exception(type(e), e, e.__traceback__),
|
|
|
|
)
|
|
|
|
finally:
|
2023-02-27 02:09:42 +00:00
|
|
|
context.put_finished(name)
|
2023-02-26 21:06:40 +00:00
|
|
|
logger.info("finished job: %s", name)
|