chore: import upstream snapshot with attribution

python/ray/air/tests/execution/test_resource_manager_fixed.py

@@ -0,0 +1,178 @@
+import pytest
+
+import ray
+from ray.air.execution.resources.fixed import FixedResourceManager
+from ray.air.execution.resources.request import ResourceRequest
+from ray.util.scheduling_strategies import PlacementGroupSchedulingStrategy
+
+REQUEST_2_CPU = ResourceRequest([{"CPU": 2}])
+REQUEST_4_CPU = ResourceRequest([{"CPU": 4}])
+REQUEST_1_2_CPU = ResourceRequest([{"CPU": 1}, {"CPU": 2}])
+REQUEST_0_2_CPU = ResourceRequest([{"CPU": 0}, {"CPU": 2}])
+
+
+@pytest.fixture(scope="module")
+def ray_start_4_cpus():
+    address_info = ray.init(num_cpus=4)
+    yield address_info
+    # The code after the yield will run as teardown code.
+    ray.shutdown()
+
+
+def test_acquire_return_resources(ray_start_4_cpus):
+    manager = FixedResourceManager(total_resources={"CPU": 4})
+
+    assert not manager.has_resources_ready(REQUEST_2_CPU)
+    assert not manager.has_resources_ready(REQUEST_4_CPU)
+
+    manager.request_resources(REQUEST_2_CPU)
+    manager.request_resources(REQUEST_4_CPU)
+
+    assert manager.has_resources_ready(REQUEST_4_CPU)
+
+    ready_2 = manager.acquire_resources(REQUEST_2_CPU)
+
+    assert manager.has_resources_ready(REQUEST_2_CPU)
+    assert not manager.has_resources_ready(REQUEST_4_CPU)
+
+    manager.free_resources(ready_2)
+
+    assert manager.has_resources_ready(REQUEST_4_CPU)
+
+
+def test_numerical_error(ray_start_4_cpus):
+    """Make sure we don't run into numerical errors when using fractional resources.
+
+    Legacy test: test_trial_runner::TrialRunnerTest::testResourceNumericalError
+    """
+    manager = FixedResourceManager(
+        total_resources={"CPU": 0.99, "GPU": 0.99, "a": 0.99}
+    )
+    resource_request = ResourceRequest([{"CPU": 0.33, "GPU": 0.33, "a": 0.33}])
+
+    for i in range(3):
+        manager.request_resources(resource_request)
+        assert manager.acquire_resources(
+            resource_request=resource_request
+        ), manager._available_resources
+
+    assert manager._available_resources["CPU"] == 0
+    assert manager._available_resources["GPU"] == 0
+    assert manager._available_resources["a"] == 0
+
+
+def test_bind_two_bundles(ray_start_4_cpus):
+    """Test that binding two remote objects to a ready resource works.
+
+    - Request resources with 2 bundles (1 CPU and 2 CPUs)
+    - Bind two remote tasks to these bundles, execute
+    - Assert that resource allocation returns the correct resources: 1 CPU and 2 CPUs
+    """
+    manager = FixedResourceManager()
+    manager.request_resources(REQUEST_1_2_CPU)
+
+    assert manager.has_resources_ready(REQUEST_1_2_CPU)
+
+    @ray.remote
+    def get_assigned_resources():
+        return ray.get_runtime_context().get_assigned_resources()
+
+    acq = manager.acquire_resources(REQUEST_1_2_CPU)
+    [av1] = acq.annotate_remote_entities([get_assigned_resources])
+
+    res1 = ray.get(av1.remote())
+
+    assert sum(v for k, v in res1.items() if k.startswith("CPU")) == 1
+
+    [av1, av2] = acq.annotate_remote_entities(
+        [get_assigned_resources, get_assigned_resources]
+    )
+
+    res1, res2 = ray.get([av1.remote(), av2.remote()])
+    assert sum(v for k, v in res1.items() if k.startswith("CPU")) == 1
+    assert sum(v for k, v in res2.items() if k.startswith("CPU")) == 2
+
+
+def test_bind_empty_head_bundle(ray_start_4_cpus):
+    """Test that binding two remote objects to a ready resource works with empty head.
+
+    - Request resources with 2 bundles (0 CPU and 2 CPUs)
+    - Bind two remote tasks to these bundles, execute
+    - Assert that resource allocation returns the correct resources: 0 CPU and 2 CPUs
+    """
+    manager = FixedResourceManager()
+    assert REQUEST_0_2_CPU.head_bundle_is_empty
+    manager.request_resources(REQUEST_0_2_CPU)
+    ray.wait(manager.get_resource_futures(), num_returns=1)
+
+    assert manager.has_resources_ready(REQUEST_0_2_CPU)
+
+    @ray.remote
+    def get_assigned_resources():
+        return ray.get_runtime_context().get_assigned_resources()
+
+    acq = manager.acquire_resources(REQUEST_0_2_CPU)
+    [av1] = acq.annotate_remote_entities([get_assigned_resources])
+
+    res1 = ray.get(av1.remote())
+
+    assert sum(v for k, v in res1.items() if k.startswith("CPU")) == 0
+
+    [av1, av2] = acq.annotate_remote_entities(
+        [get_assigned_resources, get_assigned_resources]
+    )
+
+    res1, res2 = ray.get([av1.remote(), av2.remote()])
+    assert sum(v for k, v in res1.items() if k.startswith("CPU")) == 0
+    assert sum(v for k, v in res2.items() if k.startswith("CPU")) == 2
+
+
+@pytest.mark.parametrize("strategy", ["STRICT_PACK", "PACK", "SPREAD", "STRICT_SPREAD"])
+def test_strategy(ray_start_4_cpus, strategy):
+    """The fixed resoure manager does not support STRICT placement strategies."""
+    manager = FixedResourceManager()
+
+    req = ResourceRequest([{"CPU": 2}], strategy=strategy)
+
+    if strategy.startswith("STRICT_"):
+        with pytest.raises(RuntimeError):
+            manager.request_resources(req)
+    else:
+        manager.request_resources(req)
+
+
+@pytest.mark.parametrize("strategy", ["STRICT_PACK", "PACK", "SPREAD", "STRICT_SPREAD"])
+def test_strategy_nested(ray_start_4_cpus, strategy):
+    """The fixed resoure manager does not support STRICT_SPREAD within a PG."""
+
+    @ray.remote
+    def nested_test():
+        manager = FixedResourceManager()
+
+        req = ResourceRequest([{"CPU": 2}], strategy=strategy)
+
+        if strategy == "STRICT_SPREAD":
+            with pytest.raises(RuntimeError):
+                manager.request_resources(req)
+        else:
+            manager.request_resources(req)
+
+    pg = ray.util.placement_group([{"CPU": 2}])
+    ray.wait([pg.ready()])
+
+    try:
+        ray.get(
+            nested_test.options(
+                scheduling_strategy=PlacementGroupSchedulingStrategy(
+                    placement_group=pg, placement_group_capture_child_tasks=True
+                )
+            ).remote()
+        )
+    finally:
+        ray.util.remove_placement_group(pg)
+
+
+if __name__ == "__main__":
+    import sys
+
+    sys.exit(pytest.main(["-v", __file__]))