fix metadata for optimized models

Author: jan-elastic

x-pack/plugin/inference/qa/inference-service-tests/src/javaRestTest/java/org/elasticsearch/xpack/inference/MlModelServer.java

@@ -16,15 +16,21 @@
 import org.elasticsearch.logging.LogManager;
 import org.elasticsearch.logging.Logger;
 import org.elasticsearch.test.fixture.HttpHeaderParser;
+import org.elasticsearch.xcontent.XContentParser;
+import org.elasticsearch.xcontent.XContentParserConfiguration;
+import org.elasticsearch.xcontent.XContentType;
 import org.elasticsearch.xpack.core.XPackSettings;
+import org.elasticsearch.xpack.core.ml.inference.trainedmodel.ModelPackageConfig;
 import org.junit.rules.TestRule;
 import org.junit.runner.Description;
 import org.junit.runners.model.Statement;
 
+import java.io.ByteArrayInputStream;
 import java.io.IOException;
 import java.io.InputStream;
 import java.io.OutputStream;
 import java.net.InetSocketAddress;
+import java.nio.charset.StandardCharsets;
 import java.util.Random;
 import java.util.concurrent.ExecutorService;
 import java.util.concurrent.Executors;
@@ -46,41 +52,26 @@ public String getUrl() {
         return new URIBuilder().setScheme("http").setHost(HOST).setPort(port).toString();
     }
 
-    private static String getFileName(HttpExchange exchange) {
-        // Strip the leading slash
-        String fileName = exchange.getRequestURI().getPath().substring(1);
-        // If a model specifically optimized for some platform is requested,
-        // serve the default non-optimized model instead, which is compatible.
-        for (String platform : XPackSettings.ML_NATIVE_CODE_PLATFORMS) {
-            fileName = fileName.replace("_" + platform, "");
-        }
-        return fileName;
-    }
-
-    private static void handle(HttpExchange exchange) throws IOException {
-        String fileName = getFileName(exchange);
+    private void handle(HttpExchange exchange) throws IOException {
         String rangeHeader = exchange.getRequestHeaders().getFirst(HttpHeaders.RANGE);
         HttpHeaderParser.Range range = rangeHeader != null ? HttpHeaderParser.parseRangeHeader(rangeHeader) : null;
-        logger.info("Request: {} {}", fileName, range == null ? "" : range);
+        logger.info("request: {} range={}", exchange.getRequestURI().getPath(), range);
 
-        ClassLoader classloader = Thread.currentThread().getContextClassLoader();
-        try (InputStream is = classloader.getResourceAsStream(fileName)) {
+        try (InputStream is = getInputStream(exchange)) {
             int httpStatus;
             long numBytes;
             if (is == null) {
                 httpStatus = HttpStatus.SC_NOT_FOUND;
                 numBytes = 0;
+            } else if (range == null) {
+                httpStatus = HttpStatus.SC_OK;
+                numBytes = is.available();
             } else {
-                if (range == null) {
-                    httpStatus = HttpStatus.SC_OK;
-                    numBytes = is.available();
-                } else {
-                    httpStatus = HttpStatus.SC_PARTIAL_CONTENT;
-                    is.skipNBytes(range.start());
-                    numBytes = range.end() - range.start() + 1;
-                }
+                httpStatus = HttpStatus.SC_PARTIAL_CONTENT;
+                is.skipNBytes(range.start());
+                numBytes = range.end() - range.start() + 1;
             }
-            logger.info("Response: {} {}", fileName, httpStatus);
+            logger.info("response: {} {}", exchange.getRequestURI().getPath(), httpStatus);
             exchange.sendResponseHeaders(httpStatus, numBytes);
             try (OutputStream os = exchange.getResponseBody()) {
                 while (numBytes > 0) {
@@ -92,6 +83,33 @@ private static void handle(HttpExchange exchange) throws IOException {
         }
     }
 
+    private InputStream getInputStream(HttpExchange exchange) throws IOException {
+        String path = exchange.getRequestURI().getPath().substring(1);  // Strip leading slash
+        String modelId = path.substring(0, path.indexOf('.'));
+        String extension = path.substring(path.indexOf('.') + 1);
+
+        // If a model specifically optimized for some platform is requested,
+        // serve the default non-optimized model instead, which is compatible.
+        String defaultModelId = modelId;
+        for (String platform : XPackSettings.ML_NATIVE_CODE_PLATFORMS) {
+            defaultModelId = defaultModelId.replace("_" + platform, "");
+        }
+
+        ClassLoader classloader = Thread.currentThread().getContextClassLoader();
+        InputStream is = classloader.getResourceAsStream(defaultModelId + "." + extension);
+        if (is != null && modelId.equals(defaultModelId) == false && extension.equals("metadata.json")) {
+            // When an optimized version is requested, fix the default metadata,
+            // so that it contains the correct model ID.
+            try (XContentParser parser = XContentType.JSON.xContent().createParser(XContentParserConfiguration.EMPTY, is.readAllBytes())) {
+                is.close();
+                ModelPackageConfig packageConfig = ModelPackageConfig.fromXContentLenient(parser);
+                packageConfig = new ModelPackageConfig.Builder(packageConfig).setPackedModelId(modelId).build();
+                is = new ByteArrayInputStream(packageConfig.toString().getBytes(StandardCharsets.UTF_8));
+            }
+        }
+        return is;
+    }
+
     @Override
     public Statement apply(Statement statement, Description description) {
         return new Statement() {
@@ -112,7 +130,7 @@ public void evaluate() throws Throwable {
 
                 ExecutorService executor = Executors.newCachedThreadPool();
                 server.setExecutor(executor);
-                server.createContext("/", MlModelServer::handle);
+                server.createContext("/", MlModelServer.this::handle);
                 server.start();
 
                 try {