Merge pull request #24 from Yoctol/list_models

List models

serving_utils/client.py

@@ -9,7 +9,7 @@
 import grpc
 import tensorflow as tf
 
-from .protos import predict_pb2, prediction_service_pb2_grpc
+from .protos import predict_pb2, prediction_service_pb2_grpc, list_models_pb2, list_models_pb2_grpc
 
 
 def copy_message(src, dst):
@@ -123,6 +123,11 @@ def parse_predict_response(response):
             results[key] = nd_array
         return results
 
+    def list_models(self):
+        stub = list_models_pb2_grpc.ListModelsStub(self._channel)
+        response = stub.ListModels(list_models_pb2.ListModelsRequest())
+        return response.models
+
     def predict(
             self,
             data: List[PredictInput],

serving_utils/protos/list_models.proto

@@ -0,0 +1,12 @@
+syntax = "proto3";
+
+service ListModels {
+  rpc ListModels (ListModelsRequest) returns (ListModelsResponse) {}
+}
+
+message ListModelsRequest {
+}
+
+message ListModelsResponse {
+  repeated string models = 1;
+}

serving_utils/tests/test_saver.py

@@ -55,7 +55,7 @@ def test_fit_n_save_serving(self):
                 ],
             ),
             set(
-                glob(self.model.output_dir + '/*/*') +  # noqa:W504
+                glob(self.model.output_dir + '/*/*') +  # noqa: W504
                 glob(self.model.output_dir + '/*/*/*'),
             ),
         )