Model Dowload Buttons

central/src/main/java/ai/djl/serving/central/HttpStaticFileServerInitializer.java

@@ -14,6 +14,7 @@
 
 import ai.djl.serving.central.handler.HttpStaticFileServerHandler;
 import ai.djl.serving.central.handler.ModelMetaDataHandler;
+import ai.djl.serving.central.handler.ModelDownloadHandler;
 import io.netty.channel.ChannelInitializer;
 import io.netty.channel.ChannelPipeline;
 import io.netty.channel.socket.SocketChannel;
@@ -54,6 +55,7 @@ public void initChannel(SocketChannel ch) {
         pipeline.addLast(new HttpServerCodec());
         pipeline.addLast(new HttpObjectAggregator(65536));
         pipeline.addLast(new ChunkedWriteHandler());
+        pipeline.addLast(new ModelDownloadHandler());
         pipeline.addLast(new ModelMetaDataHandler());
         pipeline.addLast(new HttpStaticFileServerHandler());
     }

central/src/main/java/ai/djl/serving/central/handler/ModelDownloadHandler.java

@@ -0,0 +1,118 @@
+/*
+ * Copyright 2021 Amazon.com, Inc. or its affiliates. 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. A copy of the License is located at
+ *
+ * http://aws.amazon.com/apache2.0/
+ *
+ * or in the "license" file accompanying this file. This file 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 ai.djl.serving.central.handler;
+
+import ai.djl.Application;
+import ai.djl.repository.Artifact;
+import ai.djl.repository.zoo.ModelNotFoundException;
+import ai.djl.repository.zoo.ModelZoo;
+import ai.djl.serving.central.responseencoder.HttpRequestResponse;
+import ai.djl.serving.central.utils.NettyUtils;
+import ai.djl.serving.http.BadRequestException;
+import io.netty.channel.ChannelHandlerContext;
+import io.netty.channel.SimpleChannelInboundHandler;
+import io.netty.handler.codec.http.FullHttpRequest;
+import io.netty.handler.codec.http.QueryStringDecoder;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+import java.io.IOException;
+import java.net.URI;
+import java.util.*;
+import java.util.concurrent.CompletableFuture;
+
+final class ModelLink {
+
+    private static final Logger logger = LoggerFactory.getLogger(ModelLink.class);
+    private static Map<String, URI> links = new HashMap<String,URI>();
+    private ModelLink() {}
+
+    private static URI BASE_URI = URI.create("https://mlrepo.djl.ai/");
+
+    public static Map<String, URI> linkFinder(String modelName) throws IOException, ModelNotFoundException {
+        Map<Application, List<Artifact>> models = ModelZoo.listModels();
+        models.forEach(
+                (app, list) -> {
+                    list.forEach(
+                            artifact -> {
+                                if (artifact.getName().equals(modelName)){
+                                    for (Map.Entry<String, Artifact.Item> entry :
+                                            artifact.getFiles().entrySet()) {
+                                        URI fileUri = URI.create(entry.getValue().getUri());
+                                        URI baseUri = artifact.getMetadata().getRepositoryUri();
+                                        if (!fileUri.isAbsolute()) {
+                                            fileUri = BASE_URI.resolve(baseUri).resolve(fileUri);
+                                        }
+                                        try {
+                                            links.put(entry.getKey(),fileUri);
+                                        } catch(Exception e){
+                                            logger.info(String.valueOf(e));
+                                        }
+                                    }
+                                }});
+                });
+        return links;
+    }
+}
+
+
+/**
+ * A handler to handle download requests from the UI
+ * @author anfee1@morgan.edu
+ *
+ */
+public class ModelDownloadHandler extends SimpleChannelInboundHandler<FullHttpRequest> {
+
+    HttpRequestResponse jsonResponse;
+    public ModelDownloadHandler() { jsonResponse = new HttpRequestResponse(); }
+
+    /**
+     * handle the deployment request by forwarding the request to the serving-instance.
+     *
+     * @param ctx the context
+     * @param request the full request
+     */
+    @Override
+    protected void channelRead0(ChannelHandlerContext ctx, FullHttpRequest request) throws IOException, ModelNotFoundException {
+        final Logger logger = LoggerFactory.getLogger(ModelDownloadHandler.class);
+        QueryStringDecoder decoder = new QueryStringDecoder(request.uri());
+        String modelName=NettyUtils.getParameter(decoder, "modelName", null);
+        CompletableFuture.supplyAsync(
+                () -> {
+                    try {
+                        if (modelName!=null) {
+                            logger.info(String.valueOf(ModelLink.linkFinder(modelName)));
+                            return ModelLink.linkFinder(modelName);
+                        } else {
+                            throw new BadRequestException("modelName is mandatory.");
+                        }
+
+                    } catch (IOException | ModelNotFoundException ex) {
+                        throw new IllegalArgumentException(ex.getMessage(), ex);
+                    }
+                })
+                .exceptionally((ex) -> Collections.emptyMap())
+                .thenAccept(linksMap -> jsonResponse.sendAsJson(ctx, request, linksMap));
+    }
+
+
+    /** {@inheritDoc} */
+    @Override
+    public boolean acceptInboundMessage(Object msg) {
+        FullHttpRequest request = (FullHttpRequest) msg;
+
+        String uri = request.uri();
+        return uri.startsWith("/serving/models?");
+    }
+
+}

