def test_create_docker_with_requirements(self):
self.setup()
req_file = os.path.join(tempfile.mkdtemp(), "requirements.txt")
with open(req_file, "w") as f:
f.writelines(["tensorflow-datasets"])
lcb = containerize.LocalContainerBuilder(
self.entry_point,
None,
self.chief_config,
self.worker_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(_TF_VERSION),
"WORKDIR /app/\n",
"COPY /app/requirements.txt /app/requirements.txt\n",
"RUN if [ -e requirements.txt ]; "
"then pip install --no-cache -r requirements.txt; fi\n",
"COPY /app/ /app/\n",
'ENTRYPOINT ["python", "sample.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_get_file_path_map_with_requirements(self):
self.setup()
req_file = os.path.join(tempfile.mkdtemp(), "requirements.txt")
with open(req_file, "w") as f:
f.writelines(["tensorflow-datasets"])
lcb = containerize.LocalContainerBuilder(
self.entry_point,
None,
self.chief_config,
self.worker_config,
self.mock_registry,
self.project_id,
requirements_txt=req_file,
)
lcb._create_docker_file()
file_map = lcb._get_file_path_map()
self.assertDictEqual(
file_map,
{
lcb.docker_file_path: "Dockerfile",
req_file: "/app/requirements.txt",
self.entry_point_dir: "/app/",
},
)
os.remove(req_file)
self.cleanup(lcb.docker_file_path)
def test_get_file_path_map_with_wrapped_entry_point(self):
self.setup()
lcb = containerize.LocalContainerBuilder(
self.entry_point,
self.entry_point,
self.chief_config,
self.worker_config,
self.mock_registry,
self.project_id,
destination_dir="/my_app/temp/",
)
lcb._create_docker_file()
file_map = lcb._get_file_path_map()
self.assertDictEqual(
file_map,
{
lcb.docker_file_path: "Dockerfile",
self.entry_point_dir: "/my_app/temp/",
self.entry_point: "/my_app/temp/sample.py",
},
)
self.cleanup(lcb.docker_file_path)
def test_check_nonexistent_docker_image(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:21",
)
# Verify the docker base image is identified as nonexistent
self.assertTrue(not lcb._base_image_exist())
# Verify that the dockerfile fetches the latest stable image
lcb._create_docker_file()
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,
lcb.docker_file_path)
self.cleanup(lcb.docker_file_path)
def test_get_tar_file_path(self):
self.setup()
req_file = "requirements.txt"
with open(req_file, "w") as f:
f.writelines(["tensorflow-datasets"])
lcb = containerize.LocalContainerBuilder(
self.entry_point,
self.entry_point,
self.chief_config,
self.worker_config,
self.mock_registry,
self.project_id,
requirements_txt=req_file,
)
lcb._get_tar_file_path()
assert tarfile.is_tarfile(lcb.tar_file_path)
tar_file = tarfile.open(lcb.tar_file_path)
tar_file_names = [m.name for m in tar_file.getmembers()]
self.assertIn("app/mnist_example_using_fit.py", tar_file_names)
self.assertIn("app/requirements.txt", tar_file_names)
self.assertIn("Dockerfile", tar_file_names)
os.remove(req_file)
self.cleanup(lcb.docker_file_path)
def test_check_docker_base_image_nightly(self):
self.setup(requests_get_return_value=False)
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 nonexistent
self.assertFalse(lcb._base_image_exist())
# Verify that the dockerfile fetches the latest tfnightly
lcb._create_docker_file()
expected_docker_file_lines = [
"FROM tensorflow/tensorflow:nightly\n",
"WORKDIR /app/\n",
"COPY /app/ /app/\n",
'ENTRYPOINT ["python", "sample.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.worker_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, {"decode": True, "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_get_docker_image_with_custom_image_uri(
self, mock_api_client, mock_logger):
self.setup()
# Verify mocking is correct and mock img tag.
assert mock_api_client is docker.APIClient
assert mock_logger is containerize.logger
docker_client = mock_api_client.return_value
lcb = containerize.LocalContainerBuilder(
self.entry_point,
self.entry_point,
self.chief_config,
self.worker_config,
destination_dir="/my_app/temp/",
docker_config=docker_config.DockerConfig(image="gcr.io/test-name"),
)
lcb._get_file_path_map = mock.Mock()
lcb._get_file_path_map.return_value = {}
lcb.tar_file_path = ""
with mock.patch("builtins.open", mock.mock_open(read_data="")):
img_tag = lcb.get_docker_image()
self.assertEqual(img_tag, "gcr.io/test-name")
# Verify docker APIClient is invoked as expected.
self.assertEqual(mock_api_client.call_count, 1)
_, kwargs = mock_api_client.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, {"decode": True, "stream": True})
# Verify logger info calls.
self.assertEqual(mock_logger.info.call_count, 2)
mock_logger.info.assert_has_calls(
[
mock.call("Building Docker image: %s", img_tag),
mock.call("Publishing Docker image: %s", img_tag),
]
)
self.cleanup(lcb.docker_file_path)
def test_get_file_path_map_defaults(self):
self.setup()
lcb = containerize.LocalContainerBuilder(
self.entry_point,
None,
self.chief_config,
self.worker_config,
)
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.worker_config,
)
lcb._create_docker_file()
expected_docker_file_lines = [
"FROM tensorflow/tensorflow:{}-gpu\n".format(_TF_VERSION),
"WORKDIR /app/\n",
"COPY /app/ /app/\n",
'ENTRYPOINT ["python", "sample.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.worker_config,
)
lcb._create_docker_file()
expected_docker_file_lines = [
"FROM tensorflow/tensorflow:{}\n".format(_TF_VERSION),
"WORKDIR /app/\n",
"COPY /app/ /app/\n",
'ENTRYPOINT ["python", "sample.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_kaggle_base_image(self):
self.setup()
os.environ["KAGGLE_DOCKER_IMAGE"] = "gcr.io/kaggle-images/python"
lcb = containerize.LocalContainerBuilder(
self.entry_point,
None,
self.chief_config,
self.worker_config,
)
lcb._create_docker_file()
expected_docker_file_lines = [
"FROM gcr.io/kaggle-images/python\n",
"WORKDIR /app/\n",
"COPY /app/ /app/\n",
'ENTRYPOINT ["python", "sample.py"]',
]
self.assert_docker_file(expected_docker_file_lines,
lcb.docker_file_path)
del os.environ["KAGGLE_DOCKER_IMAGE"]
self.cleanup(lcb.docker_file_path)
def test_create_docker_file_with_docker_parent_image(self):
self.setup()
lcb = containerize.LocalContainerBuilder(
self.entry_point,
None,
self.chief_config,
self.worker_config,
docker_config=docker_config.DockerConfig(
parent_image="tensorflow/tensorflow:latest"),
)
lcb._create_docker_file()
expected_docker_file_lines = [
"FROM tensorflow/tensorflow:latest\n",
"WORKDIR /app/\n",
"COPY /app/ /app/\n",
'ENTRYPOINT ["python", "sample.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_destination_dir(self):
self.setup()
lcb = containerize.LocalContainerBuilder(
self.entry_point,
None,
self.chief_config,
self.worker_config,
self.mock_registry,
self.project_id,
destination_dir="/my_app/temp/",
)
lcb._create_docker_file()
expected_docker_file_lines = [
"FROM tensorflow/tensorflow:{}-gpu\n".format(_TF_VERSION),
"WORKDIR /my_app/temp/\n",
"COPY /my_app/temp/ /my_app/temp/\n",
'ENTRYPOINT ["python", "sample.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_tpu_config(self):
self.setup()
lcb = containerize.LocalContainerBuilder(
self.entry_point,
None,
machine_config.COMMON_MACHINE_CONFIGS["CPU"],
machine_config.COMMON_MACHINE_CONFIGS["TPU"],
self.mock_registry,
self.project_id,
)
lcb._create_docker_file()
expected_docker_file_lines = [
"FROM tensorflow/tensorflow:{}\n".format(VERSION),
"WORKDIR /app/\n",
"RUN pip install cloud-tpu-client\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_docker_base_image(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:latest",
)
lcb._create_docker_file()
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,
lcb.docker_file_path)
self.cleanup(lcb.docker_file_path)
