Allow training without S3 (#39)

This makes the container not download the training script if it is
not an S3 uri. In this case we assume it was mounted to the container
directly.

Also fix the model collection when the checkpoints happen in local
files.

Author: iquintero

src/tf_container/serve.py

@@ -75,11 +75,25 @@ def export_saved_model(checkpoint_dir, model_path, s3=boto3.client('s3')):
             s3.download_file(bucket_name, key, target)
     else:
         if os.path.exists(checkpoint_dir):
-            shutil.copy2(checkpoint_dir, model_path)
+            _recursive_copy(checkpoint_dir, model_path)
         else:
             logger.error("Failed to copy saved model. File does not exist in {}".format(checkpoint_dir))
 
 
+def _recursive_copy(src, dst):
+    for root, dirs, files in os.walk(src):
+        root = os.path.relpath(root, src)
+        current_path = os.path.join(src, root)
+        target_path = os.path.join(dst, root)
+
+        for file in files:
+            shutil.copy(os.path.join(current_path, file), os.path.join(target_path, file))
+        for dir in dirs:
+            new_dir = os.path.join(target_path, dir)
+            if not os.path.exists(new_dir):
+                os.mkdir(os.path.join(target_path, dir))
+
+
 def transformer(user_module):
     grpc_proxy_client = proxy_client.GRPCProxyClient(TF_SERVING_PORT)
     _wait_model_to_load(grpc_proxy_client, TF_SERVING_MAXIMUM_LOAD_MODEL_TIME_IN_SECONDS)

src/tf_container/train_entry_point.py

@@ -136,7 +136,8 @@ def train():
     # saving checkpoints of larger sizes.
     os.environ['S3_REQUEST_TIMEOUT_MSEC'] = str(env.hyperparameters.get('s3_checkpoint_save_timeout', 60000))
 
-    env.download_user_module()
+    if env.user_script_archive.lower().startswith('s3://'):
+        env.download_user_module()
     env.pip_install_requirements()
 
     customer_script = env.import_user_module()

test/unit/test_serve.py

@@ -80,7 +80,7 @@ def test_export_saved_model_from_filesystem(mock_exists, mock_makedirs, serve):
     checkpoint_dir = 'a/dir'
     model_path = 'possible/another/dir'
 
-    with patch('shutil.copy2') as mock_copy:
+    with patch('tf_container.serve._recursive_copy') as mock_copy:
         serve.export_saved_model(checkpoint_dir, model_path)
         mock_copy.assert_called_once_with(checkpoint_dir, model_path)