Skip to content

Commit

Permalink
[RPC] Fix tuning on macOS and Windows (apache#15771)
Browse files Browse the repository at this point in the history
Fix regression in (apache#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'.
  • Loading branch information
gmeeker committed Jan 6, 2024
1 parent 51bdaec commit 6e61c39
Showing 1 changed file with 13 additions and 8 deletions.
21 changes: 13 additions & 8 deletions python/tvm/rpc/server.py
Original file line number Diff line number Diff line change
Expand Up @@ -119,6 +119,15 @@ def download_linked_module(file_name):
return temp


def _serve_loop(sock, load_library, work_path=None):
"""Server loop"""
sockfd = sock.fileno()
temp = _server_env(load_library, work_path)
_ffi_api.ServerLoop(sockfd)
if not work_path:
temp.remove()


def _parse_server_opt(opts):
# parse client options
ret = {}
Expand All @@ -128,25 +137,21 @@ def _parse_server_opt(opts):
return ret


def _serving(sock, addr, opts, load_library):
def _serving(conn, addr, opts, load_library):
logger.info(f"connected from {addr}")
work_path = utils.tempdir()
old_cwd = os.getcwd()
os.chdir(work_path.path) # Avoiding file name conflict between sessions.
logger.info(f"start serving at {work_path.path}")

def _serve_loop():
_server_env(load_library, work_path)
_ffi_api.ServerLoop(sock.fileno())

server_proc = multiprocessing.Process(target=_serve_loop)
server_proc = multiprocessing.Process(target=_serve_loop, args=(conn, load_library, work_path))
server_proc.start()
server_proc.join(opts.get("timeout", None)) # Wait until finish or timeout.

if server_proc.is_alive():
logger.info("timeout in RPC session, kill..")
_ffi_api.ReturnException(
sock.fileno(),
conn.fileno(),
f'RPCSessionTimeoutError: Your {opts["timeout"]}s session has expired, '
f'try to increase the "session_timeout" value.',
)
Expand All @@ -166,7 +171,7 @@ def _serve_loop():
logger.info(f"finish serving {addr}")
os.chdir(old_cwd)
work_path.remove()
sock.close()
conn.close()


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

0 comments on commit 6e61c39

Please sign in to comment.