nano-vllm/task_plan.md

Task Plan: Fix Torch Distributed Port Conflict

Goal

支持多卡环境下同时启动多个独立的 nanovllm 进程进行测试，无需手动管理端口。

Problem Analysis

核心问题

当前：所有 nanovllm 实例默认使用端口 2333
     └── 多个独立进程同时运行时会冲突！

CUDA_VISIBLE_DEVICES=0 python test1.py  # 绑定端口 2333 ✓
CUDA_VISIBLE_DEVICES=1 python test2.py  # 尝试绑定 2333 → EADDRINUSE ❌

根本原因

• 端口是系统级资源，与 GPU 无关
• 即使使用不同 GPU，端口仍会冲突
• 当前硬编码默认端口 2333

---

Solution: Dynamic Port Allocation

核心方案

def _find_free_port() -> int:
    """让系统自动分配一个空闲端口"""
    with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
        s.bind(('', 0))
        return s.getsockname()[1]

# 优先使用环境变量，否则自动分配
port = os.environ.get("NANOVLLM_DIST_PORT")
if port is None:
    port = _find_free_port()
else:
    port = int(port)

效果

# 无需手动指定端口，可以同时运行多个测试
CUDA_VISIBLE_DEVICES=0 python test1.py &  # 自动端口 54321
CUDA_VISIBLE_DEVICES=1 python test2.py &  # 自动端口 54322
CUDA_VISIBLE_DEVICES=2 python test3.py &  # 自动端口 54323

# 仍然支持手动指定（向后兼容）
NANOVLLM_DIST_PORT=2333 python test.py

---

Implementation Phases

Phase 1: ModelRunner 动态端口 [pending]

File: nanovllm/engine/model_runner.py

import socket

def _find_free_port() -> int:
    """Find a free port for distributed communication."""
    with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
        s.bind(('', 0))
        s.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
        return s.getsockname()[1]

class ModelRunner:
    def __init__(self, config: Config, rank: int, event: Event | list[Event]):
        # ... existing code ...

        import os
        port = os.environ.get("NANOVLLM_DIST_PORT")
        if port is None:
            port = _find_free_port()
            logger.info(f"Auto-assigned distributed port: {port}")
        else:
            port = int(port)

        dist.init_process_group("nccl", f"tcp://localhost:{port}", ...)

Phase 2: LLMEngine 资源清理增强 [pending]

File: nanovllm/engine/llm_engine.py

添加 close() 方法和 context manager 支持，确保资源正确释放：

class LLMEngine:
    def __init__(self, model, **kwargs):
        # ... existing code ...
        self._closed = False
        atexit.register(self._atexit_handler)

    def _atexit_handler(self):
        if not self._closed:
            self.close()

    def close(self):
        """Explicitly close the engine and release all resources."""
        if self._closed:
            return
        self._closed = True
        try:
            atexit.unregister(self._atexit_handler)
        except Exception:
            pass
        self.model_runner.call("exit")
        del self.model_runner
        for p in self.ps:
            p.join()

    def exit(self):
        """Alias for close() - backward compatibility."""
        self.close()

    def __del__(self):
        try:
            self.close()
        except Exception:
            pass

    def __enter__(self):
        return self

    def __exit__(self, *args):
        self.close()
        return False

Phase 3: 测试验证 [pending]

File: tests/test_multiple_processes.py (新建)

"""Test multiple independent nanovllm processes."""
import subprocess
import sys
import time

def test_parallel_processes():
    """Test running multiple nanovllm processes in parallel."""
    script = '''
import sys
sys.path.insert(0, ".")
from nanovllm import LLM, SamplingParams
import os

gpu = os.environ.get("CUDA_VISIBLE_DEVICES", "0")
print(f"[GPU {gpu}] Starting LLM")
llm = LLM("path/to/model", enable_cpu_offload=True)
outputs = llm.generate(["Hello"], SamplingParams(max_tokens=10))
print(f"[GPU {gpu}] Output: {outputs[0]['text'][:50]}")
llm.close()
print(f"[GPU {gpu}] Done")
'''

    # Start 2 processes on different GPUs
    procs = []
    for gpu in [0, 1]:
        env = {"CUDA_VISIBLE_DEVICES": str(gpu)}
        p = subprocess.Popen(
            [sys.executable, "-c", script],
            env={**os.environ, **env}
        )
        procs.append(p)
        time.sleep(1)  # Stagger start slightly

    # Wait for all
    for p in procs:
        assert p.wait() == 0, f"Process failed with code {p.returncode}"

    print("PASSED: test_parallel_processes")

if __name__ == "__main__":
    test_parallel_processes()

Phase 4: 文档更新 [pending]

File: docs/torch_distributed_port_issue.md

更新文档标记问题已通过动态端口分配解决。

---

Usage After Fix

场景 1: 多进程并行测试（主要场景）

# 无需任何额外配置，直接运行
CUDA_VISIBLE_DEVICES=0 python test_group1.py &
CUDA_VISIBLE_DEVICES=1 python test_group2.py &
CUDA_VISIBLE_DEVICES=2 python test_group3.py &
wait

场景 2: 同一进程顺序创建（也支持）

for i in range(3):
    with LLM(model_path) as llm:
        outputs = llm.generate(prompts, params)
    # 自动清理，下一个可以使用新的随机端口

场景 3: 手动指定端口（向后兼容）

NANOVLLM_DIST_PORT=2333 python test.py

---

Success Criteria

• 多个独立进程可以同时运行（不同 GPU）
• 无需手动指定端口
• 向后兼容（环境变量仍有效）
• 同一进程顺序创建也能工作
• 资源正确清理

---

Files to Modify

File	Action	Status
nanovllm/engine/model_runner.py	Add _find_free_port()	pending
nanovllm/engine/llm_engine.py	Add close(), context manager	pending
tests/test_multiple_processes.py	Create	pending
docs/torch_distributed_port_issue.md	Update	pending