wehub-resource-sync
93 lines
2.5 KiB
Python
93 lines
2.5 KiB
Python
import os
|
|
import time
|
|
from contextlib import contextmanager
|
|
from typing import List, Optional, Tuple
|
|
|
|
import numpy as np
|
|
|
|
import ray
|
|
|
|
# Only run tests matching this filter pattern.
|
|
|
|
filter_pattern = os.environ.get("TESTS_TO_RUN", "")
|
|
skip_pattern = os.environ.get("TESTS_TO_SKIP", "")
|
|
|
|
|
|
def timeit(
|
|
name, fn, multiplier=1, warmup_time_sec=10
|
|
) -> List[Optional[Tuple[str, float, float]]]:
|
|
if filter_pattern and filter_pattern not in name:
|
|
return [None]
|
|
if skip_pattern and skip_pattern in name:
|
|
return [None]
|
|
# sleep for a while to avoid noisy neigbhors.
|
|
# related issue: https://github.com/ray-project/ray/issues/22045
|
|
time.sleep(warmup_time_sec)
|
|
# warmup
|
|
start = time.perf_counter()
|
|
count = 0
|
|
while time.perf_counter() - start < 1:
|
|
fn()
|
|
count += 1
|
|
# real run
|
|
step = count // 10 + 1
|
|
stats = []
|
|
for _ in range(4):
|
|
start = time.perf_counter()
|
|
count = 0
|
|
while time.perf_counter() - start < 2:
|
|
for _ in range(step):
|
|
fn()
|
|
count += step
|
|
end = time.perf_counter()
|
|
stats.append(multiplier * count / (end - start))
|
|
|
|
mean = np.mean(stats)
|
|
sd = np.std(stats)
|
|
print(name, "per second", round(mean, 2), "+-", round(sd, 2))
|
|
return [(name, mean, sd)]
|
|
|
|
|
|
async def asyncio_timeit(
|
|
name, async_fn, multiplier=1, warmup_time_sec=10
|
|
) -> List[Optional[Tuple[str, float, float]]]:
|
|
if filter_pattern and filter_pattern not in name:
|
|
return [None]
|
|
if skip_pattern and skip_pattern in name:
|
|
return [None]
|
|
# sleep for a while to avoid noisy neigbhors.
|
|
# related issue: https://github.com/ray-project/ray/issues/22045
|
|
time.sleep(warmup_time_sec)
|
|
# warmup
|
|
start = time.perf_counter()
|
|
count = 0
|
|
while time.perf_counter() - start < 1:
|
|
await async_fn()
|
|
count += 1
|
|
# real run
|
|
step = count // 10 + 1
|
|
stats = []
|
|
for _ in range(4):
|
|
start = time.perf_counter()
|
|
count = 0
|
|
while time.perf_counter() - start < 2:
|
|
for _ in range(step):
|
|
await async_fn()
|
|
count += step
|
|
end = time.perf_counter()
|
|
stats.append(multiplier * count / (end - start))
|
|
|
|
mean = np.mean(stats)
|
|
sd = np.std(stats)
|
|
print(name, "per second", round(mean, 2), "+-", round(sd, 2))
|
|
return [(name, mean, sd)]
|
|
|
|
|
|
@contextmanager
|
|
def ray_setup_and_teardown(**init_args):
|
|
ray.init(**init_args)
|
|
try:
|
|
yield None
|
|
finally:
|
|
ray.shutdown()
|