maga_transformer/distribute/gang_info.py

import json
import logging
import os
import socket
from typing import NamedTuple, List, Any, Dict, Optional

from maga_transformer.distribute.worker_info import g_worker_info, g_parallel_info, WorkerInfo

CONFIG_FILE_ENV = 'DISTRIBUTE_CONFIG_FILE'

def members_from_json(gang_info_json: Dict[str, Any]) -> List[WorkerInfo]:
    members: List[WorkerInfo] = []
    # here is only the fake ip
    for name, info in gang_info_json.items():
        server_port = info['port'] if 'port' in info else -1
        members.append(WorkerInfo(
            server_port=server_port,
            gang_hb_port=-1,
            http_port=-1,
            rpc_server_port=-1,
            backend_server_port=-1,
            remote_rpc_server_port=-1,
            cache_store_listen_port=-1,
            cache_store_connect_port=-1,
            cache_store_rdma_connect_port=-1,
            cache_store_rdma_listen_port=-1,
            local_rank=0,
            world_rank=0,
            name=info['name'], ip=info['ip'], info=info))
    zone_name = os.environ.get("ZONE_NAME", "")
    if zone_name:
        members = [member for member in members if member.name.split('_')[-2] == zone_name]
    masters = [member for member in members if member.name.endswith('part0')]
    if len(masters) != 1:
        raise Exception(f"gang master should contains 1 but got {len(masters)}")
    return sorted(members, key=lambda x:x.name)

'''
test env example:
name:smoke_part0,ip:127.0.0.1,port:13045;name:smoke_part1,ip:127.0.0.1,port:12053
'''
def members_from_test_env(env_str: str) -> List[WorkerInfo]:
    members: List[WorkerInfo] = []
    for member_str in env_str.split(';'):
        member_info = {}
        for item in member_str.split(','):
            key, value = item.split(':')
            member_info[key] = value
        members.append(WorkerInfo(
            server_port=int(member_info['port']),
            gang_hb_port=-1,
            http_port=-1,
            rpc_server_port=-1,
            backend_server_port=-1,
            remote_rpc_server_port=-1,
            cache_store_listen_port=-1,
            cache_store_connect_port=-1,
            cache_store_rdma_connect_port=-1,
            cache_store_rdma_listen_port=-1,
            local_rank=0,
            world_rank=0,
            name=member_info['name'], ip=member_info['ip'], info=member_info))
    masters = [member for member in members if member.name.endswith('part0')]
    if len(masters) != 1:
        raise Exception(f"gang master should contains 1 but got {len(masters)}")
    sorted_members = sorted(members, key=lambda x:x.name)
    if masters[0].name != sorted_members[0].name:
        raise Exception(f"gang master should be the first one but got {sorted_members[0].name}")
    return sorted_members

'''
raw gang info example:
app.c2.io/biz-detail-ganginfo="{\"llama13B_2A10_PCIE_1_inference_part0\":{\"name\":\"llama13B_2A10_PCIE_1_inference_part0\",\"ip\":\"33.76.194.173\"},\"llama13B_2A10_PCIE_1_inference_part1\":{\"name\":\"llama13B_2A10_PCIE_1_inference_part1\",\"ip\":\"33.76.194.182\"}}"
'''
def get_c2_members():
    file_name = os.environ.get("GANG_ANNOCATION_PATH", "/etc/podinfo/annotations")
    if not os.path.exists(file_name):
        raise Exception(f"not found file: {file_name}")

    with open(file_name, 'r') as reader:
        content = reader.read()

    infos = [x for x in content.split("\n") if "app.c2.io/biz-detail-ganginfo" in x]
    if len(infos) != 1:
        raise Exception(f"ganginfo length is not equal to 1, actual: {infos}")

    gang_info = infos[0].replace("\\", "")
    logging.info(f"gang info: {gang_info[gang_info.index('=') + 2: -1]}")
    gang_info_json = json.loads(gang_info[gang_info.index('=') + 2: -1])
    logging.info(f"gang info json: {gang_info_json}")
    return members_from_json(gang_info_json)

def get_members_from_file():
    file = os.environ[CONFIG_FILE_ENV]
    with open(file, 'r') as reader:
        config_json = json.loads(reader.read())
    return members_from_json(config_json)

class GangInfo(NamedTuple):
    members: List[WorkerInfo]
    master: WorkerInfo
    self: WorkerInfo
    num_nodes: int

    def workers(self) -> List[WorkerInfo]:
        return [member for member in self.members if not member.equals(self.master)]

def get_gang_info() -> GangInfo:
    if g_parallel_info.local_world_size < g_parallel_info.world_size:
        # from config file
        if os.environ.get(CONFIG_FILE_ENV):
            members = get_members_from_file()
        # for distributed test
        elif os.environ.get("GANG_CONFIG_STRING"):
            logging.info(f"use GANG_CONFIG_STRING: {os.environ['GANG_CONFIG_STRING']}")
            members = members_from_test_env(os.environ['GANG_CONFIG_STRING'])
        # from c2 annotation
        else:
            members = get_c2_members()
    else:
        members = [WorkerInfo(socket.gethostbyname(socket.gethostname()),
                              -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, 0, 0, 'local', None)]

    # 假设 GPU 均匀分布，可以整除
    # member 是按 part 排序的
    self: Optional[WorkerInfo] = None
    master: Optional[WorkerInfo] = None
    all_members: List[WorkerInfo] = []
    for part_rank, member in enumerate(members):
        for local_rank in range(g_parallel_info.local_world_size):
            new_member = WorkerInfo(
                ip=member.ip,
                server_port=WorkerInfo.server_port_offset(local_rank, member.server_port),
                gang_hb_port=WorkerInfo.gang_hb_port_offset(local_rank, member.server_port),
                http_port=WorkerInfo.http_port_offset(local_rank, member.server_port),
                rpc_server_port=WorkerInfo.rpc_server_port_offset(local_rank, member.server_port),
                backend_server_port=WorkerInfo.backend_server_port_offset(local_rank, member.server_port),
                cache_store_listen_port=WorkerInfo.cache_store_listen_port_offset(local_rank, member.server_port),
                cache_store_rdma_listen_port=WorkerInfo.cache_store_rdma_listen_port_offset(local_rank, member.server_port),
                remote_rpc_server_port=WorkerInfo.rpc_server_port_offset(local_rank, int(os.environ.get("REMOTE_SERVER_PORT", 0))),
                cache_store_connect_port=WorkerInfo.cache_store_listen_port_offset(local_rank, int(os.environ.get("REMOTE_SERVER_PORT", 0))),
                cache_store_rdma_connect_port=WorkerInfo.cache_store_rdma_listen_port_offset(local_rank, int(os.environ.get("REMOTE_SERVER_PORT", 0))),
                local_rank=local_rank,
                world_rank=part_rank * g_parallel_info.local_world_size + local_rank,
                name=member.name + '_' + str(local_rank),
                info=member.info)
            all_members.append(new_member)
            logging.info(f"local rank {local_rank} vs {g_parallel_info.local_rank}, \
                         new_member: {new_member.ip} vs {g_worker_info.ip}, \
                         server port {new_member.server_port} vs {g_worker_info.server_port}")
            if (local_rank == g_parallel_info.local_rank and
                new_member.ip == g_worker_info.ip and new_member.server_port == g_worker_info.server_port):
                self = new_member
            if part_rank == 0 and local_rank == 0:
                master = new_member
    # not check master and self empty here for ut
    return GangInfo(all_members, master, self, len(members))