SavedModelBundle leak fix (#335) (#350)

tensorflow · Jul 8, 2021 · e0e4545 · e0e4545
1 parent e786c78
commit e0e4545
Show file tree

Hide file tree

Showing 13 changed files with 73 additions and 58 deletions.
diff --git a/README.md b/README.md
@@ -61,12 +61,12 @@ systems, you should add the following dependencies:
 <dependency>
   <groupId>org.tensorflow</groupId>
   <artifactId>tensorflow-core-api</artifactId>
-  <version>0.3.1</version>
+  <version>0.3.2</version>
 </dependency>
 <dependency>
   <groupId>org.tensorflow</groupId>
   <artifactId>tensorflow-core-api</artifactId>
-  <version>0.3.1</version>
+  <version>0.3.2</version>
   <classifier>linux-x86_64${javacpp.platform.extension}</classifier>
 </dependency>
 ```
@@ -77,24 +77,24 @@ native dependencies as follows:
 <dependency>
   <groupId>org.tensorflow</groupId>
   <artifactId>tensorflow-core-api</artifactId>
-  <version>0.3.1</version>
+  <version>0.3.2</version>
 </dependency>
 <dependency>
   <groupId>org.tensorflow</groupId>
   <artifactId>tensorflow-core-api</artifactId>
-  <version>0.3.1</version>
+  <version>0.3.2</version>
   <classifier>linux-x86_64${javacpp.platform.extension}</classifier>
 </dependency>
 <dependency>
   <groupId>org.tensorflow</groupId>
   <artifactId>tensorflow-core-api</artifactId>
-  <version>0.3.1</version>
+  <version>0.3.2</version>
   <classifier>macosx-x86_64${javacpp.platform.extension}</classifier>
 </dependency>
 <dependency>
   <groupId>org.tensorflow</groupId>
   <artifactId>tensorflow-core-api</artifactId>
-  <version>0.3.1</version>
+  <version>0.3.2</version>
   <classifier>windows-x86_64${javacpp.platform.extension}</classifier>
 </dependency>
 ```
@@ -107,7 +107,7 @@ artifact includes transitively all the artifacts above as a single dependency:
 <dependency>
   <groupId>org.tensorflow</groupId>
   <artifactId>tensorflow-core-platform${javacpp.platform.extension}</artifactId>
-  <version>0.3.1</version>
+  <version>0.3.2</version>
 </dependency>
 ```
 
@@ -137,7 +137,7 @@ to add Sonatype OSS repository in your pom.xml, like the following
     <dependency>
         <groupId>org.tensorflow</groupId>
         <artifactId>tensorflow-core-platform</artifactId>
-        <version>0.3.0-SNAPSHOT</version>
+        <version>0.4.0-SNAPSHOT</version>
     </dependency>
 </dependencies>
 ```
@@ -151,6 +151,7 @@ This table shows the mapping between different version of TensorFlow for Java an
 | 0.2.0  | 2.3.1  |
 | 0.3.0  | 2.4.1  |
 | 0.3.1  | 2.4.1  |
+| 0.3.2  | 2.4.1  |
 
 ## How to Contribute?
 

diff --git a/ndarray/pom.xml b/ndarray/pom.xml
@@ -22,7 +22,7 @@
   <parent>
     <groupId>org.tensorflow</groupId>
     <artifactId>tensorflow-java</artifactId>
-    <version>0.3.1</version>
+    <version>0.3.2</version>
   </parent>
   <artifactId>ndarray</artifactId>
   <packaging>jar</packaging>

diff --git a/pom.xml b/pom.xml
@@ -5,7 +5,7 @@
 
   <groupId>org.tensorflow</groupId>
   <artifactId>tensorflow-java</artifactId>
-  <version>0.3.1</version>
+  <version>0.3.2</version>
   <packaging>pom</packaging>
 
   <name>TensorFlow Java Parent</name>

diff --git a/tensorflow-core/pom.xml b/tensorflow-core/pom.xml
@@ -22,7 +22,7 @@
   <parent>
     <groupId>org.tensorflow</groupId>
     <artifactId>tensorflow-java</artifactId>
-    <version>0.3.1</version>
+    <version>0.3.2</version>
   </parent>
   <artifactId>tensorflow-core</artifactId>
   <packaging>pom</packaging>

diff --git a/tensorflow-core/tensorflow-core-api/pom.xml b/tensorflow-core/tensorflow-core-api/pom.xml
@@ -6,7 +6,7 @@
   <parent>
     <groupId>org.tensorflow</groupId>
     <artifactId>tensorflow-core</artifactId>
-    <version>0.3.1</version>
+    <version>0.3.2</version>
   </parent>
   <artifactId>tensorflow-core-api</artifactId>
   <packaging>jar</packaging>

diff --git a/tensorflow-core/tensorflow-core-api/src/main/java/org/tensorflow/SavedModelBundle.java b/tensorflow-core/tensorflow-core-api/src/main/java/org/tensorflow/SavedModelBundle.java
@@ -11,11 +11,10 @@
 WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 See the License for the specific language governing permissions and
 limitations under the License.
-==============================================================================*/
-
+=======================================================================
+*/
 package org.tensorflow;
 
-import static org.tensorflow.internal.c_api.global.tensorflow.TF_LoadSessionFromSavedModel;
 import static org.tensorflow.internal.c_api.global.tensorflow.TF_NewGraph;
 import static org.tensorflow.internal.c_api.global.tensorflow.TF_SetConfig;
 
@@ -32,7 +31,6 @@
 import java.util.Map;
 import java.util.stream.Collectors;
 import org.bytedeco.javacpp.BytePointer;
-import org.bytedeco.javacpp.PointerPointer;
 import org.bytedeco.javacpp.PointerScope;
 import org.tensorflow.exceptions.TensorFlowException;
 import org.tensorflow.internal.c_api.TF_Buffer;
@@ -412,14 +410,19 @@ private static SavedModelBundle load(
       // load the session
       TF_Graph graph = TF_NewGraph();
       TF_Buffer metagraphDef = TF_Buffer.newBuffer();
-      TF_Session session = TF_LoadSessionFromSavedModel(
-          opts, runOpts, new BytePointer(exportDir), new PointerPointer(tags),
-          tags.length, graph, metagraphDef, status);
+      TF_Session session =
+          TF_Session.loadSessionFromSavedModel(
+              opts, runOpts, exportDir, tags, graph, metagraphDef, status);
       status.throwExceptionIfNotOK();
 
       // handle the result
       try {
-        bundle = fromHandle(graph, session, MetaGraphDef.parseFrom(metagraphDef.dataAsByteBuffer()));
+        bundle =
+            fromHandle(graph, session, MetaGraphDef.parseFrom(metagraphDef.dataAsByteBuffer()));
+        // Only retain the references if the metagraphdef parses correctly,
+        // otherwise allow the pointer scope to clean them up
+        graph.retainReference();
+        session.retainReference();
       } catch (InvalidProtocolBufferException e) {
         throw new TensorFlowException("Cannot parse MetaGraphDef protocol buffer", e);
       }

diff --git a/...ore/tensorflow-core-api/src/main/java/org/tensorflow/internal/c_api/AbstractTF_Graph.java b/...ore/tensorflow-core-api/src/main/java/org/tensorflow/internal/c_api/AbstractTF_Graph.java
@@ -1,19 +1,19 @@
 /*
- Copyright 2019 The TensorFlow Authors. All Rights Reserved.
+Copyright 2019 The TensorFlow Authors. All Rights Reserved.
 
- Licensed under the Apache License, Version 2.0 (the "License");
- you may not use this file except in compliance with the License.
- You may obtain a copy of the License at
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
 
-     http://www.apache.org/licenses/LICENSE-2.0
+    http://www.apache.org/licenses/LICENSE-2.0
 
- Unless required by applicable law or agreed to in writing, software
- distributed under the License is distributed on an "AS IS" BASIS,
- WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
- See the License for the specific language governing permissions and
- limitations under the License.
- =======================================================================
- */
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+=======================================================================
+*/
 
 package org.tensorflow.internal.c_api;
 
@@ -25,29 +25,40 @@
 
 @Properties(inherit = org.tensorflow.internal.c_api.presets.tensorflow.class)
 public abstract class AbstractTF_Graph extends Pointer {
-    protected static class DeleteDeallocator extends TF_Graph implements Pointer.Deallocator {
-        DeleteDeallocator(TF_Graph s) { super(s); }
-        @Override public void deallocate() { if (!isNull()) TF_DeleteGraph(this); setNull(); }
+  protected static class DeleteDeallocator extends TF_Graph implements Pointer.Deallocator {
+    DeleteDeallocator(TF_Graph s) {
+      super(s);
     }
 
-    public AbstractTF_Graph(Pointer p) { super(p); }
-
-    /**
-     * Calls TF_NewGraph(), and registers a deallocator.
-     * @return TF_Graph created. Do not call TF_DeleteGraph() on it.
-     */
-    public static TF_Graph newGraph() {
-        TF_Graph g = TF_NewGraph();
-        if (g != null) {
-            g.deallocator(new DeleteDeallocator(g));
-        }
-        return g;
+    @Override
+    public void deallocate() {
+      if (!isNull()) TF_DeleteGraph(this);
+      setNull();
     }
+  }
 
-    /**
-     * Calls the deallocator, if registered, otherwise has no effect.
-     */
-    public void delete() {
-        deallocate();
+  public AbstractTF_Graph(Pointer p) {
+    super(p);
+  }
+
+  /**
+   * Calls TF_NewGraph(), and registers a deallocator.
+   *
+   * <p>Note {@link org.tensorflow.Graph} will call TF_DeleteGraph on close, so do not use this
+   * method when constructing a reference for use inside a {@code Graph} object.
+   *
+   * @return TF_Graph created. Do not call TF_DeleteGraph() on it.
+   */
+  public static TF_Graph newGraph() {
+    TF_Graph g = TF_NewGraph();
+    if (g != null) {
+      g.deallocator(new DeleteDeallocator(g));
     }
+    return g;
+  }
+
+  /** Calls the deallocator, if registered, otherwise has no effect. */
+  public void delete() {
+    deallocate();
+  }
 }
diff --git a/tensorflow-core/tensorflow-core-generator/pom.xml b/tensorflow-core/tensorflow-core-generator/pom.xml
@@ -5,7 +5,7 @@
   <parent>
     <groupId>org.tensorflow</groupId>
     <artifactId>tensorflow-core</artifactId>
-    <version>0.3.1</version>
+    <version>0.3.2</version>
   </parent>
   <artifactId>tensorflow-core-generator</artifactId>
   <packaging>jar</packaging>

diff --git a/tensorflow-core/tensorflow-core-platform-gpu/pom.xml b/tensorflow-core/tensorflow-core-platform-gpu/pom.xml
@@ -22,7 +22,7 @@
   <parent>
     <groupId>org.tensorflow</groupId>
     <artifactId>tensorflow-core</artifactId>
-    <version>0.3.1</version>
+    <version>0.3.2</version>
   </parent>
   <artifactId>tensorflow-core-platform-gpu</artifactId>
   <name>TensorFlow Core API Library Platform GPU</name>

diff --git a/tensorflow-core/tensorflow-core-platform-mkl-gpu/pom.xml b/tensorflow-core/tensorflow-core-platform-mkl-gpu/pom.xml
@@ -22,7 +22,7 @@
   <parent>
     <groupId>org.tensorflow</groupId>
     <artifactId>tensorflow-core</artifactId>
-    <version>0.3.1</version>
+    <version>0.3.2</version>
   </parent>
   <artifactId>tensorflow-core-platform-mkl-gpu</artifactId>
   <name>TensorFlow Core API Library Platform MKL GPU</name>

diff --git a/tensorflow-core/tensorflow-core-platform-mkl/pom.xml b/tensorflow-core/tensorflow-core-platform-mkl/pom.xml
@@ -22,7 +22,7 @@
   <parent>
     <groupId>org.tensorflow</groupId>
     <artifactId>tensorflow-core</artifactId>
-    <version>0.3.1</version>
+    <version>0.3.2</version>
   </parent>
   <artifactId>tensorflow-core-platform-mkl</artifactId>
   <name>TensorFlow Core API Library Platform MKL</name>

diff --git a/tensorflow-core/tensorflow-core-platform/pom.xml b/tensorflow-core/tensorflow-core-platform/pom.xml
@@ -22,7 +22,7 @@
   <parent>
     <groupId>org.tensorflow</groupId>
     <artifactId>tensorflow-core</artifactId>
-    <version>0.3.1</version>
+    <version>0.3.2</version>
   </parent>
   <artifactId>tensorflow-core-platform</artifactId>
   <name>TensorFlow Core API Library Platform</name>

diff --git a/tensorflow-framework/pom.xml b/tensorflow-framework/pom.xml
@@ -22,7 +22,7 @@
   <parent>
     <groupId>org.tensorflow</groupId>
     <artifactId>tensorflow-java</artifactId>
-    <version>0.3.1</version>
+    <version>0.3.2</version>
   </parent>
   <artifactId>tensorflow-framework</artifactId>
   <packaging>jar</packaging>