Use fileAcquirer

4 files changed, 29 insertions, 13 deletions

diff --git a/model-evaluation/src/main/java/ai/vespa/models/evaluation/ModelsEvaluator.java b/model-evaluation/src/main/java/ai/vespa/models/evaluation/ModelsEvaluator.java
index b507c5d7f08..a0b859bf930 100644
--- a/model-evaluation/src/main/java/ai/vespa/models/evaluation/ModelsEvaluator.java
+++ b/model-evaluation/src/main/java/ai/vespa/models/evaluation/ModelsEvaluator.java
@@ -5,6 +5,7 @@ import com.google.common.annotations.Beta;
 import com.google.common.collect.ImmutableMap;
 import com.google.inject.Inject;
 import com.yahoo.component.AbstractComponent;
+import com.yahoo.filedistribution.fileacquirer.FileAcquirer;
 import com.yahoo.vespa.config.search.RankProfilesConfig;
 import com.yahoo.vespa.config.search.core.RankingConstantsConfig;
 
@@ -24,8 +25,10 @@ public class ModelsEvaluator extends AbstractComponent {
     private final ImmutableMap<String, Model> models;
 
     @Inject
-    public ModelsEvaluator(RankProfilesConfig config, RankingConstantsConfig constantsConfig) {
-        this(new RankProfilesConfigImporter().importFrom(config, constantsConfig));
+    public ModelsEvaluator(RankProfilesConfig config,
+                           RankingConstantsConfig constantsConfig,
+                           FileAcquirer fileAcquirer) {
+        this(new RankProfilesConfigImporter(fileAcquirer).importFrom(config, constantsConfig));
     }
 
     public ModelsEvaluator(Map<String, Model> models) {
diff --git a/model-evaluation/src/main/java/ai/vespa/models/evaluation/RankProfilesConfigImporter.java b/model-evaluation/src/main/java/ai/vespa/models/evaluation/RankProfilesConfigImporter.java
index 00f60654318..87ac53488db 100644
--- a/model-evaluation/src/main/java/ai/vespa/models/evaluation/RankProfilesConfigImporter.java
+++ b/model-evaluation/src/main/java/ai/vespa/models/evaluation/RankProfilesConfigImporter.java
@@ -1,6 +1,8 @@
 // Copyright 2018 Yahoo Holdings. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.
 package ai.vespa.models.evaluation;
 
+import com.yahoo.config.FileReference;
+import com.yahoo.filedistribution.fileacquirer.FileAcquirer;
 import com.yahoo.io.GrowableByteBuffer;
 import com.yahoo.io.IOUtils;
 import com.yahoo.searchlib.rankingexpression.ExpressionFunction;
@@ -11,7 +13,6 @@ import com.yahoo.tensor.TensorType;
 import com.yahoo.tensor.serialization.TypedBinaryFormat;
 import com.yahoo.vespa.config.search.RankProfilesConfig;
 import com.yahoo.vespa.config.search.core.RankingConstantsConfig;
-import com.yahoo.vespa.defaults.Defaults;
 
 import java.io.File;
 import java.io.IOException;
@@ -21,7 +22,7 @@ import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
 import java.util.Optional;
-import java.util.logging.Logger;
+import java.util.concurrent.TimeUnit;
 
 /**
  * Converts RankProfilesConfig instances to RankingExpressions for evaluation.
@@ -31,6 +32,12 @@ import java.util.logging.Logger;
  */
 public class RankProfilesConfigImporter {
 
+    private final FileAcquirer fileAcquirer;
+
+    public RankProfilesConfigImporter(FileAcquirer fileAcquirer) {
+        this.fileAcquirer = fileAcquirer;
+    }
+
     /**
      * Returns a map of the models contained in this config, indexed on name.
      * The map is modifiable and owned by the caller.
@@ -107,16 +114,14 @@ public class RankProfilesConfigImporter {
             constants.add(new Constant(constantConfig.name(),
                                        readTensorFromFile(constantConfig.name(),
                                                           TensorType.fromSpec(constantConfig.type()),
-                                                          constantConfig.fileref().value())));
+                                                          constantConfig.fileref())));
         }
         return constants;
     }
 
-    protected Tensor readTensorFromFile(String name, TensorType type, String fileReference) {
+    protected Tensor readTensorFromFile(String name, TensorType type, FileReference fileReference) {
         try {
-            File dir = new File(Defaults.getDefaults().underVespaHome("var/db/vespa/filedistribution"), fileReference);
-            File file = dir.listFiles()[0]; // directory contains one file having the original name
-
+            File file = fileAcquirer.waitFor(fileReference, 7, TimeUnit.DAYS);
             if (file.getName().endsWith(".tbf"))
                 return TypedBinaryFormat.decode(Optional.of(type),
                                                 GrowableByteBuffer.wrap(IOUtils.readFileBytes(file)));
@@ -125,6 +130,9 @@ public class RankProfilesConfigImporter {
                                                    file + " for constant " + name);
             // TODO: Support json and json.lz4
         }
+        catch (InterruptedException e) {
+            throw new IllegalStateException("Gave up waiting for constant " + name);
+        }
         catch (IOException e) {
             throw new UncheckedIOException(e);
         }
diff --git a/model-evaluation/src/test/java/ai/vespa/models/evaluation/ModelTester.java b/model-evaluation/src/test/java/ai/vespa/models/evaluation/ModelTester.java
index 0a400f7126f..0aceaccc3e0 100644
--- a/model-evaluation/src/test/java/ai/vespa/models/evaluation/ModelTester.java
+++ b/model-evaluation/src/test/java/ai/vespa/models/evaluation/ModelTester.java
@@ -1,8 +1,11 @@
 // Copyright 2018 Yahoo Holdings. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.
 package ai.vespa.models.evaluation;
 
+import com.yahoo.config.FileReference;
 import com.yahoo.config.subscription.ConfigGetter;
 import com.yahoo.config.subscription.FileSource;
+import com.yahoo.filedistribution.fileacquirer.FileAcquirer;
+import com.yahoo.filedistribution.fileacquirer.MockFileAcquirer;
 import com.yahoo.io.GrowableByteBuffer;
 import com.yahoo.io.IOUtils;
 import com.yahoo.path.Path;
@@ -42,7 +45,7 @@ public class ModelTester {
                                                        RankProfilesConfig.class).getConfig("");
         RankingConstantsConfig constantsConfig = new ConfigGetter<>(new FileSource(configDir.append("ranking-constants.cfg").toFile()),
                                                                     RankingConstantsConfig.class).getConfig("");
-        return new RankProfilesConfigImporterWithMockedConstants(Path.fromString(path).append("constants"))
+        return new RankProfilesConfigImporterWithMockedConstants(Path.fromString(path).append("constants"), MockFileAcquirer.returnFile(null))
                        .importFrom(config, constantsConfig);
     }
 
@@ -68,12 +71,13 @@ public class ModelTester {
 
         private final Path constantsPath;
 
-        public RankProfilesConfigImporterWithMockedConstants(Path constantsPath) {
+        public RankProfilesConfigImporterWithMockedConstants(Path constantsPath, FileAcquirer fileAcquirer) {
+            super(fileAcquirer);
             this.constantsPath = constantsPath;
         }
 
         @Override
-        protected Tensor readTensorFromFile(String name, TensorType type, String fileReference) {
+        protected Tensor readTensorFromFile(String name, TensorType type, FileReference fileReference) {
             try {
                 return TypedBinaryFormat.decode(Optional.of(type),
                                                 GrowableByteBuffer.wrap(IOUtils.readFileBytes(constantsPath.append(name).toFile())));
diff --git a/model-evaluation/src/test/java/ai/vespa/models/evaluation/ModelsEvaluatorTest.java b/model-evaluation/src/test/java/ai/vespa/models/evaluation/ModelsEvaluatorTest.java
index 0708d872ff6..bd1ff6b8ed7 100644
--- a/model-evaluation/src/test/java/ai/vespa/models/evaluation/ModelsEvaluatorTest.java
+++ b/model-evaluation/src/test/java/ai/vespa/models/evaluation/ModelsEvaluatorTest.java
@@ -3,6 +3,7 @@ package ai.vespa.models.evaluation;
 
 import com.yahoo.config.subscription.ConfigGetter;
 import com.yahoo.config.subscription.FileSource;
+import com.yahoo.filedistribution.fileacquirer.MockFileAcquirer;
 import com.yahoo.path.Path;
 import com.yahoo.tensor.Tensor;
 import com.yahoo.vespa.config.search.RankProfilesConfig;
@@ -47,7 +48,7 @@ public class ModelsEvaluatorTest {
                                                        RankProfilesConfig.class).getConfig("");
         RankingConstantsConfig constantsConfig = new ConfigGetter<>(new FileSource(configDir.append("ranking-constants.cfg").toFile()),
                                                                     RankingConstantsConfig.class).getConfig("");
-        return new ModelsEvaluator(config, constantsConfig);
+        return new ModelsEvaluator(config, constantsConfig, MockFileAcquirer.returnFile(null));
     }
 
 }