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