Only store large constants under global models

5 files changed, 103 insertions, 61 deletions

diff --git a/config-model/src/main/java/com/yahoo/searchdefinition/expressiontransforms/ConvertedModel.java b/config-model/src/main/java/com/yahoo/searchdefinition/expressiontransforms/ConvertedModel.java
index 852972803f3..050a5226324 100644
--- a/config-model/src/main/java/com/yahoo/searchdefinition/expressiontransforms/ConvertedModel.java
+++ b/config-model/src/main/java/com/yahoo/searchdefinition/expressiontransforms/ConvertedModel.java
@@ -240,9 +240,12 @@ public class ConvertedModel {
         profile.addConstant(constantName, asValue(constantValue));
     }
 
-    private static void transformLargeConstant(ModelStore store, RankProfile profile, QueryProfileRegistry queryProfiles,
-                                        Set<String> constantsReplacedByMacros,
-                                        String constantName, Tensor constantValue) {
+    private static void transformLargeConstant(ModelStore store,
+                                               RankProfile profile,
+                                               QueryProfileRegistry queryProfiles,
+                                               Set<String> constantsReplacedByMacros,
+                                               String constantName,
+                                               Tensor constantValue) {
         RankProfile.Macro macroOverridingConstant = profile.getMacros().get(constantName);
         if (macroOverridingConstant != null) {
             TensorType macroType = macroOverridingConstant.getRankingExpression().type(profile.typeContext(queryProfiles));
@@ -255,7 +258,7 @@ public class ConvertedModel {
             Path constantPath = store.writeLargeConstant(constantName, constantValue);
             if ( ! profile.rankingConstants().asMap().containsKey(constantName)) {
                 profile.rankingConstants().add(new RankingConstant(constantName, constantValue.type(),
-                                                           constantPath.toString()));
+                                                                   constantPath.toString()));
             }
         }
     }
@@ -612,8 +615,12 @@ public class ConvertedModel {
                        .writeFile(new StringReader(name + ":" + constant.type() + ":" + correct(constantPath)));
 
             // Write content explicitly as a file on the file system as this is distributed using file distribution
-            createIfNeeded(constantsPath);
-            IOUtils.writeFile(application.getFileReference(constantPath), TypedBinaryFormat.encode(constant));
+            // - but only if this is a global model to avoid writing the same constants for each rank profile
+            //   where they are used
+            if (modelFiles.modelName.isGlobal()) {
+                createIfNeeded(constantsPath);
+                IOUtils.writeFile(application.getFileReference(constantPath), TypedBinaryFormat.encode(constant));
+            }
             return correct(constantPath);
         }
 
@@ -683,9 +690,13 @@ public class ConvertedModel {
             return ApplicationPackage.MODELS_GENERATED_REPLICATED_DIR.append(modelName.fullName());
         }
 
