refactor context socket manager

jina-ai · Dec 15, 2018 · 1803281 · 1803281
1 parent b91684b
commit 1803281
Showing 1 changed file with 4 additions and 1 deletion.
diff --git a/server/bert_serving/server/__init__.py b/server/bert_serving/server/__init__.py
@@ -337,9 +337,12 @@ def model_fn(features, labels, mode, params):
 
         return Estimator(model_fn=model_fn, config=RunConfig(session_config=config))
 
+    def run(self):
+        self._run()
+
     @zmqd.socket(zmq.PULL)
     @zmqd.socket(zmq.PUSH)
-    def run(self, receiver, sink):
+    def _run(self, receiver, sink):
         self.logger.info('use device %s, load graph from %s' %
                          ('cpu' if self.device_id < 0 else ('gpu: %d' % self.device_id), self.graph_path))
         os.environ['CUDA_VISIBLE_DEVICES'] = str(self.device_id)