[SPARK-51711][ML][PYTHON][CONNECT] Propagates the active remote spark session to new threads to fix CrossValidator

[xi-db]

What changes were proposed in this pull request?

In SparkML with Spark Connect, the _parallelFitTasks fails when running CrossValidator fitting, as the active remote spark session is not properly propagated to the new threads.

Before the PR, this code will fail in the line cvModel = cv.fit(data):

from pyspark.ml.classification import RandomForestClassifier
from pyspark.ml.evaluation import BinaryClassificationEvaluator
from pyspark.ml.tuning import ParamGridBuilder, CrossValidator
from pyspark.ml.linalg import Vectors

data = spark.createDataFrame([
    (Vectors.dense(1.0, 2.0), 0),
    (Vectors.dense(2.0, 3.0), 1),
    (Vectors.dense(1.5, 2.5), 0),
    (Vectors.dense(3.0, 4.0), 1),
    (Vectors.dense(1.1, 2.1), 0),
    (Vectors.dense(2.5, 3.5), 1),
], ["features", "label"])

rf = RandomForestClassifier(labelCol="label", featuresCol="features")
evaluator = BinaryClassificationEvaluator(labelCol="label")
paramGrid = (ParamGridBuilder()
             .addGrid(rf.maxDepth, [2])
             .addGrid(rf.numTrees, [5, 10])
             .build())

cv = CrossValidator(estimator=rf,
                    estimatorParamMaps=paramGrid,
                    evaluator=evaluator,
                    numFolds=3)

cvModel = cv.fit(data)

bestModel = cvModel.bestModel
print(f"Best maxDepth: {bestModel.getMaxDepth()}")
print(f"Best maxBins: {bestModel.getMaxBins()}")
print(f"Best numTrees: {bestModel.getNumTrees}")

It fails because the active remote spark session is not properly set on that thread:

File ~/spark/python/pyspark/ml/util.py:250, in try_remote_call.<locals>.wrapped(self, name, *args)
    247 from pyspark.ml.connect.serialize import serialize, deserialize
    249 session = SparkSession.getActiveSession()
--> 250 assert session is not None
    251 assert isinstance(self._java_obj, str)
    252 methods, obj_ref = _extract_id_methods(self._java_obj)

AssertionError: 

With this fix, the above code snippet works correctly.

Why are the changes needed?

It fixes a bug with CrossValidator fitting.

Does this PR introduce any user-facing change?

No.

How was this patch tested?

New test.

Was this patch authored or co-authored using generative AI tooling?

No.

[xi-db]

Hi @zhengruifeng , could you review this PR?

It fixes a bug in SparkML via SparkConnect. The bug is reproducible with the code example in the PR description.

Thanks!

[WeichenXu123]

LGTM

[WeichenXu123]

LGTM

[zhengruifeng]

Is this necessary?
it is under if not is_remote() branch

[zhengruifeng]

is this related?
if no, I think we should file a separate PR to increase the value

[zhengruifeng]

@xi-db have we tried just replacing SparkSession.getActiveSession() with SparkSession.active() ?