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)