You are viewing a plain text version of this content. The canonical link for it is here.
Posted to reviews@spark.apache.org by "zhengruifeng (via GitHub)" <gi...@apache.org> on 2025/02/20 02:53:27 UTC
[PR] [SPARK-51261][ML][PYTHON][CONNECT] Introduce model size estimation to control ml cache [spark]
zhengruifeng opened a new pull request, #50013:
URL: https://github.com/apache/spark/pull/50013
### What changes were proposed in this pull request?
Introduce model size estimation to control ml cache
### Why are the changes needed?
We should control the memory usage of cached ml models
1, When `CONNECT_SESSION_ML_CACHE_SIZE>0`,
before training a new model, estimate its size, fail the training if the estimated size > `CONNECT_SESSION_ML_CACHE_SIZE`
before caching a new model, fail if `model.estimatedSize > CONNECT_SESSION_ML_CACHE_SIZE`;
2, When `CONNECT_SESSION_ML_CACHE_TOTAL_ITEM_SIZE>0`,
set the `maximumWeight` of the cache, and apply the model size as the `weighter`, the guava cache will evict old model (because it hasn't been used recently or very often) when necessary.
3, Add config `CONNECT_SESSION_ML_CACHE_SIZE` to control the number of cached items, it only take effect when `CONNECT_SESSION_ML_CACHE_TOTAL_ITEM_SIZE<=0`, because guava cache doesn't support set `maximumWeight` and `maximumSize` together.
4, Add config `CONNECT_SESSION_ML_CACHE_TIMEOUT`, to evict cached item when a duration has elapsed after the creation or last access.
### Does this PR introduce _any_ user-facing change?
Only for ML connect
### How was this patch tested?
added tests
### Was this patch authored or co-authored using generative AI tooling?
no
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-51261][ML][PYTHON][CONNECT] Introduce model size estimation to control ml cache [spark]
Posted by "hvanhovell (via GitHub)" <gi...@apache.org>.
hvanhovell commented on code in PR #50013:
URL: https://github.com/apache/spark/pull/50013#discussion_r1968738761
##########
common/utils/src/main/resources/error/error-conditions.json:
##########
@@ -780,6 +780,11 @@
"Cannot retrieve <objectName> from the ML cache. It is probably because the entry has been evicted."
]
},
+ "CACHE_ITEM_EXCEEDED" : {
+ "message" : [
+ "Cannot <operation> object <name>, since its estimated size <estimatedSize> exceeds the limitation <maximumSize>."
Review Comment:
`...exceeds the maximum allowed size <maximumSize>.` instead of `exceeds the limitation <maximumSize>.`?
Does this use proper size formats?
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-51261][ML][PYTHON][CONNECT] Introduce model size estimation to control ml cache [spark]
Posted by "hvanhovell (via GitHub)" <gi...@apache.org>.
hvanhovell commented on code in PR #50013:
URL: https://github.com/apache/spark/pull/50013#discussion_r1968705260
##########
sql/connect/server/src/main/scala/org/apache/spark/sql/connect/ml/MLException.scala:
##########
@@ -36,3 +36,17 @@ private[spark] case class MLCacheInvalidException(objectName: String)
errorClass = "CONNECT_ML.CACHE_INVALID",
messageParameters = Map("objectName" -> objectName),
cause = null)
+
+private[spark] case class MlItemSizeExceededException(
Review Comment:
Let's make sure the client side error mappings can deal with this new Exception. cc @itholic .
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-51261][ML][PYTHON][CONNECT] Introduce model size estimation to control ml cache [spark]
Posted by "zhengruifeng (via GitHub)" <gi...@apache.org>.
zhengruifeng commented on code in PR #50013:
URL: https://github.com/apache/spark/pull/50013#discussion_r1970737739
##########
sql/connect/server/src/main/scala/org/apache/spark/sql/connect/ml/MLHandler.scala:
##########
@@ -125,6 +127,15 @@ private[connect] object MLHandler extends Logging {
val dataset = MLUtils.parseRelationProto(fitCmd.getDataset, sessionHolder)
val estimator =
MLUtils.getEstimator(sessionHolder, estimatorProto, Some(fitCmd.getParams))
+
+ // pre-training model size check
+ val maxSize = conf.getConf(Connect.CONNECT_SESSION_ML_CACHE_SINGLE_ITEM_SIZE)
+ if (maxSize > 0) {
+ val estimatedSize = estimator.estimateModelSize(dataset)
Review Comment:
we will need unit tests to check this estimation
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-51261][ML][PYTHON][CONNECT] Introduce model size estimation to control ml cache [spark]
Posted by "hvanhovell (via GitHub)" <gi...@apache.org>.
hvanhovell commented on code in PR #50013:
URL: https://github.com/apache/spark/pull/50013#discussion_r1968704565
##########
sql/connect/server/src/main/scala/org/apache/spark/sql/connect/ml/MLHandler.scala:
##########
@@ -125,6 +127,15 @@ private[connect] object MLHandler extends Logging {
val dataset = MLUtils.parseRelationProto(fitCmd.getDataset, sessionHolder)
val estimator =
MLUtils.getEstimator(sessionHolder, estimatorProto, Some(fitCmd.getParams))
+
+ // pre-training model size check
+ val maxSize = conf.getConf(Connect.CONNECT_SESSION_ML_CACHE_SINGLE_ITEM_SIZE)
+ if (maxSize > 0) {
+ val estimatedSize = estimator.estimateModelSize(dataset)
Review Comment:
How accurate is this? How much state is added to the model while fitting?
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-51261][ML][PYTHON][CONNECT] Introduce model size estimation to control ml cache [spark]
Posted by "hvanhovell (via GitHub)" <gi...@apache.org>.
hvanhovell commented on code in PR #50013:
URL: https://github.com/apache/spark/pull/50013#discussion_r1968711609
##########
mllib/src/main/scala/org/apache/spark/ml/util/Summary.scala:
##########
@@ -18,11 +18,21 @@
package org.apache.spark.ml.util
import org.apache.spark.annotation.Since
+import org.apache.spark.util.KnownSizeEstimation
/**
+ * For ml connect only.
* Trait for the Summary
* All the summaries should extend from this Summary in order to
* support connect.
*/
@Since("4.0.0")
-private[spark] trait Summary
+private[spark] trait Summary extends KnownSizeEstimation {
+
+ // A summary is normally a small object, with several RDDs or DataFrame.
+ // The SizeEstimator is likely to overestimate the size of the summary,
+ // because it will also count the underlying SparkSession and/or SparkContext,
Review Comment:
This is why I don't like the SizeEstimator.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-51261][ML][PYTHON][CONNECT] Introduce model size estimation to control ml cache [spark]
Posted by "hvanhovell (via GitHub)" <gi...@apache.org>.
hvanhovell commented on code in PR #50013:
URL: https://github.com/apache/spark/pull/50013#discussion_r1968734181
##########
mllib/src/main/scala/org/apache/spark/ml/classification/LogisticRegression.scala:
##########
@@ -1248,6 +1263,11 @@ class LogisticRegressionModel private[spark] (
}
}
+ private[spark] override def estimatedSize: Long = {
+ SizeEstimator.estimate((this.params, this.uid,
Review Comment:
NIT: Can we just make the SizeEstimator take varargs instead of creating a tuple here.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-51261][ML][PYTHON][CONNECT] Introduce model size estimation to control ml cache [spark]
Posted by "zhengruifeng (via GitHub)" <gi...@apache.org>.
zhengruifeng commented on code in PR #50013:
URL: https://github.com/apache/spark/pull/50013#discussion_r1970737739
##########
sql/connect/server/src/main/scala/org/apache/spark/sql/connect/ml/MLHandler.scala:
##########
@@ -125,6 +127,15 @@ private[connect] object MLHandler extends Logging {
val dataset = MLUtils.parseRelationProto(fitCmd.getDataset, sessionHolder)
val estimator =
MLUtils.getEstimator(sessionHolder, estimatorProto, Some(fitCmd.getParams))
+
+ // pre-training model size check
+ val maxSize = conf.getConf(Connect.CONNECT_SESSION_ML_CACHE_SINGLE_ITEM_SIZE)
+ if (maxSize > 0) {
+ val estimatedSize = estimator.estimateModelSize(dataset)
Review Comment:
we will need unit tests to guard this estimation
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-51261][ML][PYTHON][CONNECT] Introduce model size estimation to control ml cache [spark]
Posted by "zhengruifeng (via GitHub)" <gi...@apache.org>.
zhengruifeng commented on PR #50013:
URL: https://github.com/apache/spark/pull/50013#issuecomment-2670881333
> I didn't check how to estimate the size of every model/summary, which needs ML experts to do. The ml connect part looks good to me.
let me add more ut for this estimation
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-51261][ML][PYTHON][CONNECT] Introduce model size estimation to control ml cache [spark]
Posted by "zhengruifeng (via GitHub)" <gi...@apache.org>.
zhengruifeng commented on code in PR #50013:
URL: https://github.com/apache/spark/pull/50013#discussion_r1970735576
##########
mllib/src/main/scala/org/apache/spark/ml/classification/FMClassifier.scala:
##########
@@ -235,6 +236,13 @@ class FMClassifier @Since("3.0.0") (
model.setSummary(Some(summary))
}
+ override def estimateModelSize(dataset: Dataset[_]): Long = {
+ val numFeatures = DatasetUtils.getNumFeatures(dataset, $(featuresCol))
Review Comment:
DatasetUtils.getNumFeatures is quite cheap, it will try to fetch `numFeatures` from the metadata, and if there is no such metadata, it just infer the `numFeatures` from the first row.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-51261][ML][PYTHON][CONNECT] Introduce model size estimation to control ml cache [spark]
Posted by "hvanhovell (via GitHub)" <gi...@apache.org>.
hvanhovell commented on code in PR #50013:
URL: https://github.com/apache/spark/pull/50013#discussion_r1968733126
##########
mllib/src/main/scala/org/apache/spark/ml/classification/FMClassifier.scala:
##########
@@ -235,6 +236,13 @@ class FMClassifier @Since("3.0.0") (
model.setSummary(Some(summary))
}
+ override def estimateModelSize(dataset: Dataset[_]): Long = {
+ val numFeatures = DatasetUtils.getNumFeatures(dataset, $(featuresCol))
Review Comment:
I am worried that executing the input query multiple times will be considered wasteful. I am wondering if we should do the check while fitting instead; we fail as soon as the model gets too large.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-51261][ML][PYTHON][CONNECT] Introduce model size estimation to control ml cache [spark]
Posted by "wbo4958 (via GitHub)" <gi...@apache.org>.
wbo4958 commented on code in PR #50013:
URL: https://github.com/apache/spark/pull/50013#discussion_r1962817999
##########
mllib/src/main/scala/org/apache/spark/ml/Model.scala:
##########
@@ -18,6 +18,7 @@
package org.apache.spark.ml
import org.apache.spark.ml.param.ParamMap
+import org.apache.spark.util.SizeEstimator
Review Comment:
This is my first time knowing SizeEstimator to estimating the java object size. That's cool.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-51261][ML][PYTHON][CONNECT] Introduce model size estimation to control ml cache [spark]
Posted by "hvanhovell (via GitHub)" <gi...@apache.org>.
hvanhovell commented on code in PR #50013:
URL: https://github.com/apache/spark/pull/50013#discussion_r1968717393
##########
mllib-local/src/main/scala/org/apache/spark/ml/linalg/Vectors.scala:
##########
@@ -504,6 +506,10 @@ object Vectors {
/** Max number of nonzero entries used in computing hash code. */
private[linalg] val MAX_HASH_NNZ = 128
+
+ private[ml] def getSparseSize(nnz: Long): Long = nnz * 12 + 20
Review Comment:
Let's try to avoid hard coding constants.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-51261][ML][PYTHON][CONNECT] Introduce model size estimation to control ml cache [spark]
Posted by "zhengruifeng (via GitHub)" <gi...@apache.org>.
zhengruifeng commented on code in PR #50013:
URL: https://github.com/apache/spark/pull/50013#discussion_r1970735576
##########
mllib/src/main/scala/org/apache/spark/ml/classification/FMClassifier.scala:
##########
@@ -235,6 +236,13 @@ class FMClassifier @Since("3.0.0") (
model.setSummary(Some(summary))
}
+ override def estimateModelSize(dataset: Dataset[_]): Long = {
+ val numFeatures = DatasetUtils.getNumFeatures(dataset, $(featuresCol))
Review Comment:
DatasetUtils.getNumFeatures is quite cheap, it will try to fetch `numFeatures` from the metadata, and if there is no such metadata, it just infer the `numFeatures` from the first row.
https://github.com/apache/spark/blob/9cf6dc873ff34412df6256cdc7613eed40716570/mllib/src/main/scala/org/apache/spark/ml/util/DatasetUtils.scala#L206-L214
##########
mllib-local/src/main/scala/org/apache/spark/ml/linalg/Vectors.scala:
##########
@@ -504,6 +506,10 @@ object Vectors {
/** Max number of nonzero entries used in computing hash code. */
private[linalg] val MAX_HASH_NNZ = 128
+
+ private[ml] def getSparseSize(nnz: Long): Long = nnz * 12 + 20
Review Comment:
SG
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-51261][ML][PYTHON][CONNECT] Introduce model size estimation to control ml cache [spark]
Posted by "zhengruifeng (via GitHub)" <gi...@apache.org>.
zhengruifeng commented on code in PR #50013:
URL: https://github.com/apache/spark/pull/50013#discussion_r1964614803
##########
sql/connect/server/src/main/scala/org/apache/spark/sql/connect/ml/MLHandler.scala:
##########
@@ -125,6 +127,15 @@ private[connect] object MLHandler extends Logging {
val dataset = MLUtils.parseRelationProto(fitCmd.getDataset, sessionHolder)
val estimator =
MLUtils.getEstimator(sessionHolder, estimatorProto, Some(fitCmd.getParams))
+
+ // pre-training model size check
+ val maxSize = conf.getConf(Connect.CONNECT_SESSION_ML_CACHE_SINGLE_ITEM_SIZE)
+ if (maxSize > 0) {
+ val estimatedSize = estimator.estimateModelSize(dataset)
+ if (estimatedSize > maxSize) {
+ throw MlItemSizeExceededException("fit", estimator.uid, estimatedSize, maxSize)
Review Comment:
I think we should not stop the session in this case
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-51261][ML][PYTHON][CONNECT] Introduce model size estimation to control ml cache [spark]
Posted by "zhengruifeng (via GitHub)" <gi...@apache.org>.
zhengruifeng commented on PR #50013:
URL: https://github.com/apache/spark/pull/50013#issuecomment-2673148757
currently, we have a pre-training model size check, if the model to be trained seems larger, the ml server will directly fail the training.
regarding model loading, it first load the model into memory, then throw error if its size is large. Probably we should also use a pre-loading model size check.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-51261][ML][PYTHON][CONNECT] Introduce model size estimation to control ml cache [spark]
Posted by "wbo4958 (via GitHub)" <gi...@apache.org>.
wbo4958 commented on code in PR #50013:
URL: https://github.com/apache/spark/pull/50013#discussion_r1962815802
##########
sql/connect/server/src/main/scala/org/apache/spark/sql/connect/ml/MLHandler.scala:
##########
@@ -125,6 +127,15 @@ private[connect] object MLHandler extends Logging {
val dataset = MLUtils.parseRelationProto(fitCmd.getDataset, sessionHolder)
val estimator =
MLUtils.getEstimator(sessionHolder, estimatorProto, Some(fitCmd.getParams))
+
+ // pre-training model size check
+ val maxSize = conf.getConf(Connect.CONNECT_SESSION_ML_CACHE_SINGLE_ITEM_SIZE)
+ if (maxSize > 0) {
+ val estimatedSize = estimator.estimateModelSize(dataset)
+ if (estimatedSize > maxSize) {
+ throw MlItemSizeExceededException("fit", estimator.uid, estimatedSize, maxSize)
Review Comment:
similar comment as before.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-51261][ML][PYTHON][CONNECT] Introduce model size estimation to control ml cache [spark]
Posted by "wbo4958 (via GitHub)" <gi...@apache.org>.
wbo4958 commented on PR #50013:
URL: https://github.com/apache/spark/pull/50013#issuecomment-2670403064
I didn't check how to estimate the size of every model/summary, which needs ML experts to do. The ml connect part looks good to me.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-51261][ML][PYTHON][CONNECT] Introduce model size estimation to control ml cache [spark]
Posted by "hvanhovell (via GitHub)" <gi...@apache.org>.
hvanhovell commented on code in PR #50013:
URL: https://github.com/apache/spark/pull/50013#discussion_r1968703858
##########
sql/connect/server/src/main/scala/org/apache/spark/sql/connect/ml/MLHandler.scala:
##########
@@ -125,6 +127,15 @@ private[connect] object MLHandler extends Logging {
val dataset = MLUtils.parseRelationProto(fitCmd.getDataset, sessionHolder)
val estimator =
MLUtils.getEstimator(sessionHolder, estimatorProto, Some(fitCmd.getParams))
+
+ // pre-training model size check
+ val maxSize = conf.getConf(Connect.CONNECT_SESSION_ML_CACHE_SINGLE_ITEM_SIZE)
+ if (maxSize > 0) {
+ val estimatedSize = estimator.estimateModelSize(dataset)
+ if (estimatedSize > maxSize) {
+ throw MlItemSizeExceededException("fit", estimator.uid, estimatedSize, maxSize)
Review Comment:
No we should not stop the session. The client can technically recover by increasing the limits.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-51261][ML][PYTHON][CONNECT] Introduce model size estimation to control ml cache [spark]
Posted by "hvanhovell (via GitHub)" <gi...@apache.org>.
hvanhovell commented on code in PR #50013:
URL: https://github.com/apache/spark/pull/50013#discussion_r1968708896
##########
sql/connect/server/src/main/scala/org/apache/spark/sql/connect/ml/MLCache.scala:
##########
@@ -21,23 +21,52 @@ import java.util.concurrent.{ConcurrentMap, TimeUnit}
import com.google.common.cache.CacheBuilder
-import org.apache.spark.internal.Logging
+import org.apache.spark.internal.{Logging, LogKeys, MDC}
+import org.apache.spark.ml.Model
import org.apache.spark.ml.util.ConnectHelper
+import org.apache.spark.sql.classic.SparkSession
+import org.apache.spark.sql.connect.config.Connect._
+import org.apache.spark.util.SizeEstimator
/**
* MLCache is for caching ML objects, typically for models and summaries evaluated by a model.
*/
-private[connect] class MLCache extends Logging {
+private[connect] class MLCache(session: SparkSession) extends Logging {
private val helper = new ConnectHelper()
private val helperID = "______ML_CONNECT_HELPER______"
+ private def conf = session.sessionState.conf
- private val cachedModel: ConcurrentMap[String, Object] = CacheBuilder
- .newBuilder()
- .softValues()
- .maximumSize(MLCache.MAX_CACHED_ITEMS)
- .expireAfterAccess(MLCache.CACHE_TIMEOUT_MINUTE, TimeUnit.MINUTES)
- .build[String, Object]()
- .asMap()
+ private val cachedModel: ConcurrentMap[String, (Object, Long)] = {
+ val builder = CacheBuilder.newBuilder().softValues()
+
+ val cacheWeight = conf.getConf(CONNECT_SESSION_ML_CACHE_TOTAL_ITEM_SIZE)
+ val cacheSize = conf.getConf(CONNECT_SESSION_ML_CACHE_SIZE)
+ val timeOut = conf.getConf(CONNECT_SESSION_ML_CACHE_TIMEOUT)
+
+ if (cacheWeight > 0) {
+ builder
Review Comment:
What is the eviction policy here? LRU?
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-51261][ML][PYTHON][CONNECT] Introduce model size estimation to control ml cache [spark]
Posted by "zhengruifeng (via GitHub)" <gi...@apache.org>.
zhengruifeng commented on code in PR #50013:
URL: https://github.com/apache/spark/pull/50013#discussion_r1964613438
##########
sql/connect/server/src/main/scala/org/apache/spark/sql/connect/ml/MLCache.scala:
##########
@@ -47,8 +76,18 @@ private[connect] class MLCache extends Logging {
* the key
*/
def register(obj: Object): String = {
+ val (estimatedSize, name) = obj match {
+ case m: Model[_] => (m.estimatedSize, m.uid)
+ case o => (SizeEstimator.estimate(o), Option(o).map(_.getClass.getName).getOrElse("NULL"))
+ }
+
+ val maxSize = conf.getConf(Connect.CONNECT_SESSION_ML_CACHE_SINGLE_ITEM_SIZE)
+ if (maxSize > 0 && estimatedSize > maxSize) {
+ throw MlItemSizeExceededException("cache", name, estimatedSize, maxSize)
Review Comment:
I think we should not stop the whole sessionholder here.
Just throwing an error seems a right behavior
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-51261][ML][PYTHON][CONNECT] Introduce model size estimation to control ml cache [spark]
Posted by "hvanhovell (via GitHub)" <gi...@apache.org>.
hvanhovell commented on code in PR #50013:
URL: https://github.com/apache/spark/pull/50013#discussion_r1968701252
##########
sql/connect/server/src/main/scala/org/apache/spark/sql/connect/config/Connect.scala:
##########
@@ -313,4 +313,49 @@ object Connect {
.internal()
.booleanConf
.createWithDefault(true)
+
+ val CONNECT_SESSION_ML_CACHE_TIMEOUT =
+ buildConf("spark.connect.session.mlCache.timeout")
Review Comment:
I am not sure if we can do this. This is a non-recoverable situation for a client. The client itself should tell us they are done with a model. I would suggest instead that we either fail and give them a list of models they can evict, or that we store old models on disk.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-51261][ML][PYTHON][CONNECT] Introduce model size estimation to control ml cache [spark]
Posted by "zhengruifeng (via GitHub)" <gi...@apache.org>.
zhengruifeng commented on code in PR #50013:
URL: https://github.com/apache/spark/pull/50013#discussion_r1970742537
##########
sql/connect/server/src/main/scala/org/apache/spark/sql/connect/ml/MLCache.scala:
##########
@@ -21,23 +21,52 @@ import java.util.concurrent.{ConcurrentMap, TimeUnit}
import com.google.common.cache.CacheBuilder
-import org.apache.spark.internal.Logging
+import org.apache.spark.internal.{Logging, LogKeys, MDC}
+import org.apache.spark.ml.Model
import org.apache.spark.ml.util.ConnectHelper
+import org.apache.spark.sql.classic.SparkSession
+import org.apache.spark.sql.connect.config.Connect._
+import org.apache.spark.util.SizeEstimator
/**
* MLCache is for caching ML objects, typically for models and summaries evaluated by a model.
*/
-private[connect] class MLCache extends Logging {
+private[connect] class MLCache(session: SparkSession) extends Logging {
private val helper = new ConnectHelper()
private val helperID = "______ML_CONNECT_HELPER______"
+ private def conf = session.sessionState.conf
- private val cachedModel: ConcurrentMap[String, Object] = CacheBuilder
- .newBuilder()
- .softValues()
- .maximumSize(MLCache.MAX_CACHED_ITEMS)
- .expireAfterAccess(MLCache.CACHE_TIMEOUT_MINUTE, TimeUnit.MINUTES)
- .build[String, Object]()
- .asMap()
+ private val cachedModel: ConcurrentMap[String, (Object, Long)] = {
+ val builder = CacheBuilder.newBuilder().softValues()
+
+ val cacheWeight = conf.getConf(CONNECT_SESSION_ML_CACHE_TOTAL_ITEM_SIZE)
+ val cacheSize = conf.getConf(CONNECT_SESSION_ML_CACHE_SIZE)
+ val timeOut = conf.getConf(CONNECT_SESSION_ML_CACHE_TIMEOUT)
+
+ if (cacheWeight > 0) {
+ builder
Review Comment:
it said
```
*
* <p>When eviction is necessary, the cache evicts entries that are less likely to be used again.
* For example, the cache may evict an entry because it hasn't been used recently or very often.
*
```
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-51261][ML][PYTHON][CONNECT] Introduce model size estimation to control ml cache [spark]
Posted by "zhengruifeng (via GitHub)" <gi...@apache.org>.
zhengruifeng commented on code in PR #50013:
URL: https://github.com/apache/spark/pull/50013#discussion_r1970737074
##########
sql/connect/server/src/main/scala/org/apache/spark/sql/connect/ml/MLHandler.scala:
##########
@@ -125,6 +127,15 @@ private[connect] object MLHandler extends Logging {
val dataset = MLUtils.parseRelationProto(fitCmd.getDataset, sessionHolder)
val estimator =
MLUtils.getEstimator(sessionHolder, estimatorProto, Some(fitCmd.getParams))
+
+ // pre-training model size check
+ val maxSize = conf.getConf(Connect.CONNECT_SESSION_ML_CACHE_SINGLE_ITEM_SIZE)
+ if (maxSize > 0) {
+ val estimatedSize = estimator.estimateModelSize(dataset)
Review Comment:
That depends, we are trying to estimate an upper bound of this size. E.g.
> - Given a LogisticRegression estimator, assume the coefficients are dense, even though the actual fitted model might be sparse (by L1 penalty).
> - Given a tree model, assume all underlying trees are complete binary trees, even though some branches might be pruned or truncated.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-51261][ML][PYTHON][CONNECT] Introduce model size estimation to control ml cache [spark]
Posted by "zhengruifeng (via GitHub)" <gi...@apache.org>.
zhengruifeng closed pull request #50013: [SPARK-51261][ML][PYTHON][CONNECT] Introduce model size estimation to control ml cache
URL: https://github.com/apache/spark/pull/50013
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-51261][ML][PYTHON][CONNECT] Introduce model size estimation to control ml cache [spark]
Posted by "wbo4958 (via GitHub)" <gi...@apache.org>.
wbo4958 commented on code in PR #50013:
URL: https://github.com/apache/spark/pull/50013#discussion_r1962815424
##########
sql/connect/server/src/main/scala/org/apache/spark/sql/connect/ml/MLCache.scala:
##########
@@ -47,8 +76,18 @@ private[connect] class MLCache extends Logging {
* the key
*/
def register(obj: Object): String = {
+ val (estimatedSize, name) = obj match {
+ case m: Model[_] => (m.estimatedSize, m.uid)
+ case o => (SizeEstimator.estimate(o), Option(o).map(_.getClass.getName).getOrElse("NULL"))
+ }
+
+ val maxSize = conf.getConf(Connect.CONNECT_SESSION_ML_CACHE_SINGLE_ITEM_SIZE)
+ if (maxSize > 0 && estimatedSize > maxSize) {
+ throw MlItemSizeExceededException("cache", name, estimatedSize, maxSize)
Review Comment:
Do we need to remove the existing cache before throwing exception? or the whole sessionholder will be closed?
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-51261][ML][PYTHON][CONNECT] Introduce model size estimation to control ml cache [spark]
Posted by "wbo4958 (via GitHub)" <gi...@apache.org>.
wbo4958 commented on code in PR #50013:
URL: https://github.com/apache/spark/pull/50013#discussion_r1962815424
##########
sql/connect/server/src/main/scala/org/apache/spark/sql/connect/ml/MLCache.scala:
##########
@@ -47,8 +76,18 @@ private[connect] class MLCache extends Logging {
* the key
*/
def register(obj: Object): String = {
+ val (estimatedSize, name) = obj match {
+ case m: Model[_] => (m.estimatedSize, m.uid)
+ case o => (SizeEstimator.estimate(o), Option(o).map(_.getClass.getName).getOrElse("NULL"))
+ }
+
+ val maxSize = conf.getConf(Connect.CONNECT_SESSION_ML_CACHE_SINGLE_ITEM_SIZE)
+ if (maxSize > 0 && estimatedSize > maxSize) {
+ throw MlItemSizeExceededException("cache", name, estimatedSize, maxSize)
Review Comment:
Do we need to remove the existing cache before throwing exception?
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org