chore: import upstream snapshot with attribution

python/ray/tests/test_tqdm.py

@@ -0,0 +1,112 @@
+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__]))