zhengruifeng commented on code in PR #57754:
URL: https://github.com/apache/spark/pull/57754#discussion_r3757953371
##########
sql/connect/server/src/main/scala/org/apache/spark/sql/connect/ui/SparkConnectServerPage.scala:
##########
@@ -431,8 +521,152 @@ 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 || mlCacheStatuses.isEmpty) {
+ <span>N/A</span>
+ } else {
+ mlCacheStatuses
+ .flatMap(_.get(SessionKey(session.userId, session.sessionId)))
+ .filter(_.models.nonEmpty)
+ .map { status =>
+ if (status.memoryControlEnabled) {
+ val inMemoryModels = status.models.count(_.inMemory)
+ val objectLabel = if (inMemoryModels == 1) "object" else "objects"
+ <span>
+ {s"$inMemoryModels $objectLabel 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 objectLabel = if (cachedModels == 1) "object" else "objects"
+ <span>
+ Memory control disabled<br/>
+ {s"$cachedModels cached $objectLabel"}
+ </span>
+ }
+ }
+ .getOrElse(<span>Not used</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.isDefined,
row.model.estimatedSizeBytes.getOrElse(0L)))
Review Comment:
Addressed in c5c91ba9ac2. Estimated Size now orders directly by
Option[Long], preserving None-before-Some without tuple allocation.
##########
sql/connect/server/src/main/scala/org/apache/spark/sql/connect/ui/SparkConnectServerPage.scala:
##########
@@ -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)
Review Comment:
Addressed in c5c91ba9ac2. The manager now returns every active session with
an optional snapshot, so an uninitialized or empty cache renders Not used,
while unavailable or removed live state renders N/A.
##########
sql/connect/server/src/main/scala/org/apache/spark/sql/connect/ml/MLCache.scala:
##########
@@ -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))
Review Comment:
Addressed in c5c91ba9ac2. getStatus now returns the metadata iteration
directly without per-session sorting; the UI data source remains responsible
for display ordering.
##########
sql/connect/server/src/test/scala/org/apache/spark/sql/connect/ui/SparkConnectServerPageSuite.scala:
##########
@@ -134,6 +140,118 @@ 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.hasLiveMLCacheStatus).thenReturn(true)
+ when(tab.getMLCacheStatuses).thenReturn(
Review Comment:
Addressed in c5c91ba9ac2. Added a focused test using the production
SessionManager -> SessionHolder -> MLCache -> SparkConnectServerTab -> page
path. It verifies that rendering neither initializes an unused cache nor
updates last-access time, both before and after explicit cache initialization.
--
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: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
