wehub-resource-sync
138 lines
5.4 KiB
Python
138 lines
5.4 KiB
Python
# Copyright (c) ModelScope Contributors. All rights reserved.
|
|
# Some code borrowed from ROLL: https://github.com/alibaba/ROLL
|
|
import math
|
|
import os
|
|
from dataclasses import dataclass, field
|
|
from typing import Any, Dict, List
|
|
|
|
|
|
@dataclass
|
|
class NodeGroup:
|
|
device_count: int
|
|
nodes: List[Any] = field(default_factory=list)
|
|
|
|
|
|
def get_node_rank():
|
|
return int(os.environ.get('NODE_RANK', '0'))
|
|
|
|
|
|
class ResourceManager:
|
|
|
|
possible_keys = ['nproc_per_node', 'nnodes']
|
|
|
|
def __init__(self, groups: Dict[str, Any]):
|
|
import ray
|
|
from ray.util.placement_group import PlacementGroup
|
|
nproc_per_node = int(groups['nproc_per_node'])
|
|
device_types = set([group['device'].upper()
|
|
for group in groups.values() if hasattr(group, '__getitem__')]) - {'CPU'}
|
|
assert len(device_types) == 1
|
|
device_type = next(iter(device_types))
|
|
all_ranks = []
|
|
last_rank = -1
|
|
cpu_proc_count = 0
|
|
for group_name, group in groups.items():
|
|
if group_name in self.possible_keys:
|
|
continue
|
|
ranks = group['ranks']
|
|
device = group['device'].upper()
|
|
if device == 'CPU':
|
|
assert isinstance(ranks, int), 'CPU group only supports integer ranks'
|
|
cpu_proc_count += ranks
|
|
continue
|
|
try:
|
|
ranks = int(ranks) # int type
|
|
ranks = list(range(last_rank + 1, last_rank + 1 + ranks))
|
|
except Exception: # noqa
|
|
if isinstance(ranks, str):
|
|
ranks = eval(ranks, {'__builtins__': {'list': list, 'range': range}})
|
|
finally:
|
|
all_ranks.extend(ranks)
|
|
group['ranks'] = ranks
|
|
last_rank = ranks[-1]
|
|
|
|
assert len(set(all_ranks)) == len(all_ranks)
|
|
groups['nnodes'] = math.ceil(len(all_ranks) / nproc_per_node)
|
|
|
|
self.nodes = []
|
|
for node in ray.nodes():
|
|
resource = node['Resources']
|
|
node_gpu_num = int(resource.get(device_type, 0))
|
|
if node_gpu_num >= nproc_per_node:
|
|
self.nodes.append(node)
|
|
|
|
bundles = []
|
|
cpu_bundles = []
|
|
for i in range(groups['nnodes']):
|
|
node = self.nodes[i]
|
|
node_cpu = int(node['Resources']['CPU'])
|
|
bundles.append({device_type: nproc_per_node, 'CPU': node_cpu // 2 + 1})
|
|
cpu_bundles.append({'CPU': node_cpu // 4 + 1}) # TODO dynamic scheduling
|
|
|
|
nproc_cpu_per_node = cpu_proc_count // len(cpu_bundles) + 1
|
|
self.cpu_node_map = {}
|
|
for i in range(cpu_proc_count):
|
|
node_idx = i // nproc_cpu_per_node
|
|
cpu_cnt = cpu_bundles[node_idx]['CPU']
|
|
self.cpu_node_map[i] = (node_idx, cpu_cnt // nproc_cpu_per_node)
|
|
|
|
self.placement_groups = [ray.util.placement_group([bundle]) for bundle in bundles]
|
|
self.cpu_placement_groups = [ray.util.placement_group([bundle]) for bundle in cpu_bundles]
|
|
cpu_bundles.sort(key=lambda bundle: bundle['CPU'], reverse=True)
|
|
ray.get([pg.ready() for pg in self.placement_groups])
|
|
ray.get([pg.ready() for pg in self.cpu_placement_groups])
|
|
|
|
self.node_ranks = ray.get(
|
|
[ray.remote(get_node_rank).options(placement_group=pg).remote() for pg in self.placement_groups])
|
|
if self.node_ranks.count(0) > 1:
|
|
self.node_ranks = list(range(len(self.placement_groups)))
|
|
|
|
self.node2pg: Dict[int, PlacementGroup] = {}
|
|
for node_rank, placement_group in zip(self.node_ranks, self.placement_groups):
|
|
self.node2pg[node_rank] = placement_group
|
|
|
|
self.device_groups = {}
|
|
ray_address = str(ray.get_runtime_context().gcs_address)
|
|
for group_name, group in groups.items():
|
|
if group_name in self.possible_keys:
|
|
continue
|
|
|
|
if group['device'] != 'CPU':
|
|
ranks = group['ranks']
|
|
local_device_groups = []
|
|
for rank in ranks:
|
|
node_rank = rank // nproc_per_node
|
|
gpu_rank = rank % nproc_per_node
|
|
local_device_groups.append(
|
|
dict(
|
|
node_rank=node_rank,
|
|
gpu_rank=[gpu_rank],
|
|
placement_group=self.node2pg[node_rank],
|
|
ray_address=ray_address))
|
|
for worker in group['workers']:
|
|
self.device_groups[worker] = local_device_groups
|
|
else:
|
|
ranks = group['ranks']
|
|
local_device_groups = []
|
|
global_cpu_proc_idx = 0
|
|
for _ in range(ranks):
|
|
local_device_groups.append(
|
|
dict(
|
|
placement_group=self.cpu_placement_groups[self.cpu_node_map[global_cpu_proc_idx][0]],
|
|
ray_address=ray_address))
|
|
global_cpu_proc_idx += 1
|
|
for worker in group['workers']:
|
|
self.device_groups[worker] = local_device_groups
|
|
|
|
self.groups = groups
|
|
|
|
def resource(self, worker):
|
|
return self.device_groups[worker]
|
|
|
|
def destroy_placement_group(self):
|
|
import ray
|
|
for pg in self.placement_groups:
|
|
ray.util.remove_placement_group(pg)
|
|
for pg in self.cpu_placement_groups:
|
|
ray.util.remove_placement_group(pg)
|