fix: race condition issue initializing result db (#5369)

* fix: race condition issue initializing result db

Signed-off-by: Frost Ming <me@frostming.com>

* fix: only init db when needed

Signed-off-by: Frost Ming <me@frostming.com>

Author: frostming

src/_bentoml_impl/server/app.py

@@ -306,9 +306,10 @@ async def create_instance(self, app: Starlette) -> None:
         self._service_instance.to_async = proxy.to_async  # type: ignore[attr-defined]
         self._service_instance.to_sync = proxy.to_sync  # type: ignore[attr-defined]
         set_current_service(self._service_instance)
-        store_path = BentoMLContainer.result_store_file.get()
-        self._result_store = Sqlite3Store(store_path)
-        await self._result_store.__aenter__()
+        if self.service.needs_task_db():
+            store_path = BentoMLContainer.result_store_file.get()
+            self._result_store = Sqlite3Store(store_path)
+            await self._result_store.__aenter__()
 
     @inject
     def _add_response_headers(
@@ -356,7 +357,8 @@ async def destroy_instance(self, _: Starlette) -> None:
         logger.info("Service instance cleanup finalized")
         self._service_instance = None
         set_current_service(None)
-        await self._result_store.__aexit__(None, None, None)
+        if hasattr(self, "_result_store"):
+            await self._result_store.__aexit__(None, None, None)
 
     async def livez(self, _: Request) -> Response:
         from starlette.exceptions import HTTPException

src/_bentoml_impl/server/serving.py

@@ -140,6 +140,8 @@ def server_on_deployment(
     svc: AnyService, result_file: str = Provide[BentoMLContainer.result_store_file]
 ) -> None:
     # Resolve models before server starts.
+    from ..tasks.result import Sqlite3Store
+
     if bento := svc.bento:
         for model in bento.info.all_models:
             model.to_model().resolve()
@@ -150,8 +152,10 @@ def server_on_deployment(
         member = getattr(svc.inner, name)
         if callable(member) and getattr(member, "__bentoml_deployment_hook__", False):
             member()
-    if os.path.exists(result_file):
-        os.remove(result_file)
+    if svc.needs_task_db():
+        if os.path.exists(result_file):
+            os.remove(result_file)
+        Sqlite3Store.init_db(result_file)
 
 
 @inject(squeeze_none=True)

src/_bentoml_impl/tasks/result.py

@@ -109,28 +109,33 @@ def _connect(self, db_file: str) -> Connection:
 
     async def __aenter__(self) -> "t.Self":
         self._conn = await self._conn
-        await self._init_db()
         return self
 
     async def __aexit__(self, *_: t.Any) -> None:
         await self._conn.close()
 
-    async def _init_db(self) -> None:
-        await self._conn.execute(
-            textwrap.dedent("""\
-            CREATE TABLE IF NOT EXISTS result (
-                task_id TEXT PRIMARY KEY,
-                name TEXT,
-                input BLOB,
-                status TEXT,
-                result BLOB,
-                created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
-                executed_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
-                completed_at TIMESTAMP DEFAULT NULL
+    @classmethod
+    def init_db(cls, db_file: str) -> None:
+        import sqlite3
+
+        with sqlite3.connect(
+            db_file, detect_types=sqlite3.PARSE_DECLTYPES | sqlite3.PARSE_COLNAMES
+        ) as conn:
+            conn.execute(
+                textwrap.dedent("""\
+                CREATE TABLE IF NOT EXISTS result (
+                    task_id TEXT PRIMARY KEY,
+                    name TEXT,
+                    input BLOB,
+                    status TEXT,
+                    result BLOB,
+                    created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
+                    executed_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
+                    completed_at TIMESTAMP DEFAULT NULL
+                )
+                """)
             )
-            """)
-        )
-        await self._conn.commit()
+            conn.commit()
 
     async def new_entry(self, name: str, input: Request) -> str:
         task_id = uuid.uuid4().hex

src/_bentoml_sdk/service/factory.py

@@ -473,6 +473,11 @@ def on_load_bento(self, bento: Bento) -> None:
                 model.revision = revision
             svc.bento = bento
 
+    def needs_task_db(self) -> bool:
+        if "BENTOCLOUD_DEPLOYMENT_URL" in os.environ:
+            return False
+        return any(method.is_task for method in self.apis.values())
+
 
 @t.overload
 def service(inner: type[T], /) -> Service[T]: ...