You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@spark.apache.org by ru...@apache.org on 2025/04/18 03:47:55 UTC
(spark) branch branch-4.0 updated: [SPARK-51829][PYTHON][ML] Client side should update `client.thread_local.ml_caches` after deletion
This is an automated email from the ASF dual-hosted git repository.
ruifengz pushed a commit to branch branch-4.0
in repository https://gitbox.apache.org/repos/asf/spark.git
The following commit(s) were added to refs/heads/branch-4.0 by this push:
new 10643890ffaa [SPARK-51829][PYTHON][ML] Client side should update `client.thread_local.ml_caches` after deletion
10643890ffaa is described below
commit 10643890ffaa5c9ad8456a225ca57c8ec535e6e9
Author: Ruifeng Zheng <ru...@apache.org>
AuthorDate: Fri Apr 18 11:47:21 2025 +0800
[SPARK-51829][PYTHON][ML] Client side should update `client.thread_local.ml_caches` after deletion
### What changes were proposed in this pull request?
Client side should update `client.thread_local.ml_caches` after deletion:
- server side returns the list of ids of objects which have been successfully deleted from ml cache;
- client side also remove them from `client.thread_local.ml_caches`
### Why are the changes needed?
`client.thread_local.ml_caches` should not keep deleted objects, but
```
In [1]: from pyspark.ml.linalg import Vectors
...: from pyspark.ml.classification import *
...: from pyspark.ml.regression import *
...:
...: df = spark.createDataFrame([
...: (1.0, 2.0, Vectors.dense(1.0)),
...: (0.0, 2.0, Vectors.sparse(1, [], []))], ["label", "weight", "features"])
...:
...: lr = LinearRegression(regParam=0.0, solver="normal")
...: model = lr.fit(df)
25/04/17 13:23:16 WARN Instrumentation: [15c70922] regParam is zero, which might cause numerical instability and overfitting.
[--------------------------------------------------------------------------------] 0.00% Complete (0 Tasks running, 2s, Scanned 0.0[--------------------------------------------------------------------------------] 0.00% Complete (0 Tasks running, 2s, Scanned 0.025/04/17 13:23:18 WARN InstanceBuilder: Failed to load implementation from:dev.ludovic.netlib.lapack.JNILAPACK
[********************************************************************************] 100.00% Complete (0 Tasks running, 2s, Scanned 0[********************************************************************************] 100.00% Complete (0 Tasks running, 2s, Scanned 0
In [2]: spark.client.thread_local.ml_caches
Out[2]: {'e531e646-6f41-42ad-a9b1-1ce78e35c79c'}
In [3]: del model
In [4]: spark.client.thread_local.ml_caches
Out[4]: {'e531e646-6f41-42ad-a9b1-1ce78e35c79c'}
```
### Does this PR introduce _any_ user-facing change?
no
### How was this patch tested?
added test
### Was this patch authored or co-authored using generative AI tooling?
no
Closes #50618 from zhengruifeng/ml_connect_del_py_cache.
Authored-by: Ruifeng Zheng <ru...@apache.org>
Signed-off-by: Ruifeng Zheng <ru...@apache.org>
(cherry picked from commit 0d3d3a47d3a5829a7acc61bab5024005296a6ab1)
Signed-off-by: Ruifeng Zheng <ru...@apache.org>
---
dev/sparktestsupport/modules.py | 1 +
.../pyspark/ml/tests/connect/test_connect_cache.py | 104 +++++++++++++++++++++
python/pyspark/sql/connect/client/core.py | 22 ++++-
.../org/apache/spark/sql/connect/ml/MLCache.scala | 6 +-
.../apache/spark/sql/connect/ml/MLHandler.scala | 17 ++--
5 files changed, 136 insertions(+), 14 deletions(-)
diff --git a/dev/sparktestsupport/modules.py b/dev/sparktestsupport/modules.py
index b4681574c4bd..d4fd054b7977 100644
--- a/dev/sparktestsupport/modules.py
+++ b/dev/sparktestsupport/modules.py
@@ -1113,6 +1113,7 @@ pyspark_ml_connect = Module(
# ml doctests
"pyspark.ml.connect.functions",
# ml unittests
+ "pyspark.ml.tests.connect.test_connect_cache",
"pyspark.ml.tests.connect.test_connect_function",
"pyspark.ml.tests.connect.test_parity_torch_distributor",
"pyspark.ml.tests.connect.test_parity_torch_data_loader",
diff --git a/python/pyspark/ml/tests/connect/test_connect_cache.py b/python/pyspark/ml/tests/connect/test_connect_cache.py
new file mode 100644
index 000000000000..9d907f8c9fdf
--- /dev/null
+++ b/python/pyspark/ml/tests/connect/test_connect_cache.py
@@ -0,0 +1,104 @@
+#
+# Licensed to the Apache Software Foundation (ASF) under one or more
+# contributor license agreements. See the NOTICE file distributed with
+# this work for additional information regarding copyright ownership.
+# The ASF licenses this file to You under the Apache License, Version 2.0
+# (the "License"); you may not use this file except in compliance with
+# the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+
+import unittest
+
+from pyspark.ml.linalg import Vectors
+from pyspark.ml.classification import LinearSVC
+from pyspark.testing.connectutils import ReusedConnectTestCase
+
+
+class MLConnectCacheTests(ReusedConnectTestCase):
+ def test_delete_model(self):
+ spark = self.spark
+ spark.client._cleanup_ml()
+
+ df = (
+ spark.createDataFrame(
+ [
+ (1.0, 1.0, Vectors.dense(0.0, 5.0)),
+ (0.0, 2.0, Vectors.dense(1.0, 2.0)),
+ (1.0, 3.0, Vectors.dense(2.0, 1.0)),
+ (0.0, 4.0, Vectors.dense(3.0, 3.0)),
+ ],
+ ["label", "weight", "features"],
+ )
+ .coalesce(1)
+ .sortWithinPartitions("weight")
+ )
+ svc = LinearSVC(maxIter=1, regParam=1.0)
+
+ model = svc.fit(df)
+
+ # model is cached in python side
+ self.assertEqual(len(spark.client.thread_local.ml_caches), 1)
+
+ # explicitly delete the model
+ del model
+
+ # model is removed in python side
+ self.assertEqual(len(spark.client.thread_local.ml_caches), 0)
+
+ def test_cleanup_ml(self):
+ spark = self.spark
+ spark.client._cleanup_ml()
+
+ df = (
+ spark.createDataFrame(
+ [
+ (1.0, 1.0, Vectors.dense(0.0, 5.0)),
+ (0.0, 2.0, Vectors.dense(1.0, 2.0)),
+ (1.0, 3.0, Vectors.dense(2.0, 1.0)),
+ (0.0, 4.0, Vectors.dense(3.0, 3.0)),
+ ],
+ ["label", "weight", "features"],
+ )
+ .coalesce(1)
+ .sortWithinPartitions("weight")
+ )
+
+ svc = LinearSVC(maxIter=1, regParam=1.0)
+ model1 = svc.fit(df)
+ model2 = svc.fit(df)
+ model3 = svc.fit(df)
+ self.assertEqual(len([model1, model2, model3]), 3)
+
+ # all 3 models are cached in python side
+ self.assertEqual(len(spark.client.thread_local.ml_caches), 3)
+
+ # explicitly delete the model1
+ del model1
+
+ # model1 is removed in python side
+ self.assertEqual(len(spark.client.thread_local.ml_caches), 2)
+
+ spark.client._cleanup_ml()
+
+ # All models are removed in python side
+ self.assertEqual(len(spark.client.thread_local.ml_caches), 0)
+
+
+if __name__ == "__main__":
+ from pyspark.ml.tests.connect.test_connect_cache import * # noqa: F401
+
+ try:
+ import xmlrunner # type: ignore[import]
+
+ testRunner = xmlrunner.XMLTestRunner(output="target/test-reports", verbosity=2)
+ except ImportError:
+ testRunner = None
+ unittest.main(testRunner=testRunner, verbosity=2)
diff --git a/python/pyspark/sql/connect/client/core.py b/python/pyspark/sql/connect/client/core.py
index 4c5a262ebdc8..7eb24cfa083f 100644
--- a/python/pyspark/sql/connect/client/core.py
+++ b/python/pyspark/sql/connect/client/core.py
@@ -1971,9 +1971,11 @@ class SparkConnectClient(object):
self.thread_local.ml_caches = set()
if cache_id in self.thread_local.ml_caches:
- self._delete_ml_cache([cache_id])
+ deleted = self._delete_ml_cache([cache_id])
+ for obj_id in deleted:
+ self.thread_local.ml_caches.remove(obj_id)
- def _delete_ml_cache(self, cache_ids: List[str]) -> None:
+ def _delete_ml_cache(self, cache_ids: List[str]) -> List[str]:
# try best to delete the cache
try:
if len(cache_ids) > 0:
@@ -1981,13 +1983,23 @@ class SparkConnectClient(object):
command.ml_command.delete.obj_refs.extend(
[pb2.ObjectRef(id=cache_id) for cache_id in cache_ids]
)
- self.execute_command(command)
+ (_, properties, _) = self.execute_command(command)
+
+ assert properties is not None
+
+ if properties is not None and "ml_command_result" in properties:
+ ml_command_result = properties["ml_command_result"]
+ deleted = ml_command_result.operator_info.obj_ref.id.split(",")
+ return cast(List[str], deleted)
+ return []
except Exception:
- pass
+ return []
def _cleanup_ml(self) -> None:
if not hasattr(self.thread_local, "ml_caches"):
self.thread_local.ml_caches = set()
self.disable_reattachable_execute()
- self._delete_ml_cache([model_id for model_id in self.thread_local.ml_caches])
+ deleted = self._delete_ml_cache(list(self.thread_local.ml_caches))
+ for obj_id in deleted:
+ self.thread_local.ml_caches.remove(obj_id)
diff --git a/sql/connect/server/src/main/scala/org/apache/spark/sql/connect/ml/MLCache.scala b/sql/connect/server/src/main/scala/org/apache/spark/sql/connect/ml/MLCache.scala
index e8d858502072..311da95c6299 100644
--- a/sql/connect/server/src/main/scala/org/apache/spark/sql/connect/ml/MLCache.scala
+++ b/sql/connect/server/src/main/scala/org/apache/spark/sql/connect/ml/MLCache.scala
@@ -72,8 +72,10 @@ private[connect] class MLCache extends Logging {
* @param refId
* the key used to look up the corresponding object
*/
- def remove(refId: String): Unit = {
- cachedModel.remove(refId)
+ def remove(refId: String): Boolean = {
+ val removed = cachedModel.remove(refId)
+ // remove returns null if the key is not present
+ removed != null
}
/**
diff --git a/sql/connect/server/src/main/scala/org/apache/spark/sql/connect/ml/MLHandler.scala b/sql/connect/server/src/main/scala/org/apache/spark/sql/connect/ml/MLHandler.scala
index 9e3472ba8181..0615fd4afb8e 100644
--- a/sql/connect/server/src/main/scala/org/apache/spark/sql/connect/ml/MLHandler.scala
+++ b/sql/connect/server/src/main/scala/org/apache/spark/sql/connect/ml/MLHandler.scala
@@ -17,6 +17,7 @@
package org.apache.spark.sql.connect.ml
+import scala.collection.mutable
import scala.jdk.CollectionConverters.CollectionHasAsScala
import org.apache.spark.connect.proto
@@ -155,9 +156,7 @@ private[connect] object MLHandler extends Logging {
.setObjRef(proto.ObjectRef.newBuilder().setId(id)))
.build()
case a: Array[_] if a.nonEmpty && a.forall(_.isInstanceOf[Model[_]]) =>
- val ids = a.map { m =>
- mlCache.register(m.asInstanceOf[Model[_]])
- }
+ val ids = a.map(m => mlCache.register(m.asInstanceOf[Model[_]]))
proto.MlCommandResult
.newBuilder()
.setOperatorInfo(
@@ -171,16 +170,20 @@ private[connect] object MLHandler extends Logging {
}
case proto.MlCommand.CommandCase.DELETE =>
- var result = false
+ val ids = mutable.ArrayBuilder.make[String]
mlCommand.getDelete.getObjRefsList.asScala.toArray.foreach { objId =>
if (!objId.getId.contains(".")) {
- mlCache.remove(objId.getId)
- result = true
+ if (mlCache.remove(objId.getId)) {
+ ids += objId.getId
+ }
}
}
proto.MlCommandResult
.newBuilder()
- .setParam(LiteralValueProtoConverter.toLiteralProto(result))
+ .setOperatorInfo(
+ proto.MlCommandResult.MlOperatorInfo
+ .newBuilder()
+ .setObjRef(proto.ObjectRef.newBuilder().setId(ids.result().mkString(","))))
.build()
case proto.MlCommand.CommandCase.WRITE =>
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@spark.apache.org
For additional commands, e-mail: commits-help@spark.apache.org