WeichenXu123 commented on code in PR #40607:
URL: https://github.com/apache/spark/pull/40607#discussion_r1155428117
##########
python/pyspark/ml/torch/distributor.py:
##########
@@ -32,54 +32,65 @@
from pyspark import cloudpickle
from pyspark.sql import SparkSession
+from pyspark.taskcontext import BarrierTaskContext
from pyspark.ml.torch.log_communication import ( # type: ignore
- get_driver_host,
LogStreamingClient,
LogStreamingServer,
)
-from pyspark.context import SparkContext
-from pyspark.taskcontext import BarrierTaskContext
-# TODO(SPARK-41589): will move the functions and tests to an external file
-# once we are in agreement about which functions should be in utils.py
-def get_conf_boolean(sc: SparkContext, key: str, default_value: str) -> bool:
- """Get the conf "key" from the given spark context,
+def _get_active_session() -> SparkSession:
+ from pyspark.sql.utils import is_remote
+
+ if not is_remote():
+ spark = SparkSession.getActiveSession()
+ else:
+ from pyspark.sql.connect.session import _active_spark_session
+
+ spark = _active_spark_session # type: ignore[assignment]
Review Comment:
```suggestion
import pyspark.sql.connect.session
spark = pyspark.sql.connect.session._active_spark_session # type:
ignore[assignment]
```
--
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]
