xi-db commented on code in PR #50752:
URL: https://github.com/apache/spark/pull/50752#discussion_r2071378989
##########
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:
Do you think it is risky that mlcache has no memory limit when offloading is
disabled? Probably it makes sense to keep the old behaviour in this case:
mlcache also has memory limit and retention time when offloading is disabled,
but once a model is evicted, future access to it will throw an error
CACHE_INVALID. WDYT?
--
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
