From d1961afdbc9dcf89c2c68c1967c223be8d07d1e6 Mon Sep 17 00:00:00 2001 From: Sean Sube Date: Sun, 26 Feb 2023 14:36:32 -0600 Subject: [PATCH] re-implement cancellation --- api/onnx_web/worker/pool.py | 15 +++++++++++++-- 1 file changed, 13 insertions(+), 2 deletions(-) diff --git a/api/onnx_web/worker/pool.py b/api/onnx_web/worker/pool.py index 810bb91a..a58fd3c0 100644 --- a/api/onnx_web/worker/pool.py +++ b/api/onnx_web/worker/pool.py @@ -82,11 +82,22 @@ class DevicePoolExecutor: the future and never execute it. If the job has been started, it should be cancelled on the next progress callback. """ - raise NotImplementedError() + if key not in self.jobs: + logger.warn("attempting to cancel unknown job: %s", key) + return False + + device = self.jobs[key] + cancel = self.context[device].cancel + logger.info("cancelling job %s on device %s", key, device) + + if cancel.get_lock(): + cancel.value = True + + return True def done(self, key: str) -> Tuple[Optional[bool], int]: if key not in self.jobs: - logger.warn("checking status for unknown key: %s", key) + logger.warn("checking status for unknown job: %s", key) return (None, 0) device = self.jobs[key]