Merge pull request #218 from christianversloot:master

PiperOrigin-RevId: 345099633

Author: Tensorflow Cloud maintainers

src/python/tensorflow_cloud/core/containerize.py

@@ -103,6 +103,8 @@ def __init__(
         # Those will be populated lazily.
         self.tar_file_path = None
         self.docker_client = None
+        self.tar_file_descriptor = None
+        self.docker_file_descriptor = None
 
     def get_docker_image(
         self, max_status_check_attempts=None, delay_between_status_checks=None
@@ -117,15 +119,30 @@ def get_docker_image(
         """
         raise NotImplementedError
 
-    def get_generated_files(self):
-        return [self.docker_file_path, self.tar_file_path]
+    def get_generated_files(self, return_descriptors=False):
+        """Get generated file paths and/or descriptors for generated files.
+
+        Args:
+            return_descriptors: Whether to return descriptors as well.
+
+        Returns:
+          Docker and tar file paths. Depending on return_descriptors, possibly
+          their file descriptors as well.
+        """
+        if return_descriptors:
+            return [
+                (self.docker_file_path, self.docker_file_descriptor),
+                (self.tar_file_path, self.tar_file_descriptor)
+            ]
+        else:
+            return [self.docker_file_path, self.tar_file_path]
 
     def _get_tar_file_path(self):
         """Packages files into a tarball."""
         self._create_docker_file()
         file_path_map = self._get_file_path_map()
 
-        _, self.tar_file_path = tempfile.mkstemp()
+        self.tar_file_descriptor, self.tar_file_path = tempfile.mkstemp()
         with tarfile.open(self.tar_file_path, "w:gz", dereference=True) as tar:
             for source, destination in file_path_map.items():
                 tar.add(source, arcname=destination)
@@ -233,7 +250,7 @@ def _create_docker_file(self):
         )
 
         content = "\n".join(lines)
-        _, self.docker_file_path = tempfile.mkstemp()
+        self.docker_file_descriptor, self.docker_file_path = tempfile.mkstemp()
         with open(self.docker_file_path, "w") as f:
             f.write(content)
 

src/python/tensorflow_cloud/core/preprocess.py

@@ -47,6 +47,7 @@ def get_preprocessed_entry_point(
     worker_count,
     distribution_strategy,
     called_from_notebook=False,
+    return_file_descriptor=False
 ):
     """Creates python script for distribution based on the given `entry_point`.
 
@@ -107,13 +108,15 @@ def get_preprocessed_entry_point(
             `worker_count` params.
         called_from_notebook: Boolean. True if the API is run in a
             notebook environment.
+        return_file_descriptor: Boolean. True if the file descriptor for the
+            temporary file is also returned.
 
     Returns:
         The `preprocessed_entry_point` file path.
 
     Raises:
         RuntimeError: If invoked from Notebook but unable to access it.
-                      Typically, this is due to missing the `nbconvert` package.
+            Typically, this is due to missing the `nbconvert` package.
     """
 
     # Set `TF_KERAS_RUNNING_REMOTELY` env variable. This is required in order
@@ -198,10 +201,15 @@ def get_preprocessed_entry_point(
                 script_lines.append(line)
 
     # Create a tmp wrapped entry point script file.
-    _, output_file = tempfile.mkstemp(suffix=".py")
+    file_descriptor, output_file = tempfile.mkstemp(suffix=".py")
     with open(output_file, "w") as f:
         f.writelines(script_lines)
-    return output_file
+
+    # Returning file descriptor could be necessary for some os.close calls
+    if return_file_descriptor:
+      return (output_file, file_descriptor)
+    else:
+      return output_file
 
 
 def _get_colab_notebook_content():

src/python/tensorflow_cloud/core/run.py

@@ -194,14 +194,16 @@ def run(
     if (distribution_strategy == "auto"
         or entry_point.endswith("ipynb")
         or entry_point is None):
-        preprocessed_entry_point = preprocess.get_preprocessed_entry_point(
-            entry_point,
-            chief_config,
-            worker_config,
-            worker_count,
-            distribution_strategy,
-            called_from_notebook=called_from_notebook,
-        )
+        preprocessed_entry_point, \
+          pep_file_descriptor = preprocess.get_preprocessed_entry_point(
+              entry_point,
+              chief_config,
+              worker_config,
+              worker_count,
+              distribution_strategy,
+              called_from_notebook=called_from_notebook,
+              return_file_descriptor=True,
+          )
 
     # Create Docker file, generate a tarball, build and push Docker
     # image using the tarball.
@@ -227,9 +229,12 @@ def run(
 
     # Delete all the temporary files we created.
     if preprocessed_entry_point is not None:
+        os.close(pep_file_descriptor)
         os.remove(preprocessed_entry_point)
-    for f in container_builder.get_generated_files():
-        os.remove(f)
+    for file_path, file_descriptor \
+        in container_builder.get_generated_files(return_descriptors=True):
+        os.close(file_descriptor)
+        os.remove(file_path)
 
     # Deploy Docker image on the cloud.
     job_id = deploy.deploy_job(