diff --git a/python/ray/experimental/state.py b/python/ray/experimental/state.py index ce887245f..b08b4e70e 100644 --- a/python/ray/experimental/state.py +++ b/python/ray/experimental/state.py @@ -3,6 +3,7 @@ from __future__ import division from __future__ import print_function import copy +from collections import defaultdict import heapq import json import redis @@ -949,3 +950,24 @@ class GlobalState(object): if num_tasks is 0: return 0, 0, 0 return overall_smallest, overall_largest, num_tasks + + def cluster_resources(self): + """Get the current total cluster resources. + + Note that this information can grow stale as nodes are added to or + removed from the cluster. + + Returns: + A dictionary mapping resource name to the total quantity of that + resource in the cluster. + """ + local_schedulers = self.local_schedulers() + resources = defaultdict(lambda: 0) + + for local_scheduler in local_schedulers: + for key, value in local_scheduler.items(): + if key not in ["ClientType", "Deleted", "DBClientID", + "AuxAddress", "LocalSchedulerSocketName"]: + resources[key] += value + + return dict(resources) diff --git a/test/runtest.py b/test/runtest.py index 757d27bbc..6eef49bec 100644 --- a/test/runtest.py +++ b/test/runtest.py @@ -1840,7 +1840,10 @@ class GlobalStateAPI(unittest.TestCase): with self.assertRaises(Exception): ray.global_state.log_files() - ray.init() + ray.init(num_cpus=5, num_gpus=3, resources={"CustomResource": 1}) + + resources = {"CPU": 5, "GPU": 3, "CustomResource": 1} + assert ray.global_state.cluster_resources() == resources self.assertEqual(ray.global_state.object_table(), dict())