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)