-        /** Files stored below this path will not be replicated in zookeeper */
-        public Path storedModelPath() {
-            return ApplicationPackage.MODELS_GENERATED_DIR.append(modelName.fullName());
+        /**
+         * Files stored below this path will not be replicated in zookeeper.
+         * Large constants are only stored under the global (not rank-profile-specific)
+         * path to avoid storing the same large constant multiple times.
+         */
+        public Path storedGlobalModelPath() {
+            return ApplicationPackage.MODELS_GENERATED_DIR.append(modelName.localName());
         }
 
         public Path expressionPath(String name) {
@@ -702,7 +713,7 @@ public class ConvertedModel {
 
         /** Path to the large (ranking) constants directory */
         public Path largeConstantsContentPath() {
-            return storedModelPath().append("constants");
+            return storedGlobalModelPath().append("constants");
         }
 
         /** Path to the large (ranking) constants directory */
@@ -790,10 +801,15 @@ public class ConvertedModel {
             this.fullName = (namespace != null ? namespace + "." : "") + name;
         }
 
+        /** Returns true if the local name of this is not in a namespace */
+        public boolean isGlobal() { return namespace == null; }
+
+        /** Returns the namespace, or null if this is global */
         public String namespace() { return namespace; }
         public String localName() { return name; }
         public String fullName() { return fullName; }
 
+
         @Override
         public boolean equals(Object o) {
             if (o == this) return true;
diff --git a/config-model/src/test/integration/onnx/services.xml b/config-model/src/test/integration/onnx/services.xml
new file mode 100644
index 00000000000..f623b2464fc
--- /dev/null
+++ b/config-model/src/test/integration/onnx/services.xml
@@ -0,0 +1,5 @@
+<services>
+    <container version="1.0">
+
+    </container>
+</services>
\ No newline at end of file
diff --git a/config-model/src/test/integration/tensorflow/services.xml b/config-model/src/test/integration/tensorflow/services.xml
new file mode 100644
index 00000000000..f623b2464fc
--- /dev/null
+++ b/config-model/src/test/integration/tensorflow/services.xml
@@ -0,0 +1,5 @@
+<services>
+    <container version="1.0">
+
+    </container>
+</services>
\ No newline at end of file
diff --git a/config-model/src/test/java/com/yahoo/searchdefinition/processing/RankingExpressionWithOnnxTestCase.java b/config-model/src/test/java/com/yahoo/searchdefinition/processing/RankingExpressionWithOnnxTestCase.java
index 414a77e9164..3ca4cdaf547 100644
--- a/config-model/src/test/java/com/yahoo/searchdefinition/processing/RankingExpressionWithOnnxTestCase.java
+++ b/config-model/src/test/java/com/yahoo/searchdefinition/processing/RankingExpressionWithOnnxTestCase.java
@@ -3,14 +3,17 @@
 package com.yahoo.searchdefinition.processing;
 
 import com.yahoo.config.application.api.ApplicationPackage;
+import com.yahoo.config.model.ApplicationPackageTester;
 import com.yahoo.io.GrowableByteBuffer;
 import com.yahoo.io.IOUtils;
 import com.yahoo.path.Path;
 import com.yahoo.search.query.profile.QueryProfileRegistry;
 import com.yahoo.searchdefinition.RankingConstant;
 import com.yahoo.searchdefinition.parser.ParseException;
+import com.yahoo.searchlib.rankingexpression.integration.ml.ImportedModel;
 import com.yahoo.tensor.Tensor;
 import com.yahoo.tensor.serialization.TypedBinaryFormat;
+import com.yahoo.vespa.model.VespaModel;
 import com.yahoo.yolean.Exceptions;
 import org.junit.After;
 import org.junit.Test;
@@ -20,6 +23,7 @@ import java.io.UncheckedIOException;
 import java.util.Optional;
 
 import com.yahoo.searchdefinition.processing.RankingExpressionWithTensorFlowTestCase.StoringApplicationPackage;
+import org.xml.sax.SAXException;
 
 import static junit.framework.TestCase.assertTrue;
 import static org.junit.Assert.assertEquals;
@@ -41,14 +45,36 @@ public class RankingExpressionWithOnnxTestCase {
     }
 
     @Test
+    public void testGlobalOnnxModel() throws SAXException, IOException {
+        ApplicationPackageTester tester = ApplicationPackageTester.create(applicationDir.toString());
+        VespaModel model = new VespaModel(tester.app());
+        assertLargeConstant(name + "_Variable_1", model, Optional.of(10L));
+        assertLargeConstant(name + "_Variable", model, Optional.of(7840L));
+
+        // At this point the expression is stored - copy application to another location which do not have a models dir
+        Path storedAppDir = applicationDir.append("copy");
+        try {
+            storedAppDir.toFile().mkdirs();
+            IOUtils.copy(applicationDir.append("services.xml").toString(), storedAppDir.append("services.xml").toString());
+            IOUtils.copyDirectory(applicationDir.append(ApplicationPackage.MODELS_GENERATED_DIR).toFile(),
+                                  storedAppDir.append(ApplicationPackage.MODELS_GENERATED_DIR).toFile());
+            ApplicationPackageTester storedTester = ApplicationPackageTester.create(storedAppDir.toString());
+            VespaModel storedModel = new VespaModel(storedTester.app());
+            assertLargeConstant(name + "_Variable_1", storedModel, Optional.of(10L));
+            assertLargeConstant(name + "_Variable", storedModel, Optional.of(7840L));
+        }
+        finally {
+            IOUtils.recursiveDeleteDir(storedAppDir.toFile());
+        }
+    }
+
+    @Test
     public void testOnnxReferenceWithConstantFeature() {
         RankProfileSearchFixture search = fixtureWith("constant(mytensor)",
                 "onnx('mnist_softmax.onnx')",
                 "constant mytensor { file: ignored\ntype: tensor(d0[7],d1[784]) }",
                 null);
         search.assertFirstPhaseExpression(vespaExpression, "my_profile");
-        assertLargeConstant(name + "_Variable_1", search, Optional.of(10L));
-        assertLargeConstant(name + "_Variable", search, Optional.of(7840L));
     }
 
     @Test
@@ -68,8 +94,6 @@ public class RankingExpressionWithOnnxTestCase {
                 "Placeholder",
                 application);
         search.assertFirstPhaseExpression(vespaExpression, "my_profile");
-        assertLargeConstant(name + "_Variable_1", search, Optional.of(10L));
-        assertLargeConstant(name + "_Variable", search, Optional.of(7840L));
     }
 
     @Test
@@ -82,8 +106,6 @@ public class RankingExpressionWithOnnxTestCase {
                 "Placeholder",
                 application);
         search.assertFirstPhaseExpression(vespaExpression, "my_profile");
-        assertLargeConstant( name + "_Variable_1", search, Optional.of(10L));
-        assertLargeConstant( name + "_Variable", search, Optional.of(7840L));
     }
 
 
@@ -104,8 +126,6 @@ public class RankingExpressionWithOnnxTestCase {
                 "Placeholder",
                 application);
         search.assertFirstPhaseExpression(vespaExpression, "my_profile");
-        assertLargeConstant( name + "_Variable_1", search, Optional.of(10L));
-        assertLargeConstant( name + "_Variable", search, Optional.of(7840L));
     }
 
 
@@ -114,8 +134,6 @@ public class RankingExpressionWithOnnxTestCase {
         RankProfileSearchFixture search = fixtureWith("tensor(d0[2],d1[784])(0.0)",
                 "5 + sum(onnx('mnist_softmax.onnx'))");
         search.assertFirstPhaseExpression("5 + reduce(" + vespaExpression + ", sum)", "my_profile");
-        assertLargeConstant( name + "_Variable_1", search, Optional.of(10L));
-        assertLargeConstant( name + "_Variable", search, Optional.of(7840L));
     }
 
     @Test
@@ -181,9 +199,6 @@ public class RankingExpressionWithOnnxTestCase {
                 "onnx('mnist_softmax.onnx')");
         search.assertFirstPhaseExpression(vespaExpression, "my_profile");
 
-        assertLargeConstant( name + "_Variable_1", search, Optional.of(10L));
-        assertLargeConstant( name + "_Variable", search, Optional.of(7840L));
-
         // At this point the expression is stored - copy application to another location which do not have a models dir
         Path storedApplicationDirectory = applicationDir.getParentPath().append("copy");
         try {
@@ -200,8 +215,6 @@ public class RankingExpressionWithOnnxTestCase {
             searchFromStored.assertFirstPhaseExpression(vespaExpression, "my_profile");
             // Verify that the constants exists, but don't verify the content as we are not
             // simulating file distribution in this test
-            assertLargeConstant( name + "_Variable_1", searchFromStored, Optional.empty());
-            assertLargeConstant( name + "_Variable", searchFromStored, Optional.empty());
         }
         finally {
             IOUtils.recursiveDeleteDir(storedApplicationDirectory.toFile());
@@ -232,7 +245,6 @@ public class RankingExpressionWithOnnxTestCase {
 
         assertNull("Constant overridden by macro is not added",
                 search.search().rankingConstants().get( name + "_Variable"));
-        assertLargeConstant( name + "_Variable_1", search, Optional.of(10L));
 
         // At this point the expression is stored - copy application to another location which do not have a models dir
         Path storedApplicationDirectory = applicationDir.getParentPath().append("copy");
@@ -245,8 +257,7 @@ public class RankingExpressionWithOnnxTestCase {
             searchFromStored.compileRankProfile("my_profile", applicationDir.append("models"));
             searchFromStored.assertFirstPhaseExpression(vespaExpressionWithoutConstant, "my_profile");
             assertNull("Constant overridden by macro is not added",
-                    searchFromStored.search().rankingConstants().get( name + "_Variable"));
-            assertLargeConstant( name + "_Variable_1", searchFromStored, Optional.of(10L));
+                       searchFromStored.search().rankingConstants().get( name + "_Variable"));
         } finally {
             IOUtils.recursiveDeleteDir(storedApplicationDirectory.toFile());
         }
@@ -256,19 +267,19 @@ public class RankingExpressionWithOnnxTestCase {
      * Verifies that the constant with the given name exists, and - only if an expected size is given -
      * that the content of the constant is available and has the expected size.
      */
-    private void assertLargeConstant(String name, RankProfileSearchFixture search, Optional<Long> expectedSize) {
+    private void assertLargeConstant(String constantName, VespaModel model, Optional<Long> expectedSize) {
         try {
-            Path constantApplicationPackagePath = Path.fromString("models.generated/my_profile.mnist_softmax.onnx/constants").append(name + ".tbf");
-            RankingConstant rankingConstant = search.search().rankingConstants().get(name);
-            assertEquals(name, rankingConstant.getName());
+            Path constantApplicationPackagePath = Path.fromString("models.generated/" + name + "/constants").append(constantName + ".tbf");
+            RankingConstant rankingConstant = model.rankingConstants().get(constantName);
+            assertEquals(constantName, rankingConstant.getName());
             assertTrue(rankingConstant.getFileName().endsWith(constantApplicationPackagePath.toString()));
 
             if (expectedSize.isPresent()) {
                 Path constantPath = applicationDir.append(constantApplicationPackagePath);
                 assertTrue("Constant file '" + constantPath + "' has been written",
-                        constantPath.toFile().exists());
+                           constantPath.toFile().exists());
                 Tensor deserializedConstant = TypedBinaryFormat.decode(Optional.empty(),
-                        GrowableByteBuffer.wrap(IOUtils.readFileBytes(constantPath.toFile())));
+                                                                       GrowableByteBuffer.wrap(IOUtils.readFileBytes(constantPath.toFile())));
                 assertEquals(expectedSize.get().longValue(), deserializedConstant.size());
             }
         }
diff --git a/config-model/src/test/java/com/yahoo/searchdefinition/processing/RankingExpressionWithTensorFlowTestCase.java b/config-model/src/test/java/com/yahoo/searchdefinition/processing/RankingExpressionWithTensorFlowTestCase.java
index 450c66e04ef..71b45ec628a 100644
--- a/config-model/src/test/java/com/yahoo/searchdefinition/processing/RankingExpressionWithTensorFlowTestCase.java
+++ b/config-model/src/test/java/com/yahoo/searchdefinition/processing/RankingExpressionWithTensorFlowTestCase.java
@@ -3,6 +3,7 @@ package com.yahoo.searchdefinition.processing;
 
 import com.yahoo.config.application.api.ApplicationFile;
 import com.yahoo.config.application.api.ApplicationPackage;
+import com.yahoo.config.model.ApplicationPackageTester;
 import com.yahoo.config.model.test.MockApplicationPackage;
 import com.yahoo.io.GrowableByteBuffer;
 import com.yahoo.io.IOUtils;
@@ -15,9 +16,11 @@ import com.yahoo.searchlib.rankingexpression.evaluation.Value;
 import com.yahoo.tensor.Tensor;
 import com.yahoo.tensor.TensorType;
 import com.yahoo.tensor.serialization.TypedBinaryFormat;
+import com.yahoo.vespa.model.VespaModel;
 import com.yahoo.yolean.Exceptions;
 import org.junit.After;
 import org.junit.Test;
+import org.xml.sax.SAXException;
 
 import java.io.BufferedInputStream;
 import java.io.File;
@@ -36,6 +39,7 @@ import java.util.List;
 import java.util.Optional;
 import java.util.stream.Collectors;
 
+import static junit.framework.TestCase.assertTrue;
 import static org.junit.Assert.*;
 
 /**
@@ -56,12 +60,34 @@ public class RankingExpressionWithTensorFlowTestCase {
     }
 
     @Test
+    public void testGlobalTensorFlowModel() throws SAXException, IOException {
+        ApplicationPackageTester tester = ApplicationPackageTester.create(applicationDir.toString());
+        VespaModel model = new VespaModel(tester.app());
+        assertLargeConstant(name + "_layer_Variable_1_read", model, Optional.of(10L));
+        assertLargeConstant(name + "_layer_Variable_read", model, Optional.of(7840L));
+
+        // At this point the expression is stored - copy application to another location which do not have a models dir
+        Path storedAppDir = applicationDir.append("copy");
+        try {
+            storedAppDir.toFile().mkdirs();
+            IOUtils.copy(applicationDir.append("services.xml").toString(), storedAppDir.append("services.xml").toString());
+            IOUtils.copyDirectory(applicationDir.append(ApplicationPackage.MODELS_GENERATED_DIR).toFile(),
+                                  storedAppDir.append(ApplicationPackage.MODELS_GENERATED_DIR).toFile());
+            ApplicationPackageTester storedTester = ApplicationPackageTester.create(storedAppDir.toString());
+            VespaModel storedModel = new VespaModel(storedTester.app());
+            assertLargeConstant(name + "_layer_Variable_1_read", storedModel, Optional.of(10L));
+            assertLargeConstant(name + "_layer_Variable_read", storedModel, Optional.of(7840L));
+        }
+        finally {
+            IOUtils.recursiveDeleteDir(storedAppDir.toFile());
+        }
+    }
+
+    @Test
     public void testTensorFlowReference() {
         RankProfileSearchFixture search = fixtureWith("tensor(d0[2],d1[784])(0.0)",
                                                       "tensorflow('mnist_softmax/saved')");
         search.assertFirstPhaseExpression(vespaExpression, "my_profile");
-        assertLargeConstant(name + "_layer_Variable_1_read", search, Optional.of(10L));
-        assertLargeConstant(name + "_layer_Variable_read", search, Optional.of(7840L));
     }
 
     @Test
@@ -71,8 +97,6 @@ public class RankingExpressionWithTensorFlowTestCase {
                                                       "constant mytensor { file: ignored\ntype: tensor(d0[7],d1[784]) }",
                                                       null);
         search.assertFirstPhaseExpression(vespaExpression, "my_profile");
-        assertLargeConstant(name + "_layer_Variable_1_read", search, Optional.of(10L));
-        assertLargeConstant(name + "_layer_Variable_read", search, Optional.of(7840L));
     }
 
     @Test
@@ -91,8 +115,6 @@ public class RankingExpressionWithTensorFlowTestCase {
                                                       "Placeholder",
                                                       application);
         search.assertFirstPhaseExpression(vespaExpression, "my_profile");
-        assertLargeConstant(name + "_layer_Variable_1_read", search, Optional.of(10L));
-        assertLargeConstant(name + "_layer_Variable_read", search, Optional.of(7840L));
     }
 
     @Test
@@ -105,8 +127,6 @@ public class RankingExpressionWithTensorFlowTestCase {
                                                       "Placeholder",
                                                       application);
         search.assertFirstPhaseExpression(vespaExpression, "my_profile");
-        assertLargeConstant(name + "_layer_Variable_1_read", search, Optional.of(10L));
-        assertLargeConstant(name + "_layer_Variable_read", search, Optional.of(7840L));
     }
 
     @Test
@@ -125,8 +145,6 @@ public class RankingExpressionWithTensorFlowTestCase {
                                                       "Placeholder",
                                                       application);
         search.assertFirstPhaseExpression(vespaExpression, "my_profile");
-        assertLargeConstant(name + "_layer_Variable_1_read", search, Optional.of(10L));
-        assertLargeConstant(name + "_layer_Variable_read", search, Optional.of(7840L));
     }
 
     @Test
@@ -134,8 +152,6 @@ public class RankingExpressionWithTensorFlowTestCase {
         RankProfileSearchFixture search = fixtureWith("tensor(d0[2],d1[784])(0.0)",
                                                       "5 + sum(tensorflow('mnist_softmax/saved'))");
         search.assertFirstPhaseExpression("5 + reduce(" + vespaExpression + ", sum)", "my_profile");
-        assertLargeConstant(name + "_layer_Variable_1_read", search, Optional.of(10L));
-        assertLargeConstant(name + "_layer_Variable_read", search, Optional.of(7840L));
     }
 
     @Test
@@ -233,9 +249,6 @@ public class RankingExpressionWithTensorFlowTestCase {
                                                       "tensorflow('mnist_softmax/saved')");
         search.assertFirstPhaseExpression(vespaExpression, "my_profile");
 
-        assertLargeConstant(name + "_layer_Variable_1_read", search, Optional.of(10L));
-        assertLargeConstant(name + "_layer_Variable_read", search, Optional.of(7840L));
-
         // At this point the expression is stored - copy application to another location which do not have a models dir
         Path storedApplicationDirectory = applicationDir.getParentPath().append("copy");
         try {
@@ -250,10 +263,6 @@ public class RankingExpressionWithTensorFlowTestCase {
                                                                     "Placeholder",
                                                                     storedApplication);
             searchFromStored.assertFirstPhaseExpression(vespaExpression, "my_profile");
-            // Verify that the constants exists, but don't verify the content as we are not
-            // simulating file distribution in this test
-            assertLargeConstant(name + "_layer_Variable_1_read", searchFromStored, Optional.empty());
-            assertLargeConstant(name + "_layer_Variable_read", searchFromStored, Optional.empty());
         }
         finally {
             IOUtils.recursiveDeleteDir(storedApplicationDirectory.toFile());
@@ -287,7 +296,6 @@ public class RankingExpressionWithTensorFlowTestCase {
 
         assertNull("Constant overridden by macro is not added",
                    search.search().rankingConstants().get("mnist_softmax_saved_layer_Variable_read"));
-        assertLargeConstant(name + "_layer_Variable_1_read", search, Optional.of(10L));
 
         // At this point the expression is stored - copy application to another location which do not have a models dir
         Path storedApplicationDirectory = applicationDir.getParentPath().append("copy");
@@ -303,7 +311,6 @@ public class RankingExpressionWithTensorFlowTestCase {
             searchFromStored.assertFirstPhaseExpression(vespaExpressionWithoutConstant, "my_profile_child");
             assertNull("Constant overridden by macro is not added",
                        searchFromStored.search().rankingConstants().get("mnist_softmax_saved_layer_Variable_read"));
-            assertLargeConstant(name + "_layer_Variable_1_read", searchFromStored, Optional.of(10L));
         }
         finally {
             IOUtils.recursiveDeleteDir(storedApplicationDirectory.toFile());
@@ -316,8 +323,6 @@ public class RankingExpressionWithTensorFlowTestCase {
         RankProfileSearchFixture search = fixtureWith("tensor(d0[1],d1[784])(0.0)",
                                                       "tensorflow('mnist_softmax/saved')");
         search.assertFirstPhaseExpression(expression, "my_profile");
-        assertLargeConstant(name + "_layer_Variable_1_read", search, Optional.of(10L));
-        assertLargeConstant(name + "_layer_Variable_read", search, Optional.of(7840L));
     }
 
     @Test
@@ -401,11 +406,11 @@ public class RankingExpressionWithTensorFlowTestCase {
      * Verifies that the constant with the given name exists, and - only if an expected size is given -
      * that the content of the constant is available and has the expected size.
      */
-    private void assertLargeConstant(String name, RankProfileSearchFixture search, Optional<Long> expectedSize) {
+    private void assertLargeConstant(String constantName, VespaModel model, Optional<Long> expectedSize) {
         try {
-            Path constantApplicationPackagePath = Path.fromString("models.generated/my_profile.mnist_softmax_saved/constants").append(name + ".tbf");
-            RankingConstant rankingConstant = search.search().rankingConstants().get(name);
-            assertEquals(name, rankingConstant.getName());
+            Path constantApplicationPackagePath = Path.fromString("models.generated/" + name + "/constants").append(constantName + ".tbf");
+            RankingConstant rankingConstant = model.rankingConstants().get(constantName);
+            assertEquals(constantName, rankingConstant.getName());
             assertTrue(rankingConstant.getFileName().endsWith(constantApplicationPackagePath.toString()));
 
             if (expectedSize.isPresent()) {