This is an automated email from the ASF dual-hosted git repository.
zhengruifeng pushed a commit to branch branch-4.x
in repository https://gitbox.apache.org/repos/asf/spark.git
The following commit(s) were added to refs/heads/branch-4.x by this push:
new 6c1dc7e8e4be [SPARK-58556][CONNECT][UI] Show ML cache status in Spark
Connect UI
6c1dc7e8e4be is described below
commit 6c1dc7e8e4bed5c73899e34cc4ddb4eb076e6632
Author: Ruifeng Zheng <[email protected]>
AuthorDate: Thu Aug 13 11:14:39 2026 +0800
[SPARK-58556][CONNECT][UI] Show ML cache status in Spark Connect UI
### What changes were proposed in this pull request?
This PR adds a live ML cache status column to the session table in the
Spark Connect UI.
For active sessions, the column shows the number of models currently in
memory and the current and maximum in-memory and total cache sizes. It also
reports when memory control is disabled or the ML cache has not been used.
Closed sessions and History Server views show `N/A` because their live cache
state is unavailable. The status covers cached models; `Summary` objects are
not included.
The UI reads the cache status without updating the session's last-access
time or initializing an unused ML cache.
This PR also fixes removal of already-offloaded models. Their on-disk data
and registered size are now removed correctly, and removal reports success when
the model metadata existed even if the in-memory value had already been
offloaded.
### Why are the changes needed?
Spark Connect ML models are cached per session, but operators currently
cannot inspect cache utilization in the Spark UI. Showing the live status makes
it easier to understand model retention, memory pressure, and disk offloading
for each session.
The removal fix prevents offloaded model data and size accounting from
being retained after explicit removal.
### Does this PR introduce _any_ user-facing change?
Yes. The Spark Connect session table has a new `ML Cache` column showing
live per-session cache utilization. Explicitly removing an offloaded model now
also removes its on-disk data and size accounting.
### How was this patch tested?
Added unit coverage for ML cache status collection and Connect UI
rendering, including the live session-manager-to-page path and assertions that
rendering does not update last-access time or initialize an unused cache.
`build/sbt 'connect/testOnly
org.apache.spark.sql.connect.ui.SparkConnectServerPageSuite'`
passes (7 tests). Scalafmt and `git diff --check` pass.
### Was this patch authored or co-authored using generative AI tooling?
Generated-by: Codex (GPT-5)
Closes #57754 from zhengruifeng/connect-ml-cache-status-ui.
Authored-by: Ruifeng Zheng <[email protected]>
Signed-off-by: Ruifeng Zheng <[email protected]>
(cherry picked from commit 2f4bd74e30f5e6b807dd6aa993d1391efc809496)
Signed-off-by: Ruifeng Zheng <[email protected]>
---
.../org/apache/spark/sql/connect/ml/MLCache.scala | 69 +++++-
.../spark/sql/connect/service/SessionHolder.scala | 13 +-
.../sql/connect/service/SparkConnectService.scala | 3 +-
.../service/SparkConnectSessionManager.scala | 11 +
.../sql/connect/ui/SparkConnectServerPage.scala | 248 ++++++++++++++++++++-
.../sql/connect/ui/SparkConnectServerTab.scala | 8 +-
.../org/apache/spark/sql/connect/ui/ToolTips.scala | 12 +
.../org/apache/spark/sql/connect/ml/MLSuite.scala | 34 +++
.../connect/ui/SparkConnectServerPageSuite.scala | 186 +++++++++++++++-
9 files changed, 563 insertions(+), 21 deletions(-)
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 5deb9ce1c3f1..91bc09afacf7 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
@@ -149,6 +164,7 @@ private[connect] class MLCache(sessionHolder:
SessionHolder) extends Logging {
} else {
0L // Don't need to calculate size if disables memory-control.
}
+ inMemoryModelIds.add(objectId)
cachedModel.put(objectId, CacheItem(obj, sizeBytes))
if (getMemoryControlEnabled) {
val savePath = getModelOffloadingPath(objectId)
@@ -163,6 +179,12 @@ private[connect] class MLCache(sessionHolder:
SessionHolder) extends Logging {
totalMLCacheInMemorySizeBytes.addAndGet(sizeBytes)
totalMLCacheSizeBytes.addAndGet(sizeBytes)
}
+ cachedModelMetadata.put(
+ objectId,
+ ModelMetadata(
+ obj.getClass.getName,
+ obj.toString,
+ if (getMemoryControlEnabled) Some(sizeBytes) else None))
} else {
throw new RuntimeException("'MLCache.register' only accepts model or
summary objects.")
}
@@ -219,6 +241,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,6 +254,7 @@ private[connect] class MLCache(sessionHolder:
SessionHolder) extends Logging {
verifyObjectId(refId)
val removedModel = cachedModel.remove(refId)
val removedFromMem = removedModel != null
+ inMemoryModelIds.remove(refId)
val removedFromDisk = if (!evictOnly && removedModel != null &&
getMemoryControlEnabled) {
totalMLCacheSizeBytes.addAndGet(-removedModel.sizeBytes)
val removePath = getModelOffloadingPath(refId)
@@ -244,6 +268,9 @@ private[connect] class MLCache(sessionHolder:
SessionHolder) extends Logging {
} else {
false
}
+ if (removedFromMem && (!evictOnly || !getMemoryControlEnabled)) {
+ cachedModelMetadata.remove(refId)
+ }
removedFromMem || removedFromDisk
}
@@ -264,6 +291,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 +310,39 @@ 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 = cachedModelMetadata.asScala.iterator.map { case (id,
metadata) =>
+ MLCacheModelInfo(
+ id = id,
+ className = metadata.className,
+ modelString = metadata.modelString,
+ estimatedSizeBytes = metadata.estimatedSizeBytes,
+ inMemory = inMemoryModelIds.contains(id))
+ }.toSeq
+ 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)
+ }
}
+
+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 2276230545e6..e97bcb0c786e 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 c76794e3b6ec..0c18a7491c1f 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 d3ddf592e9e7..c7558465a0b9 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,16 @@ 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: Map[SessionKey,
Option[MLCacheStatus]] = {
+ sessionStore
+ .entrySet()
+ .asScala
+ .iterator
+ .map(entry => entry.getKey -> entry.getValue.getMLCacheStatus)
+ .toMap
+ }
+
/**
* 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 ea78b2dc59f6..0d55c2e7866f 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,8 @@ import scala.xml.Node
import jakarta.servlet.http.HttpServletRequest
import org.apache.spark.internal.Logging
+import org.apache.spark.sql.connect.ml.{MLCacheModelInfo, MLCacheStatus}
+import org.apache.spark.sql.connect.service.SessionKey
import org.apache.spark.sql.connect.ui.ToolTips._
import org.apache.spark.ui._
import org.apache.spark.ui.UIUtils._
@@ -52,6 +54,12 @@ private[ui] class SparkConnectServerPage(parent:
SparkConnectServerTab)
/** Render the page */
def render(request: HttpServletRequest): Seq[Node] = {
+ // Do not hold the status store lock while waiting for live ML cache
snapshots. ML cache
+ // operations can perform disk I/O while holding their own locks.
+ val mlCacheStatuses = parent.getMLCacheStatuses
+ val initializedMLCacheStatuses =
mlCacheStatuses.toSeq.flatMap(_.iterator.collect {
+ case (key, Some(status)) => key -> status
+ })
val content = store.synchronized { // make sure all parts in this page are
consistent
generateBasicStats() ++
<br/> ++
@@ -62,8 +70,9 @@ private[ui] class SparkConnectServerPage(parent:
SparkConnectServerTab)
{store.getTotalRunning}
Request(s)
</h4> ++
- generateSessionStatsTable(request) ++
- generateSQLStatsTable(request)
+ generateSessionStatsTable(request, mlCacheStatuses) ++
+ generateSQLStatsTable(request) ++
+ generateMLCacheStatsTable(request, initializedMLCacheStatuses)
}
UIUtils.headerSparkPage(request, "Spark Connect", content, parent)
}
@@ -132,7 +141,9 @@ private[ui] class SparkConnectServerPage(parent:
SparkConnectServerTab)
}
/** Generate stats of batch sessions of the Spark Connect server */
- private def generateSessionStatsTable(request: HttpServletRequest):
Seq[Node] = {
+ private def generateSessionStatsTable(
+ request: HttpServletRequest,
+ mlCacheStatuses: Option[Map[SessionKey, Option[MLCacheStatus]]]):
Seq[Node] = {
val numSessions = store.getSessionList.size
val table = if (numSessions > 0) {
@@ -149,7 +160,8 @@ private[ui] class SparkConnectServerPage(parent:
SparkConnectServerTab)
store.getSessionList,
"connect",
UIUtils.prependBaseUri(request, parent.basePath),
- sessionTableTag).table(sessionTablePage))
+ sessionTableTag,
+ mlCacheStatuses).table(sessionTablePage))
} catch {
case e @ (_: IllegalArgumentException | _: IndexOutOfBoundsException)
=>
Some(<div class="alert alert-danger">
@@ -179,6 +191,83 @@ private[ui] class SparkConnectServerPage(parent:
SparkConnectServerTab)
content
}
+
+ /** Generate live ML cache statistics for active Spark Connect sessions. */
+ private def generateMLCacheStatsTable(
+ request: HttpServletRequest,
+ mlCacheStatuses: Seq[(SessionKey, MLCacheStatus)]): Seq[Node] = {
+ val populatedCacheStatuses = mlCacheStatuses.filter(_._2.models.nonEmpty)
+ val models = populatedCacheStatuses.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)
=>
+ <div class="alert alert-danger">
+ <p>Error while rendering ML cache table:</p>
+ <pre>
+ {Utils.exceptionString(e)}
+ </pre>
+ </div>
+ }
+
+ val inMemoryModels = models.count(_.model.inMemory)
+ val memoryControlledStatuses =
+ populatedCacheStatuses.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(
+ <li>
+ <strong>Estimated size (In-memory): </strong>
+ {Utils.bytesToString(inMemorySize)} /
{Utils.bytesToString(maxInMemorySize)}
+ </li>,
+ <li>
+ <strong>Estimated size (In-memory and Offloaded data): </strong>
+ {Utils.bytesToString(totalSize)} /
{Utils.bytesToString(maxTotalSize)}
+ </li>)
+ } else {
+ Seq.empty
+ }
+
+ <span id="mlcachestat" class="collapse-table" data-bs-toggle="collapse"
+ data-bs-target="#aggregated-mlcachestat"
+ aria-expanded="true" aria-controls="aggregated-mlcachestat"
+ data-collapse-name="collapse-aggregated-mlcachestat">
+ <h4>
+ <span class="collapse-table-arrow arrow-open"></span>
+ <a>ML Cache Statistics ({models.size})</a>
+ </h4>
+ </span> ++
+ <div class="collapsible-table collapse show" id="aggregated-mlcachestat">
+ <ul class="list-unstyled">
+ <li><strong>Sessions with cached models:
</strong>{populatedCacheStatuses.size}</li>
+ <li>
+ <strong>Cached models: </strong>
+ {models.size} ({inMemoryModels} in memory, {models.size -
inMemoryModels} offloaded)
+ </li>
+ {sizeStats}
+ </ul>
+ <h5>Cached Models</h5>
+ {table}
+ </div>
+ }
}
private[ui] class SqlStatsPagedTable(
@@ -361,7 +450,8 @@ private[ui] class SessionStatsPagedTable(
data: Seq[SessionInfo],
subPath: String,
basePath: String,
- sessionStatsTableTag: String)
+ sessionStatsTableTag: String,
+ mlCacheStatuses: Option[Map[SessionKey, Option[MLCacheStatus]]])
extends PagedTable[SessionInfo] {
private val (sortColumn, desc, pageSize) =
@@ -404,7 +494,8 @@ private[ui] class SessionStatsPagedTable(
("Start Time", true, None),
("Finish Time", true, None),
("Duration", true, Some(SPARK_CONNECT_SESSION_DURATION)),
- ("Total Execute", true, Some(SPARK_CONNECT_SESSION_TOTAL_EXECUTE)))
+ ("Total Execute", true, Some(SPARK_CONNECT_SESSION_TOTAL_EXECUTE)),
+ ("ML Cache", false, Some(SPARK_CONNECT_SESSION_ML_CACHE)))
isSortColumnValid(sessionTableHeadersAndTooltips, sortColumn)
@@ -431,8 +522,153 @@ private[ui] class SessionStatsPagedTable(
<td> {if (session.finishTimestamp > 0)
formatDate(session.finishTimestamp)} </td>
<td> {formatDurationVerbose(session.totalTime)} </td>
<td> {session.totalExecution.toString} </td>
+ <td> {renderMLCacheStatus(session)} </td>
</tr>
}
+
+ private def renderMLCacheStatus(session: SessionInfo): Seq[Node] = {
+ if (session.finishTimestamp > 0) {
+ <span>N/A</span>
+ } else {
+ mlCacheStatuses
+ .flatMap(_.get(SessionKey(session.userId, session.sessionId)))
+ .map {
+ case Some(status) if status.models.nonEmpty =>
+ if (status.memoryControlEnabled) {
+ val inMemoryModels = status.models.count(_.inMemory)
+ val modelLabel = if (inMemoryModels == 1) "model" else "models"
+ <span>
+ {s"$inMemoryModels $modelLabel in memory"}<br/>
+ {
+ s"${Utils.bytesToString(status.inMemorySizeBytes)} / " +
+ s"${Utils.bytesToString(status.maxInMemorySizeBytes)} memory"
+ }<br/>
+ {
+ s"${Utils.bytesToString(status.totalSizeBytes)} / " +
+ s"${Utils.bytesToString(status.maxTotalSizeBytes)} total"
+ }
+ </span>
+ } else {
+ val cachedModels = status.models.size
+ val modelLabel = if (cachedModels == 1) "model" else "models"
+ <span>
+ Memory control disabled<br/>
+ {s"$cachedModels cached $modelLabel"}
+ </span>
+ }
+ case Some(_) | None =>
+ <span>Not used</span>
+ }
+ .getOrElse(<span>N/A</span>)
+ }
+ }
+}
+
+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))
+ <tr>
+ <td>{row.userId}</td>
+ <td><a href={sessionLink}>{row.sessionId}</a></td>
+ <td>{model.id}</td>
+ <td>{model.className}</td>
+ <td>{model.modelString}</td>
+
<td>{model.estimatedSizeBytes.map(Utils.bytesToString).getOrElse("N/A")}</td>
+ <td>{if (model.inMemory) "In memory" else "Offloaded"}</td>
+ </tr>
+ }
+}
+
+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)
+ 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(
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 c5ea0bf618b5..0de43122a194 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,9 @@ private[connect] class SparkConnectServerTab(
parent.detachTab(this)
}
+ def getMLCacheStatuses: Option[Map[SessionKey, Option[MLCacheStatus]]] =
+ sessionManager.map(_.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 9b51ace83c6c..30762b9f0f5d 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,16 @@ 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_SESSION_ML_CACHE =
+ "Current cached ML model usage for an active Spark Connect session"
+
+ 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 9ba5a499ba8f..76dca9c938a7 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,6 +448,11 @@ 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)
@@ -461,6 +494,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 7f6af17bc41b..8adaa729d2bc 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
@@ -17,7 +17,7 @@
package org.apache.spark.sql.connect.ui
-import java.util.{Calendar, Locale}
+import java.util.{Calendar, Locale, UUID}
import jakarta.servlet.http.HttpServletRequest
import org.apache.commons.text.StringEscapeUtils
@@ -26,8 +26,11 @@ import org.scalatest.BeforeAndAfter
import org.apache.spark.{SharedSparkContext, SparkConf, SparkFunSuite}
import org.apache.spark.scheduler.SparkListenerJobStart
+import org.apache.spark.sql.classic.SparkSession
+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.ui.SparkUI
import org.apache.spark.util.kvstore.InMemoryStore
class SparkConnectServerPageSuite
@@ -47,7 +50,10 @@ class SparkConnectServerPageSuite
/**
* Run a dummy session and return the store
*/
- private def getStatusStore: SparkConnectServerAppStatusStore = {
+ private def getStatusStore(
+ closeSession: Boolean = true,
+ sessionId: String = "sessionId",
+ userId: String = "userId"): SparkConnectServerAppStatusStore = {
kvstore = new ElementTrackingStore(new InMemoryStore, new SparkConf())
// val server = mock(classOf[SparkConnectServer], RETURNS_SMART_NULLS)
val sparkConf = new SparkConf
@@ -56,14 +62,14 @@ class SparkConnectServerPageSuite
val statusStore = new SparkConnectServerAppStatusStore(kvstore)
listener.onOtherEvent(
- SparkListenerConnectSessionStarted("sessionId", "userId",
System.currentTimeMillis()))
+ SparkListenerConnectSessionStarted(sessionId, userId,
System.currentTimeMillis()))
listener.onOtherEvent(
SparkListenerConnectOperationStarted(
"jobTag",
"operationId",
System.currentTimeMillis(),
- "sessionId",
- "userId",
+ sessionId,
+ userId,
"userName",
"dummy query",
Set()))
@@ -74,14 +80,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 +97,7 @@ class SparkConnectServerPageSuite
when(tab.store).thenReturn(store)
when(tab.appName).thenReturn("testing")
when(tab.headerTabs).thenReturn(Seq.empty)
+ when(tab.getMLCacheStatuses).thenReturn(None)
val page = new SparkConnectServerPage(tab)
val html = page.render(request).toString().toLowerCase(Locale.ROOT)
@@ -96,6 +105,9 @@ class SparkConnectServerPageSuite
assert(html.contains("session statistics (1)"))
assert(html.contains("request statistics (1)"))
assert(html.contains("dummy query"))
+ assert(html.contains("ml cache"))
+ assert(html.contains("<span>n/a</span>"))
+ assert(!html.contains("ml cache statistics"))
// Pagination support
assert(html.contains("<label class=\"text-nowrap\">1 pages. jump
to</label>"))
@@ -107,7 +119,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 +146,162 @@ 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)
+ val status = 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)))
+ when(tab.getMLCacheStatuses).thenReturn(
+ Some(Map(SessionKey("userId", "sessionId") -> Some(status))))
+
+ 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("1 model in memory"))
+ assert(html.contains("1024.0 b / 4.0 kib memory"))
+ assert(html.contains("2.0 kib / 8.0 kib total"))
+ 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 Connect Server page should show unused and unavailable ML cache
states") {
+ 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(None)
+ val unavailableHtml =
+ new
SparkConnectServerPage(tab).render(request).toString().toLowerCase(Locale.ROOT)
+ assert(unavailableHtml.contains("<span>n/a</span>"))
+
+ when(tab.getMLCacheStatuses).thenReturn(Some(Map(SessionKey("userId",
"sessionId") -> None)))
+ val uninitializedHtml =
+ new
SparkConnectServerPage(tab).render(request).toString().toLowerCase(Locale.ROOT)
+ assert(uninitializedHtml.contains("<span>not used</span>"))
+
+ val clearedStatus = MLCacheStatus(
+ memoryControlEnabled = true,
+ inMemorySizeBytes = 0,
+ maxInMemorySizeBytes = 4096,
+ totalSizeBytes = 0,
+ maxTotalSizeBytes = 8192,
+ models = Seq.empty)
+ when(tab.getMLCacheStatuses).thenReturn(
+ Some(Map(SessionKey("userId", "sessionId") -> Some(clearedStatus))))
+ val clearedHtml =
+ new
SparkConnectServerPage(tab).render(request).toString().toLowerCase(Locale.ROOT)
+ assert(clearedHtml.contains("<span>not used</span>"))
+ assert(!clearedHtml.contains("ml cache statistics"))
+
+ when(tab.getMLCacheStatuses).thenReturn(Some(Map.empty))
+ val removedHtml =
+ new
SparkConnectServerPage(tab).render(request).toString().toLowerCase(Locale.ROOT)
+ assert(removedHtml.contains("<span>n/a</span>"))
+
+ val disabledStatus = MLCacheStatus(
+ memoryControlEnabled = false,
+ inMemorySizeBytes = 0,
+ maxInMemorySizeBytes = 0,
+ totalSizeBytes = 0,
+ maxTotalSizeBytes = 0,
+ models = Seq(
+ MLCacheModelInfo(
+ id = "model-id",
+ className = "model-class",
+ modelString = "model-details",
+ estimatedSizeBytes = None,
+ inMemory = true)))
+ when(tab.getMLCacheStatuses).thenReturn(
+ Some(Map(SessionKey("userId", "sessionId") -> Some(disabledStatus))))
+ val disabledHtml =
+ new
SparkConnectServerPage(tab).render(request).toString().toLowerCase(Locale.ROOT)
+ assert(disabledHtml.contains("memory control disabled"))
+ assert(disabledHtml.contains("1 cached model"))
+ }
+
+ test("Spark Connect Server page should read ML cache status without touching
the session") {
+ val key = SessionKey("userId", UUID.randomUUID().toString)
+ val store =
+ getStatusStore(closeSession = false, sessionId = key.sessionId, userId =
key.userId)
+ val sessionManager = new SparkConnectSessionManager()
+ sessionManager.initializeBaseSession(() => new SparkSession(sc))
+ val sessionHolder = sessionManager.getOrCreateIsolatedSession(key, None)
+
+ val request = mock(classOf[HttpServletRequest])
+ val sparkUI = SparkUI.create(
+ Some(sc),
+ sc.statusStore,
+ sc.conf,
+ sc.env.securityManager,
+ sc.appName,
+ "",
+ sc.startTime)
+ val tab =
+ new SparkConnectServerTab(store, sparkUI, Some(sessionManager))
+
+ try {
+ val lastAccessTime = sessionHolder.getSessionHolderInfo.lastAccessTimeMs
+ Thread.sleep(10)
+ val unusedHtml =
+ new
SparkConnectServerPage(tab).render(request).toString().toLowerCase(Locale.ROOT)
+ assert(unusedHtml.contains("<span>not used</span>"))
+ assert(sessionHolder.getMLCacheStatus.isEmpty)
+ assert(sessionHolder.getSessionHolderInfo.lastAccessTimeMs ===
lastAccessTime)
+
+ sessionHolder.mlCache
+ val initializedLastAccessTime =
sessionHolder.getSessionHolderInfo.lastAccessTimeMs
+ Thread.sleep(10)
+ val initializedHtml =
+ new
SparkConnectServerPage(tab).render(request).toString().toLowerCase(Locale.ROOT)
+ assert(initializedHtml.contains("<span>not used</span>"))
+ assert(sessionHolder.getMLCacheStatus.exists(_.models.isEmpty))
+ assert(sessionHolder.getSessionHolderInfo.lastAccessTimeMs ===
initializedLastAccessTime)
+ } finally {
+ tab.detach()
+ sessionManager.closeSession(key)
+ sessionManager.shutdown()
+ }
+ }
+
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())
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
