record metrics after the last byte was sent

bentoml · Sep 15, 2022 · 0b716c6 · 0b716c6
1 parent 0a72c8e
commit 0b716c6
Show file tree

Hide file tree

Showing 2 changed files with 27 additions and 21 deletions.
diff --git a/bentoml/_internal/server/instruments.py b/bentoml/_internal/server/instruments.py
@@ -19,6 +19,7 @@
 START_TIME_VAR: "contextvars.ContextVar[float]" = contextvars.ContextVar(
     "START_TIME_VAR"
 )
+STATUS_VAR: "contextvars.ContextVar[int]" = contextvars.ContextVar("STATUS_VAR")
 from ..context import component_context
 
 
@@ -251,28 +252,33 @@ async def __call__(
 
         async def wrapped_send(message: "ext.ASGIMessage") -> None:
             if message["type"] == "http.response.start":
-                status_code = message["status"]
+                STATUS_VAR.set(message["status"])
+            elif message["type"] == "http.response.body":
+                if ("more_body" not in message) or not message["more_body"]:
+                    assert START_TIME_VAR.get() != 0
+                    assert STATUS_VAR.get() != 0
 
-                # instrument request total count
-                self.metrics_request_total.labels(
-                    endpoint=endpoint,
-                    service_name=component_context.bento_name,
-                    service_version=component_context.bento_version,
-                    http_response_code=status_code,
-                    runner_name=component_context.component_name,
-                ).inc()
+                    # instrument request total count
+                    self.metrics_request_total.labels(
+                        endpoint=endpoint,
+                        service_name=component_context.bento_name,
+                        service_version=component_context.bento_version,
+                        http_response_code=STATUS_VAR.get(),
+                        runner_name=component_context.component_name,
+                    ).inc()
 
-                # instrument request duration
-                assert START_TIME_VAR.get() != 0
-                total_time = max(default_timer() - START_TIME_VAR.get(), 0)
-                self.metrics_request_duration.labels(  # type: ignore
-                    endpoint=endpoint,
-                    service_name=component_context.bento_name,
-                    service_version=component_context.bento_version,
-                    http_response_code=status_code,
-                    runner_name=component_context.component_name,
-                ).observe(total_time)
-            START_TIME_VAR.set(0)
+                    # instrument request duration
+                    total_time = max(default_timer() - START_TIME_VAR.get(), 0)
+                    self.metrics_request_duration.labels(  # type: ignore
+                        endpoint=endpoint,
+                        service_name=component_context.bento_name,
+                        service_version=component_context.bento_version,
+                        http_response_code=STATUS_VAR.get(),
+                        runner_name=component_context.component_name,
+                    ).observe(total_time)
+
+                    START_TIME_VAR.set(0)
+                    STATUS_VAR.set(0)
             await send(message)
 
         with self.metrics_request_in_progress.labels(

diff --git a/bentoml_cli/worker/runner.py b/bentoml_cli/worker/runner.py
@@ -114,7 +114,7 @@ def main(
         component_context.bento_version = "not available"
     else:
         component_context.bento_name = service.tag.name
-        component_context.bento_version = service.tag.version
+        component_context.bento_version = service.tag.version or ""
 
     for runner in service.runners:
         if runner.name == runner_name: