Add validations on signatures and saved models

Author: karllessard

tensorflow-core/tensorflow-core-api/src/main/java/org/tensorflow/ConcreteFunction.java

@@ -231,11 +231,11 @@ public Tensor<?> call(Tensor<?> tensor) throws IllegalArgumentException {
    *
    * <p>This method is convenient shortcut equivalent to
    * {@code SavedModel.exporter(exportDir).withFunction(this).export()}
+   *
+   * @throws IOException if saved model or variable state cannot be written on disk
    */
   public void save(String exportDir) throws IOException {
-    SavedModelBundle.exporter(exportDir)
-        .withFunction(this)
-        .export();
+    SavedModelBundle.exporter(exportDir).withFunction(this).export();
   }
 
   /**

tensorflow-core/tensorflow-core-api/src/main/java/org/tensorflow/SavedModelBundle.java

@@ -25,9 +25,9 @@
 import java.io.OutputStream;
 import java.nio.file.Path;
 import java.nio.file.Paths;
-import java.util.ArrayList;
 import java.util.Arrays;
 import java.util.HashMap;
+import java.util.LinkedHashMap;
 import java.util.List;
 import java.util.Map;
 import java.util.stream.Collectors;
@@ -45,6 +45,7 @@
 import org.tensorflow.proto.framework.MetaGraphDef.MetaInfoDef;
 import org.tensorflow.proto.framework.RunOptions;
 import org.tensorflow.proto.framework.SavedModel;
+import org.tensorflow.proto.util.SaverDef;
 
 /**
  * SavedModelBundle represents a model loaded from storage.
@@ -73,6 +74,7 @@ public SavedModelBundle load() {
      * @param options A <a
      *     href="https://www.tensorflow.org/code/tensorflow/core/protobuf/config.proto">RunOptions
      *     protocol buffer</a>.
+     * @return this object
      */
     public Loader withRunOptions(RunOptions options) {
       this.runOptions = options;
@@ -85,6 +87,7 @@ public Loader withRunOptions(RunOptions options) {
      * @param configProto A <a
      *     href="https://www.tensorflow.org/code/tensorflow/core/protobuf/config.proto">ConfigProto
      *     protocol buffer</a>.
+     * @return this object
      */
     public Loader withConfigProto(ConfigProto configProto) {
       this.configProto = configProto;
@@ -97,11 +100,12 @@ public Loader withConfigProto(ConfigProto configProto) {
      * <p>Has no effect if {@code tags} is null or empty
      *
      * @param tags the tags identifying the specific MetaGraphDef to load.
+     * @return this object
+     * @throws IllegalArgumentException if tags are invalid
      */
     public Loader withTags(String... tags) {
-      if (tags != null && tags.length > 0) {
-        this.tags = tags;
-      }
+      validateTags(tags);
+      this.tags = tags;
       return this;
     }
 
@@ -110,7 +114,7 @@ private Loader(String exportDir) {
     }
 
     private String exportDir = null;
-    private String[] tags = {DEFAULT_TAG};
+    private String[] tags = { DEFAULT_TAG };
     private ConfigProto configProto = null;
     private RunOptions runOptions = null;
   }
@@ -125,9 +129,11 @@ public static final class Exporter {
      *
      * @param tags the tags identifying the specific MetaGraphDef to save.
      * @return this object
+     * @throws IllegalArgumentException if tags are invalid
      */
     public Exporter withTags(String... tags) {
-      this.tags.addAll(Arrays.asList(tags));
+      validateTags(tags);
+      this.tags = tags;
       return this;
     }
 
@@ -148,6 +154,8 @@ public Exporter withTags(String... tags) {
      * @param function a function carrying a signature and a valid session to the graph to be saved
      * @return this object
      * @throws IllegalArgumentException if a function with the same name has already been added to the model
+     * @throws UnsupportedOperationException if this function does not share the same session with the other
+     *                                       functions added to this model
      */
     public Exporter withFunction(ConcreteFunction function) {
       Signature signature = function.signature();
@@ -167,22 +175,22 @@ public Exporter withFunction(ConcreteFunction function) {
     /**
      * Save the model into the export directory.
      *
-     * @throws IOException if saved model or variable state can be written on disk
+     * @throws IOException if saved model or variable state cannot be written on disk
      */
     public void export() throws IOException {
       if (functions.isEmpty() || session == null) {
         throw new IllegalStateException("Model should contain at least one valid function");
       }
-      if (tags.isEmpty()) {
-        tags.add(DEFAULT_TAG);
-      }
+      Graph graph = session.graph();
+
       // It is imperative to retrieve the graphDef after the saverDef, as the former might add
       // new ops to the graph for saving and restoring the variables.
-      Graph graph = session.graph();
+      SaverDef saverDef = graph.saverDef();
+
       MetaGraphDef.Builder metaGraphDef = metaGraphDefBuilder
-          .setSaverDef(graph.saverDef())
+          .setSaverDef(saverDef)
           .setGraphDef(graph.toGraphDef())
-          .setMetaInfoDef(MetaInfoDef.newBuilder().addAllTags(tags));
+          .setMetaInfoDef(MetaInfoDef.newBuilder().addAllTags(Arrays.asList(tags)));
       functions.forEach((k, f) -> metaGraphDef.putSignatureDef(k, f.signature().asSignatureDef()));
 
       // Make sure saved model directories exist
@@ -205,9 +213,9 @@ public void export() throws IOException {
     }
 
     private final String exportDir;
-    private final List<String> tags = new ArrayList<>();
+    private String[] tags = { DEFAULT_TAG };
     private final MetaGraphDef.Builder metaGraphDefBuilder = MetaGraphDef.newBuilder();
-    private final Map<String, ConcreteFunction> functions = new HashMap<>();
+    private final Map<String, ConcreteFunction> functions = new LinkedHashMap<>();
     private Session session;
   }
 
@@ -227,7 +235,11 @@ public void export() throws IOException {
    * @return a bundle containing the graph and associated session.
    */
   public static SavedModelBundle load(String exportDir, String... tags) {
-    return loader(exportDir).withTags(tags).load();
+    Loader loader = loader(exportDir);
+    if (tags != null && tags.length > 0) {
+      loader.withTags(tags);
+    }
+    return loader.load();
   }
 
   /**
@@ -416,6 +428,12 @@ opts, runOpts, new BytePointer(exportDir), new PointerPointer(tags),
     return bundle;
   }
 
+  private static void validateTags(String[] tags) {
+    if (tags == null || tags.length == 0 || Arrays.stream(tags).anyMatch(t -> t == null || t.isEmpty())) {
+      throw new IllegalArgumentException("Invalid tags: " + Arrays.toString(tags));
+    }
+  }
+
   static {
     TensorFlow.init();
   }

tensorflow-core/tensorflow-core-api/src/main/java/org/tensorflow/Session.java

@@ -454,7 +454,10 @@ public void run(Op op) {
    * <i>mymodel/myvariables/variables</i>, then the generated files will be located under
    * <i>mymodel/myvariables</i> and named <i>variables.data-*-of-*</i>
    *
-   * @param prefix
+   * <p/>Note that this method might alter the underlying graph if it is the first time that one
+   * of its session is saved, see {@link Graph#saverDef()} for more details.
+   *
+   * @param prefix prefix to the variable files to save
    */
   public void save(String prefix) {
     SaverDef saverDef = graph.saverDef();

tensorflow-core/tensorflow-core-api/src/main/java/org/tensorflow/Signature.java

@@ -45,8 +45,12 @@ public static class Builder {
      *
      * @param key signature key
      * @return this builder
+     * @throws IllegalArgumentException if the key is invalid
      */
     public Builder key(String key) {
+      if (key == null || key.isEmpty()) {
+        throw new IllegalArgumentException("Invalid key: " + key);
+      }
       this.key = key;
       return this;
     }
@@ -57,8 +61,12 @@ public Builder key(String key) {
      * @param inputName user-friendly name for this input tensor
      * @param input input tensor
      * @return this builder
+     * @throws IllegalArgumentException if {@code inputName} is already mapped to another input
      */
     public Builder input(String inputName, Operand<?> input) {
+      if (signatureBuilder.containsInputs(inputName)) {
+        throw new IllegalArgumentException("\"" + inputName + "\" is already being mapped to another input");
+      }
       signatureBuilder.putInputs(inputName, toTensorInfo(input.asOutput()));
       return this;
     }
@@ -69,8 +77,12 @@ public Builder input(String inputName, Operand<?> input) {
      * @param inputName user-friendly name for this input tensor
      * @param input input tensor
      * @return this builder
+     * @throws IllegalArgumentException if {@code outputName} is already mapped to another output
      */
     public Builder output(String outputName, Operand<?> output) {
+      if (signatureBuilder.containsOutputs(outputName)) {
+        throw new IllegalArgumentException("\"" + outputName + "\" is already being mapped to another output");
+      }
       signatureBuilder.putOutputs(outputName, toTensorInfo(output.asOutput()));
       return this;
     }
@@ -79,11 +91,11 @@ public Builder output(String outputName, Operand<?> output) {
      * Provide extensible name information enabling third-party users to mark a signature as
      * supporting a particular method
      *
-     * @param methodName method name
+     * @param methodName method name or null for none (default)
      * @return this builder
      */
     public Builder methodName(String methodName) {
-      signatureBuilder.setMethodName(methodName);
+      signatureBuilder.setMethodName(methodName == null ? "" : methodName);
       return this;
     }
 
@@ -126,10 +138,10 @@ public String key() {
   }
 
   /**
-   * Returns the method name of this signature (e.g. as exposed by TF serving)
+   * Returns the method name of this signature (e.g. as exposed by TF serving) or null if none
    */
   public String methodName() {
-    return signatureDef.getMethodName();
+    return signatureDef.getMethodName().isEmpty() ? null : signatureDef.getMethodName();
   }
 
   /**

tensorflow-core/tensorflow-core-api/src/test/java/org/tensorflow/ConcreteFunctionTest.java

@@ -1,6 +1,21 @@
+/* Copyright 2020 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
+
+    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.
+==============================================================================*/
 package org.tensorflow;
 
 import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.assertThrows;
 import static org.junit.jupiter.api.Assertions.fail;
 
 import org.junit.jupiter.api.Test;
@@ -75,18 +90,8 @@ public void closingFunctionReleaseAllResourcesItOwns() {
       g = f.graph();
       s = f.session();
     }
-    try {
-      s.run("Add");
-      fail();
-    } catch (IllegalStateException e) {
-      // as expected
-    }
-    try {
-      g.toGraphDef();
-      fail();
-    } catch (IllegalStateException e) {
-      // as expected
-    }
+    assertThrows(IllegalStateException.class, () -> s.run("Add"));
+    assertThrows(IllegalStateException.class, () -> g.toGraphDef());
   }
 
   @Test
@@ -97,12 +102,7 @@ public void closingFunctionCreatedFromGraphOnlyReleaseResourcesItOwns() {
       try (ConcreteFunction f = ConcreteFunction.create(signature, g)) {
         s = f.session();
       }
-      try {
-        s.run(Init.DEFAULT_NAME);
-        fail();
-      } catch (IllegalStateException e) {
-        // as expected
-      }
+      assertThrows(IllegalStateException.class, () -> s.run(Init.DEFAULT_NAME));
       g.toGraphDef();  // check that graph is still valid
     }
   }

tensorflow-core/tensorflow-core-api/src/test/java/org/tensorflow/DeviceSpecTest.java

@@ -1,3 +1,17 @@
+/* Copyright 2020 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
+
+    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.
+==============================================================================*/
 package org.tensorflow;
 
 import org.junit.jupiter.api.Test;

tensorflow-core/tensorflow-core-api/src/test/java/org/tensorflow/SavedModelBundleTest.java

@@ -17,6 +17,7 @@
 
 import static org.junit.jupiter.api.Assertions.assertEquals;
 import static org.junit.jupiter.api.Assertions.assertNotNull;
+import static org.junit.jupiter.api.Assertions.assertThrows;
 import static org.junit.jupiter.api.Assertions.assertTrue;
 import static org.junit.jupiter.api.Assertions.fail;
 
@@ -258,6 +259,34 @@ public void cannotExportMultipleFunctionsWithSameSignatureKey() throws IOExcepti
     }
   }
 
+  @Test
+  public void cannotExportOrImportInvalidTags() {
+    assertThrows(IllegalArgumentException.class, () ->
+        SavedModelBundle.loader("/").withTags()
+    );
+    assertThrows(IllegalArgumentException.class, () ->
+        SavedModelBundle.loader("/").withTags(new String[]{})
+    );
+    assertThrows(IllegalArgumentException.class, () ->
+        SavedModelBundle.loader("/").withTags(new String[]{"tag", null})
+    );
+    assertThrows(IllegalArgumentException.class, () ->
+        SavedModelBundle.loader("/").withTags(new String[]{"tag", ""})
+    );
+    assertThrows(IllegalArgumentException.class, () ->
+        SavedModelBundle.exporter("/").withTags()
+    );
+    assertThrows(IllegalArgumentException.class, () ->
+        SavedModelBundle.exporter("/").withTags(new String[]{})
+    );
+    assertThrows(IllegalArgumentException.class, () ->
+        SavedModelBundle.exporter("/").withTags(new String[]{"tag", null})
+    );
+    assertThrows(IllegalArgumentException.class, () ->
+        SavedModelBundle.exporter("/").withTags(new String[]{"tag", ""})
+    );
+  }
+
   @Test
   public void pythonTfFunction() {
     // ConcreteFunctions on models saved using python