feat(router): add device and dtype info (#215)

343437c7 · OlivierDehaene · GitHub · ac8c0f6f · 343437c7
Unverified Commit 343437c7 authored Apr 21, 2023 by OlivierDehaene Committed by GitHub Apr 21, 2023
Hide whitespace changes
Inline Side-by-side

Showing with 3 additions and 0 deletions

server/text_generation_server/server.py server/text_generation_server/server.py +3 -0

No files found.
--- a/server/text_generation_server/server.py
+++ b/server/text_generation_server/server.py
@@ -26,6 +26,9 @@ class TextGenerationService(generate_pb2_grpc.TextGenerationServiceServicer):
            # Force inference mode for the lifetime of TextGenerationService
            self._inference_mode_raii_guard = torch._C._InferenceMode(True)

+    async def Info(self, request, context):
+        return self.model.info
+
    async def ServiceDiscovery(self, request, context):
        return generate_pb2.ServiceDiscoveryResponse(urls=self.server_urls)