chore: import upstream snapshot with attribution

python/ray/dashboard/modules/node/datacenter.py

@@ -0,0 +1,248 @@
+import copy
+from typing import List, Optional
+
+import ray.dashboard.consts as dashboard_consts
+from ray._common.utils import (
+    get_or_create_event_loop,
+)
+from ray._private.utils import (
+    parse_pg_formatted_resources_to_original,
+)
+from ray.dashboard.utils import (
+    async_loop_forever,
+    compose_state_message,
+)
+
+
+class DataSource:
+    # {node id hex(str): node stats(dict of GetNodeStatsReply
+    # in node_manager.proto)}
+    node_stats = {}
+    # {node id hex(str): node physical stats(dict from reporter_agent.py)}
+    node_physical_stats = {}
+    # {actor id hex(str): actor table data(dict of ActorTableData
+    # in gcs.proto)}
+    actors = {}
+    # {node id hex(str): gcs node info(dict of GcsNodeInfo in gcs.proto)}
+    nodes = {}
+    # {node id hex(str): worker list}
+    node_workers = {}
+    # {node id hex(str): {actor id hex(str): actor table data}}
+    node_actors = {}
+    # {worker id(str): core worker stats}
+    core_worker_stats = {}
+
+
+class DataOrganizer:
+    @staticmethod
+    @async_loop_forever(dashboard_consts.RAY_DASHBOARD_STATS_PURGING_INTERVAL)
+    async def purge():
+        # Purge data that is out of date.
+        # These data sources are maintained by DashboardHead,
+        # we do not needs to purge them:
+        #   * agents
+        #   * nodes
+        alive_nodes = {
+            node_id
+            for node_id, node_info in DataSource.nodes.items()
+            if node_info["state"] == "ALIVE"
+        }
+        for key in DataSource.node_stats.keys() - alive_nodes:
+            DataSource.node_stats.pop(key)
+
+        for key in DataSource.node_physical_stats.keys() - alive_nodes:
+            DataSource.node_physical_stats.pop(key)
+
+    @classmethod
+    @async_loop_forever(dashboard_consts.RAY_DASHBOARD_STATS_UPDATING_INTERVAL)
+    async def organize(cls, thread_pool_executor):
+        """
+        Organizes data: read from (node_physical_stats, node_stats) and updates
+        (node_workers, node_worker_stats).
+
+        This methods is not really async, but DataSource is not thread safe so we need
+        to make sure it's on the main event loop thread. To avoid blocking the main
+        event loop, we yield after each node processed.
+        """
+        loop = get_or_create_event_loop()
+
+        node_workers = {}
+        core_worker_stats = {}
+
+        # NOTE: We copy keys of the `DataSource.nodes` to make sure
+        #       it doesn't change during the iteration (since its being updated
+        #       from another async task)
+        for node_id in list(DataSource.nodes.keys()):
+            node_physical_stats = DataSource.node_physical_stats.get(node_id, {})
+            node_stats = DataSource.node_stats.get(node_id, {})
+            # Offloads the blocking operation to a thread pool executor. This also
+            # yields to the event loop.
+            workers = await loop.run_in_executor(
+                thread_pool_executor,
+                cls._extract_workers_for_node,
+                node_physical_stats,
+                node_stats,
+            )
+
+            for worker in workers:
+                for stats in worker.get("coreWorkerStats", []):
+                    worker_id = stats["workerId"]
+                    core_worker_stats[worker_id] = stats
+
+            node_workers[node_id] = workers
+
+        DataSource.node_workers = node_workers
+        DataSource.core_worker_stats = core_worker_stats
+
+    @classmethod
+    def _extract_workers_for_node(cls, node_physical_stats, node_stats):
+        workers = []
+        # Merge coreWorkerStats (node stats) to workers (node physical stats)
+        pid_to_worker_stats = {}
+        pid_to_language = {}
+        pid_to_job_id = {}
+
+        for core_worker_stats in node_stats.get("coreWorkersStats", []):
+            pid = core_worker_stats["pid"]
+
+            pid_to_worker_stats[pid] = core_worker_stats
+            pid_to_language[pid] = core_worker_stats["language"]
+            pid_to_job_id[pid] = core_worker_stats["jobId"]
+
+        for worker in node_physical_stats.get("workers", []):
+            worker = dict(worker)
+            pid = worker["pid"]
+
+            core_worker_stats = pid_to_worker_stats.get(pid)
+            # Empty list means core worker stats is not available.
+            worker["coreWorkerStats"] = [core_worker_stats] if core_worker_stats else []
+            worker["language"] = pid_to_language.get(
+                pid, dashboard_consts.DEFAULT_LANGUAGE
+            )
+            worker["jobId"] = pid_to_job_id.get(pid, dashboard_consts.DEFAULT_JOB_ID)
+
+            workers.append(worker)
+
+        return workers
+
+    @classmethod
+    async def get_node_info(cls, node_id, get_summary=False):
+        node_physical_stats = dict(DataSource.node_physical_stats.get(node_id, {}))
+        node_stats = dict(DataSource.node_stats.get(node_id, {}))
+        node = DataSource.nodes.get(node_id, {})
+
+        if get_summary:
+            node_physical_stats.pop("workers", None)
+            node_stats.pop("workersStats", None)
+        else:
+            node_stats.pop("coreWorkersStats", None)
+        store_stats = node_stats.get("storeStats", {})
+        used = int(store_stats.get("objectStoreBytesUsed", 0))
+        # objectStoreBytesAvail == total in the object_manager.cc definition.
+        total = int(store_stats.get("objectStoreBytesAvail", 0))
+        ray_stats = {
+            "object_store_used_memory": used,
+            "object_store_available_memory": total - used,
+        }
+
+        node_info = node_physical_stats
+
+        # Merge node stats to node physical stats under raylet
+        node_info["raylet"] = node_stats
+        node_info["raylet"].update(ray_stats)
+
+        # Merge GcsNodeInfo to node physical stats
+        node_info["raylet"].update(node)
+        death_info = node.get("deathInfo", {})
+        node_info["raylet"]["stateMessage"] = compose_state_message(
+            death_info.get("reason", None), death_info.get("reasonMessage", None)
+        )
+
+        # TODO(spencer-p): TPU process linking is currently prone to over-counting
+        # as it attaches all TPU workers to every chip. This will be addressed
+        # with a more robust mapping in a future update.
+        node_info["tpus"] = copy.deepcopy(node_info.get("tpus", []))
+
+        if not get_summary:
+            actor_table_entries = DataSource.node_actors.get(node_id, {})
+
+            # Merge actors to node physical stats
+            node_info["actors"] = {
+                actor_id: await DataOrganizer._get_actor_info(actor_table_entry)
+                for actor_id, actor_table_entry in actor_table_entries.items()
+            }
+
+            # Update workers to node physical stats
+            node_info["workers"] = DataSource.node_workers.get(node_id, [])
+
+        return node_info
+
+    @classmethod
+    async def get_all_node_summary(cls):
+        return [
+            # NOTE: We're intentionally awaiting in a loop to avoid excessive
+            #       concurrency spinning up excessive # of tasks for large clusters
+            await DataOrganizer.get_node_info(node_id, get_summary=True)
+            for node_id in DataSource.nodes.keys()
+        ]
+
+    @classmethod
+    async def get_actor_infos(cls, actor_ids: Optional[List[str]] = None):
+        target_actor_table_entries: dict[str, Optional[dict]]
+        if actor_ids is not None:
+            target_actor_table_entries = {
+                actor_id: DataSource.actors.get(actor_id) for actor_id in actor_ids
+            }
+        else:
+            target_actor_table_entries = DataSource.actors
+
+        return {
+            actor_id: await DataOrganizer._get_actor_info(actor_table_entry)
+            for actor_id, actor_table_entry in target_actor_table_entries.items()
+        }
+
+    @staticmethod
+    async def _get_actor_info(actor: Optional[dict]) -> Optional[dict]:
+        if actor is None:
+            return None
+
+        actor = actor.copy()
+        worker_id = actor["address"]["workerId"]
+        core_worker_stats = DataSource.core_worker_stats.get(worker_id, {})
+        actor.update(core_worker_stats)
+
+        # TODO(fyrestone): remove this, give a link from actor
+        # info to worker info in front-end.
+        node_id = actor["address"]["nodeId"]
+        pid = core_worker_stats.get("pid")
+        node_physical_stats = DataSource.node_physical_stats.get(node_id, {})
+        actor_process_stats = None
+        actor_process_gpu_stats = []
+        if pid:
+            for process_stats in node_physical_stats.get("workers", []):
+                if process_stats["pid"] == pid:
+                    actor_process_stats = process_stats
+                    break
+
+            for gpu_stats in node_physical_stats.get("gpus", []):
+                # gpu_stats.get("processesPids") can be None, an empty list or a
+                # list of dictionaries.
+                for process in gpu_stats.get("processesPids") or []:
+                    if process["pid"] == pid:
+                        actor_process_gpu_stats.append(gpu_stats)
+                        break
+
+        actor["gpus"] = actor_process_gpu_stats
+        actor["processStats"] = actor_process_stats
+        actor["mem"] = node_physical_stats.get("mem", [])
+
+        required_resources = parse_pg_formatted_resources_to_original(
+            actor["requiredResources"]
+        )
+        actor["requiredResources"] = required_resources
+
+        # TODO(spencer-p): TPU process linking is currently prone to over-counting.
+        # This will be addressed with a more robust mapping in a future update.
+        actor["tpus"] = []
+
+        return actor