From 1d577a5f628dda4f35f1bc3e195a5f17b5e2f4c1 Mon Sep 17 00:00:00 2001 From: Leon-Sander Date: Sun, 3 Nov 2024 00:51:27 +0100 Subject: [PATCH 1/2] Add task_id to ResponseContext and set it in _run_task under server/app --- src/_bentoml_impl/server/app.py | 1 + src/bentoml/_internal/context.py | 1 + 2 files changed, 2 insertions(+) diff --git a/src/_bentoml_impl/server/app.py b/src/_bentoml_impl/server/app.py index e0e229a4786..19b35f7d8ef 100644 --- a/src/_bentoml_impl/server/app.py +++ b/src/_bentoml_impl/server/app.py @@ -440,6 +440,7 @@ async def cancel_task(self, request: Request) -> Response: async def _run_task(self, task_id: str, name: str, request: Request) -> None: try: + self.service.context.response.task_id = task_id resp = await self.api_endpoint_wrapper(name, request) await self._result_store.set_result( task_id, diff --git a/src/bentoml/_internal/context.py b/src/bentoml/_internal/context.py index 9fbdd3401f0..35ec0b39f23 100644 --- a/src/bentoml/_internal/context.py +++ b/src/bentoml/_internal/context.py @@ -133,6 +133,7 @@ class ResponseContext: cookies: list[Cookie] = attr.field(factory=list) status_code: int = 200 background: BackgroundTasks = attr.field(factory=BackgroundTasks) + task_id: str | None = None @property def headers(self) -> Metadata: From 910a4d6afe7e4b787ff9a5ec4ef488aabf98910f Mon Sep 17 00:00:00 2001 From: Leon-Sander Date: Sun, 3 Nov 2024 16:58:25 +0100 Subject: [PATCH 2/2] feat: Refactor task_id placement in Inference context Previously, task_id was set in the response context, but after further consideration, it made more sense to move it to the request context. Since task_id is a fixed value and the response context is for modifications, placing it in the request context better reflects its purpose and usage. --- src/_bentoml_impl/server/app.py | 2 +- src/bentoml/_internal/context.py | 1 - 2 files changed, 1 insertion(+), 2 deletions(-) diff --git a/src/_bentoml_impl/server/app.py b/src/_bentoml_impl/server/app.py index 19b35f7d8ef..571a08978d2 100644 --- a/src/_bentoml_impl/server/app.py +++ b/src/_bentoml_impl/server/app.py @@ -440,7 +440,7 @@ async def cancel_task(self, request: Request) -> Response: async def _run_task(self, task_id: str, name: str, request: Request) -> None: try: - self.service.context.response.task_id = task_id + self.service.context.request.state.task_id = task_id resp = await self.api_endpoint_wrapper(name, request) await self._result_store.set_result( task_id, diff --git a/src/bentoml/_internal/context.py b/src/bentoml/_internal/context.py index 35ec0b39f23..9fbdd3401f0 100644 --- a/src/bentoml/_internal/context.py +++ b/src/bentoml/_internal/context.py @@ -133,7 +133,6 @@ class ResponseContext: cookies: list[Cookie] = attr.field(factory=list) status_code: int = 200 background: BackgroundTasks = attr.field(factory=BackgroundTasks) - task_id: str | None = None @property def headers(self) -> Metadata: