Skip to content

Commit 6e61c39

Browse files
committed
[RPC] Fix tuning on macOS and Windows (#15771)
Fix regression in (#15187) when multiprocessing start method is not 'fork', which prevented tuning from working. This affects macOS and Windows. Also in python 3.14 the default start method will be 'spawn'.
1 parent 51bdaec commit 6e61c39

File tree

1 file changed

+13
-8
lines changed

1 file changed

+13
-8
lines changed

python/tvm/rpc/server.py

Lines changed: 13 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -119,6 +119,15 @@ def download_linked_module(file_name):
119119
return temp
120120

121121

122+
def _serve_loop(sock, load_library, work_path=None):
123+
"""Server loop"""
124+
sockfd = sock.fileno()
125+
temp = _server_env(load_library, work_path)
126+
_ffi_api.ServerLoop(sockfd)
127+
if not work_path:
128+
temp.remove()
129+
130+
122131
def _parse_server_opt(opts):
123132
# parse client options
124133
ret = {}
@@ -128,25 +137,21 @@ def _parse_server_opt(opts):
128137
return ret
129138

130139

131-
def _serving(sock, addr, opts, load_library):
140+
def _serving(conn, addr, opts, load_library):
132141
logger.info(f"connected from {addr}")
133142
work_path = utils.tempdir()
134143
old_cwd = os.getcwd()
135144
os.chdir(work_path.path) # Avoiding file name conflict between sessions.
136145
logger.info(f"start serving at {work_path.path}")
137146

138-
def _serve_loop():
139-
_server_env(load_library, work_path)
140-
_ffi_api.ServerLoop(sock.fileno())
141-
142-
server_proc = multiprocessing.Process(target=_serve_loop)
147+
server_proc = multiprocessing.Process(target=_serve_loop, args=(conn, load_library, work_path))
143148
server_proc.start()
144149
server_proc.join(opts.get("timeout", None)) # Wait until finish or timeout.
145150

146151
if server_proc.is_alive():
147152
logger.info("timeout in RPC session, kill..")
148153
_ffi_api.ReturnException(
149-
sock.fileno(),
154+
conn.fileno(),
150155
f'RPCSessionTimeoutError: Your {opts["timeout"]}s session has expired, '
151156
f'try to increase the "session_timeout" value.',
152157
)
@@ -166,7 +171,7 @@ def _serve_loop():
166171
logger.info(f"finish serving {addr}")
167172
os.chdir(old_cwd)
168173
work_path.remove()
169-
sock.close()
174+
conn.close()
170175

171176

172177
def _listen_loop(sock, port, rpc_key, tracker_addr, load_library, custom_addr):

0 commit comments

Comments
 (0)