mkincaid commented on code in PR #56584:
URL: https://github.com/apache/spark/pull/56584#discussion_r3454908288
##########
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:
Hi @uros-b, thanks for the quick review and input. I pushed a change that
adds `synchronized` so that we wouldn't have two concurrent `estimatedSize`
calls from here. However, I'm realizing this doesn't address the second part of
your comment (a concurrent reader from elsewhere would still observe parent ==
null).
As I looked into this further, the truly non-mutating approaches I came up
with were:
- Create a `copy` of the `Model` with empty `parent`, then size that. But
this depends on the implementation of the `copy` method which is model-specific
(so not sure if it can be relied on to faithfully copy everything we care about
sizing).
- Make the `Model` object `Cloneable`, then `clone()`, clear `parent`, and
size. But changing an interface of `Model` itself seems less conservative and
beyond the scope I was intending for this original fix.
- Serialize and deserialize the object before estimating its size. Since
the `parent` is `@transient` it would be gone in the serialized copy. This
seems conceptually appealing (it seems like, in principle, the data the model
keeps and serializes is the state we care about sizing) but not sure if it
might be expensive for large models and, like the `copy` option, should I worry
about the possibility something relevant doesn’t survive the round trip.
- Target the fix elsewhere, e.g., perhaps `SizeEstimator` itself should
skip walking through `SparkSession` objects (the same way as there are existing
exclusions there for `ClassLoader` and `scala.reflect`). This also seems less
conservative since other users of `SizeEstimator` might not want the behavior
to change.
Or I may be missing something easier/cleaner. It is probably pretty obvious
that I'm new to this code base, so I want to be thoughtful about design and get
more input before proceeding. Appreciate your patience with me and looking
forward to your thoughts :)
--
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]
