import threading
import time
import torch
import torch_npu
import torch.distributed as dist


torch.npu.config.allow_internal_format = True

dist.init_process_group(backend="hccl")
rank = dist.get_rank()
world_size = dist.get_world_size()

torch.npu.set_device(rank)


new_pg = dist.new_group(ranks=[0, 1], backend="hccl")

a = torch.ones(2, 3).npu()
b  = torch.rand(3).npu()

src_tensor = torch.zeros(3, 2 ).npu()
src_tensor1 = torch.rand(3).npu()

recv_tensor  = torch.ones_like(src_tensor)
recv_tensor1 = torch.ones_like(src_tensor1)

ds  = torch.npu.current_stream()
print(f"rank {rank} ds {ds}", flush=True)


def f():
    if rank == 0:
        time.sleep(50)
        dist.send(src_tensor1, 1, group=new_pg)
        print(f"rank {rank} send", flush=True)

    else:
        dist.recv(recv_tensor1, 0, group=new_pg)
        print(f"rank {rank} recv", flush=True)

def f2():
    if rank == 1:
        time.sleep(80)
    dist.all_reduce(b)
    print(f"rank {rank} allreduce 2", flush=True)


c = torch.rand(2).npu()


def f3():
    global c
    time.sleep(20)
    torch.npu.set_device(rank)
    s = torch.npu.Stream()
    s.wait_stream(ds)
    with torch.npu.stream(s):
        for i in range(5000):
            c = c + 1
    print(f"rank {rank} f3", flush=True)

t = threading.Thread(target=f, args=())
t2 = threading.Thread(target=f2, args=())
t3 = threading.Thread(target=f3, args=())


if rank == 0:
    dist.send(src_tensor, 1, group=new_pg)
    dist.all_reduce(a)
else:
    dist.recv(recv_tensor, 0, group=new_pg)
    dist.all_reduce(a)

t.start()
t2.start()
t3.start()

t.join()
t2.join()
t3.join()

print(f"rank {rank} ok", flush=True)    

脚本运行 torchrun --nproc-per-node=2 test.py

这个脚本会卡住，可以配置export HCCL_EXEC_TIMEOUT=180让其快速超时退出

可以通过export ASCEND_PROCESS_LOG_PATH=xxx指定plog日志目录
可以通过export HCCL_ENTRY_LOG_ENABLE=1开启hccl enrty日志，对应在plog目录下查看Entry看到hccl的下发记录

这个脚本在默认开启taskqueue时，虽然多线程多流下发，但是最终都是一个taskqueue串行下发，导致hccl有依赖的任务会卡住，可以在entry中看到rank 0,1上最后一个hccl任务没有下发
如果关闭taskqueue或者配置PER_STREAM_QUEUE=1就可以看到都下发了（虽然下发，由于下发顺序不一致hccl依旧会超时报错）

推荐做法是保证有依赖的任务的下发顺序，脚本中去掉time.sleep()，让rank 0,1 send/recv以及allreduce都能配对下发即可