WeichenXu123 commented on code in PR #50752:
URL: https://github.com/apache/spark/pull/50752#discussion_r2073193261
##########
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:
This leads to non-deterministic error , CACHE_INVALID error might occur
unpredictably, this will cause user hard to debug their code, I don't suggest
this kind of 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-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
