wehub-resource-sync
77 lines
3.2 KiB
Python
77 lines
3.2 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
|
"""Unit tests for OmniDiffusionConfig master_port resolution (issue #3794)."""
|
|
|
|
import socket
|
|
|
|
import pytest
|
|
|
|
from vllm_omni.diffusion.data import OmniDiffusionConfig
|
|
|
|
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
|
|
|
|
|
|
def _get_available_port() -> int:
|
|
with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as sock:
|
|
sock.bind(("", 0))
|
|
return sock.getsockname()[1]
|
|
|
|
|
|
def _bind_available_port_below(max_port: int, port_inc: int) -> socket.socket:
|
|
for _ in range(100):
|
|
sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
|
|
try:
|
|
sock.bind(("", 0))
|
|
port = sock.getsockname()[1]
|
|
if port >= max_port - port_inc:
|
|
sock.close()
|
|
continue
|
|
with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as next_sock:
|
|
next_sock.bind(("", port + port_inc))
|
|
return sock
|
|
except OSError:
|
|
sock.close()
|
|
raise RuntimeError(f"Failed to reserve an available port below {max_port - port_inc}")
|
|
|
|
|
|
class TestOmniDiffusionConfigMasterPort:
|
|
def test_honors_master_port_env(self, monkeypatch: pytest.MonkeyPatch) -> None:
|
|
port = _get_available_port()
|
|
monkeypatch.setenv("MASTER_PORT", str(port))
|
|
config = OmniDiffusionConfig(model="test")
|
|
assert config.master_port == port
|
|
|
|
def test_honors_explicit_master_port(self, monkeypatch: pytest.MonkeyPatch) -> None:
|
|
monkeypatch.delenv("MASTER_PORT", raising=False)
|
|
port = _get_available_port()
|
|
config = OmniDiffusionConfig(model="test", master_port=port)
|
|
assert config.master_port == port
|
|
|
|
def test_master_port_env_takes_precedence_over_kwarg(self, monkeypatch: pytest.MonkeyPatch) -> None:
|
|
env_port = _get_available_port()
|
|
kwarg_port = _get_available_port()
|
|
monkeypatch.setenv("MASTER_PORT", str(env_port))
|
|
config = OmniDiffusionConfig(model="test", master_port=kwarg_port)
|
|
assert config.master_port == env_port
|
|
|
|
def test_explicit_master_port_is_stable_without_jitter(self, monkeypatch: pytest.MonkeyPatch) -> None:
|
|
monkeypatch.delenv("MASTER_PORT", raising=False)
|
|
port = _get_available_port()
|
|
ports = {OmniDiffusionConfig(model="test", master_port=port).master_port for _ in range(5)}
|
|
assert ports == {port}
|
|
|
|
def test_default_master_port_is_available(self, monkeypatch: pytest.MonkeyPatch) -> None:
|
|
monkeypatch.delenv("MASTER_PORT", raising=False)
|
|
config = OmniDiffusionConfig(model="test")
|
|
assert config.master_port is not None
|
|
with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as sock:
|
|
sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
|
|
sock.bind(("", config.master_port))
|
|
|
|
def test_settle_port_steps_when_port_is_busy(self, monkeypatch: pytest.MonkeyPatch) -> None:
|
|
monkeypatch.delenv("MASTER_PORT", raising=False)
|
|
with _bind_available_port_below(max_port=60000, port_inc=37) as sock:
|
|
busy_port = sock.getsockname()[1]
|
|
config = OmniDiffusionConfig(model="test", master_port=busy_port)
|
|
assert config.master_port == busy_port + 37
|