wehub-resource-sync
113 lines
3.1 KiB
Python
113 lines
3.1 KiB
Python
import sys
|
|
import time
|
|
|
|
import pytest
|
|
|
|
import ray
|
|
from ray._common.test_utils import wait_for_condition
|
|
from ray.experimental import tqdm_ray
|
|
|
|
|
|
def test_distributed_tqdm_remote():
|
|
@ray.remote
|
|
class Actor:
|
|
def __init__(self):
|
|
try:
|
|
self.bar = tqdm_ray.tqdm(
|
|
desc="foo", total=100, position=0, flush_interval_s=0
|
|
)
|
|
self.bar.update(42)
|
|
except Exception as e:
|
|
print(e)
|
|
|
|
def print_something(self):
|
|
print("hello there")
|
|
|
|
def update(self):
|
|
self.bar.update(1)
|
|
|
|
a = Actor.remote()
|
|
mgr = tqdm_ray.instance()
|
|
wait_for_condition(lambda: len(mgr.bar_groups) == 1)
|
|
bar_group = list(mgr.bar_groups.values())[0]
|
|
assert len(bar_group.bars_by_uuid) == 1
|
|
bar = list(bar_group.bars_by_uuid.values())[0]
|
|
assert bar.bar.n == 42, bar.bar.n
|
|
assert "foo" in bar.bar.desc, bar.bar.desc
|
|
assert not mgr.in_hidden_state
|
|
|
|
# Test stdout save/restore clearing.
|
|
assert mgr.num_hides == 0
|
|
ray.get(a.print_something.remote())
|
|
wait_for_condition(lambda: mgr.num_hides == 1)
|
|
wait_for_condition(lambda: not mgr.in_hidden_state)
|
|
|
|
|
|
def test_distributed_tqdm_local():
|
|
mgr = tqdm_ray.instance()
|
|
mgr.bar_groups.clear()
|
|
|
|
bar = tqdm_ray.tqdm(desc="bar", total=100, position=0, flush_interval_s=0)
|
|
bar.update(42)
|
|
wait_for_condition(lambda: len(mgr.bar_groups) == 1)
|
|
assert len(mgr.bar_groups) == 1
|
|
bar_group = list(mgr.bar_groups.values())[0]
|
|
assert len(bar_group.bars_by_uuid) == 1
|
|
bar = list(bar_group.bars_by_uuid.values())[0]
|
|
assert bar.bar.n == 42, bar.bar.n
|
|
assert "bar" in bar.bar.desc, bar.bar.desc
|
|
|
|
|
|
def test_distributed_tqdm_iterator():
|
|
mgr = tqdm_ray.instance()
|
|
mgr.bar_groups.clear()
|
|
|
|
assert sum(tqdm_ray.tqdm(range(100), desc="baz", flush_interval_s=0)) == sum(
|
|
range(100)
|
|
)
|
|
wait_for_condition(lambda: len(mgr.bar_groups) == 1)
|
|
assert len(mgr.bar_groups) == 1
|
|
bar_group = list(mgr.bar_groups.values())[0]
|
|
assert len(bar_group.bars_by_uuid) == 1
|
|
bar = list(bar_group.bars_by_uuid.values())[0]
|
|
assert bar.bar.n == 100, bar.bar.n
|
|
assert "baz" in bar.bar.desc, bar.bar.desc
|
|
|
|
|
|
def test_flush_interval():
|
|
mgr = tqdm_ray.instance()
|
|
mgr.bar_groups.clear()
|
|
|
|
FLUSH_INTERVAL_S = 0.1
|
|
|
|
def check_value(expected_value):
|
|
bar_group = list(mgr.bar_groups.values())[0]
|
|
assert len(bar_group.bars_by_uuid) == 1
|
|
bar = list(bar_group.bars_by_uuid.values())[0]
|
|
assert bar.bar.n == expected_value
|
|
|
|
bar = tqdm_ray.tqdm(
|
|
desc="bar",
|
|
total=100,
|
|
position=0,
|
|
flush_interval_s=FLUSH_INTERVAL_S,
|
|
)
|
|
# The first update should trigger flush.
|
|
bar.update(1)
|
|
check_value(1)
|
|
# Quickly calling update multiple times
|
|
# should not trigger flush.
|
|
for _ in range(10):
|
|
bar.update(1)
|
|
check_value(1)
|
|
# Wait for flush interval and call update again.
|
|
# This should trigger flush.
|
|
time.sleep(FLUSH_INTERVAL_S)
|
|
bar.update(1)
|
|
check_value(12)
|
|
bar.close()
|
|
|
|
|
|
if __name__ == "__main__":
|
|
sys.exit(pytest.main(["-sv", __file__]))
|