diff --git a/sql/connect/server/src/main/scala/org/apache/spark/sql/connect/ml/MLCache.scala b/sql/connect/server/src/main/scala/org/apache/spark/sql/connect/ml/MLCache.scala index 5deb9ce1c3f19..5c3db4b99e307 100644 --- a/sql/connect/server/src/main/scala/org/apache/spark/sql/connect/ml/MLCache.scala +++ b/sql/connect/server/src/main/scala/org/apache/spark/sql/connect/ml/MLCache.scala @@ -23,6 +23,7 @@ import java.util.concurrent.{ConcurrentHashMap, ConcurrentMap, TimeUnit} import java.util.concurrent.atomic.{AtomicBoolean, AtomicLong} import scala.collection.mutable +import scala.jdk.CollectionConverters._ import scala.util.control.NonFatal import com.google.common.cache.{CacheBuilder, RemovalNotification} @@ -74,14 +75,28 @@ private[connect] class MLCache(sessionHolder: SessionHolder) extends Logging { Connect.CONNECT_SESSION_CONNECT_ML_CACHE_MEMORY_CONTROL_OFFLOADING_TIMEOUT) } + private case class ModelMetadata( + className: String, + modelString: String, + estimatedSizeBytes: Option[Long]) + + // Keep lightweight metadata after a model is evicted from memory so the UI can report + // offloaded models without loading them back into memory. + private val cachedModelMetadata = new ConcurrentHashMap[String, ModelMetadata]() + private val inMemoryModelIds = ConcurrentHashMap.newKeySet[String]() + private[ml] case class CacheItem(obj: Object, sizeBytes: Long) private[ml] val cachedModel: ConcurrentMap[String, CacheItem] = { if (getMemoryControlEnabled) { CacheBuilder .newBuilder() .softValues() - .removalListener((removed: RemovalNotification[String, CacheItem]) => - totalMLCacheInMemorySizeBytes.addAndGet(-removed.getValue.sizeBytes)) + .removalListener((removed: RemovalNotification[String, CacheItem]) => { + Option(removed.getValue).foreach { value => + totalMLCacheInMemorySizeBytes.addAndGet(-value.sizeBytes) + } + inMemoryModelIds.remove(removed.getKey) + }) .maximumWeight(getMaxInMemoryCacheSizeKB) .weigher((key: String, value: CacheItem) => { Math.ceil(value.sizeBytes.toDouble / 1024).toInt @@ -142,24 +157,32 @@ private[connect] class MLCache(sessionHolder: SessionHolder) extends Logging { if (obj.isInstanceOf[Summary]) { cachedModel.put(objectId, CacheItem(obj, 0)) } else if (obj.isInstanceOf[Model[_]]) { - val sizeBytes = if (getMemoryControlEnabled) { - val _sizeBytes = estimateObjectSize(obj) - checkModelSize(_sizeBytes) - _sizeBytes + val model = obj.asInstanceOf[Model[_]] + val estimatedSizeBytes = if (getMemoryControlEnabled) { + val sizeBytes = estimateObjectSize(model) + checkModelSize(sizeBytes) + Some(sizeBytes) } else { - 0L // Don't need to calculate size if disables memory-control. + // Avoid adding model-size estimation overhead when memory control is disabled. + None } - cachedModel.put(objectId, CacheItem(obj, sizeBytes)) if (getMemoryControlEnabled) { val savePath = getModelOffloadingPath(objectId) - obj.asInstanceOf[MLWritable].write.saveToLocal(savePath.toString) - if (obj.isInstanceOf[HasTrainingSummary[_]] - && obj.asInstanceOf[HasTrainingSummary[_]].hasSummary) { - obj + model.asInstanceOf[MLWritable].write.saveToLocal(savePath.toString) + if (model.isInstanceOf[HasTrainingSummary[_]] + && model.asInstanceOf[HasTrainingSummary[_]].hasSummary) { + model .asInstanceOf[HasTrainingSummary[_]] .saveSummary(savePath.resolve("summary").toString) } - Files.writeString(savePath.resolve(modelClassNameFile), obj.getClass.getName) + Files.writeString(savePath.resolve(modelClassNameFile), model.getClass.getName) + } + cachedModelMetadata.put( + objectId, + ModelMetadata(model.getClass.getName, model.toString, estimatedSizeBytes)) + inMemoryModelIds.add(objectId) + cachedModel.put(objectId, CacheItem(model, estimatedSizeBytes.getOrElse(0L))) + estimatedSizeBytes.foreach { sizeBytes => totalMLCacheInMemorySizeBytes.addAndGet(sizeBytes) totalMLCacheSizeBytes.addAndGet(sizeBytes) } @@ -219,6 +242,7 @@ private[connect] class MLCache(sessionHolder: SessionHolder) extends Logging { loadPath.toString, loadFromLocal = true) val sizeBytes = estimateObjectSize(obj) + inMemoryModelIds.add(refId) cachedModel.put(refId, CacheItem(obj, sizeBytes)) totalMLCacheInMemorySizeBytes.addAndGet(sizeBytes) } @@ -231,8 +255,12 @@ private[connect] class MLCache(sessionHolder: SessionHolder) extends Logging { verifyObjectId(refId) val removedModel = cachedModel.remove(refId) val removedFromMem = removedModel != null - val removedFromDisk = if (!evictOnly && removedModel != null && getMemoryControlEnabled) { - totalMLCacheSizeBytes.addAndGet(-removedModel.sizeBytes) + inMemoryModelIds.remove(refId) + val metadata = Option(cachedModelMetadata.get(refId)) + val removedFromDisk = if (!evictOnly && metadata.nonEmpty && getMemoryControlEnabled) { + metadata.get.estimatedSizeBytes.foreach { sizeBytes => + totalMLCacheSizeBytes.addAndGet(-sizeBytes) + } val removePath = getModelOffloadingPath(refId) val offloadingPath = new File(removePath.toString) if (offloadingPath.exists()) { @@ -244,7 +272,9 @@ private[connect] class MLCache(sessionHolder: SessionHolder) extends Logging { } else { false } - removedFromMem || removedFromDisk + val removeMetadata = !evictOnly || !getMemoryControlEnabled + val removedMetadata = removeMetadata && cachedModelMetadata.remove(refId) != null + removedFromMem || removedFromDisk || removedMetadata } /** @@ -264,6 +294,9 @@ private[connect] class MLCache(sessionHolder: SessionHolder) extends Logging { def clear(): Int = this.synchronized { val size = cachedModel.size() cachedModel.clear() + cachedModelMetadata.clear() + inMemoryModelIds.clear() + totalMLCacheInMemorySizeBytes.set(0) totalMLCacheSizeBytes.set(0) if (getMemoryControlEnabled) { SparkFileUtils.cleanDirectory(new File(offloadedModelsDir.toString)) @@ -280,4 +313,40 @@ private[connect] class MLCache(sessionHolder: SessionHolder) extends Logging { } info.result() } + + /** Returns a cache snapshot without loading or touching any cached model. */ + def getStatus: MLCacheStatus = this.synchronized { + val models = mutable.ArrayBuilder.make[MLCacheModelInfo] + cachedModelMetadata.asScala.foreach { case (id, metadata) => + models += MLCacheModelInfo( + id = id, + className = metadata.className, + modelString = metadata.modelString, + estimatedSizeBytes = metadata.estimatedSizeBytes, + inMemory = inMemoryModelIds.contains(id)) + } + MLCacheStatus( + memoryControlEnabled = getMemoryControlEnabled, + inMemorySizeBytes = totalMLCacheInMemorySizeBytes.get(), + maxInMemorySizeBytes = sessionHolder.session.conf.get( + Connect.CONNECT_SESSION_CONNECT_ML_CACHE_MEMORY_CONTROL_MAX_IN_MEMORY_SIZE), + totalSizeBytes = totalMLCacheSizeBytes.get(), + maxTotalSizeBytes = getMLCacheMaxSize, + models = models.result().toIndexedSeq.sortBy(_.id)) + } } + +private[connect] case class MLCacheModelInfo( + id: String, + className: String, + modelString: String, + estimatedSizeBytes: Option[Long], + inMemory: Boolean) + +private[connect] case class MLCacheStatus( + memoryControlEnabled: Boolean, + inMemorySizeBytes: Long, + maxInMemorySizeBytes: Long, + totalSizeBytes: Long, + maxTotalSizeBytes: Long, + models: Seq[MLCacheModelInfo]) diff --git a/sql/connect/server/src/main/scala/org/apache/spark/sql/connect/service/SessionHolder.scala b/sql/connect/server/src/main/scala/org/apache/spark/sql/connect/service/SessionHolder.scala index 2276230545e67..e97bcb0c786e3 100644 --- a/sql/connect/server/src/main/scala/org/apache/spark/sql/connect/service/SessionHolder.scala +++ b/sql/connect/server/src/main/scala/org/apache/spark/sql/connect/service/SessionHolder.scala @@ -38,7 +38,7 @@ import org.apache.spark.sql.classic.SparkSession import org.apache.spark.sql.connect.IllegalStateErrors import org.apache.spark.sql.connect.common.InvalidPlanInput import org.apache.spark.sql.connect.config.Connect -import org.apache.spark.sql.connect.ml.MLCache +import org.apache.spark.sql.connect.ml.{MLCache, MLCacheStatus} import org.apache.spark.sql.connect.pipelines.DataflowGraphRegistry import org.apache.spark.sql.connect.planner.PythonStreamingQueryListener import org.apache.spark.sql.connect.planner.StreamingForeachBatchHelper @@ -132,7 +132,16 @@ case class SessionHolder(userId: String, sessionId: String, session: SparkSessio new ConcurrentHashMap() // ML model cache - private[connect] lazy val mlCache = new MLCache(this) + @volatile private var mlCacheInitialized = false + private[connect] lazy val mlCache = { + val cache = new MLCache(this) + mlCacheInitialized = true + cache + } + + private[connect] def getMLCacheStatus: Option[MLCacheStatus] = { + if (mlCacheInitialized) Some(mlCache.getStatus) else None + } // Mapping from id to StreamingQueryListener. Used for methods like removeListener() in // StreamingQueryManager. diff --git a/sql/connect/server/src/main/scala/org/apache/spark/sql/connect/service/SparkConnectService.scala b/sql/connect/server/src/main/scala/org/apache/spark/sql/connect/service/SparkConnectService.scala index c76794e3b6ec1..0c18a7491c1f2 100644 --- a/sql/connect/server/src/main/scala/org/apache/spark/sql/connect/service/SparkConnectService.scala +++ b/sql/connect/server/src/main/scala/org/apache/spark/sql/connect/service/SparkConnectService.scala @@ -386,7 +386,8 @@ object SparkConnectService extends Logging { Some( new SparkConnectServerTab( new SparkConnectServerAppStatusStore(kvStore), - SparkConnectServerTab.getSparkUI(sc))) + SparkConnectServerTab.getSparkUI(sc), + Some(sessionManager))) } else { None } diff --git a/sql/connect/server/src/main/scala/org/apache/spark/sql/connect/service/SparkConnectSessionManager.scala b/sql/connect/server/src/main/scala/org/apache/spark/sql/connect/service/SparkConnectSessionManager.scala index d3ddf592e9e7d..ca124ba9693fa 100644 --- a/sql/connect/server/src/main/scala/org/apache/spark/sql/connect/service/SparkConnectSessionManager.scala +++ b/sql/connect/server/src/main/scala/org/apache/spark/sql/connect/service/SparkConnectSessionManager.scala @@ -32,6 +32,7 @@ import org.apache.spark.internal.Logging import org.apache.spark.internal.LogKeys.{INTERVAL, SESSION_HOLD_INFO} import org.apache.spark.sql.classic.SparkSession import org.apache.spark.sql.connect.config.Connect.{CONNECT_SESSION_MANAGER_CLOSED_SESSIONS_TOMBSTONES_SIZE, CONNECT_SESSION_MANAGER_DEFAULT_SESSION_TIMEOUT, CONNECT_SESSION_MANAGER_MAINTENANCE_INTERVAL} +import org.apache.spark.sql.connect.ml.MLCacheStatus import org.apache.spark.util.ThreadUtils /** @@ -284,6 +285,18 @@ class SparkConnectSessionManager extends Logging { closedSessionsCache.asMap.asScala.values.toSeq } + // Read live cache state directly without updating the sessions' last-access times. + private[connect] def getMLCacheStatuses: Seq[(SessionKey, MLCacheStatus)] = { + sessionStore + .entrySet() + .asScala + .flatMap { entry => + entry.getValue.getMLCacheStatus.map(entry.getKey -> _) + } + .toSeq + .sortBy { case (key, _) => (key.userId, key.sessionId) } + } + /** * Schedules periodic maintenance checks if it is not already scheduled. * diff --git a/sql/connect/server/src/main/scala/org/apache/spark/sql/connect/ui/SparkConnectServerPage.scala b/sql/connect/server/src/main/scala/org/apache/spark/sql/connect/ui/SparkConnectServerPage.scala index ea78b2dc59f6a..b3af78f62e361 100644 --- a/sql/connect/server/src/main/scala/org/apache/spark/sql/connect/ui/SparkConnectServerPage.scala +++ b/sql/connect/server/src/main/scala/org/apache/spark/sql/connect/ui/SparkConnectServerPage.scala @@ -26,6 +26,7 @@ import scala.xml.Node import jakarta.servlet.http.HttpServletRequest import org.apache.spark.internal.Logging +import org.apache.spark.sql.connect.ml.MLCacheModelInfo import org.apache.spark.sql.connect.ui.ToolTips._ import org.apache.spark.ui._ import org.apache.spark.ui.UIUtils._ @@ -63,7 +64,8 @@ private[ui] class SparkConnectServerPage(parent: SparkConnectServerTab) Request(s) ++ generateSessionStatsTable(request) ++ - generateSQLStatsTable(request) + generateSQLStatsTable(request) ++ + generateMLCacheStatsTable(request) } UIUtils.headerSparkPage(request, "Spark Connect", content, parent) } @@ -179,6 +181,80 @@ private[ui] class SparkConnectServerPage(parent: SparkConnectServerTab) content } + + /** Generate live ML cache statistics for active Spark Connect sessions. */ + private def generateMLCacheStatsTable(request: HttpServletRequest): Seq[Node] = { + val cacheStatuses = parent.getMLCacheStatuses.filter(_._2.models.nonEmpty) + val models = cacheStatuses.flatMap { case (key, status) => + status.models.map(MLCacheModelTableRow(key.userId, key.sessionId, _)) + } + if (models.isEmpty) { + return Seq.empty + } + + val tableTag = "mlcachemodels" + val tablePage = Option(request.getParameter(s"$tableTag.page")).map(_.toInt).getOrElse(1) + val table = + try { + new MLCacheModelStatsPagedTable( + request, + parent, + models, + "connect", + UIUtils.prependBaseUri(request, parent.basePath), + tableTag).table(tablePage) + } catch { + case e @ (_: IllegalArgumentException | _: IndexOutOfBoundsException) => +
+

