polish code

Author: jan-elastic

x-pack/plugin/inference/qa/inference-service-tests/build.gradle

@@ -1,6 +1,7 @@
 apply plugin: 'elasticsearch.internal-java-rest-test'
 
 dependencies {
+  javaRestTestImplementation project(path: xpackModule('core'))
   javaRestTestImplementation project(path: xpackModule('inference'))
   clusterPlugins project(':x-pack:plugin:inference:qa:test-service-plugin')
 }

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

@@ -58,14 +58,15 @@ public class InferenceBaseRestTest extends ESRestTestCase {
 
     @Before
     public void setMlModelRepository() throws IOException {
+        logger.info("setting ML model repository");
         var request = new Request("PUT", "/_cluster/settings");
         request.setJsonEntity(Strings.format("""
             {
               "persistent": {
-                "xpack.ml.model_repository": "http://localhost:%d"
+                "xpack.ml.model_repository": "%s"
               }
-            }""", mlModelServer.getPort()));
-        client().performRequest(request);
+            }""", mlModelServer.getUrl()));
+        assertOK(client().performRequest(request));
     }
 
     @Override

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

@@ -10,9 +10,13 @@
 import com.sun.net.httpserver.HttpExchange;
 import com.sun.net.httpserver.HttpServer;
 
+import org.apache.http.HttpHeaders;
 import org.apache.http.HttpStatus;
+import org.apache.http.client.utils.URIBuilder;
 import org.elasticsearch.logging.LogManager;
 import org.elasticsearch.logging.Logger;
+import org.elasticsearch.test.fixture.HttpHeaderParser;
+import org.elasticsearch.xpack.core.XPackSettings;
 import org.junit.rules.TestRule;
 import org.junit.runner.Description;
 import org.junit.runners.model.Statement;
@@ -27,58 +31,62 @@
 
 /**
  * Simple model server to serve ML models.
- * The URL path corresponds to file name in this class's resources.
+ * The URL path corresponds to a file name in this class's resources.
  * If the file is found, its content is returned, otherwise 404.
  * Respects a range header to serve partial content.
  */
 public class MlModelServer implements TestRule {
 
+    private static final String HOST = "localhost";
     private static final Logger logger = LogManager.getLogger(MlModelServer.class);
 
     private int port;
 
-    int getPort() {
-        return port;
+    public String getUrl() {
+        return new URIBuilder().setScheme("http").setHost(HOST).setPort(port).toString();
     }
 
-    private static void handle(HttpExchange exchange) throws IOException {
+    private static String getFileName(HttpExchange exchange) {
+        // Strip the leading slash
         String fileName = exchange.getRequestURI().getPath().substring(1);
-        // If this architecture is requested, serve the default model instead.
-        fileName = fileName.replace("_linux-x86_64", "");
-        String range = exchange.getRequestHeaders().getFirst("Range");
-        Integer rangeFrom = null;
-        Integer rangeTo = null;
-        if (range != null) {
-            assert range.startsWith("bytes=");
-            assert range.contains("-");
-            rangeFrom = Integer.parseInt(range.substring("bytes=".length(), range.indexOf('-')));
-            rangeTo = Integer.parseInt(range.substring(range.indexOf('-') + 1)) + 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, "");
         }
-        logger.info("Request: {} range=[{},{})", fileName, rangeFrom, rangeTo);
+        return fileName;
+    }
+
+    private static void handle(HttpExchange exchange) throws IOException {
+        String fileName = getFileName(exchange);
+        String rangeHeader = exchange.getRequestHeaders().getFirst(HttpHeaders.RANGE);
+        HttpHeaderParser.Range range = rangeHeader != null ? HttpHeaderParser.parseRangeHeader(rangeHeader) : null;
+        logger.info("Request: {} {}", fileName, range == null ? "" : range);
+
         ClassLoader classloader = Thread.currentThread().getContextClassLoader();
         try (InputStream is = classloader.getResourceAsStream(fileName)) {
+            int httpStatus;
+            long numBytes;
             if (is == null) {
-                logger.info("Response: {} 404", fileName);
-                exchange.sendResponseHeaders(HttpStatus.SC_NOT_FOUND, 0);
+                httpStatus = HttpStatus.SC_NOT_FOUND;
+                numBytes = 0;
             } else {
-                try (OutputStream os = exchange.getResponseBody()) {
-                    int httpStatus;
-                    int numBytes;
-                    if (range == null) {
-                        httpStatus = HttpStatus.SC_OK;
-                        numBytes = is.available();
-                    } else {
-                        httpStatus = HttpStatus.SC_PARTIAL_CONTENT;
-                        is.skipNBytes(rangeFrom);
-                        numBytes = rangeTo - rangeFrom;
-                    }
-                    logger.info("Response: {} {}", fileName, httpStatus);
-                    exchange.sendResponseHeaders(httpStatus, numBytes);
-                    while (numBytes > 0) {
-                        byte[] bytes = is.readNBytes(Math.min(1 << 20, numBytes));
-                        os.write(bytes);
-                        numBytes -= bytes.length;
-                    }
+                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;
+                }
+            }
+            logger.info("Response: {} {}", fileName, httpStatus);
+            exchange.sendResponseHeaders(httpStatus, numBytes);
+            try (OutputStream os = exchange.getResponseBody()) {
+                while (numBytes > 0) {
+                    byte[] bytes = is.readNBytes((int) Math.min(1 << 20, numBytes));
+                    os.write(bytes);
+                    numBytes -= bytes.length;
                 }
             }
         }
@@ -91,11 +99,10 @@ public Statement apply(Statement statement, Description description) {
             public void evaluate() throws Throwable {
                 logger.info("Starting ML model server");
                 HttpServer server = HttpServer.create();
-                server.createContext("/", MlModelServer::handle);
                 while (true) {
                     port = new Random().nextInt(10000, 65536);
                     try {
-                        server.bind(new InetSocketAddress("localhost", port), 1);
+                        server.bind(new InetSocketAddress(HOST, port), 1);
                     } catch (Exception e) {
                         continue;
                     }
@@ -105,12 +112,13 @@ public void evaluate() throws Throwable {
 
                 ExecutorService executor = Executors.newCachedThreadPool();
                 server.setExecutor(executor);
+                server.createContext("/", MlModelServer::handle);
                 server.start();
 
                 try {
                     statement.evaluate();
                 } finally {
-                    logger.info("Stopping ML model server in port {}", port);
+                    logger.info("Stopping ML model server on port {}", port);
                     server.stop(1);
                     executor.shutdown();
                 }