Increase cloudbuild timeout to 20 minutes.

PiperOrigin-RevId: 344908082

Author: rosbo

src/python/tensorflow_cloud/core/containerize.py

@@ -396,7 +396,7 @@ class CloudContainerBuilder(ContainerBuilder):
     """Container builder that uses Google cloud build."""
 
     def get_docker_image(
-        self, max_status_check_attempts=20, delay_between_status_checks=30
+        self, max_status_check_attempts=40, delay_between_status_checks=30
     ):
         """Builds, publishes and returns a Docker image.
 
@@ -427,7 +427,9 @@ def get_docker_image(
             requestBuilder=google_api_client.TFCloudHttpRequest,
         )
         request_dict = self._create_cloud_build_request_dict(
-            image_uri, storage_object_name
+            image_uri,
+            storage_object_name,
+            max_status_check_attempts*delay_between_status_checks
         )
 
         try:
@@ -494,7 +496,9 @@ def _upload_tar_to_gcs(self):
         blob.upload_from_filename(self.tar_file_path)
         return storage_object_name
 
-    def _create_cloud_build_request_dict(self, image_uri, storage_object_name):
+    def _create_cloud_build_request_dict(
+        self, image_uri, storage_object_name, timeout_sec
+    ):
         """Creates request body for cloud build JSON API call.
 
         `create` body should be a `Build` object
@@ -503,6 +507,7 @@ def _create_cloud_build_request_dict(self, image_uri, storage_object_name):
         Args:
             image_uri: GCR Docker image URI.
             storage_object_name: Name of the tarfile object in GCS.
+            timeout_sec: timeout for the CloudBuild in seconds.
 
         Returns:
             Build request dictionary.
@@ -511,6 +516,7 @@ def _create_cloud_build_request_dict(self, image_uri, storage_object_name):
         request_dict["projectId"] = self.project_id
         request_dict["images"] = [[image_uri]]
         request_dict["steps"] = []
+        request_dict["timeout"] = "{}s".format(timeout_sec)
         build_args = ["build", "-t", image_uri, "."]
 
         if self.docker_config:

src/python/tensorflow_cloud/core/tests/unit/containerize_test.py

@@ -444,6 +444,7 @@ def test_get_docker_image_cloud_build(
 
         request_dict = {}
         request_dict["projectId"] = self.project_id
+        request_dict["timeout"] = "1200s"
         request_dict["images"] = [["gcr.io/my-project/tf_cloud_train:abcde"]]
         request_dict["steps"] = [{
             "name": "gcr.io/cloud-builders/docker",
@@ -612,6 +613,7 @@ def test_get_docker_image_cloud_build_with_custom_image_uri(
 
         request_dict = {}
         request_dict["projectId"] = self.project_id
+        request_dict["timeout"] = "1200s"
         request_dict["images"] = [[img_tag]]
         request_dict["steps"] = [{
             "name": "gcr.io/cloud-builders/docker",
@@ -732,6 +734,7 @@ def test_get_docker_image_cloud_build_with_cache(
 
         request_dict = {}
         request_dict["projectId"] = self.project_id
+        request_dict["timeout"] = "1200s"
         request_dict["images"] = [[img_tag]]
         request_dict["steps"] = [{
             "name": "gcr.io/cloud-builders/docker",