Merge branch 'master' into leewyang_update_1.12

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,6 @@
 from pyspark.sql import Row, SparkSession
 
 import tensorflow as tf
-
 from tensorflow.contrib.saved_model.python.saved_model import reader
 from tensorflow.python.saved_model import loader
 from . import TFCluster, gpu_info, dfutil

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()