apextoaster
/
onnx-web
1
0
Fork 0
Code Actions

set queue timeouts

This commit is contained in:
Sean Sube 2023-02-27 22:37:43 -06:00
parent 0793b61c3a
commit 136759285d
Signed by: ssube
GPG Key ID: 3EED7B957D362AF1
3 changed files with 9 additions and 6 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
api/onnx_web/worker/context.py
View File

@ -63,7 +63,7 @@ class WorkerContext:
self.cancel.value = cancel
def set_progress(self, progress: int) -> None:
self.progress.put((self.job, self.device.device, progress))
self.progress.put((self.job, self.device.device, progress), block=False)
def set_finished(self) -> None:
self.finished.put((self.job, self.device.device))

8
api/onnx_web/worker/pool.py
View File

@ -88,7 +88,7 @@ class DevicePoolExecutor:
while True:
try:
job = logs.get()
job = logs.get(timeout=(self.join_timeout / 2))
with open("worker.log", "w") as f:
logger.info("got log: %s", job)
f.write(str(job) + "\n\n")
@ -106,7 +106,7 @@ class DevicePoolExecutor:
logger.info("checking in from progress worker thread")
while True:
try:
job, device, value = progress.get()
job, device, value = progress.get(timeout=(self.join_timeout / 2))
logger.info("progress update for job: %s to %s", job, value)
self.active_jobs[job] = (device, value)
if job in self.cancelled_jobs:
@ -128,7 +128,7 @@ class DevicePoolExecutor:
logger.info("checking in from finished worker thread")
while True:
try:
job, device = finished.get()
job, device = finished.get(timeout=(self.join_timeout / 2))
logger.info("job has been finished: %s", job)
context = self.context[device]
_device, progress = self.active_jobs[job]
@ -256,7 +256,7 @@ class DevicePoolExecutor:
)
device = self.devices[device_idx].device
self.pending[device].put((key, fn, args, kwargs))
self.pending[device].put((key, fn, args, kwargs), block=False)
def status(self) -> List[Tuple[str, int, bool, bool]]:
history = [

5
api/onnx_web/worker/worker.py
View File

@ -1,4 +1,5 @@
from logging import getLogger
from queue import Empty
from traceback import format_exception
from setproctitle import setproctitle
@ -18,7 +19,7 @@ def worker_main(context: WorkerContext, server: ServerContext):
logger.info("checking in from worker, %s", get_available_providers())
while True:
name, fn, args, kwargs = context.pending.get()
name, fn, args, kwargs = context.pending.get(timeout=1.0)
logger.info("worker for %s got job: %s", context.device.device, name)
try:
@ -27,6 +28,8 @@ def worker_main(context: WorkerContext, server: ServerContext):
logger.info("starting job: %s", name)
fn(context, *args, **kwargs)
logger.info("job succeeded: %s", name)
except Empty:
pass
except Exception as e:
logger.error(
"error while running job: %s",