[core] Report resource load by shape (#9806)

* Report and aggregate resource load by shape

* python test

* python test

* x

* update

python/ray/tests/test_global_state.py

@@ -1,3 +1,4 @@
+import json
 import pytest
 try:
     import pytest_timeout
@@ -6,6 +7,8 @@ except ImportError:
 import time
 
 import ray
+import ray.ray_constants
+import ray.test_utils
 
 
 # TODO(rliaw): The proper way to do this is to have the pytest config setup.
@@ -130,6 +133,88 @@ def test_global_state_actor_entry(ray_start_regular):
         actor_id=b_actor_id)["State"] == ray.gcs_utils.ActorTableData.ALIVE
 
 
+@pytest.mark.parametrize("max_shapes", [0, 2, -1])
+def test_load_report(shutdown_only, max_shapes):
+    resource1 = "A"
+    resource2 = "B"
+    cluster = ray.init(
+        num_cpus=1,
+        resources={resource1: 1},
+        _internal_config=json.dumps({
+            "max_resource_shapes_per_load_report": max_shapes,
+        }))
+    redis = ray.services.create_redis_client(
+        cluster["redis_address"],
+        password=ray.ray_constants.REDIS_DEFAULT_PASSWORD)
+    client = redis.pubsub(ignore_subscribe_messages=True)
+    client.psubscribe(ray.gcs_utils.XRAY_HEARTBEAT_BATCH_PATTERN)
+
+    @ray.remote
+    def sleep():
+        time.sleep(1000)
+
+    sleep.remote()
+    for _ in range(3):
+        sleep.remote()
+        sleep.options(resources={resource1: 1}).remote()
+        sleep.options(resources={resource2: 1}).remote()
+
+    class Checker:
+        def __init__(self):
+            self.report = None
+
+        def check_load_report(self):
+            try:
+                message = client.get_message()
+            except redis.exceptions.ConnectionError:
+                pass
+            if message is None:
+                return False
+
+            pattern = message["pattern"]
+            data = message["data"]
+            if pattern != ray.gcs_utils.XRAY_HEARTBEAT_BATCH_PATTERN:
+                return False
+
+            pub_message = ray.gcs_utils.PubSubMessage.FromString(data)
+            heartbeat_data = pub_message.data
+            heartbeat = ray.gcs_utils.HeartbeatBatchTableData.FromString(
+                heartbeat_data)
+            self.report = heartbeat.resource_load_by_shape.resource_demands
+            if max_shapes == 0:
+                return True
+            elif max_shapes == 2:
+                return len(self.report) >= 2
+            else:
+                return len(self.report) >= 3
+
+    # Wait for load information to arrive.
+    checker = Checker()
+    ray.test_utils.wait_for_condition(checker.check_load_report)
+
+    # Check that we respect the max shapes limit.
+    if max_shapes != -1:
+        assert len(checker.report) <= max_shapes
+
+    if max_shapes > 0:
+        # Check that we always include the 1-CPU resource shape.
+        one_cpu_shape = {"CPU": 1}
+        one_cpu_found = False
+        for demand in checker.report:
+            if demand.shape == one_cpu_shape:
+                one_cpu_found = True
+        assert one_cpu_found
+
+        # Check that we differentiate between infeasible and ready tasks.
+        for demand in checker.report:
+            if resource2 in demand.shape:
+                assert demand.num_infeasible_requests_queued > 0
+                assert demand.num_ready_requests_queued == 0
+            else:
+                assert demand.num_ready_requests_queued > 0
+                assert demand.num_infeasible_requests_queued == 0
+
+
 if __name__ == "__main__":
     import pytest
     import sys