diff --git a/deeplearning4j/deeplearning4j-ui-parent/deeplearning4j-vertx/src/main/java/org/deeplearning4j/ui/VertxUIServer.java b/deeplearning4j/deeplearning4j-ui-parent/deeplearning4j-vertx/src/main/java/org/deeplearning4j/ui/VertxUIServer.java index 7c6a27f4e..ff59dd2d8 100644 --- a/deeplearning4j/deeplearning4j-ui-parent/deeplearning4j-vertx/src/main/java/org/deeplearning4j/ui/VertxUIServer.java +++ b/deeplearning4j/deeplearning4j-ui-parent/deeplearning4j-vertx/src/main/java/org/deeplearning4j/ui/VertxUIServer.java @@ -75,7 +75,8 @@ public class VertxUIServer extends AbstractVerticle implements UIServer { private static Function statsStorageProvider; private static Integer instancePort; - private static Thread autoStopThread; + @Getter + private static Thread shutdownHook; /** * Get (and, initialize if necessary) the UI server. This synchronous function will wait until the server started. @@ -188,19 +189,17 @@ public class VertxUIServer extends AbstractVerticle implements UIServer { Vertx vertx = Vertx.vertx(); vertx.deployVerticle(VertxUIServer.class.getName(), promise); - Thread currentThread = Thread.currentThread(); - VertxUIServer.autoStopThread = new Thread(() -> { - try { - currentThread.join(); - if (VertxUIServer.instance != null && !VertxUIServer.instance.isStopped()) { - log.info("Deeplearning4j UI server is auto-stopping after thread (name: {}) died.", - currentThread.getName()); + VertxUIServer.shutdownHook = new Thread(() -> { + if (VertxUIServer.instance != null && !VertxUIServer.instance.isStopped()) { + log.info("Deeplearning4j UI server is auto-stopping in shutdown hook."); + try { instance.stop(); + } catch (InterruptedException e) { + log.error("Interrupted stopping of Deeplearning4j UI server in shutdown hook.", e); } - } catch (InterruptedException e) { - log.error("Deeplearning4j UI server auto-stop thread was interrupted.", e); } }); + Runtime.getRuntime().addShutdownHook(shutdownHook); } @@ -385,7 +384,6 @@ public class VertxUIServer extends AbstractVerticle implements UIServer { + server.actualPort(), result.cause())); } }); - VertxUIServer.autoStopThread.start(); } private List extractArgsFromRoute(String path, RoutingContext rc) { diff --git a/deeplearning4j/deeplearning4j-ui-parent/deeplearning4j-vertx/src/main/java/org/deeplearning4j/ui/api/UIServer.java b/deeplearning4j/deeplearning4j-ui-parent/deeplearning4j-vertx/src/main/java/org/deeplearning4j/ui/api/UIServer.java index 2bc98c20c..56b049f6f 100644 --- a/deeplearning4j/deeplearning4j-ui-parent/deeplearning4j-vertx/src/main/java/org/deeplearning4j/ui/api/UIServer.java +++ b/deeplearning4j/deeplearning4j-ui-parent/deeplearning4j-vertx/src/main/java/org/deeplearning4j/ui/api/UIServer.java @@ -168,4 +168,14 @@ public interface UIServer { * @param stopCallback callback {@link Promise} to notify on completion */ void stopAsync(Promise stopCallback); + + /** + * Get shutdown hook of UI server, that will stop the server when the Runtime is stopped. + * You may de-register this shutdown hook with {@link Runtime#removeShutdownHook(Thread)}, + * and add your own hook with {@link Runtime#addShutdownHook(Thread)} + * @return shutdown hook + */ + static Thread getShutdownHook() { + return VertxUIServer.getShutdownHook(); + }; } diff --git a/deeplearning4j/deeplearning4j-ui-parent/deeplearning4j-vertx/src/test/java/org/deeplearning4j/ui/TestVertxUI.java b/deeplearning4j/deeplearning4j-ui-parent/deeplearning4j-vertx/src/test/java/org/deeplearning4j/ui/TestVertxUI.java index f7dc6f079..2ed7e1b74 100644 --- a/deeplearning4j/deeplearning4j-ui-parent/deeplearning4j-vertx/src/test/java/org/deeplearning4j/ui/TestVertxUI.java +++ b/deeplearning4j/deeplearning4j-ui-parent/deeplearning4j-vertx/src/test/java/org/deeplearning4j/ui/TestVertxUI.java @@ -346,12 +346,16 @@ public class TestVertxUI extends BaseDL4JTest { } @Test - public void testUIAutoStopOnThreadExit() throws InterruptedException { - AtomicReference uiServer = new AtomicReference<>(); - Thread thread = new Thread(() -> uiServer.set(UIServer.getInstance())); - thread.start(); - thread.join(); - Thread.sleep(1_000); - assertTrue(uiServer.get().isStopped()); + public void testUIShutdownHook() throws InterruptedException { + UIServer uIServer = UIServer.getInstance(); + Thread shutdownHook = UIServer.getShutdownHook(); + shutdownHook.start(); + shutdownHook.join(); + /* + * running the shutdown hook thread before the Runtime is terminated + * enables us to check if the UI server has been shut down or not + */ + assertTrue(uIServer.isStopped()); + log.info("Deeplearning4j UI server stopped in shutdown hook."); } }