export_saved_model: copy asset files (#100)

* export_saved_model: copy asset files

Author: lukasz-artificial

src/tf_container/serve.py

@@ -12,6 +12,7 @@
 #  permissions and limitations under the License.
 
 import json
+import re
 import shutil
 import subprocess
 import boto3
@@ -54,34 +55,33 @@ def export_saved_model(checkpoint_dir, model_path, s3=boto3.client('s3', region_
             raise e
         # Select most recent saved_model.pb
         saved_model_path = saved_model_path_array[-1]
-
-        variables_path = [x['Key'] for x in contents if 'variables/variables' in x['Key']]
-        variable_names_to_paths = {v.split('/').pop(): v for v in variables_path}
+        saved_model_base_path = os.path.dirname(saved_model_path)
 
         prefixes = key_prefix.split('/')
-        folders = saved_model_path.split('/')[len(prefixes):]
-        saved_model_filename = folders.pop()
+        folders = saved_model_path.split('/')[len(prefixes):-1]
         path_to_save_model = os.path.join(model_path, *folders)
 
-        path_to_variables = os.path.join(path_to_save_model, 'variables')
-
-        os.makedirs(path_to_variables)
-
-        target = os.path.join(path_to_save_model, saved_model_filename)
-        s3.download_file(bucket_name, saved_model_path, target)
-        logger.info("Downloaded saved model at {}".format(target))
+        def file_filter(x): return x['Key'].startswith(saved_model_base_path) and not x['Key'].endswith("/")
+        paths_to_copy = [x['Key'] for x in contents if file_filter(x)]
 
-        for filename, full_path in variable_names_to_paths.items():
-            key = full_path
-            target = os.path.join(path_to_variables, filename)
+        for key in paths_to_copy:
+            target = re.sub(r"^"+saved_model_base_path, path_to_save_model, key)
+            _makedirs_for_file(target)
             s3.download_file(bucket_name, key, target)
+        logger.info("Downloaded saved model at {}".format(path_to_save_model))
     else:
         if os.path.exists(checkpoint_dir):
             _recursive_copy(checkpoint_dir, model_path)
         else:
             logger.error("Failed to copy saved model. File does not exist in {}".format(checkpoint_dir))
 
 
+def _makedirs_for_file(file_path):
+    directory = os.path.dirname(file_path)
+    if not os.path.exists(directory):
+        os.makedirs(directory)
+
+
 def _recursive_copy(src, dst):
     for root, dirs, files in os.walk(src):
         root = os.path.relpath(root, src)

test/integ/container_tests/sentiment_classification.py

@@ -0,0 +1,32 @@
+import json
+import os
+import requests
+
+
+def test_json_request():
+    data = ["great", "fantastic", "movie"]
+
+    serialized_output = requests.post("http://localhost:8080/invocations",
+                                      data=json.dumps(data),
+                                      headers={'Content-type': 'application/json',
+                                               'Accept': 'application/json'}).content
+
+    prediction_result = json.loads(serialized_output)
+    classes = prediction_result['result']['classifications'][0]['classes']
+    assert len(classes) == 2
+    assert classes[0].keys() == [u'score', u'label']
+
+
+def test_assets_were_restored():
+    found_vocabulary_files = _find_files("/opt/ml/model/export/Servo",
+                                         dir_predicate=lambda d: d.endswith("/assets"),
+                                         file_predicate=lambda f: f == "vocabulary.txt")
+    assert len(found_vocabulary_files) > 0, 'At least one "vocabulary.txt" asset file is expected'
+
+
+def _find_files(root_dir, dir_predicate=lambda _: True, file_predicate=lambda _: True):
+    found_files = []
+    for root, _, files in os.walk(root_dir):
+        if dir_predicate(root):
+            found_files += [root + "/" + file for file in files if file_predicate(file)]
+    return found_files

test/integ/test_save_restore_assets.py

@@ -0,0 +1,37 @@
+import os
+
+from sagemaker import fw_utils
+
+from test.integ.docker_utils import train, HostingContainer
+from test.integ.utils import copy_resource, create_config_files
+from test.integ.conftest import SCRIPT_PATH
+
+import uuid
+
+
+def test_save_restore_assets(docker_image, sagemaker_session, opt_ml, processor):
+    resource_path = os.path.join(SCRIPT_PATH, '../resources/sentiment')
+
+    default_bucket = sagemaker_session.default_bucket()
+
+    copy_resource(resource_path, opt_ml, 'data', 'input/data')
+
+    s3_source_archive = fw_utils.tar_and_upload_dir(session=sagemaker_session.boto_session,
+                                bucket=sagemaker_session.default_bucket(),
+                                s3_key_prefix='test_job',
+                                script='sentiment.py',
+                                directory=os.path.join(resource_path, 'code'))
+
+    checkpoint_s3_path = 's3://{}/save_restore_assets/output-{}'.format(default_bucket, uuid.uuid4())
+
+    additional_hyperparameters = dict(
+        training_steps=1000,
+        evaluation_steps=100,
+        checkpoint_path=checkpoint_s3_path)
+    create_config_files('sentiment.py', s3_source_archive.s3_prefix, opt_ml, additional_hyperparameters)
+    os.makedirs(os.path.join(opt_ml, 'model'))
+
+    train(docker_image, opt_ml, processor)
+
+    with HostingContainer(opt_ml=opt_ml, image=docker_image, script_name='sentiment.py', processor=processor) as c:
+        c.execute_pytest('test/integ/container_tests/sentiment_classification.py')

test/resources/sentiment/__init__.py

@@ -0,0 +1,48 @@
+from __future__ import print_function
+
+import tensorflow as tf
+
+INPUT_TENSOR_NAME = 'inputs'
+CHANNEL_BASEDIR = '/opt/ml/input/data/'
+
+
+def estimator_fn(run_config, hyperparameters):
+    feature_columns = [tf.feature_column.embedding_column(
+        tf.feature_column.categorical_column_with_vocabulary_file(
+            INPUT_TENSOR_NAME, CHANNEL_BASEDIR + "vocabulary.txt"),
+        combiner="mean",
+        dimension=32)]
+
+    return tf.estimator.DNNClassifier(feature_columns=feature_columns,
+                                      hidden_units=[10],
+                                      n_classes=2,
+                                      config=run_config)
+
+
+def train_input_fn(training_dir, hyperparameters):
+    return _generate_input_fn("train")
+    
+    
+def eval_input_fn(training_dir, hyperparameters):
+    return _generate_input_fn("test")
+
+
+def serving_input_fn(hyperparameters):
+    feature_spec = {INPUT_TENSOR_NAME: tf.VarLenFeature(dtype=tf.string)}
+    return tf.estimator.export.build_parsing_serving_input_receiver_fn(feature_spec)()
+
+
+def _generate_input_fn(data_dir):
+    dataset = tf.data.Dataset.from_generator(lambda: _input_data_reader(data_dir), (tf.string, tf.int32)).repeat(3)
+    doc, label = dataset.make_one_shot_iterator().get_next()
+    return {INPUT_TENSOR_NAME: doc}, label
+
+
+def _input_data_reader(subdir):
+    data_file = CHANNEL_BASEDIR + subdir + ".txt"
+    with open(data_file, "rt") as f:
+        for line in f.readlines():
+            if line.startswith("#"):
+                continue
+            label, doc = line.strip().split("\t")
+            yield [doc.lower().split(" ")], [[int(label)]]

test/resources/sentiment/code/__init__.py

@@ -0,0 +1,31 @@
+# samples from Large Movie Review Dataset (http://ai.stanford.edu/~amaas/data/sentiment/ , http://www.aclweb.org/anthology/P11-1015)
+1	One of the funniest movies made in recent years. Good characterization, plot and exceptional chemistry make this one a classic
+0	Hated it with all my being. Worst movie ever. Mentally- scarred. Help me. It was that bad.TRUST ME!!!
+0	This is quite possibly the worst sequel ever made. The script is unfunny and the acting stinks. The exact opposite of the original.
+0	This movie is terrible. It's about some no brain surfin dude that inherits some company. Does Carrot Top have no shame?<br /><br />
+1	This is the greatest movie ever. If you have written it off with out ever seeing it. You must give it a second try.
+1	Just love the interplay between two great characters of stage & screen - Veidt & Barrymore
+0	Don't waste your time and money on it. It's not quite as bad as "Adrenalin", by the same director but that's not saying much.
+1	Absolutely fantastic! Whatever I say wouldn't do this underrated movie the justice it deserves. Watch it now! FANTASTIC!
+0	This movie is so bad it's almost good. Bad story, bad acting, bad music, you name it. O.K., who are the jokers that gave this flick a '10'?
+1	Add this little gem to your list of holiday regulars. It is<br /><br />sweet, funny, and endearing
+0	Ten minutes of people spewing gallons of pink vomit. Recurring scenes of enormous piles of dog excrement - need one say more???
+0	Primary plot!Primary direction!Poor interpretation.
+0	Read the book, forget the movie!
+1	This is a good film. This is very funny. Yet after this film there were no good Ernest films!
+1	If you like Pauly Shore, you'll love Son in Law. If you hate Pauly Shore, then, well...I liked it!
+1	Brilliant and moving performances by Tom Courtenay and Peter Finch.
+1	A touching movie. It is full of emotions and wonderful acting. I could have sat through it a second time.
+1	Without a doubt, one of Tobe Hoppor's best! Epic storytellng, great special effects, and The Spacegirl (vamp me baby!).
+0	Widow hires a psychopath as a handyman. Sloppy film noir thriller which doesn't make much of its tension promising set-up. (3/10)
+1	For pure gothic vampire cheese nothing can compare to the Subspecies films. I highly recommend each and every one of them.
+0	More suspenseful, more subtle, much, much more disturbing....
+1	a mesmerizing film that certainly keeps your attention... Ben Daniels is fascinating (and courageous) to watch.
+0	What a script, what a story, what a mess!
+1	If you've ever had a mad week-end out with your mates then you'll appreciate this film. Excellent fun and a laugh a minute.
+0	I caught this film late at night on HBO. Talk about wooden acting, unbelievable plot, et al. Very little going in its favor. Skip it.
+1	Brilliant. Ranks along with Citizen Kane, The Matrix and Godfathers. Must see, at least for basset in her early days. Watch it.
+0	The plot was really weak and confused. This is a true Oprah flick. (In Oprah's world, all men are evil and all women are victims.)
+0	I hope this group of film-makers never re-unites.
+1	This is a great movie. Too bad it is not available on home video.
+0	This is a terrible movie, don't waste your money on it. Don't even watch it for free. That's all I have to say.