fix flaky tests

src/sagemaker/fw_utils.py

@@ -219,9 +219,11 @@ def framework_name_from_image(image_name):
     else:
         # extract framework, python version and image tag
         # We must support both the legacy and current image name format.
-        name_pattern = \
-            re.compile('^sagemaker(?:-rl)?-(tensorflow|mxnet|chainer|pytorch|scikit-learn):(.*)-(.*?)-(py2|py3)$')
-        legacy_name_pattern = re.compile('^sagemaker-(tensorflow|mxnet)-(py2|py3)-(cpu|gpu):(.*)$')
+        name_pattern = re.compile(
+            r'^sagemaker(?:-rl)?-(tensorflow|mxnet|chainer|pytorch|scikit-learn):(.*)-(.*?)-(py2|py3)$')
+        legacy_name_pattern = re.compile(
+            r'^sagemaker-(tensorflow|mxnet)-(py2|py3)-(cpu|gpu):(.*)$')
+
         name_match = name_pattern.match(sagemaker_match.group(8))
         legacy_match = legacy_name_pattern.match(sagemaker_match.group(8))
 

src/sagemaker/tensorflow/serving.py

@@ -13,7 +13,6 @@
 from __future__ import absolute_import
 
 import logging
-
 import sagemaker
 from sagemaker.content_types import CONTENT_TYPE_JSON
 from sagemaker.fw_utils import create_image_uri
@@ -144,7 +143,6 @@ def _get_image_uri(self, instance_type):
         if self.image:
             return self.image
 
-        # reuse standard image uri function, then strip unwanted python component
         region_name = self.sagemaker_session.boto_region_name
         return create_image_uri(region_name, Model.FRAMEWORK_NAME, instance_type,
                                 self._framework_version)

src/sagemaker/utils.py

@@ -14,6 +14,7 @@
 
 import errno
 import os
+import random
 import re
 import sys
 import tarfile
@@ -64,6 +65,14 @@ def name_from_base(base, max_length=63, short=False):
     return '{}-{}'.format(trimmed_base, timestamp)
 
 
+def unique_name_from_base(base, max_length=63):
+    unique = '%04x' % random.randrange(16**4)  # 4-digit hex
+    ts = str(int(time.time()))
+    available_length = max_length - 2 - len(ts) - len(unique)
+    trimmed = base[:available_length]
+    return '{}-{}-{}'.format(trimmed, ts, unique)
+
+
 def airflow_name_from_base(base, max_length=63, short=False):
     """Append airflow execution_date macro (https://airflow.apache.org/code.html?#macros)
     to the provided string. The macro will beevaluated in Airflow operator runtime.

tests/integ/local_mode_utils.py

@@ -0,0 +1,38 @@
+# Copyright 2017-2018 Amazon.com, Inc. or its affiliates. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License"). You
+# may not use this file except in compliance with the License. A copy of
+# the License is located at
+#
+#     http://aws.amazon.com/apache2.0/
+#
+# or in the "license" file accompanying this file. This file is
+# distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF
+# ANY KIND, either express or implied. See the License for the specific
+# language governing permissions and limitations under the License.
+from __future__ import absolute_import
+
+import fcntl
+import os
+import time
+from contextlib import contextmanager
+
+import tests.integ
+
+LOCK_PATH = os.path.join(tests.integ.DATA_DIR, 'local_mode_lock')
+
+
+@contextmanager
+def lock():
+    # Since Local Mode uses the same port for serving, we need a lock in order
+    # to allow concurrent test execution.
+    local_mode_lock_fd = open(LOCK_PATH, 'w')
+    local_mode_lock = local_mode_lock_fd.fileno()
+
+    fcntl.lockf(local_mode_lock, fcntl.LOCK_EX)
+
+    try:
+        yield
+    finally:
+        time.sleep(5)
+        fcntl.lockf(local_mode_lock, fcntl.LOCK_UN)