feat(client): add a client for benchmarking and testing

gnes/cli/api.py

@@ -46,7 +46,7 @@ def route(args):
 
 
 def frontend(args):
-    from gnes.service.frontend import FrontendService
+    from ..service.frontend import FrontendService
     import threading
     with FrontendService(args):
         forever = threading.Event()
@@ -61,7 +61,7 @@ def client(args):
     elif args.client == 'benchmark':
         return _client_bm(args)
     else:
-        raise ValueError('gnes client must follow a client type from {"http", "cli", "benchmark"...}\n'
+        raise ValueError('gnes client must follow with a client type from {http, cli, benchmark...}\n'
                          'see "gnes client --help" for details')
 
 

gnes/cli/parser.py

@@ -290,12 +290,13 @@ def set_benchmark_client_parser(parser=None):
     if not parser:
         parser = set_base_parser()
     _set_grpc_parser(parser)
-
+    parser.add_argument('--batch_size', type=int, default=64,
+                        help='the size of the request to split')
     parser.add_argument('--request_length', type=int,
                         default=1024,
                         help='binary string length of each request')
     parser.add_argument('--num_requests', type=int,
-                        default=1024,
+                        default=128,
                         help='number of total requests')
     return parser
 

gnes/client/benchmark.py

@@ -14,6 +14,8 @@
 #  limitations under the License.
 
 
+import time
+
 import grpc
 
 from ..helper import TimeContext
@@ -32,16 +34,22 @@ def __init__(self, args):
             stub = gnes_pb2_grpc.GnesRPCStub(channel)
 
             id = 0
-
-            with TimeContext('StreamCall num_requests=%d request_length=%d' % (args.num_requests, args.request_length)):
-                resp = list(stub.StreamCall(RequestGenerator.index(all_bytes, args.batch_size)))[-1]
+            with TimeContext('StreamCall') as tc:
+                resp = stub.StreamCall(RequestGenerator.index(all_bytes, args.batch_size))
                 for r in resp:
                     assert r.request_id == str(id)
                     id += 1
+            stream_call_el = tc.duration
 
             id = 0
-            with TimeContext('Call num_requests=%d request_length=%d' % (args.num_requests, args.request_length)):
+            with TimeContext('Call') as tc:
                 for req in RequestGenerator.index(all_bytes, 1):
                     r = stub.Call(req)
                     assert r.request_id == str(id)
                     id += 1
+            call_el = tc.duration
+
+        print('num_requests     %d\n'
+              'request_length   %d' % (args.num_requests, args.request_length))
+        print('StreamCall       %3.3f s\n'
+              'Call             %3.3f s\n' % (stream_call_el, call_el))

gnes/helper.py

@@ -14,8 +14,6 @@
 #  limitations under the License.
 
 
-
-
 import fcntl
 import logging
 import os
@@ -256,10 +254,12 @@ def warning(self, msg, **kwargs):
 class TimeContext:
     def __init__(self, msg):
         self._msg = msg
+        self.duration = 0
 
     def __enter__(self):
         self.start = time.perf_counter()
         print(self._msg, end=' ...\t', flush=True)
+        return self
 
     def __exit__(self, typ, value, traceback):
         self.duration = time.perf_counter() - self.start

tests/test_stream_grpc.py

@@ -4,7 +4,8 @@
 
 import grpc
 
-from gnes.cli.parser import set_frontend_parser, set_router_service_parser
+from gnes.cli.parser import set_frontend_parser, set_router_service_parser, set_benchmark_client_parser
+from gnes.client.benchmark import BenchmarkClient
 from gnes.helper import TimeContext
 from gnes.proto import RequestGenerator, gnes_pb2_grpc
 from gnes.service.base import SocketType, MessageHandler, BaseService as BS
@@ -40,6 +41,26 @@ def setUp(self):
         os.unsetenv('http_proxy')
         os.unsetenv('https_proxy')
 
+    def test_bm_frontend(self):
+        args = set_frontend_parser().parse_args([
+            '--grpc_host', '127.0.0.1',
+        ])
+
+        p_args = set_router_service_parser().parse_args([
+            '--port_in', str(args.port_out),
+            '--port_out', str(args.port_in),
+            '--socket_in', str(SocketType.PULL_CONNECT),
+            '--socket_out', str(SocketType.PUSH_CONNECT),
+            '--yaml_path', 'BaseRouter'
+        ])
+
+        b_args = set_benchmark_client_parser().parse_args([
+            '--num_requests', '10',
+            '--request_length', '65536'
+        ])
+        with RouterService(p_args), FrontendService(args):
+            BenchmarkClient(b_args)
+
     def test_grpc_frontend(self):
         args = set_frontend_parser().parse_args([
             '--grpc_host', '127.0.0.1',