You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@spark.apache.org by ho...@apache.org on 2026/09/03 12:52:48 UTC
(spark) branch master updated: [PYTHON][MINOR] Improve TorchDistributor log server provisioning
This is an automated email from the ASF dual-hosted git repository.
asf-gitbox-commits 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 77f1e05768f3 [PYTHON][MINOR] Improve TorchDistributor log server provisioning
77f1e05768f3 is described below
commit 77f1e05768f356a613466211db5cd99b5959fe05
Author: Holden Karau <ho...@snowflake.com>
AuthorDate: Mon Aug 31 22:47:52 2026 +0000
[PYTHON][MINOR] Improve TorchDistributor log server provisioning
The log streaming server binds the driver address that workers connect to,
takes a per-run token as its first frame, and bounds the frame length it
accepts.
Co-authored-by: Cursor <cu...@cursor.com>
Co-Authored-By: Claude Opus 5 (1M context) <no...@anthropic.com>
Co-Authored-By: Holden Karau <ho...@pigscanfly.ca>
---
python/pyspark/ml/torch/distributor.py | 7 ++-
python/pyspark/ml/torch/log_communication.py | 40 ++++++++++++-
.../ml/torch/tests/test_log_communication.py | 68 ++++++++++++++++++++--
3 files changed, 105 insertions(+), 10 deletions(-)
diff --git a/python/pyspark/ml/torch/distributor.py b/python/pyspark/ml/torch/distributor.py
index 697272909853..d8934d660819 100644
--- a/python/pyspark/ml/torch/distributor.py
+++ b/python/pyspark/ml/torch/distributor.py
@@ -654,6 +654,7 @@ class TorchDistributor(Distributor):
input_params = self.input_params
driver_address = self.driver_address
log_streaming_server_port = self.log_streaming_server_port
+ log_streaming_auth_secret = self.log_streaming_auth_secret
is_spark_local_master = self.is_spark_local_master
driver_owned_gpus: List[str] = []
if is_spark_local_master and use_gpu:
@@ -732,7 +733,9 @@ class TorchDistributor(Distributor):
os.environ[CUDA_VISIBLE_DEVICES] = ""
set_torch_config(context)
- log_streaming_client = LogStreamingClient(driver_address, log_streaming_server_port)
+ log_streaming_client = LogStreamingClient(
+ driver_address, log_streaming_server_port, auth_secret=log_streaming_auth_secret
+ )
input_params["log_streaming_client"] = log_streaming_client
try:
with TorchDistributor._setup_spark_partition_data(iterator, schema_json):
@@ -785,10 +788,12 @@ class TorchDistributor(Distributor):
log_streaming_server.start(spark_host_address=self.driver_address)
time.sleep(1) # wait for the server to start
self.log_streaming_server_port = log_streaming_server.port
+ self.log_streaming_auth_secret = log_streaming_server.auth_secret
except Exception as e:
# If starting log streaming server failed, we don't need to break
# the distributor training but emit a warning instead.
self.log_streaming_server_port = -1
+ self.log_streaming_auth_secret = None
self.logger.warning(
"Start torch distributor log streaming server failed, "
"You cannot receive logs sent from distributor workers, ",
diff --git a/python/pyspark/ml/torch/log_communication.py b/python/pyspark/ml/torch/log_communication.py
index d6c19af1553e..bffd409f49a7 100644
--- a/python/pyspark/ml/torch/log_communication.py
+++ b/python/pyspark/ml/torch/log_communication.py
@@ -16,6 +16,8 @@
#
# type: ignore
+import hmac
+import secrets
import socket
import socketserver
import sys
@@ -25,11 +27,14 @@ import traceback
import warnings
from contextlib import closing
from struct import pack, unpack
-from typing import Generator
+from typing import Generator, Optional
# Use b'\x00' as separator instead of b'\n', because the bytes are encoded in utf-8
_SERVER_POLL_INTERVAL = 0.1
_TRUNCATE_MSG_LEN = 4000
+# Clients truncate messages to _TRUNCATE_MSG_LEN characters before sending,
+# so a frame declaring a much larger length is malformed.
+_MAX_MSG_BYTES = 65536
_log_print_lock = threading.Lock() # pylint: disable=invalid-name
@@ -46,12 +51,24 @@ class WriteLogToStdout(socketserver.StreamRequestHandler):
time.sleep(_SERVER_POLL_INTERVAL)
continue
number_bytes = unpack(">i", packed_number_bytes)[0]
+ if number_bytes < 0 or number_bytes > _MAX_MSG_BYTES:
+ # Malformed or oversized frame: drop the connection rather than
+ # allocate a buffer of the declared size.
+ return
message = self.rfile.read(number_bytes)
yield message
def handle(self) -> None:
self.request.setblocking(0) # non-blocking mode
+ authenticated = False
for bline in self._read_bline():
+ if not authenticated:
+ # The first frame carries the per-run token; the
+ # connection is closed if it does not match.
+ if hmac.compare_digest(bline, self.server.auth_secret.encode("utf-8")):
+ authenticated = True
+ continue
+ return
with _get_log_print_lock():
sys.stderr.write(bline.decode("utf-8") + "\n")
sys.stderr.flush()
@@ -63,6 +80,7 @@ class LogStreamingServer:
self.server = None
self.serve_thread = None
self.port = None
+ self.auth_secret = secrets.token_urlsafe(32)
@staticmethod
def _get_free_port(spark_host_address: str = "") -> int:
@@ -76,7 +94,11 @@ class LogStreamingServer:
raise RuntimeError("Cannot start the server twice.")
def serve_task(port: int) -> None:
- with socketserver.ThreadingTCPServer(("0.0.0.0", port), WriteLogToStdout) as server:
+ # Bind the address workers connect to (the driver address).
+ with socketserver.ThreadingTCPServer(
+ (spark_host_address, port), WriteLogToStdout
+ ) as server:
+ server.auth_secret = self.auth_secret
self.server = server
server.is_active = True
server.serve_forever(poll_interval=_SERVER_POLL_INTERVAL)
@@ -135,7 +157,9 @@ class LogStreamingClient(LogStreamingClientBase):
if LogStreamingClient._log_callback_client is not None:
LogStreamingClient._log_callback_client.close()
- def __init__(self, address: str, port: int, timeout: int = 10):
+ def __init__(
+ self, address: str, port: int, timeout: int = 10, auth_secret: Optional[str] = None
+ ):
"""
Creates a connection to the logging server and authenticates.This client is best effort,
if authentication or sending a message fails, the client will be marked as not alive and
@@ -143,9 +167,14 @@ class LogStreamingClient(LogStreamingClientBase):
:param address: Address where the service is running.
:param port: Port where the service is listening for new connections.
+ :param timeout: Timeout in seconds for connecting to the server.
+ :param auth_secret: Shared secret generated by the server
+ (:class:`LogStreamingServer` ``auth_secret``), sent as the first
+ frame after connecting; without it the server ignores this client.
"""
self.address = address
self.port = port
+ self.auth_secret = auth_secret
self.timeout = timeout
self.sock = None
self.failed = True
@@ -164,6 +193,11 @@ class LogStreamingClient(LogStreamingClientBase):
sock.settimeout(self.timeout)
sock.connect((self.address, self.port))
sock.setsockopt(socket.IPPROTO_TCP, socket.TCP_NODELAY, 1)
+ if self.auth_secret is not None:
+ # The server drops connections whose first frame is not the
+ # shared secret, so it goes first.
+ binary_secret = self.auth_secret.encode("utf-8")
+ sock.sendall(pack(">i", len(binary_secret)) + binary_secret)
self.sock = sock
self.failed = False
except (OSError, IOError): # pylint: disable=broad-except
diff --git a/python/pyspark/ml/torch/tests/test_log_communication.py b/python/pyspark/ml/torch/tests/test_log_communication.py
index 6485bc9c6401..8f65775df4d5 100644
--- a/python/pyspark/ml/torch/tests/test_log_communication.py
+++ b/python/pyspark/ml/torch/tests/test_log_communication.py
@@ -22,6 +22,7 @@ import sys
import time
import unittest
from io import StringIO
+from struct import pack
from typing import Any, Callable
import pyspark.ml.torch.log_communication
@@ -57,7 +58,7 @@ class LogStreamingServiceTestCase(unittest.TestCase):
server = LogStreamingServer()
server.start()
time.sleep(1)
- client = LogStreamingClient("localhost", server.port)
+ client = LogStreamingClient("localhost", server.port, auth_secret=server.auth_secret)
with patch_stderr() as output:
client.send("msg 001")
client.send("msg 002")
@@ -77,8 +78,8 @@ class LogStreamingServiceTestCase(unittest.TestCase):
server = LogStreamingServer()
server.start()
time.sleep(1)
- client1 = LogStreamingClient("localhost", server.port)
- client2 = LogStreamingClient("localhost", server.port)
+ client1 = LogStreamingClient("localhost", server.port, auth_secret=server.auth_secret)
+ client2 = LogStreamingClient("localhost", server.port, auth_secret=server.auth_secret)
with patch_stderr() as output:
client1.send("c1 msg1")
time.sleep(_SERVER_POLL_INTERVAL + 1)
@@ -98,7 +99,7 @@ class LogStreamingServiceTestCase(unittest.TestCase):
server = LogStreamingServer()
server.start()
time.sleep(1)
- client = LogStreamingClient("localhost", server.port)
+ client = LogStreamingClient("localhost", server.port, auth_secret=server.auth_secret)
client.send("msg 001")
server.shutdown()
for i in range(5):
@@ -111,7 +112,7 @@ class LogStreamingServiceTestCase(unittest.TestCase):
server = LogStreamingServer()
server.start()
time.sleep(1)
- client = LogStreamingClient("localhost", server.port)
+ client = LogStreamingClient("localhost", server.port, auth_secret=server.auth_secret)
with patch_stderr() as output:
client._connect()
# test client send half message first
@@ -131,7 +132,7 @@ class LogStreamingServiceTestCase(unittest.TestCase):
server = LogStreamingServer()
server.start()
time.sleep(1)
- client = LogStreamingClient("localhost", server.port)
+ client = LogStreamingClient("localhost", server.port, auth_secret=server.auth_secret)
client_ops(client)
server.shutdown()
client.close()
@@ -159,6 +160,61 @@ class LogStreamingServiceTestCase(unittest.TestCase):
run_test(client_ops_send_a_msg)
run_test(client_ops_send_a_msg_and_close)
+ def test_unauthenticated_client_is_ignored(self) -> None:
+ server = LogStreamingServer()
+ server.start()
+ time.sleep(1)
+ good_client = LogStreamingClient("localhost", server.port, auth_secret=server.auth_secret)
+ bad_client = LogStreamingClient("localhost", server.port, auth_secret="wrong-secret")
+ no_secret_client = LogStreamingClient("localhost", server.port)
+ with patch_stderr() as output:
+ bad_client.send("forged msg")
+ no_secret_client.send("noauth msg")
+ time.sleep(_SERVER_POLL_INTERVAL + 1)
+ good_client.send("real msg")
+ time.sleep(_SERVER_POLL_INTERVAL + 1)
+ output = output.getvalue()
+ self.assertIn("real msg\n", output)
+ self.assertNotIn("forged msg", output)
+ self.assertNotIn("noauth msg", output)
+ bad_client.close()
+ no_secret_client.close()
+ good_client.close()
+ server.shutdown()
+
+ def test_oversized_frame_drops_connection(self) -> None:
+ server = LogStreamingServer()
+ server.start()
+ time.sleep(1)
+ client = LogStreamingClient("localhost", server.port, auth_secret=server.auth_secret)
+ with patch_stderr() as output:
+ client._connect()
+ # Declare a frame far larger than the server-side cap: the server
+ # must drop the connection instead of buffering the payload.
+ client.sock.sendall(pack(">i", 0x7FFFFFFF))
+ time.sleep(_SERVER_POLL_INTERVAL + 1)
+ client.send("dropped")
+ time.sleep(_SERVER_POLL_INTERVAL + 1)
+ self.assertNotIn("dropped", output.getvalue())
+ client.close()
+ server.shutdown()
+
+ def test_empty_message_frame(self) -> None:
+ # Blank log lines reach the client as empty strings; the zero-length
+ # frame must not break the connection.
+ server = LogStreamingServer()
+ server.start()
+ time.sleep(1)
+ client = LogStreamingClient("localhost", server.port, auth_secret=server.auth_secret)
+ with patch_stderr() as output:
+ client.send("")
+ time.sleep(_SERVER_POLL_INTERVAL + 1)
+ client.send("alive")
+ time.sleep(_SERVER_POLL_INTERVAL + 1)
+ self.assertIn("\nalive\n", output.getvalue())
+ client.close()
+ server.shutdown()
+
if __name__ == "__main__":
from pyspark.testing import main
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@spark.apache.org
For additional commands, e-mail: commits-help@spark.apache.org