From 1faf05ea56505f66139acf5ba4816acea21ec367 Mon Sep 17 00:00:00 2001 From: aria Date: Wed, 12 Oct 2022 16:15:33 +0800 Subject: [PATCH] Ability to customize model name in response metadata (#14) Add `--model-name` flag in example server to allow customizing response metadata. --- example/go/simple/server/main.go | 8 ++++++-- example/go/simple/server/server.go | 7 ++++--- 2 files changed, 10 insertions(+), 5 deletions(-) diff --git a/example/go/simple/server/main.go b/example/go/simple/server/main.go index 294ff4a..0e93f57 100644 --- a/example/go/simple/server/main.go +++ b/example/go/simple/server/main.go @@ -6,11 +6,15 @@ import ( ) var ( - port = flag.Int("port", 50051, "The server port") + port = flag.Int("port", 50051, "The server port") + modelName = flag.String("model-name", "model-a", "Model name to be returned in response metadata") ) func main() { flag.Parse() - upiServer := UpiServer{} + upiServer := UpiServer{ + modelName: *modelName, + } + upiServer.Run(fmt.Sprintf(":%d", *port)) } diff --git a/example/go/simple/server/server.go b/example/go/simple/server/server.go index 4bc6466..839b916 100644 --- a/example/go/simple/server/server.go +++ b/example/go/simple/server/server.go @@ -12,9 +12,10 @@ import ( type UpiServer struct { upiv1.UnimplementedUniversalPredictionServiceServer + modelName string } -func (_ *UpiServer) PredictValues( +func (s *UpiServer) PredictValues( _ context.Context, req *upiv1.PredictValuesRequest, ) (*upiv1.PredictValuesResponse, error) { @@ -22,7 +23,7 @@ func (_ *UpiServer) PredictValues( Metadata: &upiv1.ResponseMetadata{ Models: []*upiv1.ModelMetadata{ { - Name: "Echo Request Table", + Name: s.modelName, Version: "1", }, }, @@ -39,7 +40,7 @@ func (us *UpiServer) Run(address string) { s := grpc.NewServer() upiv1.RegisterUniversalPredictionServiceServer(s, us) reflection.Register(s) - log.Printf("listening on port %s", address) + log.Printf("running model %s on port %s", us.modelName, address) if err := s.Serve(lis); err != nil { log.Fatalf("failed to serve: %v", err) }