#7064 Fix DL4J UIServer (temporary) memory leak

deeplearning4j · Jan 29, 2019 · fc2c9e5 · fc2c9e5
1 parent 0877205
commit fc2c9e5
Show file tree

Hide file tree

Showing 2 changed files with 17 additions and 8 deletions.
diff --git a/...ent/deeplearning4j-play/src/main/java/org/deeplearning4j/ui/module/train/TrainModule.java b/...ent/deeplearning4j-play/src/main/java/org/deeplearning4j/ui/module/train/TrainModule.java
@@ -56,6 +56,7 @@
 import java.text.DecimalFormat;
 import java.text.SimpleDateFormat;
 import java.util.*;
+import java.util.concurrent.ConcurrentHashMap;
 import java.util.concurrent.atomic.AtomicInteger;
 
 import static play.mvc.Results.ok;
@@ -85,12 +86,12 @@ private enum ModelType {
     };
 
     private final int maxChartPoints; //Technically, the way it's set up: won't exceed 2*maxChartPoints
-    private Map<String, StatsStorage> knownSessionIDs = Collections.synchronizedMap(new LinkedHashMap<>());
+    private Map<String, StatsStorage> knownSessionIDs = Collections.synchronizedMap(new WeakHashMap<>());
     private String currentSessionID;
     private int currentWorkerIdx;
-    private Map<String, AtomicInteger> workerIdxCount = Collections.synchronizedMap(new HashMap<>()); //Key: session ID
-    private Map<String, Map<Integer, String>> workerIdxToName = Collections.synchronizedMap(new HashMap<>()); //Key: session ID
-    private Map<String, Long> lastUpdateForSession = Collections.synchronizedMap(new HashMap<>());
+    private Map<String, AtomicInteger> workerIdxCount = new ConcurrentHashMap<>(); //Key: session ID
+    private Map<String, Map<Integer, String>> workerIdxToName = new ConcurrentHashMap<>(); //Key: session ID
+    private Map<String, Long> lastUpdateForSession = new ConcurrentHashMap<>();
 
     public TrainModule() {
         String maxChartPointsProp = System.getProperty(DL4JSystemProperties.CHART_MAX_POINTS_PROPERTY);
@@ -180,19 +181,25 @@ public synchronized void onAttach(StatsStorage statsStorage) {
     }
 
     @Override
-    public void onDetach(StatsStorage statsStorage) {
+    public synchronized void onDetach(StatsStorage statsStorage) {
+        Set<String> toRemove = new HashSet<>();
         for (String s : knownSessionIDs.keySet()) {
             if (knownSessionIDs.get(s) == statsStorage) {
-                knownSessionIDs.remove(s);
+//                knownSessionIDs.remove(s);
+                toRemove.add(s);
                 workerIdxCount.remove(s);
                 workerIdxToName.remove(s);
                 currentSessionID = null;
                 getDefaultSession();
             }
         }
+        for(String s : toRemove) {
+//            knownSessionIDs.put(s, null);
+            knownSessionIDs.remove(s);
+        }
     }
 
-    private void getDefaultSession() {
+    private synchronized void getDefaultSession() {
         if (currentSessionID != null)
             return;
 
@@ -376,7 +383,7 @@ private static void cleanLegacyIterationCounts(List<Integer> iterationCounts) {
     }
 
     private Result getOverviewData() {
-        Long lastUpdate = lastUpdateForSession.get(currentSessionID);
+        Long lastUpdate = (currentSessionID == null ? null : lastUpdateForSession.get(currentSessionID));
         if (lastUpdate == null)
             lastUpdate = -1L;
         I18N i18N = I18NProvider.getInstance();

diff --git a/...-ui-parent/deeplearning4j-play/src/main/java/org/deeplearning4j/ui/play/PlayUIServer.java b/...-ui-parent/deeplearning4j-play/src/main/java/org/deeplearning4j/ui/play/PlayUIServer.java
@@ -420,6 +420,8 @@ private void runHelper() throws Exception {
                     m.reportStorageEvents(out);
                 }
 
+                events.clear();
+
                 try {
                     Thread.sleep(uiProcessingDelay);
                 } catch (InterruptedException e) {