add explicit pull for local serve

Author: Dan Choi

src/sagemaker/local/image.py

@@ -668,3 +668,13 @@ def _ecr_login_if_needed(boto_session, image):
 
     cmd = "docker login -u AWS -p %s %s" % (token, ecr_url)
     subprocess.check_output(cmd, shell=True)
+
+    return True
+
+
+def _pull_image(image):
+    pull_image_command = ('docker pull %s' % image).strip()
+    print('docker command: {}'.format(pull_image_command))
+
+    subprocess.check_output(pull_image_command, shell=True)
+    print('image pulled: {}'.format(image))

tests/unit/test_image.py

@@ -513,13 +513,26 @@ def test_ecr_login_needed(check_output):
     }
     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)
+    result = 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'])
 
+    assert result
+
+
+@patch('subprocess.check_output', return_value=''.encode('utf-8'))
+def test_pull_image(check_output):
+    image = '520713654638.dkr.ecr.us-east-1.amazonaws.com/image-i-need:1.1'
+
+    sagemaker.local.image._pull_image(image)
+
+    expected_command = 'docker pull %s' % image
+
+    check_output.assert_called_once_with(expected_command, shell=True)
+
 
 def test__aws_credentials_with_long_lived_credentials():
     credentials = Credentials(access_key=_random_string(), secret_key=_random_string(), token=None)