chore: import upstream snapshot with attribution

python/ray/tests/test_basic_3.py

@@ -0,0 +1,153 @@
+# coding: utf-8
+import gc
+import logging
+import math
+import random
+import sys
+import time
+from typing import Dict
+
+import pytest
+
+import ray
+import ray.cluster_utils
+from ray._common.test_utils import wait_for_condition
+
+logger = logging.getLogger(__name__)
+
+
+def test_auto_global_gc(shutdown_only):
+    ray.init(num_cpus=1, object_store_memory=100 * 1024 * 1024)
+
+    @ray.remote
+    class Test:
+        def __init__(self):
+            self.collected = False
+            gc.disable()
+
+            def gc_called(phase, info):
+                self.collected = True
+
+            gc.callbacks.append(gc_called)
+
+        def circular_ref(self):
+            # 20MB
+            buf1 = b"0" * (10 * 1024 * 1024)
+            buf2 = b"1" * (10 * 1024 * 1024)
+            ref1 = ray.put(buf1)
+            ref2 = ray.put(buf2)
+            b = []
+            a = []
+            b.append(a)
+            a.append(b)
+            b.append(ref1)
+            a.append(ref2)
+            return a
+
+        def collected(self):
+            return self.collected
+
+    test = Test.remote()
+    # 60MB
+    for i in range(3):
+        ray.get(test.circular_ref.remote())
+    time.sleep(2)
+    assert not ray.get(test.collected.remote())
+    # 80MB
+    for _ in range(1):
+        ray.get(test.circular_ref.remote())
+    time.sleep(2)
+    assert ray.get(test.collected.remote())
+
+
+def _resource_dicts_close(d1: Dict, d2: Dict, *, abs_tol: float = 1e-4):
+    """Return if all values in the dicts are within the abs_tol."""
+
+    # A resource value of 0 is equivalent to the key not being present,
+    # so filter keys whose values are 0.
+    d1 = {k: v for k, v in d1.items() if v != 0}
+    d2 = {k: v for k, v in d2.items() if v != 0}
+
+    if d1.keys() != d2.keys():
+        return False
+
+    for k, v in d1.items():
+        if (
+            isinstance(v, float)
+            and isinstance(d2[k], float)
+            and math.isclose(v, d2[k], abs_tol=abs_tol)
+        ):
+            continue
+        if v != d2[k]:
+            return False
+    return True
+
+
+def test_many_fractional_resources(shutdown_only):
+    ray.init(num_cpus=2, num_gpus=2, resources={"Custom": 2})
+
+    def _get_available_resources() -> Dict[str, float]:
+        """Get only the resources we care about in this test."""
+        return {
+            k: v
+            for k, v in ray.available_resources().items()
+            if k in {"CPU", "GPU", "Custom"}
+        }
+
+    original_available_resources = _get_available_resources()
+
+    @ray.remote
+    def g():
+        return 1
+
+    @ray.remote
+    def check_assigned_resources(block: bool, expected_resources: Dict[str, float]):
+        assigned_resources = ray.get_runtime_context().get_assigned_resources()
+
+        # Have some tasks block to release their occupied resources to further
+        # stress the scheduler.
+        if block:
+            ray.get(g.remote())
+
+        if not _resource_dicts_close(assigned_resources, expected_resources):
+            raise RuntimeError(
+                "Mismatched resources.",
+                "Expected:",
+                expected_resources,
+                "Assigned:",
+                assigned_resources,
+            )
+
+    def _rand_resource_val() -> float:
+        return int(random.random() * 10000) / 10000
+
+    # Submit many tasks with random resource requirements and assert that they are
+    # assigned the correct resources.
+    result_ids = []
+    for i in range(10):
+        resources = {
+            "CPU": _rand_resource_val(),
+            "GPU": _rand_resource_val(),
+            "Custom": _rand_resource_val(),
+        }
+
+        for block in [False, True]:
+            result_ids.append(
+                check_assigned_resources.options(
+                    num_cpus=resources["CPU"],
+                    num_gpus=resources["GPU"],
+                    resources={"Custom": resources["Custom"]},
+                ).remote(block, resources)
+            )
+
+    # This would raise if any assigned resources don't match the expectation.
+    ray.get(result_ids)
+
+    # Check that the available resources are reset to their original values.
+    wait_for_condition(
+        lambda: _get_available_resources() == original_available_resources,
+    )
+
+
+if __name__ == "__main__":
+    sys.exit(pytest.main(["-sv", __file__]))