IGNITE-7375: Right cleanup after training.
authorArtem Malykh <amalykh@gridgain.com>
Tue, 6 Feb 2018 17:40:47 +0000 (20:40 +0300)
committerYury Babak <ybabak@gridgain.com>
Tue, 6 Feb 2018 17:40:47 +0000 (20:40 +0300)
this closes #3478

modules/ml/src/main/java/org/apache/ignite/ml/nn/trainers/distributed/MLPGroupUpdateTrainer.java

index 8e97d87..333afcc 100644 (file)
@@ -22,7 +22,9 @@ import java.util.ArrayList;
 import java.util.List;
 import java.util.Map;
 import java.util.Objects;
+import java.util.Set;
 import java.util.UUID;
+import java.util.stream.Collectors;
 import java.util.stream.Stream;
 import org.apache.ignite.Ignite;
 import org.apache.ignite.Ignition;
@@ -324,7 +326,9 @@ public class MLPGroupUpdateTrainer<U extends Serializable> extends
 
     /** {@inheritDoc} */
     @Override protected void cleanup(MLPGroupUpdateTrainerLocalContext locCtx) {
-
+        MLPGroupUpdateTrainerDataCache.getOrCreate(ignite).remove(locCtx.trainingUUID());
+        Set<GroupTrainerCacheKey<Void>> toRmv = MLPCache.allKeys(locCtx.parallelTrainingsCnt(), locCtx.trainingUUID()).collect(Collectors.toSet());
+        MLPCache.getOrCreate(ignite).removeAll(toRmv);
     }
 
     /**