central/src/main/java/ai/djl/serving/central/http/BadRequestException.java

@@ -0,0 +1,40 @@
+/*
+ * Copyright 2020 Amazon.com, Inc. or its affiliates. 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. A copy of the License is located at
+ *
+ * http://aws.amazon.com/apache2.0/
+ *
+ * or in the "license" file accompanying this file. This file 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 ai.djl.serving.http;
+
+/** Thrown when a bad HTTP request is received. */
+public class BadRequestException extends IllegalArgumentException {
+
+    static final long serialVersionUID = 1L;
+
+    /**
+     * Constructs an {@code BadRequestException} with the specified detail message.
+     *
+     * @param message The detail message (which is saved for later retrieval by the {@link
+     *     #getMessage()} method)
+     */
+    public BadRequestException(String message) {
+        super(message);
+    }
+
+    /**
+     * Constructs an {@code BadRequestException} with the specified detail message and a root cause.
+     *
+     * @param message The detail message (which is saved for later retrieval by the {@link
+     *     #getMessage()} method)
+     * @param cause root cause
+     */
+    public BadRequestException(String message, Throwable cause) {
+        super(message, cause);
+    }
+}

central/src/main/java/ai/djl/serving/central/responseencoder/HttpRequestResponse.java

@@ -0,0 +1,125 @@
+/*
+ * Copyright 2021 Amazon.com, Inc. or its affiliates. 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. A copy of the License is located at
+ *
+ * http://aws.amazon.com/apache2.0/
+ *
+ * or in the "license" file accompanying this file. This file 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 ai.djl.serving.central.responseencoder;
+
+import ai.djl.modality.Classifications;
+import ai.djl.modality.Classifications.ClassificationsSerializer;
+import ai.djl.modality.cv.output.DetectedObjects;
+import ai.djl.repository.Metadata;
+import com.google.gson.Gson;
+import com.google.gson.GsonBuilder;
+import com.google.gson.JsonPrimitive;
+import com.google.gson.JsonSerializer;
+import io.netty.buffer.ByteBuf;
+import io.netty.channel.ChannelFuture;
+import io.netty.channel.ChannelFutureListener;
+import io.netty.channel.ChannelHandlerContext;
+import io.netty.handler.codec.http.DefaultFullHttpResponse;
+import io.netty.handler.codec.http.FullHttpRequest;
+import io.netty.handler.codec.http.FullHttpResponse;
+import io.netty.handler.codec.http.HttpHeaderNames;
+import io.netty.handler.codec.http.HttpHeaderValues;
+import io.netty.handler.codec.http.HttpResponseStatus;
+import io.netty.handler.codec.http.HttpUtil;
+import io.netty.handler.codec.http.HttpVersion;
+import io.netty.util.CharsetUtil;
+import java.lang.reflect.Modifier;
+
+/**
+ * serialize to json and send the response to the client.
+ *
+ * @author erik.bamberg@web.de
+ */
+public class HttpRequestResponse {
+
+    private static final Gson GSON_WITH_TRANSIENT_FIELDS =
+            new GsonBuilder()
+                    .setDateFormat("yyyy-MM-dd'T'HH:mm:ss.SSS'Z'")
+                    .setPrettyPrinting()
+                    .excludeFieldsWithModifiers(Modifier.STATIC)
+                    .registerTypeAdapter(Classifications.class, new ClassificationsSerializer())
+                    .registerTypeAdapter(DetectedObjects.class, new ClassificationsSerializer())
+                    .registerTypeAdapter(Metadata.class, new MetaDataSerializer())
+                    .registerTypeAdapter(
+                            Double.class,
+                            (JsonSerializer<Double>)
+                                    (src, t, ctx) -> {
+                                        long v = src.longValue();
+                                        if (src.equals(Double.valueOf(String.valueOf(v)))) {
+                                            return new JsonPrimitive(v);
+                                        }
+                                        return new JsonPrimitive(src);
+                                    })
+                    .create();
+
+    /**
+     * send a response to the client.
+     *
+     * @param ctx channel context
+     * @param request full request
+     * @param entity the response
+     */
+    public void sendAsJson(ChannelHandlerContext ctx, FullHttpRequest request, Object entity) {
+
+        String serialized = GSON_WITH_TRANSIENT_FIELDS.toJson(entity);
+        ByteBuf buffer = ctx.alloc().buffer(serialized.length());
+        buffer.writeCharSequence(serialized, CharsetUtil.UTF_8);
+
+        FullHttpResponse response =
+                new DefaultFullHttpResponse(HttpVersion.HTTP_1_1, HttpResponseStatus.OK, buffer);
+        response.headers().set(HttpHeaderNames.CONTENT_TYPE, "application/json; charset=UTF-8");
+        boolean keepAlive = HttpUtil.isKeepAlive(request);
+        this.sendAndCleanupConnection(ctx, response, keepAlive);
+    }
+
+    /**
+     * send content of a ByteBuffer as
+     *  response to the client.
+     *
+     * @param ctx channel context
+     * @param request full request
+     * @param entity the response
+     */
+    public void sendByteBuffer(ChannelHandlerContext ctx, ByteBuf buffer) {
+
+        FullHttpResponse response =
+                new DefaultFullHttpResponse(HttpVersion.HTTP_1_1, HttpResponseStatus.OK, buffer);
+        response.headers().set(HttpHeaderNames.CONTENT_TYPE, "application/json; charset=UTF-8");
+        this.sendAndCleanupConnection(ctx, response, false);
+    }
+
+    /**
+     * If Keep-Alive is disabled, attaches "Connection: close" header to the response and closes the
+     * connection after the response being sent.
+     *
+     * @param ctx context
+     * @param request full request
+     * @param response full response
+     */
+    private void sendAndCleanupConnection(
+            ChannelHandlerContext ctx, FullHttpResponse response, boolean keepAlive ) {
+        HttpUtil.setContentLength(response, response.content().readableBytes());
+        if (!keepAlive) {
+            // We're going to close the connection as soon as the response is sent,
+            // so we should also make it clear for the client.
+            response.headers().set(HttpHeaderNames.CONNECTION, HttpHeaderValues.CLOSE);
+        }
+
+        ChannelFuture flushPromise = ctx.writeAndFlush(response);
+
+        if (!keepAlive) {
+            // Close the connection as soon as the response is sent.
+            flushPromise.addListener(ChannelFutureListener.CLOSE);
+        }
+    }
+}