cloud-fan commented on code in PR #57754:
URL: https://github.com/apache/spark/pull/57754#discussion_r3751227032
##########
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:
This comparator allocates a tuple on every key extraction throughout the
sort. Ordering directly by `row.model.estimatedSizeBytes` preserves the current
`None`-before-`Some` order and avoids those repeated temporary allocations.
##########
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:
Could this exercise the production `SessionManager -> SessionHolder ->
MLCache -> page` handoff instead of stubbing `getMLCacheStatuses`? That focused
integration test should also assert that rendering neither refreshes the
session's last-access time nor initializes an unused cache; the current cache
and page tests cover only the two ends independently.
--
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]
