Login to ECR if needed for Local Mode

CHANGELOG.rst

@@ -2,6 +2,10 @@
 CHANGELOG
 =========
 
+1.2.1
+========
+* bug-fix: Change Local Mode to use a sagemaker-local docker network
+
 1.2.0
 ========
 

setup.py

@@ -11,7 +11,7 @@ def read(fname):
 
 
 setup(name="sagemaker",
-      version="1.2.0",
+      version="1.2.1",
       description="Open source library for training and deploying models on Amazon SageMaker.",
       packages=find_packages('src'),
       package_dir={'': 'src'},

src/sagemaker/local/image.py

@@ -10,13 +10,16 @@
 # 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.
+import base64
 import errno
 import json
 import logging
 import os
 import platform
+import random
 import shlex
 import shutil
+import string
 import subprocess
 import sys
 import tempfile
@@ -59,7 +62,10 @@ def __init__(self, instance_type, instance_count, image, sagemaker_session=None)
         self.instance_type = instance_type
         self.instance_count = instance_count
         self.image = image
-        self.hosts = ['{}-{}'.format(CONTAINER_PREFIX, i) for i in range(1, self.instance_count + 1)]
+        # Since we are using a single docker network, Generate a random suffix to attach to the container names.
+        #  This way multiple jobs can run in parallel.
+        suffix = ''.join(random.choice(string.ascii_uppercase + string.digits) for _ in range(5))
+        self.hosts = ['{}-{}-{}'.format(CONTAINER_PREFIX, i, suffix) for i in range(1, self.instance_count + 1)]
         self.container_root = None
         self.container = None
         # set the local config. This is optional and will use reasonable defaults
@@ -110,6 +116,8 @@ def train(self, input_data_config, hyperparameters):
 
         compose_data = self._generate_compose_file('train', additional_volumes=volumes)
         compose_command = self._compose()
+
+        _ecr_login_if_needed(self.sagemaker_session.boto_session, self.image)
         _execute_and_stream_output(compose_command)
 
         s3_model_artifacts = self.retrieve_model_artifacts(compose_data)
@@ -152,6 +160,8 @@ def serve(self, primary_container):
 
         env_vars = ['{}={}'.format(k, v) for k, v in primary_container['Environment'].items()]
 
+        _ecr_login_if_needed(self.sagemaker_session.boto_session, self.image)
+
         self._generate_compose_file('serve', additional_env_vars=env_vars)
         compose_command = self._compose()
         self.container = _HostingContainer(compose_command)
@@ -296,7 +306,11 @@ def _generate_compose_file(self, command, additional_volumes=None, additional_en
         content = {
             # Some legacy hosts only support the 2.1 format.
             'version': '2.1',
-            'services': services
+            'services': services,
+            'networks': {
+                'sagemaker-local': {'name': 'sagemaker-local'}
+            }
+
         }
 
         docker_compose_path = os.path.join(self.container_root, DOCKER_COMPOSE_FILENAME)
@@ -335,7 +349,12 @@ def _create_docker_host(self, host, environment, optml_subdirs, command, volumes
             'tty': True,
             'volumes': [v.map for v in optml_volumes],
             'environment': environment,
-            'command': command
+            'command': command,
+            'networks': {
+                'sagemaker-local': {
+                    'aliases': [host]
+                }
+            }
         }
 
         serving_port = 8080 if self.local_config is None else self.local_config.get('serving_port', 8080)
@@ -390,7 +409,8 @@ def _build_optml_volumes(self, host, subdirs):
         return volumes
 
     def _cleanup(self):
-        _check_output('docker network prune -f')
+        # we don't need to cleanup anything at the moment
+        pass
 
 
 class _HostingContainer(object):
@@ -525,3 +545,24 @@ def _aws_credentials(session):
 def _write_json_file(filename, content):
     with open(filename, 'w') as f:
         json.dump(content, f)
+
+
+def _ecr_login_if_needed(boto_session, image):
+    # Only ECR images need login
+    if not ('dkr.ecr' in image and 'amazonaws.com' in image):
+        return
+
+    # do we have the image?
+    if _check_output('docker images -q %s' % image).strip():
+        return
+
+    ecr = boto_session.client('ecr')
+    auth = ecr.get_authorization_token(registryIds=[image.split('.')[0]])
+    authorization_data = auth['authorizationData'][0]
+
+    raw_token = base64.b64decode(authorization_data['authorizationToken'])
+    token = raw_token.decode('utf-8').strip('AWS:')
+    ecr_url = auth['authorizationData'][0]['proxyEndpoint']
+
+    cmd = "docker login -u AWS -p %s %s" % (token, ecr_url)
+    subprocess.check_output(cmd, shell=True)

tests/unit/test_image.py

@@ -10,6 +10,7 @@
 # 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.
+import base64
 import json
 import os
 
@@ -105,6 +106,7 @@ def test_write_config_file(LocalSession, tmpdir):
 @patch('sagemaker.local.local_session.LocalSession')
 def test_retrieve_artifacts(LocalSession, tmpdir):
     sagemaker_container = _SageMakerContainer('local', 2, 'my-image')
+    sagemaker_container.hosts = ['algo-1', 'algo-2']  # avoid any randomness
     sagemaker_container.container_root = str(tmpdir.mkdir('container-root'))
 
     volume1 = os.path.join(sagemaker_container.container_root, 'algo-1/output/')
@@ -227,3 +229,53 @@ def test_serve(up, copy, copytree, tmpdir, sagemaker_session):
             for h in sagemaker_container.hosts:
                 assert config['services'][h]['image'] == image
                 assert config['services'][h]['command'] == 'serve'
+
+
+def test_ecr_login_non_ecr():
+    session_mock = Mock()
+    sagemaker.local.image._ecr_login_if_needed(session_mock, 'ubuntu')
+
+    session_mock.assert_not_called()
+
+
+@patch('sagemaker.local.image._check_output', return_value='123451324')
+def test_ecr_login_image_exists(_check_output):
+    session_mock = Mock()
+
+    image = '520713654638.dkr.ecr.us-east-1.amazonaws.com/image-i-have:1.0'
+    sagemaker.local.image._ecr_login_if_needed(session_mock, image)
+
+    session_mock.assert_not_called()
+    _check_output.assert_called()
+
+
+@patch('subprocess.check_output', return_value=''.encode('utf-8'))
+def test_ecr_login_needed(check_output):
+    session_mock = Mock()
+
+    token = 'very-secure-token'
+    token_response = 'AWS:%s' % token
+    b64_token = base64.b64encode(token_response.encode('utf-8'))
+    response = {
+        u'authorizationData':
+            [
+                {
+                    u'authorizationToken': b64_token,
+                    u'proxyEndpoint': u'https://520713654638.dkr.ecr.us-east-1.amazonaws.com'
+                }
+            ],
+        'ResponseMetadata':
+            {
+                'RetryAttempts': 0,
+                'HTTPStatusCode': 200,
+                'RequestId': '25b2ac63-36bf-11e8-ab6a-e5dc597d2ad9',
+            }
+    }
+    session_mock.client('ecr').get_authorization_token.return_value = response
+    image = '520713654638.dkr.ecr.us-east-1.amazonaws.com/image-i-need:1.1'
+    sagemaker.local.image._ecr_login_if_needed(session_mock, image)
+
+    expected_command = 'docker login -u AWS -p %s https://520713654638.dkr.ecr.us-east-1.amazonaws.com' % token
+
+    check_output.assert_called_with(expected_command, shell=True)
+    session_mock.client('ecr').get_authorization_token.assert_called_with(registryIds=['520713654638'])