zhengruifeng commented on code in PR #50751:
URL: https://github.com/apache/spark/pull/50751#discussion_r2072807492
##########
sql/connect/server/src/main/scala/org/apache/spark/sql/connect/ml/MLCache.scala:
##########
@@ -61,6 +61,31 @@ private[connect] class MLCache(sessionHolder: SessionHolder)
extends Logging {
private[ml] val totalSizeBytes: AtomicLong = new AtomicLong(0)
+ private[ml] val totalModelCacheSizeBytes: AtomicLong = new AtomicLong(0)
+ private[spark] def getModelCacheMaxSize: Long = {
+
sessionHolder.session.conf.get(Connect.CONNECT_SESSION_CONNECT_MODEL_CACHE_MAX_SIZE)
+ }
+ private[spark] def getModelMaxSize: Long = {
+
sessionHolder.session.conf.get(Connect.CONNECT_SESSION_CONNECT_MODEL_MAX_SIZE)
+ }
+
+ def checkModelSize(estimatedModelSize: Long): Unit = {
+ if (totalModelCacheSizeBytes.get() + estimatedModelSize >
getModelCacheMaxSize) {
+ throw MLModelCacheSizeOverflowException(
+ "The model cache size in current session is about to exceed" +
+ f"$getModelCacheMaxSize bytes. " +
+ "Please delete existing cached model by executing 'del model' in
python client " +
+ "before fitting new model or loading new model, or increase " +
+ "Spark config 'spark.connect.session.connectML.model.maxSize'.")
+ }
+ if (estimatedModelSize > getModelMaxSize) {
+ throw MLModelSizeOverflowException(
+ f"The fitted or loaded model size exceeds $getModelMaxSize bytes. " +
+ f"Please fit or load a model smaller than $getModelMaxSize bytes. " +
+ f"or increase Spark config
'spark.connect.session.connectML.modelCache.maxSize'.")
+ }
Review Comment:
normally we put the formats in the definitions of error classes.
--
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-unsubscr...@spark.apache.org
For queries about this service, please contact Infrastructure at:
us...@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscr...@spark.apache.org
For additional commands, e-mail: reviews-h...@spark.apache.org