Error while rendering ML cache table:

+ 
+            {Utils.exceptionString(e)}
+
+
+ } + + val inMemoryModels = models.count(_.model.inMemory) + val memoryControlledStatuses = cacheStatuses.map(_._2).filter(_.memoryControlEnabled) + val inMemorySize = memoryControlledStatuses.map(s => BigInt(s.inMemorySizeBytes)).sum + val maxInMemorySize = memoryControlledStatuses.map(s => BigInt(s.maxInMemorySizeBytes)).sum + val totalSize = memoryControlledStatuses.map(s => BigInt(s.totalSizeBytes)).sum + val maxTotalSize = memoryControlledStatuses.map(s => BigInt(s.maxTotalSizeBytes)).sum + val sizeStats = if (memoryControlledStatuses.nonEmpty) { + Seq( +
  • + Estimated size (In-memory): + {Utils.bytesToString(inMemorySize)} / {Utils.bytesToString(maxInMemorySize)} +
    • , +
  • + Estimated size (In-memory and Offloaded data): + {Utils.bytesToString(totalSize)} / {Utils.bytesToString(maxTotalSize)} +
    • ) + } else { + Seq.empty + } + + +

    + + ML Cache Statistics ({models.size}) +

    +     ++ +
    + +
    Cached Models
    + {table} +
    + } } private[ui] class SqlStatsPagedTable( @@ -435,6 +511,114 @@ private[ui] class SessionStatsPagedTable( } } +private[ui] case class MLCacheModelTableRow( + userId: String, + sessionId: String, + model: MLCacheModelInfo) + +private[ui] class MLCacheModelStatsPagedTable( + request: HttpServletRequest, + parent: SparkConnectServerTab, + data: Seq[MLCacheModelTableRow], + subPath: String, + basePath: String, + tableTag: String) + extends PagedTable[MLCacheModelTableRow] { + + private val (sortColumn, desc, pageSize) = + getTableParameters(request, tableTag, "Estimated Size") + + private val encodedSortColumn = URLEncoder.encode(sortColumn, UTF_8.name()) + private val parameterPath = s"$basePath/$subPath/?${getParameterOtherTable(request, tableTag)}" + + override val dataSource = + new MLCacheModelTableDataSource(data, pageSize, sortColumn, desc) + + override def tableId: String = tableTag + + override def tableCssClass: String = + "table table-bordered table-sm table-striped table-head-clickable table-cell-width-limited" + + override def pageLink(page: Int): String = { + parameterPath + + s"&$pageNumberFormField=$page" + + s"&$tableTag.sort=$encodedSortColumn" + + s"&$tableTag.desc=$desc" + + s"&$pageSizeFormField=$pageSize" + + s"#$tableTag" + } + + override def pageSizeFormField: String = s"$tableTag.pageSize" + + override def pageNumberFormField: String = s"$tableTag.page" + + override def goButtonFormPath: String = + s"$parameterPath&$tableTag.sort=$encodedSortColumn" + + s"&$tableTag.desc=$desc#$tableTag" + + override def headers: Seq[Node] = { + val headersAndTooltips: Seq[(String, Boolean, Option[String])] = Seq( + ("User", true, None), + ("Session ID", true, None), + ("Model ID", true, None), + ("Model Class", true, None), + ("Model Details", true, Some(SPARK_CONNECT_ML_CACHE_MODEL_DETAILS)), + ("Estimated Size", true, Some(SPARK_CONNECT_ML_CACHE_ESTIMATED_SIZE)), + ("Storage", true, Some(SPARK_CONNECT_ML_CACHE_STORAGE))) + + isSortColumnValid(headersAndTooltips, sortColumn) + headerRow(headersAndTooltips, desc, pageSize, sortColumn, parameterPath, tableTag, tableTag) + } + + override def row(row: MLCacheModelTableRow): Seq[Node] = { + val model = row.model + val sessionLink = "%s/%s/session/?id=%s&userId=%s".format( + UIUtils.prependBaseUri(request, parent.basePath), + parent.prefix, + URLEncoder.encode(row.sessionId, UTF_8.name()), + ConnectUiUtils.encodeUserId(row.userId)) + + {row.userId} + {row.sessionId} + {model.id} + {model.className} + {model.modelString} + {model.estimatedSizeBytes.map(Utils.bytesToString).getOrElse("N/A")} + {if (model.inMemory) "In memory" else "Offloaded"} + + } +} + +private[ui] class MLCacheModelTableDataSource( + info: Seq[MLCacheModelTableRow], + pageSize: Int, + sortColumn: String, + desc: Boolean) + extends PagedDataSource[MLCacheModelTableRow](pageSize) { + + private val data = info.sorted(ordering(sortColumn, desc)) + + override def dataSize: Int = data.size + + override def sliceData(from: Int, to: Int): Seq[MLCacheModelTableRow] = data.slice(from, to) + + private def ordering(sortColumn: String, desc: Boolean): Ordering[MLCacheModelTableRow] = { + val ordering: Ordering[MLCacheModelTableRow] = sortColumn match { + case "User" => Ordering.by(_.userId) + case "Session ID" => Ordering.by(_.sessionId) + case "Model ID" => Ordering.by(_.model.id) + case "Model Class" => Ordering.by(_.model.className) + case "Model Details" => Ordering.by(_.model.modelString) + case "Estimated Size" => + Ordering.by((row: MLCacheModelTableRow) => + (row.model.estimatedSizeBytes.isDefined, row.model.estimatedSizeBytes.getOrElse(0L))) + case "Storage" => Ordering.by(_.model.inMemory) + case unknownColumn => throw new IllegalArgumentException(s"Unknown column: $unknownColumn") + } + if (desc) ordering.reverse else ordering + } +} + private[ui] class SqlStatsTableRow( val jobTag: String, val jobId: Seq[String], diff --git a/sql/connect/server/src/main/scala/org/apache/spark/sql/connect/ui/SparkConnectServerTab.scala b/sql/connect/server/src/main/scala/org/apache/spark/sql/connect/ui/SparkConnectServerTab.scala index c5ea0bf618b52..d720acb31fee2 100644 --- a/sql/connect/server/src/main/scala/org/apache/spark/sql/connect/ui/SparkConnectServerTab.scala +++ b/sql/connect/server/src/main/scala/org/apache/spark/sql/connect/ui/SparkConnectServerTab.scala @@ -21,12 +21,15 @@ import java.util.Date import org.apache.spark.SparkContext import org.apache.spark.internal.Logging +import org.apache.spark.sql.connect.ml.MLCacheStatus +import org.apache.spark.sql.connect.service.{SessionKey, SparkConnectSessionManager} import org.apache.spark.sql.errors.QueryExecutionErrors import org.apache.spark.ui.{SparkUI, SparkUITab} private[connect] class SparkConnectServerTab( val store: SparkConnectServerAppStatusStore, - sparkUI: SparkUI) + sparkUI: SparkUI, + sessionManager: Option[SparkConnectSessionManager] = None) extends SparkUITab(sparkUI, "connect") with Logging { @@ -47,6 +50,10 @@ private[connect] class SparkConnectServerTab( parent.detachTab(this) } + def getMLCacheStatuses: Seq[(SessionKey, MLCacheStatus)] = { + sessionManager.toSeq.flatMap(_.getMLCacheStatuses) + } + override def displayOrder: Int = 3 } diff --git a/sql/connect/server/src/main/scala/org/apache/spark/sql/connect/ui/ToolTips.scala b/sql/connect/server/src/main/scala/org/apache/spark/sql/connect/ui/ToolTips.scala index 9b51ace83c6c1..be2ab4c4903b2 100644 --- a/sql/connect/server/src/main/scala/org/apache/spark/sql/connect/ui/ToolTips.scala +++ b/sql/connect/server/src/main/scala/org/apache/spark/sql/connect/ui/ToolTips.scala @@ -36,4 +36,13 @@ private[ui] object ToolTips { val SPARK_CONNECT_SESSION_DURATION = "Elapsed time since session start, or until closed if the session was closed" + val SPARK_CONNECT_ML_CACHE_ESTIMATED_SIZE = + "Approximate model size recorded when it was added to the Spark Connect ML cache" + + val SPARK_CONNECT_ML_CACHE_MODEL_DETAILS = + "Output of model.toString recorded when the model was added to the cache" + + val SPARK_CONNECT_ML_CACHE_STORAGE = + "Whether the model is currently in driver memory or offloaded to driver-local disk" + } diff --git a/sql/connect/server/src/test/scala/org/apache/spark/sql/connect/ml/MLSuite.scala b/sql/connect/server/src/test/scala/org/apache/spark/sql/connect/ml/MLSuite.scala index 9ba5a499ba8fe..b78423e8687e9 100644 --- a/sql/connect/server/src/test/scala/org/apache/spark/sql/connect/ml/MLSuite.scala +++ b/sql/connect/server/src/test/scala/org/apache/spark/sql/connect/ml/MLSuite.scala @@ -387,6 +387,34 @@ class MLSuite extends MLHelper { } } + test("MLCache status") { + val sessionHolder = SparkConnectTestUtils.createDummySessionHolder(spark) + sessionHolder.session.conf + .set(Connect.CONNECT_SESSION_CONNECT_ML_CACHE_MEMORY_CONTROL_ENABLED.key, "true") + sessionHolder.session.conf + .set(Connect.CONNECT_SESSION_CONNECT_ML_CACHE_MEMORY_CONTROL_MAX_IN_MEMORY_SIZE.key, 16384) + sessionHolder.session.conf + .set(Connect.CONNECT_SESSION_CONNECT_ML_CACHE_MEMORY_CONTROL_MAX_STORAGE_SIZE.key, 65536) + + // Reading UI status should not initialize an unused cache. + assert(sessionHolder.getMLCacheStatus.isEmpty) + + val modelId = trainLogisticRegressionModel(sessionHolder) + val status = sessionHolder.getMLCacheStatus.get + assert(status.memoryControlEnabled) + assert(status.inMemorySizeBytes > 0) + assert(status.maxInMemorySizeBytes === 16384) + assert(status.totalSizeBytes === status.inMemorySizeBytes) + assert(status.maxTotalSizeBytes === 65536) + assert(status.models.size === 1) + val modelInfo = status.models.head + assert(modelInfo.id === modelId) + assert(modelInfo.className === classOf[LogisticRegressionModel].getName) + assert(modelInfo.modelString.startsWith("LogisticRegressionModel: uid=")) + assert(modelInfo.estimatedSizeBytes.contains(status.totalSizeBytes)) + assert(modelInfo.inMemory) + } + test("MLCache offloading works") { val sessionHolder = SparkConnectTestUtils.createDummySessionHolder(spark) sessionHolder.session.conf @@ -420,10 +448,24 @@ class MLSuite extends MLHelper { assert(sessionHolder.mlCache.totalMLCacheInMemorySizeBytes.get() <= memorySizeBytes) } + val status = sessionHolder.getMLCacheStatus.get + assert(status.models.size === modelIdList.size) + assert(status.models.count(_.inMemory) === maxNumModels) + assert(status.models.map(_.id).toSet === modelIdList.toSet) + // Assert all models can be loaded back from disk after they are offloaded. for (modelId <- modelIdList) { assert(sessionHolder.mlCache.get(modelId) != null) } + + val statusBeforeRemove = sessionHolder.getMLCacheStatus.get + val offloadedModel = statusBeforeRemove.models.find(!_.inMemory).get + assert(sessionHolder.mlCache.remove(offloadedModel.id)) + val statusAfterRemove = sessionHolder.getMLCacheStatus.get + assert(!statusAfterRemove.models.exists(_.id == offloadedModel.id)) + assert( + statusAfterRemove.totalSizeBytes === + statusBeforeRemove.totalSizeBytes - offloadedModel.estimatedSizeBytes.get) } test("Model size limit") { @@ -461,6 +503,7 @@ class MLSuite extends MLHelper { assert(mlCache2.get(modelId) != null) mlCache2.close() assert(mlCache2.cachedModel.isEmpty) + assert(mlCache2.getStatus.models.isEmpty) // Test 3: Edge case - register then remove model, close should still run cleanup val edgeCaseSessionHolder = SparkConnectTestUtils.createDummySessionHolder(spark) diff --git a/sql/connect/server/src/test/scala/org/apache/spark/sql/connect/ui/SparkConnectServerPageSuite.scala b/sql/connect/server/src/test/scala/org/apache/spark/sql/connect/ui/SparkConnectServerPageSuite.scala index 7f6af17bc41b3..8def63046cecf 100644 --- a/sql/connect/server/src/test/scala/org/apache/spark/sql/connect/ui/SparkConnectServerPageSuite.scala +++ b/sql/connect/server/src/test/scala/org/apache/spark/sql/connect/ui/SparkConnectServerPageSuite.scala @@ -26,6 +26,7 @@ import org.scalatest.BeforeAndAfter import org.apache.spark.{SharedSparkContext, SparkConf, SparkFunSuite} import org.apache.spark.scheduler.SparkListenerJobStart +import org.apache.spark.sql.connect.ml.{MLCacheModelInfo, MLCacheStatus} import org.apache.spark.sql.connect.service._ import org.apache.spark.status.ElementTrackingStore import org.apache.spark.util.kvstore.InMemoryStore @@ -47,7 +48,7 @@ class SparkConnectServerPageSuite /** * Run a dummy session and return the store */ - private def getStatusStore: SparkConnectServerAppStatusStore = { + private def getStatusStore(closeSession: Boolean = true): SparkConnectServerAppStatusStore = { kvstore = new ElementTrackingStore(new InMemoryStore, new SparkConf()) // val server = mock(classOf[SparkConnectServer], RETURNS_SMART_NULLS) val sparkConf = new SparkConf @@ -74,14 +75,16 @@ class SparkConnectServerPageSuite SparkListenerConnectOperationFinished("jobTag", "operationId", System.currentTimeMillis())) listener.onOtherEvent( SparkListenerConnectOperationClosed("jobTag", "operationId", System.currentTimeMillis())) - listener.onOtherEvent( - SparkListenerConnectSessionClosed("sessionId", "userId", System.currentTimeMillis())) + if (closeSession) { + listener.onOtherEvent( + SparkListenerConnectSessionClosed("sessionId", "userId", System.currentTimeMillis())) + } statusStore } test("Spark Connect Server page should load successfully") { - val store = getStatusStore + val store = getStatusStore() val request = mock(classOf[HttpServletRequest]) val tab = mock(classOf[SparkConnectServerTab], RETURNS_SMART_NULLS) @@ -89,6 +92,7 @@ class SparkConnectServerPageSuite when(tab.store).thenReturn(store) when(tab.appName).thenReturn("testing") when(tab.headerTabs).thenReturn(Seq.empty) + when(tab.getMLCacheStatuses).thenReturn(Seq.empty) val page = new SparkConnectServerPage(tab) val html = page.render(request).toString().toLowerCase(Locale.ROOT) @@ -96,6 +100,7 @@ class SparkConnectServerPageSuite assert(html.contains("session statistics (1)")) assert(html.contains("request statistics (1)")) assert(html.contains("dummy query")) + assert(!html.contains("ml cache statistics")) // Pagination support assert(html.contains("")) @@ -107,7 +112,7 @@ class SparkConnectServerPageSuite } test("Spark Connect Server session page should load successfully") { - val store = getStatusStore + val store = getStatusStore() val request = mock(classOf[HttpServletRequest]) when(request.getParameter("id")).thenReturn("sessionId") @@ -134,6 +139,56 @@ class SparkConnectServerPageSuite " data-bs-target=\"#aggregated-sqlsessionstat\"")) } + test("Spark Connect Server page should show live ML cache statistics and model details") { + val store = getStatusStore(closeSession = false) + + val request = mock(classOf[HttpServletRequest]) + val tab = mock(classOf[SparkConnectServerTab], RETURNS_SMART_NULLS) + when(tab.startTime).thenReturn(Calendar.getInstance().getTime) + when(tab.store).thenReturn(store) + when(tab.appName).thenReturn("testing") + when(tab.headerTabs).thenReturn(Seq.empty) + when(tab.getMLCacheStatuses).thenReturn( + Seq( + SessionKey("userId", "sessionId") -> + MLCacheStatus( + memoryControlEnabled = true, + inMemorySizeBytes = 1024, + maxInMemorySizeBytes = 4096, + totalSizeBytes = 2048, + maxTotalSizeBytes = 8192, + models = Seq( + MLCacheModelInfo( + id = "model-id-1", + className = "org.apache.spark.ml.classification.LogisticRegressionModel", + modelString = "LogisticRegressionModel: uid=logreg-1", + estimatedSizeBytes = Some(1024), + inMemory = true), + MLCacheModelInfo( + id = "model-id-2", + className = "org.apache.spark.ml.classification.LogisticRegressionModel", + modelString = "LogisticRegressionModel: uid=logreg-2", + estimatedSizeBytes = Some(1024), + inMemory = false))))) + + val page = new SparkConnectServerPage(tab) + val html = page.render(request).toString().toLowerCase(Locale.ROOT) + + val sessionStatsIndex = html.indexOf("session statistics") + val mlCacheStatsIndex = html.indexOf("ml cache statistics (2)") + val requestStatsIndex = html.indexOf("request statistics") + assert(sessionStatsIndex < requestStatsIndex && requestStatsIndex < mlCacheStatsIndex) + assert(html.contains("2 (1 in memory, 1 offloaded)")) + assert(html.contains("estimated size (in-memory)")) + assert(html.contains("1024.0 b / 4.0 kib")) + assert(html.contains("estimated size (in-memory and offloaded data)")) + assert(html.contains("2.0 kib / 8.0 kib")) + assert(html.contains("model-id-1")) + assert(html.contains("logisticregressionmodel: uid=logreg-1")) + assert(html.contains("in memory")) + assert(html.contains("offloaded")) + } + test("SPARK-58097: session page only shows the requested user's operations") { // Two users share the same session UUID, each running a distinct query. kvstore = new ElementTrackingStore(new InMemoryStore, new SparkConf())