Skip to content

Commit

Permalink
feat(api): make tile retries configurable
Browse files Browse the repository at this point in the history
  • Loading branch information
ssube committed Sep 11, 2023
1 parent cd06f92 commit 6fb0366
Show file tree
Hide file tree
Showing 3 changed files with 6 additions and 1 deletion.
3 changes: 3 additions & 0 deletions api/onnx_web/server/context.py
Original file line number Diff line number Diff line change
Expand Up @@ -37,6 +37,7 @@ def __init__(
memory_limit: Optional[int] = None,
admin_token: Optional[str] = None,
server_version: Optional[str] = DEFAULT_SERVER_VERSION,
worker_retries: Optional[int] = 3,
) -> None:
self.bundle_path = bundle_path
self.model_path = model_path
Expand All @@ -56,6 +57,7 @@ def __init__(
self.memory_limit = memory_limit
self.admin_token = admin_token or token_urlsafe()
self.server_version = server_version
self.worker_retries = worker_retries

self.cache = ModelCache(self.cache_limit)

Expand Down Expand Up @@ -88,6 +90,7 @@ def from_environ(cls):
server_version=environ.get(
"ONNX_WEB_SERVER_VERSION", DEFAULT_SERVER_VERSION
),
worker_retries=int(environ.get("ONNX_WEB_WORKER_RETRIES", 3)),
)

def torch_dtype(self):
Expand Down
3 changes: 2 additions & 1 deletion api/onnx_web/worker/context.py
Original file line number Diff line number Diff line change
Expand Up @@ -36,6 +36,7 @@ def __init__(
progress: "Queue[ProgressCommand]",
active_pid: "Value[int]",
idle: "Value[bool]",
retries: int,
):
self.job = None
self.name = name
Expand All @@ -47,8 +48,8 @@ def __init__(
self.active_pid = active_pid
self.last_progress = None
self.idle = idle
self.retries = retries
self.timeout = 1.0
self.retries = 3 # TODO: get from env

def start(self, job: str) -> None:
self.job = job
Expand Down
1 change: 1 addition & 0 deletions api/onnx_web/worker/pool.py
Original file line number Diff line number Diff line change
Expand Up @@ -124,6 +124,7 @@ def create_device_worker(self, device: DeviceParams) -> None:
pending=self.pending[name],
active_pid=current,
idle=self.worker_idle[name],
retries=self.server.worker_retries,
)
self.context[name] = context

Expand Down

0 comments on commit 6fb0366

Please sign in to comment.