From 55e44e8ac9842e5c514ab07811130bdd8cf58ed6 Mon Sep 17 00:00:00 2001 From: Sean Sube Date: Sun, 26 Mar 2023 11:18:27 -0500 Subject: [PATCH] fix(api): increment job counter for worker when it starts a new job (#283) --- api/onnx_web/worker/pool.py | 24 ++++++++++++------------ 1 file changed, 12 insertions(+), 12 deletions(-) diff --git a/api/onnx_web/worker/pool.py b/api/onnx_web/worker/pool.py index 21c5d7a9..2699d41b 100644 --- a/api/onnx_web/worker/pool.py +++ b/api/onnx_web/worker/pool.py @@ -355,24 +355,14 @@ class DevicePoolExecutor: **kwargs, ) -> None: device_idx = self.get_next_device(needs_device=needs_device) + device = self.devices[device_idx].device logger.info( "assigning job %s to device %s: %s", key, device_idx, - self.devices[device_idx], + device, ) - # increment job count before recycling (why tho?) - device = self.devices[device_idx].device - if device in self.total_jobs: - self.total_jobs[device] += 1 - else: - self.total_jobs[device] = 1 - - # recycle before attempting to run - logger.debug("job count for device %s: %s", device, self.total_jobs[device]) - self.rlock() - # build and queue job job = JobCommand(key, device, fn, args, kwargs) self.pending_jobs.append(job) @@ -439,6 +429,16 @@ class DevicePoolExecutor: job for job in self.pending_jobs if job.name != progress.job ] + # increment job counter if this is the start of a new job + if progress.progress == 0: + if progress.device in self.total_jobs: + self.total_jobs[progress.device] += 1 + else: + self.total_jobs[progress.device] = 1 + + logger.debug("updating job count for device %s: %s", progress.device, self.total_jobs[progress.device]) + + # check if the job has been cancelled if progress.job in self.cancelled_jobs: logger.debug( "setting flag for cancelled job: %s on %s",