Merge pull request #367 from AvihayTsayeg/master

Enable configuring TFoS server IP and PORT

Author: leewyang

.gitignore

@@ -10,3 +10,5 @@ docs/.doctrees
 target
 test-data
 dependency-reduced-pom.xml
+venv
+.idea

.travis.yml

@@ -13,6 +13,7 @@ matrix:
     - export PYTHONPATH=$(pwd)
     install:
     - pip install -r requirements.txt
+    - pip install mock
     script:
     - test/run_tests.sh
   - language: python

tensorflowonspark/pipeline.py

@@ -23,7 +23,7 @@
 from pyspark.sql import Row, SparkSession
 
 import tensorflow as tf
-from tensorflow.contrib.saved_model.python.saved_model import reader, signature_def_utils
+from tensorflow.contrib.saved_model.python.saved_model import reader
 from tensorflow.python.saved_model import loader
 from . import TFCluster, gpu_info, dfutil
 
@@ -503,7 +503,7 @@ def _run_model(iterator, args, tf_args):
     assert args.export_dir, "Inferencing with signature_def_key requires --export_dir argument"
     logging.info("===== loading meta_graph_def for tag_set ({0}) from saved_model: {1}".format(args.tag_set, args.export_dir))
     meta_graph_def = get_meta_graph_def(args.export_dir, args.tag_set)
-    signature = signature_def_utils.get_signature_def_by_key(meta_graph_def, args.signature_def_key)
+    signature = meta_graph_def.signature_def[args.signature_def_key]
     logging.debug("signature: {}".format(signature))
     inputs_tensor_info = signature.inputs
     logging.debug("inputs_tensor_info: {0}".format(inputs_tensor_info))

tensorflowonspark/reservation.py

@@ -9,6 +9,7 @@
 from __future__ import print_function
 
 import logging
+import os
 import pickle
 import select
 import socket
@@ -19,6 +20,8 @@
 
 from . import util
 
+TFOS_SERVER_PORT = "TFOS_SERVER_PORT"
+TFOS_SERVER_HOST = "TFOS_SERVER_HOST"
 BUFSIZE = 1024
 MAX_RETRIES = 3
 
@@ -146,13 +149,10 @@ def start(self):
     Returns:
       address of the Server as a tuple of (host, port)
     """
-    server_sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
-    server_sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
-    server_sock.bind(('', 0))
-    server_sock.listen(10)
+    server_sock = self.start_listening_socket()
 
     # hostname may not be resolvable but IP address probably will be
-    host = util.get_ip_address()
+    host = self.get_server_ip()
     port = server_sock.getsockname()[1]
     addr = (host, port)
     logging.info("listening for reservations at {0}".format(addr))
@@ -185,6 +185,18 @@ def _listen(self, sock):
 
     return addr
 
+  def get_server_ip(self):
+    return os.getenv(TFOS_SERVER_HOST) if os.getenv(TFOS_SERVER_HOST) else util.get_ip_address()
+
+
+  def start_listening_socket(self):
+    port_number = int(os.getenv(TFOS_SERVER_PORT)) if  os.getenv(TFOS_SERVER_PORT) else 0
+    server_sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
+    server_sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
+    server_sock.bind(('', port_number))
+    server_sock.listen(10)
+    return server_sock
+
   def stop(self):
     """Stop the Server's socket listener."""
     self.done = True

test/test_reservation.py

@@ -1,9 +1,16 @@
+import os
+import sys
 import threading
 import time
 import unittest
 
+from tensorflowonspark import util
 from tensorflowonspark.reservation import Reservations, Server, Client
 
+if sys.version_info >= (3, 3):
+    from unittest import mock
+else:
+    import mock
 
 class ReservationTest(unittest.TestCase):
   def test_reservation_class(self):
@@ -48,6 +55,25 @@ def test_reservation_server(self):
     time.sleep(1)
     self.assertEqual(s.done, True)
 
+  def test_reservation_enviroment_exists_get_server_ip_return_environment_value(self):
+      tfso_server = Server(5)
+      with mock.patch.dict(os.environ,{'TFOS_SERVER_HOST':'my_host_ip'}):
+        assert tfso_server.get_server_ip() == "my_host_ip"
+
+  def test_reservation_enviroment_not_exists_get_server_ip_return_actual_host_ip(self):
+    tfso_server = Server(5)
+    assert tfso_server.get_server_ip() == util.get_ip_address()
+
+  def test_reservation_enviroment_exists_start_listening_socket_return_socket_listening_to_environment_port_value(self):
+    tfso_server = Server(1)
+    with mock.patch.dict(os.environ, {'TFOS_SERVER_PORT': '9999'}):
+      assert tfso_server.start_listening_socket().getsockname()[1] == 9999
+
+  def test_reservation_enviroment_not_exists_start_listening_socket_return_socket(self):
+    tfso_server = Server(1)
+    print(tfso_server.start_listening_socket().getsockname()[1])
+    assert type(tfso_server.start_listening_socket().getsockname()[1]) == int
+
   def test_reservation_server_multi(self):
     """Test reservation server, expecting multiple reservations"""
     num_clients = 4
@@ -85,4 +111,4 @@ def reserve(num):
 
 
 if __name__ == '__main__':
-  unittest.main()
+  unittest.main()