PCA records final Job state in model

Richiexy · Aug 7, 2014 · 67e6ae6 · 67e6ae6
1 parent c2c6a5c
commit 67e6ae6
Show file tree

Hide file tree

Showing 2 changed files with 12 additions and 2 deletions.
diff --git a/src/main/java/hex/pca/PCA.java b/src/main/java/hex/pca/PCA.java
@@ -5,7 +5,7 @@
 import hex.FrameTask.DataInfo;
 import hex.gram.Gram.GramTask;
 import water.Job.ColumnsJob;
-import water.Key;
+import water.*;
 import water.api.DocGen;
 import water.fvec.Frame;
 import water.fvec.Vec;
@@ -83,6 +83,15 @@ public PCA(String desc, Key dest, Frame src, int max_pc, double tolerance, boole
     PCAModel myModel = buildModel(dinfo, tsk);
     myModel.delete_and_lock(self());
     myModel.unlock(self());
+    remove();                   // Close/remove job
+    final JobState state = UKV.<Job>get(self()).state;
+    new TAtomic<PCAModel>() {
+      @Override
+      public PCAModel atomic(PCAModel m) {
+        if (m != null) m.get_params().state = state;
+        return m;
+      }
+    }.invoke(dest());
   }
 
   @Override protected void init() {

diff --git a/src/test/java/hex/PCATest.java b/src/test/java/hex/PCATest.java
@@ -71,7 +71,8 @@ public void checkEigvec(double[][] expected, double[][] actual) {
     Key kpca = Key.make("basicdata.pca");
     new PCA("PCA on basic small dataset", kpca, fr, 0.0, standardize).invoke();
     model = DKV.get(kpca).get();
-    Assert.assertTrue(model.get_params().state == Job.JobState.DONE); //HEX-1817
+    Job.JobState jstate = model.get_params().state;
+    Assert.assertTrue(jstate == Job.JobState.DONE); //HEX-1817
     testHTML(model);
     } finally {
       if( fr    != null ) fr   .delete();