Skip to content

Commit

Permalink
fix RPC waiting for device (#10255)
Browse files Browse the repository at this point in the history
Co-authored-by: pfk-beta <this_email_isnot_working@gmail.com>
  • Loading branch information
pfk-beta and pfk-beta authored Feb 16, 2022
1 parent 47ca5fa commit 75122db
Showing 1 changed file with 6 additions and 1 deletion.
7 changes: 6 additions & 1 deletion python/tvm/autotvm/measure/measure_methods.py
Original file line number Diff line number Diff line change
Expand Up @@ -806,17 +806,22 @@ def check_remote(target, device_key, host=None, port=None, priority=100, timeout
"""

def _check():
logger.debug("waiting for device...")
remote = request_remote(device_key, host, port, priority)
dev = remote.device(str(target))
while not dev.exist: # wait until we get an available device
pass
logger.debug("device available")

t = threading.Thread(
target=_check,
)
t.start()
t.join(timeout)
return not t.is_alive()

remote = request_remote(device_key, host, port, priority)
dev = remote.device(str(target))
return dev.exist


def set_cuda_target_arch(arch):
Expand Down

0 comments on commit 75122db

Please sign in to comment.