def test_get_file_path_map_defaults(self):
self.setup()
self.docker_file = containerize.create_docker_file(
self.entry_point, self.chief_config)
file_map = containerize.get_file_path_map(
self.entry_point, self.docker_file)
self.assertDictEqual(file_map, {
self.docker_file: 'Dockerfile',
self.entry_point_dir: '/app/'})
self.cleanup()
def test_create_docker_file_defaults(self):
self.setup()
self.docker_file = containerize.create_docker_file(
self.entry_point, self.chief_config)
expected_docker_file_lines = [
'FROM tensorflow/tensorflow:{}-gpu\n'.format(VERSION),
'WORKDIR /app/\n',
'COPY /app/ /app/\n',
'ENTRYPOINT ["python", "mnist_example_using_fit.py"]',
]
self.assert_docker_file(expected_docker_file_lines)
self.cleanup()
def test_get_docker_image(self, MockAPIClient, MockLogger):
self.setup()
mock_registry = 'gcr.io/my-project'
mock_img_tag = mock_registry + '/tensorflow-train:abcde'
# Verify mocking is correct and mock img tag.
assert MockAPIClient is containerize.APIClient
assert MockLogger is containerize.logger
docker_client = MockAPIClient.return_value
def _mock_generate_name(docker_registry):
return mock_img_tag
containerize._generate_name = _mock_generate_name
self.docker_file = containerize.create_docker_file(
self.entry_point, self.chief_config)
file_map = containerize.get_file_path_map(
self.entry_point, self.docker_file)
tarball = package.get_tar_file_path(file_map)
img_tag = containerize.get_docker_image(mock_registry, tarball)
self.assertEqual(img_tag, mock_img_tag)
# Verify docker APIClient is invoked as expected.
self.assertEqual(MockAPIClient.call_count, 1)
_, kwargs = MockAPIClient.call_args
self.assertDictEqual(kwargs, {'version': 'auto'})
# Verify APIClient().build is invoked as expected.
self.assertEqual(docker_client.build.call_count, 1)
_, kwargs = docker_client.build.call_args
expected = {
'path': '.',
'custom_context': True,
'encoding': 'utf-8',
'tag': img_tag
}
self.assertTrue(set(expected.items()).issubset(set(kwargs.items())))
# Verify APIClient().push is invoked as expected.
self.assertEqual(docker_client.push.call_count, 1)
args, kwargs = docker_client.push.call_args
self.assertListEqual(list(args), [img_tag])
self.assertDictEqual(kwargs, {'stream': True})
# Verify logger info calls.
self.assertEqual(MockLogger.info.call_count, 2)
MockLogger.info.assert_has_calls([
call(r' Building docker image: ' + img_tag),
call(r' Publishing docker image: ' + img_tag)])
self.cleanup()
def test_create_docker_file_with_docker_base_image(self):
self.setup()
self.docker_file = containerize.create_docker_file(
self.entry_point, self.chief_config,
docker_base_image='tensorflow/tensorflow:latest')
expected_docker_file_lines = [
'FROM tensorflow/tensorflow:latest\n',
'WORKDIR /app/\n',
'COPY /app/ /app/\n',
'ENTRYPOINT ["python", "mnist_example_using_fit.py"]',
]
self.assert_docker_file(expected_docker_file_lines)
self.cleanup()
def test_get_file_path_map_with_destination_dir(self):
self.setup()
self.docker_file = containerize.create_docker_file(
self.entry_point, self.chief_config,
destination_dir='/my_app/temp/')
file_map = containerize.get_file_path_map(
self.entry_point, self.docker_file,
destination_dir='/my_app/temp/')
self.assertDictEqual(file_map, {
self.docker_file: 'Dockerfile',
self.entry_point_dir: '/my_app/temp/'})
self.cleanup()
def test_create_docker_file_with_cpu_config(self):
self.setup()
self.docker_file = containerize.create_docker_file(
self.entry_point,
machine_config.COMMON_MACHINE_CONFIGS['CPU'])
expected_docker_file_lines = [
'FROM tensorflow/tensorflow:{}\n'.format(VERSION),
'WORKDIR /app/\n',
'COPY /app/ /app/\n',
'ENTRYPOINT ["python", "mnist_example_using_fit.py"]',
]
self.assert_docker_file(expected_docker_file_lines)
self.cleanup()
def test_get_file_path_map_with_wrapped_entry_point(self):
self.setup()
self.docker_file = containerize.create_docker_file(
self.entry_point, self.chief_config,
destination_dir='/my_app/temp/')
# entry_point file mimic wrapped entry_point here.
file_map = containerize.get_file_path_map(
self.entry_point, self.docker_file,
wrapped_entry_point=self.entry_point,
destination_dir='/my_app/temp/')
self.assertDictEqual(file_map, {
self.docker_file: 'Dockerfile',
self.entry_point_dir: '/my_app/temp/',
self.entry_point: '/my_app/temp/mnist_example_using_fit.py'})
self.cleanup()
def test_get_file_path_map_with_requirements(self):
self.setup()
req_file = 'requirements.txt'
with open(req_file, 'w') as f:
f.writelines(['tensorflow-datasets'])
self.docker_file = containerize.create_docker_file(
self.entry_point, self.chief_config,
requirements_txt=req_file)
file_map = containerize.get_file_path_map(
self.entry_point, self.docker_file, requirements_txt=req_file)
self.assertDictEqual(file_map, {
self.docker_file: 'Dockerfile',
req_file: '/app/requirements.txt',
self.entry_point_dir: '/app/'})
os.remove(req_file)
self.cleanup()
def test_create_docker_with_requirements(self):
self.setup()
req_file = 'requirements.txt'
with open(req_file, 'w') as f:
f.writelines(['tensorflow-datasets'])
self.docker_file = containerize.create_docker_file(
self.entry_point, self.chief_config,
requirements_txt=req_file)
expected_docker_file_lines = [
'FROM tensorflow/tensorflow:{}-gpu\n'.format(VERSION),
'WORKDIR /app/\n',
'COPY /app/ /app/\n',
'RUN if [ -e requirements.txt ]; '
'then pip install --no-cache -r requirements.txt; fi\n',
'ENTRYPOINT ["python", "mnist_example_using_fit.py"]',
]
self.assert_docker_file(expected_docker_file_lines)
os.remove(req_file)
self.cleanup()