[SW-2345] Delete H2OFrames Produced by Algorithm Parameters

h2oai · Jul 1, 2020 · df828c9 · df828c9
1 parent 83ddc01
commit df828c9
Show file tree

Hide file tree

Showing 7 changed files with 27 additions and 3 deletions.
diff --git a/ml/src/main/scala/ai/h2o/sparkling/ml/algos/H2OAlgorithm.scala b/ml/src/main/scala/ai/h2o/sparkling/ml/algos/H2OAlgorithm.scala
@@ -76,6 +76,7 @@ abstract class H2OAlgorithm[P <: Model.Parameters: ClassTag]
         }
         .getOrElse(Map())
     val modelId = trainAndGetDestinationKey(s"/3/ModelBuilders/${parameters.algoName().toLowerCase}", params)
+    deleteRegisteredH2OFrames()
     H2OModel(modelId).toMOJOModel(
       Identifiable.randomUID(parameters.algoName()),
       H2OMOJOSettings.createFromModelParams(this),

diff --git a/ml/src/main/scala/ai/h2o/sparkling/ml/algos/H2OAutoML.scala b/ml/src/main/scala/ai/h2o/sparkling/ml/algos/H2OAutoML.scala
@@ -89,6 +89,7 @@ class H2OAutoML(override val uid: String)
     amlKeyOption = Some(autoMLId)
 
     val algoName = getLeaderboard().select("model_id").head().getString(0)
+    deleteRegisteredH2OFrames()
     H2OModel(getLeaderModelId(autoMLId))
       .toMOJOModel(Identifiable.randomUID(algoName), H2OMOJOSettings.createFromModelParams(this), internalFeatureCols)
   }

diff --git a/ml/src/main/scala/ai/h2o/sparkling/ml/algos/H2OGridSearch.scala b/ml/src/main/scala/ai/h2o/sparkling/ml/algos/H2OGridSearch.scala
@@ -143,6 +143,8 @@ class H2OGridSearch(override val uid: String)
           case _ => throw e
         }
     }
+    algo.deleteRegisteredH2OFrames()
+    deleteRegisteredH2OFrames()
     val unsortedGridModels = getGridModels(gridId, algoName, internalFeatureCols)
     if (unsortedGridModels.isEmpty) {
       throw new IllegalArgumentException("No model returned.")

diff --git a/ml/src/main/scala/ai/h2o/sparkling/ml/params/H2OAlgoParamsBase.scala b/ml/src/main/scala/ai/h2o/sparkling/ml/params/H2OAlgoParamsBase.scala
@@ -23,6 +23,8 @@ import org.apache.spark.ml.linalg.{DenseMatrix, DenseVector}
 import org.apache.spark.ml.param._
 import org.apache.spark.sql.SparkSession
 
+import scala.collection.mutable.ArrayBuffer
+
 trait H2OAlgoParamsBase extends Params {
   private[sparkling] def getH2OAlgorithmParams(trainingFrame: H2OFrame): Map[String, Any] = Map.empty
 
@@ -109,7 +111,9 @@ trait H2OAlgoParamsBase extends Params {
     convertWithH2OContext(input) { (spark, hc) =>
       import spark.implicits._
       val df = spark.sparkContext.parallelize(input).toDF()
-      hc.asH2OFrame(df).frameId
+      val frame = hc.asH2OFrame(df)
+      registerH2OFrameForDeletion(frame)
+      frame.frameId
     }
   }
 
@@ -118,7 +122,9 @@ trait H2OAlgoParamsBase extends Params {
       import spark.implicits._
       vectors.map { vector =>
         val df = spark.sparkContext.parallelize(vector.values).toDF()
-        hc.asH2OFrame(df).frameId
+        val frame = hc.asH2OFrame(df)
+        registerH2OFrameForDeletion(frame)
+        frame.frameId
       }
     }
   }
@@ -129,8 +135,19 @@ trait H2OAlgoParamsBase extends Params {
       matrix.map { matrix =>
         val rows = matrix.rowIter.map(_.toArray).toArray
         val df = spark.sparkContext.parallelize(rows).toDF()
-        hc.asH2OFrame(df).frameId
+        val frame = hc.asH2OFrame(df)
+        registerH2OFrameForDeletion(frame)
+        frame.frameId
       }
     }
   }
+
+  private val h2oFramesToBeDeleted = new ArrayBuffer[H2OFrame]()
+
+  protected sealed def registerH2OFrameForDeletion(frame: H2OFrame): Unit = h2oFramesToBeDeleted.append(frame)
+
+  protected sealed def deleteRegisteredH2OFrames(): Unit = {
+    h2oFramesToBeDeleted.foreach(_.delete())
+    h2oFramesToBeDeleted.clear()
+  }
 }
diff --git a/ml/src/main/scala/ai/h2o/sparkling/ml/params/HasBetaConstraints.scala b/ml/src/main/scala/ai/h2o/sparkling/ml/params/HasBetaConstraints.scala
@@ -41,6 +41,7 @@ trait HasBetaConstraints extends H2OAlgoParamsBase {
         s"H2OContext needs to be created in order to train the ${this.getClass.getSimpleName} model. " +
           "Please create one as H2OContext.getOrCreate().")
       val frame = hc.asH2OFrame(dataFrame)
+      registerH2OFrameForDeletion(frame)
       frame.frameId
     }
   }

diff --git a/ml/src/main/scala/ai/h2o/sparkling/ml/params/HasCalibrationDataFrame.scala b/ml/src/main/scala/ai/h2o/sparkling/ml/params/HasCalibrationDataFrame.scala
@@ -42,6 +42,7 @@ trait HasCalibrationDataFrame extends H2OAlgoParamsBase {
         s"H2OContext needs to be created in order to train the ${this.getClass.getSimpleName} model. " +
           "Please create one as H2OContext.getOrCreate().")
       val frame = hc.asH2OFrame(dataFrame)
+      registerH2OFrameForDeletion(frame)
       frame.frameId
     }
   }

diff --git a/ml/src/main/scala/ai/h2o/sparkling/ml/params/HasPlugValues.scala b/ml/src/main/scala/ai/h2o/sparkling/ml/params/HasPlugValues.scala
@@ -65,6 +65,7 @@ trait HasPlugValues extends H2OAlgoParamsBase {
       if (stringFieldsIndices.nonEmpty) {
         frame.convertColumnsToCategorical(stringFieldsIndices)
       }
+      registerH2OFrameForDeletion(frame)
       frame.frameId
     }
   }