Merge pull request #333 from gnes-ai/fix-ctrl-ipc

fix(service): use ipc for control socket

Author: mergify[bot]

docs/chapter/enviromentvars.md

@@ -38,3 +38,6 @@ Default is not set. A random port will be used.
 
 (*depreciated*) Paths of the third party components. See examples in GNES hub for latest usage.
 
+## `GNES_IPC_SOCK_TMP`
+
+Temp directory for ipc sockets, not used on Windows.

gnes/cli/parser.py

@@ -202,6 +202,8 @@ def set_service_parser(parser=None):
     parser.add_argument('--squeeze_pb', action=ActionNoYes, default=True,
                         help='sending bytes and ndarray separately apart from the protobuf message, '
                              'usually yields better network efficiency')
+    parser.add_argument('--ctrl_with_ipc', action='store_true', default=False,
+                        help='use ipc protocol for control socket')
     return parser
 
 

gnes/flow/__init__.py

@@ -703,7 +703,7 @@ def _build_graph(self, copy_flow: bool) -> 'Flow':
         op_flow._build_level = BuildLevel.GRAPH
         return op_flow
 
-    def build(self, backend: Optional[str] = 'thread', copy_flow: bool = False, *args, **kwargs) -> 'Flow':
+    def build(self, backend: Optional[str] = 'process', copy_flow: bool = False, *args, **kwargs) -> 'Flow':
         """
         Build the current flow and make it ready to use
 

gnes/service/base.py

@@ -15,10 +15,13 @@
 
 import copy
 import multiprocessing
+import os
 import random
+import tempfile
 import threading
 import time
 import types
+import uuid
 from contextlib import ExitStack
 from enum import Enum
 from typing import Tuple, List, Union, Type
@@ -113,8 +116,19 @@ class EventLoopEnd(Exception):
     pass
 
 
-def build_socket(ctx: 'zmq.Context', host: str, port: int, socket_type: 'SocketType', identity: 'str' = None) -> Tuple[
-    'zmq.Socket', str]:
+def get_random_ipc() -> str:
+    try:
+        tmp = os.environ['GNES_IPC_SOCK_TMP']
+        if not os.path.exists(tmp):
+            raise ValueError('This directory for sockets ({}) does not seems to exist.'.format(tmp))
+        tmp = os.path.join(tmp, str(uuid.uuid1())[:8])
+    except KeyError:
+        tmp = tempfile.NamedTemporaryFile().name
+    return 'ipc://%s' % tmp
+
+
+def build_socket(ctx: 'zmq.Context', host: str, port: int,
+                 socket_type: 'SocketType', identity: 'str' = None, use_ipc: bool = False) -> Tuple['zmq.Socket', str]:
     sock = {
         SocketType.PULL_BIND: lambda: ctx.socket(zmq.PULL),
         SocketType.PULL_CONNECT: lambda: ctx.socket(zmq.PULL),
@@ -129,11 +143,14 @@ def build_socket(ctx: 'zmq.Context', host: str, port: int, socket_type: 'SocketT
     }[socket_type]()
 
     if socket_type.is_bind:
-        host = BaseService.default_host
-        if port is None:
-            sock.bind_to_random_port('tcp://%s' % host)
+        if use_ipc:
+            sock.bind(host)
         else:
-            sock.bind('tcp://%s:%d' % (host, port))
+            host = BaseService.default_host
+            if port is None:
+                sock.bind_to_random_port('tcp://%s' % host)
+            else:
+                sock.bind('tcp://%s:%d' % (host, port))
     else:
         if port is None:
             sock.connect(host)
@@ -329,8 +346,12 @@ def __init__(self, args):
         self.last_dump_time = time.perf_counter()
         self._model = None
         self.use_event_loop = True
-        self.ctrl_addr = 'tcp://%s:%d' % (self.default_host, self.args.port_ctrl)
-        self.logger.info('control address: %s' % self.ctrl_addr)
+        self.ctrl_with_ipc = (os.name != 'nt') and self.args.ctrl_with_ipc
+        if self.ctrl_with_ipc:
+            self.ctrl_addr = get_random_ipc()
+        else:
+            self.ctrl_addr = 'tcp://%s:%d' % (self.default_host, self.args.port_ctrl)
+
         self.send_recv_kwargs = dict(
             check_version=self.args.check_version,
             timeout=self.args.timeout,
@@ -401,7 +422,12 @@ def _run(self, ctx):
         self.handler.service_context = self
         # print('!!!! t_id: %d service_context: %r' % (threading.get_ident(), self.handler.service_context))
         self.logger.info('bind sockets...')
-        ctrl_sock, ctrl_addr = build_socket(ctx, self.default_host, self.args.port_ctrl, SocketType.PAIR_BIND)
+        if self.ctrl_with_ipc:
+            ctrl_sock, ctrl_addr = build_socket(ctx, self.ctrl_addr, None, SocketType.PAIR_BIND,
+                                                use_ipc=self.ctrl_with_ipc)
+        else:
+            ctrl_sock, ctrl_addr = build_socket(ctx, self.default_host, self.args.port_ctrl, SocketType.PAIR_BIND)
+
         self.logger.info('control over %s' % (colored(ctrl_addr, 'yellow')))
 
         in_sock, _ = build_socket(ctx, self.args.host_in, self.args.port_in, self.args.socket_in,
@@ -412,7 +438,6 @@ def _run(self, ctx):
                                    self.args.identity)
         self.logger.info('output %s:%s' % (self.args.host_out, colored(self.args.port_out, 'yellow')))
 
-
         self.logger.info(
             'input %s:%s\t output %s:%s\t control over %s' % (
                 self.args.host_in, colored(self.args.port_in, 'yellow'),