v0.10.1rc1

examples/disaggregated_prefill_v1/gen_ranktable.py

@@ -0,0 +1,122 @@
+import argparse
+import json
+import os
+
+import torch.distributed as dist
+
+from vllm_ascend.utils import AscendSocVersion, init_ascend_soc_version, get_ascend_soc_version
+
+parser = argparse.ArgumentParser(
+    description="Arguments of rank table generator", )
+parser.add_argument("--local-host", type=str, required=True, help="local ip")
+parser.add_argument("--prefill-device-cnt",
+                    type=int,
+                    required=True,
+                    help="number of prefill devices")
+parser.add_argument("--decode-device-cnt",
+                    type=int,
+                    required=True,
+                    help="number of decode devices")
+args = parser.parse_args()
+local_host = args.local_host
+prefill_device_cnt = args.prefill_device_cnt
+decode_device_cnt = args.decode_device_cnt
+
+print("enter py")
+
+hccn_tool_path = os.environ.get("HCCN_TOOL_PATH",
+                                "/usr/local/Ascend/driver/tools/hccn_tool")
+master_addr = os.environ.get("MASTER_ADDR")
+master_port = os.environ.get("MASTER_PORT")
+rank = os.environ.get("RANK")
+local_rank = os.environ.get("LOCAL_RANK")
+# This variable is set by torchrun,
+# and is different from WORLD_SIZE in gen_rank_table.sh.
+world_size = os.environ.get("WORLD_SIZE")
+
+init_ascend_soc_version()
+soc_info = get_ascend_soc_version()
+
+
+def get_cmd_stdout(cmd):
+    import subprocess
+    return subprocess.run(cmd, capture_output=True,
+                          shell=True).stdout.decode("utf-8").strip()
+
+
+print(f"local_host: {local_host}")
+print("gen ranktable.json")
+
+num_cards = get_cmd_stdout("npu-smi info -l | grep \"Total Count\"").split(
+    ":")[1].strip()
+num_cards = int(num_cards)
+chips_per_card = get_cmd_stdout("npu-smi info -l | grep \"Chip Count\"").split(
+    "\n")[0].split(":")[1].strip()
+chips_per_card = int(chips_per_card)
+
+# generate local device list for local rank 0, and gather it to all ranks
+local_device_list: list[dict[str, str]] = list()
+if local_rank == "0":
+    super_pod_id = "0"
+    for card_id in range(num_cards):
+        for chip_id in range(chips_per_card):
+            device_id = card_id * chips_per_card + chip_id
+            if soc_info == AscendSocVersion.A3:
+                device_ip = get_cmd_stdout(
+                    f"{hccn_tool_path} -i {device_id} -vnic -g | grep ipaddr"
+                ).split(":")[1].strip()
+                super_device_id = get_cmd_stdout(
+                    f"npu-smi info -t spod-info -i {card_id} -c {chip_id} | grep SDID"
+                ).split(":")[1].strip()
+                super_pod_id = get_cmd_stdout(
+                    f"npu-smi info -t spod-info -i {card_id} -c {chip_id} | grep \"Super Pod ID\""
+                ).split(":")[1].strip()
+            else:
+                device_ip = get_cmd_stdout(
+                    f"{hccn_tool_path} -i {device_id} -ip -g | grep ipaddr"
+                ).split(":")[1].strip()
+
+            device_info = {
+                "server_id": local_host,
+                "device_id": str(device_id),
+                "device_ip": str(device_ip),
+            }
+            if soc_info == AscendSocVersion.A3:
+                device_info.update({
+                    "super_pod_id": str(super_pod_id),
+                    "super_device_id": str(super_device_id)
+                })
+            local_device_list.append(device_info)
+
+dist.init_process_group(backend=dist.Backend.GLOO)
+global_device_list = [None] * dist.get_world_size()
+dist.all_gather_object(global_device_list, local_device_list)
+global_device_list = [
+    device_info for device_list in global_device_list
+    for device_info in device_list  # type: ignore[attr-defined]
+]
+cnt = 1
+for device_info in global_device_list:  # type: ignore[assignment]
+    device_info["cluster_id"] = str(cnt)
+    cnt += 1
+assert (prefill_device_cnt + decode_device_cnt) <= len(global_device_list), \
+"prefill_device_cnt + decode_device_cnt must be less than or equal to number of all devices in cluster"
+ranktable = {
+    "version":
+    "1.2",
+    "server_count":
+    str(world_size),
+    "prefill_device_list":
+    global_device_list[:prefill_device_cnt],
+    "decode_device_list":
+    global_device_list[prefill_device_cnt:prefill_device_cnt +
+                       decode_device_cnt],
+    "status":
+    "completed"
+}
+
+if local_rank == '0':
+    with open("ranktable.json", "w") as f:
+        json.dump(ranktable, f, indent=4)
+
+    print("gen ranktable.json done")