feat(api): make tile retries configurable

api/onnx_web/server/context.py

@@ -37,6 +37,7 @@ class ServerContext:
         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
@@ -56,6 +57,7 @@ class ServerContext:
         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)
 
@@ -88,6 +90,7 @@ class ServerContext:
             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):

api/onnx_web/worker/context.py

@@ -36,6 +36,7 @@ class WorkerContext:
         progress: "Queue[ProgressCommand]",
         active_pid: "Value[int]",
         idle: "Value[bool]",
+        retries: int,
     ):
         self.job = None
         self.name = name
@@ -47,8 +48,8 @@ class WorkerContext:
         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

api/onnx_web/worker/pool.py

@@ -124,6 +124,7 @@ class DevicePoolExecutor:
             pending=self.pending[name],
             active_pid=current,
             idle=self.worker_idle[name],
+            retries=self.server.worker_retries,
         )
         self.context[name] = context