zhengruifeng commented on code in PR #56584:
URL: https://github.com/apache/spark/pull/56584#discussion_r3592040047
##########
mllib/src/main/scala/org/apache/spark/ml/Model.scala:
##########
@@ -57,5 +57,19 @@ abstract class Model[M <: Model[M]] extends Transformer {
self =>
* 4, For 3-rd extension, if external languages are used, it is recommended
to override
* this method and return a proper size.
*/
- private[spark] def estimatedSize: Long = SizeEstimator.estimate(self)
+ private[spark] def estimatedSize: Long = {
+ // SPARK-57521: Temporarily clear the parent reference during size
estimation.
+ // After fit(), the parent estimator may retain indirect references to the
SparkSession
+ // (via closures or query plan state from DataFrame operations executed
during fit).
+ // SizeEstimator traverses the entire reachable object graph, causing it
to count
+ // shared SparkSession state as part of every model's size.
+ // The parent is @transient (not persisted) and is not needed for
transform() or save().
+ val savedParent = parent
+ parent = null
Review Comment:
Sharing some analysis on the design options here, since a few of us were
pinged as ML/Connect folks.
**On the current approach:** the `synchronized` addition closes the
save/restore interleave, but as noted a concurrent reader of `parent`
(`hasParent`, or a direct `.parent` access) on a path that doesn't lock on
`self` could still observe `null` during the estimation window. In practice
this isn't an active bug today -- `estimatedSize` is only called from
`MLCache.register` (itself `synchronized`) on a freshly-fit, not-yet-shared
model, and `transform`/`save`/`write` don't read `parent`. So the
mutate-under-lock fix is acceptable as a scoped change, with the caveat that it
makes a previously side-effect-free `private[spark]` method mutating.
**A cleaner, non-mutating alternative** worth considering: address it in
`SizeEstimator` so nothing on the model has to be mutated at all. Two forms:
1. **`SparkSession extends KnownSizeEstimation` with `estimatedSize = 0`.**
In `SizeEstimator.visitSingleObject`, the `case s: KnownSizeEstimation` branch
adds the reported size and -- crucially -- does *not* enqueue the object's
pointer fields, so a `0`-returning session both contributes nothing and cuts
off the entire session subgraph. This would have to go on the concrete
`sql/core` `classic.SparkSession` (the abstract `sql/api` one has no
`spark-core` dependency and is shared with the Connect client).
2. **A class-name skip** alongside the existing `ClassLoader` /
`scala.reflect` exclusions in `visitSingleObject`. Since `SizeEstimator` (in
`spark-core`) can't import the sql-layer `SparkSession` type, this would have
to match by class-name string, which is more brittle than option 1.
Both are non-mutating and, unlike the `parent`-null approach, fix
overcounting for *any* reference path to the session rather than only the one
through `Model.parent`.
**The trade-off** is scope of review: option 1/2 change size accounting for
*every* `SizeEstimator` caller (block-manager memory accounting, `SizeTracker`,
spill thresholds, ...), not just the ML cache. That's almost certainly the
correct accounting -- a shared session isn't attributable to the object being
sized, which is exactly this PR's argument -- but it's a core-wide behavior
change that warrants core/storage owner sign-off, whereas the current
`parent`-null fix stays contained to one `private[spark]` ML method. (Note also
that the heavy shared state hangs off `SparkContext`; if the intent is "never
attribute cluster-shared state to a model," `SparkContext` may want the same
treatment.)
My take: if we want the contained fix now, the current approach is fine once
the thread-safety semantics are settled. If we're open to the broader change,
option 1 (`KnownSizeEstimation` on the classic session) is the cleanest
long-term design.
--
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]
