def test_create_docker_with_requirements(self):
self.setup()
req_file = "requirements.txt"
with open(req_file, "w") as f:
f.writelines(["tensorflow-datasets"])
lcb = containerize.LocalContainerBuilder(
self.entry_point,
None,
self.chief_config,
self.mock_registry,
self.project_id,
requirements_txt=req_file,
)
lcb._create_docker_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, lcb.docker_file_path)
os.remove(req_file)
self.cleanup(lcb.docker_file_path)
def test_check_docker_base_image_nightly(self):
self.setup()
lcb = containerize.LocalContainerBuilder(
self.entry_point,
None,
self.chief_config,
self.worker_config,
self.mock_registry,
self.project_id,
docker_base_image="tensorflow/tensorflow:2.3.0-dev20200605",
)
# Verify the docker base image is identified as nonexist
self.assertTrue(not lcb._base_image_exist())
# Verify the dockerfile ask for tfnightly
lcb._create_docker_file()
expected_docker_file_lines = [
"FROM tensorflow/tensorflow:nightly\n",
"WORKDIR /app/\n",
"COPY /app/ /app/\n",
'ENTRYPOINT ["python", "mnist_example_using_fit.py"]',
]
self.assert_docker_file(expected_docker_file_lines,
lcb.docker_file_path)
self.cleanup(lcb.docker_file_path)
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
lcb = containerize.LocalContainerBuilder(
self.entry_point,
self.entry_point,
self.chief_config,
self.mock_registry,
self.project_id,
destination_dir="/my_app/temp/",
)
def _mock_generate_name():
return mock_img_tag
lcb._generate_name = _mock_generate_name
img_tag = lcb.get_docker_image()
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(lcb.docker_file_path)
def test_check_docker_base_image_non_exist(self):
self.setup()
lcb = containerize.LocalContainerBuilder(
self.entry_point,
None,
self.chief_config,
self.worker_config,
self.mock_registry,
self.project_id,
docker_base_image="tensorflow/tensorflow:2.100.0",
)
# Verify the docker base image is identified as nonexist
self.assertTrue(not lcb._base_image_exist())
# Verify value error is raised
with self.assertRaises(ValueError) as context:
lcb._create_docker_file()
self.assertTrue("base image" in str(context.exception))
def test_get_file_path_map_defaults(self):
self.setup()
lcb = containerize.LocalContainerBuilder(
self.entry_point,
None,
self.chief_config,
self.mock_registry,
self.project_id,
)
lcb._create_docker_file()
file_map = lcb._get_file_path_map()
self.assertDictEqual(
file_map,
{lcb.docker_file_path: "Dockerfile", self.entry_point_dir: "/app/"},
)
self.cleanup(lcb.docker_file_path)
def test_create_docker_file_defaults(self):
self.setup()
lcb = containerize.LocalContainerBuilder(
self.entry_point,
None,
self.chief_config,
self.mock_registry,
self.project_id,
)
lcb._create_docker_file()
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, lcb.docker_file_path)
self.cleanup(lcb.docker_file_path)
def test_create_docker_file_with_cpu_config(self):
self.setup()
lcb = containerize.LocalContainerBuilder(
self.entry_point,
None,
machine_config.COMMON_MACHINE_CONFIGS["CPU"],
self.mock_registry,
self.project_id,
)
lcb._create_docker_file()
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, lcb.docker_file_path)
self.cleanup(lcb.docker_file_path)
