WeichenXu123 commented on code in PR #50752:
URL: https://github.com/apache/spark/pull/50752#discussion_r2073207712
##########
sql/connect/server/src/main/scala/org/apache/spark/sql/connect/ml/MLCache.scala:
##########
@@ -36,38 +39,74 @@ import org.apache.spark.sql.connect.service.SessionHolder
private[connect] class MLCache(sessionHolder: SessionHolder) extends Logging {
private val helper = new ConnectHelper()
private val helperID = "______ML_CONNECT_HELPER______"
+ private val modelClassNameFile = "__model_class_name__"
- private def getMaxCacheSizeKB: Long = {
-
sessionHolder.session.conf.get(Connect.CONNECT_SESSION_CONNECT_ML_CACHE_MAX_SIZE)
/ 1024
+ // TODO: rename it to `totalInMemorySizeBytes` because it only counts the
in-memory
+ // part data size.
+ private[ml] val totalSizeBytes: AtomicLong = new AtomicLong(0)
+
+ val offloadedModelsDir: Path = {
+ val path = Paths.get(
+ System.getProperty("java.io.tmpdir"),
+ "spark_connect_model_cache",
+ sessionHolder.sessionId)
+ Files.createDirectories(path)
+ }
+ private[spark] def getOffloadingEnabled: Boolean = {
+
sessionHolder.session.conf.get(Connect.CONNECT_SESSION_CONNECT_ML_CACHE_OFFLOADING_ENABLED)
}
- private def getTimeoutMinute: Long = {
-
sessionHolder.session.conf.get(Connect.CONNECT_SESSION_CONNECT_ML_CACHE_TIMEOUT)
+ private def getMaxInMemoryCacheSizeKB: Long = {
+ sessionHolder.session.conf.get(
+ Connect.CONNECT_SESSION_CONNECT_ML_CACHE_OFFLOADING_MAX_IN_MEMORY_SIZE)
/ 1024
+ }
+
+ private def getOffloadingTimeoutMinute: Long = {
+
sessionHolder.session.conf.get(Connect.CONNECT_SESSION_CONNECT_ML_CACHE_OFFLOADING_TIMEOUT)
}
private[ml] case class CacheItem(obj: Object, sizeBytes: Long)
- private[ml] val cachedModel: ConcurrentMap[String, CacheItem] = CacheBuilder
- .newBuilder()
- .softValues()
- .maximumWeight(getMaxCacheSizeKB)
- .expireAfterAccess(getTimeoutMinute, TimeUnit.MINUTES)
- .weigher((key: String, value: CacheItem) => {
- Math.ceil(value.sizeBytes.toDouble / 1024).toInt
- })
- .removalListener((removed: RemovalNotification[String, CacheItem]) =>
- totalSizeBytes.addAndGet(-removed.getValue.sizeBytes))
- .build[String, CacheItem]()
- .asMap()
+ private[ml] val cachedModel: ConcurrentMap[String, CacheItem] = {
+ if (getOffloadingEnabled) {
+ CacheBuilder
+ .newBuilder()
+ .softValues()
+ .removalListener((removed: RemovalNotification[String, CacheItem]) =>
+ totalSizeBytes.addAndGet(-removed.getValue.sizeBytes))
+ .maximumWeight(getMaxInMemoryCacheSizeKB)
+ .weigher((key: String, value: CacheItem) => {
+ Math.ceil(value.sizeBytes.toDouble / 1024).toInt
+ })
+ .expireAfterAccess(getOffloadingTimeoutMinute, TimeUnit.MINUTES)
+ .build[String, CacheItem]()
+ .asMap()
+ } else {
+ new ConcurrentHashMap[String, CacheItem]()
Review Comment:
If user cares memory explosion, they can enable offloading, and it is by
default enabled. so I don't see risk 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-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
