diff --git a/api/onnx_web/device_pool.py b/api/onnx_web/device_pool.py index ff243acc..59189652 100644 --- a/api/onnx_web/device_pool.py +++ b/api/onnx_web/device_pool.py @@ -1,7 +1,7 @@ from concurrent.futures import Future, ThreadPoolExecutor, ProcessPoolExecutor from logging import getLogger from multiprocessing import Value -from typing import Any, Callable, List, Optional, Tuple, Union +from typing import Any, Callable, Dict, List, Optional, Tuple, Union logger = getLogger(__name__) @@ -121,3 +121,6 @@ class DevicePoolExecutor: future = self.pool.submit(fn, context, *args, **kwargs) job = Job(key, future, context) self.jobs.append(job) + + def status(self) -> Dict[str, Tuple[bool, int]]: + return [(job.future.done(), job.get_progress()) for job in self.jobs] \ No newline at end of file diff --git a/api/onnx_web/serve.py b/api/onnx_web/serve.py index 50386b07..157fc6a6 100644 --- a/api/onnx_web/serve.py +++ b/api/onnx_web/serve.py @@ -606,6 +606,15 @@ def chain(): return jsonify(json_params(output, params, size)) +@app.route('/api/cancel', methods=['PUT']) +def cancel(): + output_file = request.args.get('output', None) + + cancel = executor.cancel(output_file) + + return ready_reply(cancel) + + @app.route('/api/ready') def ready(): output_file = request.args.get('output', None) @@ -620,13 +629,9 @@ def ready(): return ready_reply(done, progress=progress) -@app.route('/api/cancel', methods=['PUT']) -def cancel(): - output_file = request.args.get('output', None) - - cancel = executor.cancel(output_file) - - return ready_reply(cancel) +@app.route('/api/status') +def status(): + return jsonify(executor.status()) @app.route('/output/')