def test_result_returns_within_timeout_seconds():
"""Check that pl_multi_process returns within 10 seconds"""
start = time.time()
result = xla_utils.pl_multi_process(time.sleep)(
xla_utils.TPU_CHECK_TIMEOUT * 1.25)
end = time.time()
elapsed_time = int(end - start)
assert elapsed_time <= xla_utils.TPU_CHECK_TIMEOUT
assert result is False
def test_result_returns_within_timeout_seconds():
"""Check that pl_multi_process returns within 3 seconds"""
fn = xla_utils.pl_multi_process(sleep_fn)
start = time.time()
result = fn(xla_utils.TPU_CHECK_TIMEOUT * 0.5)
end = time.time()
elapsed_time = int(end - start)
assert elapsed_time <= xla_utils.TPU_CHECK_TIMEOUT
assert result