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/21 12:28:09 UTC
(spark) branch master updated: [SPARK-51845][ML][CONNECT] Add proto messages `CleanCache` and `GetCacheInfo`
This is an automated email from the ASF dual-hosted git repository.
ruifengz pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/spark.git
The following commit(s) were added to refs/heads/master by this push:
new b557311d4394 [SPARK-51845][ML][CONNECT] Add proto messages `CleanCache` and `GetCacheInfo`
b557311d4394 is described below
commit b557311d43947f1e862345ada89f4d27de74e7b1
Author: Ruifeng Zheng <ru...@apache.org>
AuthorDate: Mon Apr 21 20:27:50 2025 +0800
[SPARK-51845][ML][CONNECT] Add proto messages `CleanCache` and `GetCacheInfo`
### What changes were proposed in this pull request?
Add proto message `CleanCache` and `GetCacheInfo`
### Why are the changes needed?
to improve debuggability
### Does this PR introduce _any_ user-facing change?
no, the new protos are only used internally
```
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/21 11:12:36 WARN Instrumentation: [91fb6f9e] regParam is zero, which might cause numerical instability and overfitting.
25/04/21 11:12:37 WARN InstanceBuilder: Failed to load implementation from:dev.ludovic.netlib.lapack.JNILAPACK
[********************************************************************************] 100.00% Complete (0 Tasks running, 2s, Sca[********************************************************************************] 100.00% Complete (0 Tasks running, 2s, Sca[********************************************************************************] 100.00% Complete (0 Tasks running, 2s, Sca[********************************************************************************] 100.00% Complete (0 Tasks runni [...]
In [2]: spark.client.thread_local.ml_caches
Out[2]: {'7c7505a7-9a9a-4716-a485-d735c48608cc'}
In [3]: spark.client._get_ml_cache_info()
Out[3]: ['id: 7c7505a7-9a9a-4716-a485-d735c48608cc, obj: class org.apache.spark.ml.regression.LinearRegressionModel, size: 5924']
In [4]: del model
In [5]: spark.client.thread_local.ml_caches
Out[5]: set()
In [6]: spark.client._get_ml_cache_info()
Out[6]: []
```
### How was this patch tested?
ci
### Was this patch authored or co-authored using generative AI tooling?
no
Closes #50643 from zhengruifeng/ml_connect_clean.
Authored-by: Ruifeng Zheng <ru...@apache.org>
Signed-off-by: Ruifeng Zheng <ru...@apache.org>
---
.../pyspark/ml/tests/connect/test_connect_cache.py | 21 ++++++++++++
python/pyspark/sql/connect/client/core.py | 40 ++++++++++++++--------
python/pyspark/sql/connect/proto/ml_pb2.py | 40 ++++++++++++----------
python/pyspark/sql/connect/proto/ml_pb2.pyi | 39 ++++++++++++++++++++-
.../src/main/protobuf/spark/connect/ml.proto | 8 +++++
.../org/apache/spark/sql/connect/ml/MLCache.scala | 14 +++++++-
.../apache/spark/sql/connect/ml/MLHandler.scala | 13 +++++++
7 files changed, 141 insertions(+), 34 deletions(-)
diff --git a/python/pyspark/ml/tests/connect/test_connect_cache.py b/python/pyspark/ml/tests/connect/test_connect_cache.py
index 9d907f8c9fdf..f01a0ef047b4 100644
--- a/python/pyspark/ml/tests/connect/test_connect_cache.py
+++ b/python/pyspark/ml/tests/connect/test_connect_cache.py
@@ -46,12 +46,20 @@ class MLConnectCacheTests(ReusedConnectTestCase):
# model is cached in python side
self.assertEqual(len(spark.client.thread_local.ml_caches), 1)
+ cache_info = spark.client._get_ml_cache_info()
+ self.assertEqual(len(cache_info), 1)
+ self.assertTrue(
+ "obj: class org.apache.spark.ml.classification.LinearSVCModel" in cache_info[0],
+ cache_info,
+ )
# explicitly delete the model
del model
# model is removed in python side
self.assertEqual(len(spark.client.thread_local.ml_caches), 0)
+ cache_info = spark.client._get_ml_cache_info()
+ self.assertEqual(len(cache_info), 0)
def test_cleanup_ml(self):
spark = self.spark
@@ -79,17 +87,30 @@ class MLConnectCacheTests(ReusedConnectTestCase):
# all 3 models are cached in python side
self.assertEqual(len(spark.client.thread_local.ml_caches), 3)
+ cache_info = spark.client._get_ml_cache_info()
+ self.assertEqual(len(cache_info), 3)
+ self.assertTrue(
+ all(
+ "obj: class org.apache.spark.ml.classification.LinearSVCModel" in c
+ for c in cache_info
+ ),
+ cache_info,
+ )
# explicitly delete the model1
del model1
# model1 is removed in python side
self.assertEqual(len(spark.client.thread_local.ml_caches), 2)
+ cache_info = spark.client._get_ml_cache_info()
+ self.assertEqual(len(cache_info), 2)
spark.client._cleanup_ml()
# All models are removed in python side
self.assertEqual(len(spark.client.thread_local.ml_caches), 0)
+ cache_info = spark.client._get_ml_cache_info()
+ self.assertEqual(len(cache_info), 0)
if __name__ == "__main__":
diff --git a/python/pyspark/sql/connect/client/core.py b/python/pyspark/sql/connect/client/core.py
index 319c98178b9b..e23275512028 100644
--- a/python/pyspark/sql/connect/client/core.py
+++ b/python/pyspark/sql/connect/client/core.py
@@ -1981,13 +1981,11 @@ class SparkConnectClient(object):
self.thread_local.ml_caches.add(cache_id)
def remove_ml_cache(self, cache_id: str) -> None:
- if not hasattr(self.thread_local, "ml_caches"):
- self.thread_local.ml_caches = set()
-
- if cache_id in self.thread_local.ml_caches:
- deleted = self._delete_ml_cache([cache_id])
- for obj_id in deleted:
- self.thread_local.ml_caches.remove(obj_id)
+ if hasattr(self.thread_local, "ml_caches"):
+ if cache_id in self.thread_local.ml_caches:
+ 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]) -> List[str]:
# try best to delete the cache
@@ -2010,10 +2008,24 @@ class SparkConnectClient(object):
return []
def _cleanup_ml(self) -> None:
- if not hasattr(self.thread_local, "ml_caches"):
- self.thread_local.ml_caches = set()
-
- self.disable_reattachable_execute()
- deleted = self._delete_ml_cache(list(self.thread_local.ml_caches))
- for obj_id in deleted:
- self.thread_local.ml_caches.remove(obj_id)
+ if hasattr(self.thread_local, "ml_caches"):
+ try:
+ command = pb2.Command()
+ command.ml_command.clean_cache.SetInParent()
+ self.execute_command(command)
+ self.thread_local.ml_caches.clear()
+ except Exception:
+ pass
+
+ def _get_ml_cache_info(self) -> List[str]:
+ if hasattr(self.thread_local, "ml_caches"):
+ command = pb2.Command()
+ command.ml_command.get_cache_info.SetInParent()
+ (_, 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"]
+ return [item.string for item in ml_command_result.param.array.elements]
+ return []
diff --git a/python/pyspark/sql/connect/proto/ml_pb2.py b/python/pyspark/sql/connect/proto/ml_pb2.py
index a2475a94c447..31fa3dd5d0ec 100644
--- a/python/pyspark/sql/connect/proto/ml_pb2.py
+++ b/python/pyspark/sql/connect/proto/ml_pb2.py
@@ -40,7 +40,7 @@ from pyspark.sql.connect.proto import ml_common_pb2 as spark_dot_connect_dot_ml_
DESCRIPTOR = _descriptor_pool.Default().AddSerializedFile(
- b'\n\x16spark/connect/ml.proto\x12\rspark.connect\x1a\x1dspark/connect/relations.proto\x1a\x1fspark/connect/expressions.proto\x1a\x1dspark/connect/ml_common.proto"\xfd\t\n\tMlCommand\x12\x30\n\x03\x66it\x18\x01 \x01(\x0b\x32\x1c.spark.connect.MlCommand.FitH\x00R\x03\x66it\x12,\n\x05\x66\x65tch\x18\x02 \x01(\x0b\x32\x14.spark.connect.FetchH\x00R\x05\x66\x65tch\x12\x39\n\x06\x64\x65lete\x18\x03 \x01(\x0b\x32\x1f.spark.connect.MlCommand.DeleteH\x00R\x06\x64\x65lete\x12\x36\n\x05write\x1 [...]
+ b'\n\x16spark/connect/ml.proto\x12\rspark.connect\x1a\x1dspark/connect/relations.proto\x1a\x1fspark/connect/expressions.proto\x1a\x1dspark/connect/ml_common.proto"\xb2\x0b\n\tMlCommand\x12\x30\n\x03\x66it\x18\x01 \x01(\x0b\x32\x1c.spark.connect.MlCommand.FitH\x00R\x03\x66it\x12,\n\x05\x66\x65tch\x18\x02 \x01(\x0b\x32\x14.spark.connect.FetchH\x00R\x05\x66\x65tch\x12\x39\n\x06\x64\x65lete\x18\x03 \x01(\x0b\x32\x1f.spark.connect.MlCommand.DeleteH\x00R\x06\x64\x65lete\x12\x36\n\x05write\ [...]
)
_globals = globals()
@@ -54,21 +54,25 @@ if not _descriptor._USE_C_DESCRIPTORS:
_globals["_MLCOMMAND_WRITE_OPTIONSENTRY"]._loaded_options = None
_globals["_MLCOMMAND_WRITE_OPTIONSENTRY"]._serialized_options = b"8\001"
_globals["_MLCOMMAND"]._serialized_start = 137
- _globals["_MLCOMMAND"]._serialized_end = 1414
- _globals["_MLCOMMAND_FIT"]._serialized_start = 480
- _globals["_MLCOMMAND_FIT"]._serialized_end = 658
- _globals["_MLCOMMAND_DELETE"]._serialized_start = 660
- _globals["_MLCOMMAND_DELETE"]._serialized_end = 721
- _globals["_MLCOMMAND_WRITE"]._serialized_start = 724
- _globals["_MLCOMMAND_WRITE"]._serialized_end = 1134
- _globals["_MLCOMMAND_WRITE_OPTIONSENTRY"]._serialized_start = 1036
- _globals["_MLCOMMAND_WRITE_OPTIONSENTRY"]._serialized_end = 1094
- _globals["_MLCOMMAND_READ"]._serialized_start = 1136
- _globals["_MLCOMMAND_READ"]._serialized_end = 1217
- _globals["_MLCOMMAND_EVALUATE"]._serialized_start = 1220
- _globals["_MLCOMMAND_EVALUATE"]._serialized_end = 1403
- _globals["_MLCOMMANDRESULT"]._serialized_start = 1417
- _globals["_MLCOMMANDRESULT"]._serialized_end = 1820
- _globals["_MLCOMMANDRESULT_MLOPERATORINFO"]._serialized_start = 1610
- _globals["_MLCOMMANDRESULT_MLOPERATORINFO"]._serialized_end = 1805
+ _globals["_MLCOMMAND"]._serialized_end = 1595
+ _globals["_MLCOMMAND_FIT"]._serialized_start = 631
+ _globals["_MLCOMMAND_FIT"]._serialized_end = 809
+ _globals["_MLCOMMAND_DELETE"]._serialized_start = 811
+ _globals["_MLCOMMAND_DELETE"]._serialized_end = 872
+ _globals["_MLCOMMAND_CLEANCACHE"]._serialized_start = 874
+ _globals["_MLCOMMAND_CLEANCACHE"]._serialized_end = 886
+ _globals["_MLCOMMAND_GETCACHEINFO"]._serialized_start = 888
+ _globals["_MLCOMMAND_GETCACHEINFO"]._serialized_end = 902
+ _globals["_MLCOMMAND_WRITE"]._serialized_start = 905
+ _globals["_MLCOMMAND_WRITE"]._serialized_end = 1315
+ _globals["_MLCOMMAND_WRITE_OPTIONSENTRY"]._serialized_start = 1217
+ _globals["_MLCOMMAND_WRITE_OPTIONSENTRY"]._serialized_end = 1275
+ _globals["_MLCOMMAND_READ"]._serialized_start = 1317
+ _globals["_MLCOMMAND_READ"]._serialized_end = 1398
+ _globals["_MLCOMMAND_EVALUATE"]._serialized_start = 1401
+ _globals["_MLCOMMAND_EVALUATE"]._serialized_end = 1584
+ _globals["_MLCOMMANDRESULT"]._serialized_start = 1598
+ _globals["_MLCOMMANDRESULT"]._serialized_end = 2001
+ _globals["_MLCOMMANDRESULT_MLOPERATORINFO"]._serialized_start = 1791
+ _globals["_MLCOMMANDRESULT_MLOPERATORINFO"]._serialized_end = 1986
# @@protoc_insertion_point(module_scope)
diff --git a/python/pyspark/sql/connect/proto/ml_pb2.pyi b/python/pyspark/sql/connect/proto/ml_pb2.pyi
index 1d2e06d536e1..9f6f4c1516d8 100644
--- a/python/pyspark/sql/connect/proto/ml_pb2.pyi
+++ b/python/pyspark/sql/connect/proto/ml_pb2.pyi
@@ -134,6 +134,24 @@ class MlCommand(google.protobuf.message.Message):
self, field_name: typing_extensions.Literal["obj_refs", b"obj_refs"]
) -> None: ...
+ class CleanCache(google.protobuf.message.Message):
+ """Force to clean up all the ML cached objects"""
+
+ DESCRIPTOR: google.protobuf.descriptor.Descriptor
+
+ def __init__(
+ self,
+ ) -> None: ...
+
+ class GetCacheInfo(google.protobuf.message.Message):
+ """Get the information of all the ML cached objects"""
+
+ DESCRIPTOR: google.protobuf.descriptor.Descriptor
+
+ def __init__(
+ self,
+ ) -> None: ...
+
class Write(google.protobuf.message.Message):
"""Command to write ML operator"""
@@ -330,6 +348,8 @@ class MlCommand(google.protobuf.message.Message):
WRITE_FIELD_NUMBER: builtins.int
READ_FIELD_NUMBER: builtins.int
EVALUATE_FIELD_NUMBER: builtins.int
+ CLEAN_CACHE_FIELD_NUMBER: builtins.int
+ GET_CACHE_INFO_FIELD_NUMBER: builtins.int
@property
def fit(self) -> global___MlCommand.Fit: ...
@property
@@ -342,6 +362,10 @@ class MlCommand(google.protobuf.message.Message):
def read(self) -> global___MlCommand.Read: ...
@property
def evaluate(self) -> global___MlCommand.Evaluate: ...
+ @property
+ def clean_cache(self) -> global___MlCommand.CleanCache: ...
+ @property
+ def get_cache_info(self) -> global___MlCommand.GetCacheInfo: ...
def __init__(
self,
*,
@@ -351,10 +375,14 @@ class MlCommand(google.protobuf.message.Message):
write: global___MlCommand.Write | None = ...,
read: global___MlCommand.Read | None = ...,
evaluate: global___MlCommand.Evaluate | None = ...,
+ clean_cache: global___MlCommand.CleanCache | None = ...,
+ get_cache_info: global___MlCommand.GetCacheInfo | None = ...,
) -> None: ...
def HasField(
self,
field_name: typing_extensions.Literal[
+ "clean_cache",
+ b"clean_cache",
"command",
b"command",
"delete",
@@ -365,6 +393,8 @@ class MlCommand(google.protobuf.message.Message):
b"fetch",
"fit",
b"fit",
+ "get_cache_info",
+ b"get_cache_info",
"read",
b"read",
"write",
@@ -374,6 +404,8 @@ class MlCommand(google.protobuf.message.Message):
def ClearField(
self,
field_name: typing_extensions.Literal[
+ "clean_cache",
+ b"clean_cache",
"command",
b"command",
"delete",
@@ -384,6 +416,8 @@ class MlCommand(google.protobuf.message.Message):
b"fetch",
"fit",
b"fit",
+ "get_cache_info",
+ b"get_cache_info",
"read",
b"read",
"write",
@@ -393,7 +427,10 @@ class MlCommand(google.protobuf.message.Message):
def WhichOneof(
self, oneof_group: typing_extensions.Literal["command", b"command"]
) -> (
- typing_extensions.Literal["fit", "fetch", "delete", "write", "read", "evaluate"] | None
+ typing_extensions.Literal[
+ "fit", "fetch", "delete", "write", "read", "evaluate", "clean_cache", "get_cache_info"
+ ]
+ | None
): ...
global___MlCommand = MlCommand
diff --git a/sql/connect/common/src/main/protobuf/spark/connect/ml.proto b/sql/connect/common/src/main/protobuf/spark/connect/ml.proto
index cdb077f5055a..22c3ca7e6e90 100644
--- a/sql/connect/common/src/main/protobuf/spark/connect/ml.proto
+++ b/sql/connect/common/src/main/protobuf/spark/connect/ml.proto
@@ -36,6 +36,8 @@ message MlCommand {
Write write = 4;
Read read = 5;
Evaluate evaluate = 6;
+ CleanCache clean_cache = 7;
+ GetCacheInfo get_cache_info = 8;
}
// Command for estimator.fit(dataset)
@@ -54,6 +56,12 @@ message MlCommand {
repeated ObjectRef obj_refs = 1;
}
+ // Force to clean up all the ML cached objects
+ message CleanCache { }
+
+ // Get the information of all the ML cached objects
+ message GetCacheInfo { }
+
// Command to write ML operator
message Write {
// It could be an estimator/evaluator or the cached model
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 aceb92a4d807..05fa976b5bea 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
@@ -20,6 +20,8 @@ import java.util.UUID
import java.util.concurrent.{ConcurrentMap, TimeUnit}
import java.util.concurrent.atomic.AtomicLong
+import scala.collection.mutable
+
import com.google.common.cache.{CacheBuilder, RemovalNotification}
import org.apache.spark.internal.Logging
@@ -113,7 +115,17 @@ private[connect] class MLCache(sessionHolder: SessionHolder) extends Logging {
/**
* Clear all the caches
*/
- def clear(): Unit = {
+ def clear(): Int = {
+ val size = cachedModel.size()
cachedModel.clear()
+ size
+ }
+
+ def getInfo(): Array[String] = {
+ val info = mutable.ArrayBuilder.make[String]
+ cachedModel.forEach { case (key, value) =>
+ info += s"id: $key, obj: ${value.obj.getClass}, size: ${value.sizeBytes}"
+ }
+ info.result()
}
}
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 0615fd4afb8e..5283639e4aa2 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
@@ -186,6 +186,19 @@ private[connect] object MLHandler extends Logging {
.setObjRef(proto.ObjectRef.newBuilder().setId(ids.result().mkString(","))))
.build()
+ case proto.MlCommand.CommandCase.CLEAN_CACHE =>
+ val size = mlCache.clear()
+ proto.MlCommandResult
+ .newBuilder()
+ .setParam(LiteralValueProtoConverter.toLiteralProto(size))
+ .build()
+
+ case proto.MlCommand.CommandCase.GET_CACHE_INFO =>
+ proto.MlCommandResult
+ .newBuilder()
+ .setParam(LiteralValueProtoConverter.toLiteralProto(mlCache.getInfo()))
+ .build()
+
case proto.MlCommand.CommandCase.WRITE =>
mlCommand.getWrite.getTypeCase match {
case proto.MlCommand.Write.TypeCase.OBJ_REF => // save a model
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@spark.apache.org
For additional commands, e-mail: commits-help@spark.apache.org