Skip to content

Commit

Permalink
fix(api): default to first available device when there are no other jobs
Browse files Browse the repository at this point in the history
  • Loading branch information
ssube committed Feb 4, 2023
1 parent efee374 commit bf3f227
Show file tree
Hide file tree
Showing 2 changed files with 22 additions and 1 deletion.
9 changes: 8 additions & 1 deletion api/onnx_web/device_pool.py
Expand Up @@ -135,8 +135,15 @@ def done(self, key: str) -> Tuple[bool, int]:
return (None, 0)

def get_next_device(self):
# use the first/default device if there are no jobs
if len(self.jobs) == 0:
return 0

job_devices = [job.context.device_index.value for job in self.jobs]
queued = Counter(job_devices).most_common()
job_counts = Counter(range(len(self.devices)))
job_counts.update(job_devices)

queued = job_counts.most_common()
logger.debug('jobs queued by device: %s', queued)

return queued[-1]
Expand Down
14 changes: 14 additions & 0 deletions api/onnx_web/serve.py
Expand Up @@ -15,6 +15,7 @@
)
from flask import Flask, jsonify, make_response, request, send_from_directory, url_for
from flask_cors import CORS
from functools import cmp_to_key
from glob import glob
from io import BytesIO
from jsonschema import validate
Expand Down Expand Up @@ -333,6 +334,19 @@ def load_platforms():
available_platforms.append(DeviceParams(
potential, platform_providers[potential]))

# make sure CPU is last on the list
def cpu_last(a: DeviceParams, b: DeviceParams):
if a.device == 'cpu' and b.device == 'cpu':
return 0

if a.device == 'cpu':
return 1

return -1


available_platforms = sorted(available_platforms, key=cmp_to_key(cpu_last))

logger.info('available acceleration platforms: %s',
', '.join([str(p) for p in available_platforms]))

Expand Down

0 comments on commit bf3f227

Please sign in to comment.