475 lines
17 KiB
Python
475 lines
17 KiB
Python
# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
|
|
# SPDX-License-Identifier: Apache-2.0
|
|
|
|
"""
|
|
Test configuration and fixtures for Dynamo Python bindings tests.
|
|
|
|
TWO MODES OF OPERATION:
|
|
|
|
1. Isolated Mode (ENABLE_ISOLATED_ETCD_AND_NATS=1):
|
|
- Each test gets fresh NATS/ETCD on random ports
|
|
- Requires: pytest-forked (uv pip install pytest-forked)
|
|
- Tests using 'runtime' fixture MUST have @pytest.mark.forked
|
|
- Safer, enables parallel execution
|
|
- Run: ENABLE_ISOLATED_ETCD_AND_NATS=1 pytest tests/ -n auto
|
|
|
|
2. Default Ports Mode (ENABLE_ISOLATED_ETCD_AND_NATS=0, default):
|
|
- All tests share NATS/ETCD on default ports (4222, 2379)
|
|
- No pytest-forked required
|
|
- No @pytest.mark.forked required
|
|
- Faster for sequential runs, but NO parallel execution
|
|
- Run: pytest tests/
|
|
|
|
Performance comparison (32-core machine, 13 tests):
|
|
Default ports (ENABLE_ISOLATED_ETCD_AND_NATS=0, default): 4.06s (sequential only)
|
|
Isolated sequential (ENABLE_ISOLATED_ETCD_AND_NATS=1): 8.58s (2.1x slower, but safer)
|
|
Isolated parallel -n 8: 2.82s (1.4x faster than default)
|
|
Isolated parallel -n 16: 2.28s (1.8x faster than default, optimal)
|
|
Isolated parallel -n 32: 2.74s (overhead dominates)
|
|
|
|
Recommendation: Default mode for simplicity. Use ENABLE_ISOLATED_ETCD_AND_NATS=1
|
|
with -n 8 to -n 16 when you need parallel execution and maximum test isolation.
|
|
"""
|
|
|
|
import asyncio
|
|
import json
|
|
import os
|
|
import re
|
|
import shutil
|
|
import socket
|
|
import subprocess
|
|
import tempfile
|
|
import time
|
|
|
|
import pytest
|
|
|
|
from dynamo.runtime import DistributedRuntime
|
|
|
|
# Configuration constants
|
|
# ENABLE_ISOLATED_ETCD_AND_NATS: When True, each test gets isolated NATS/ETCD instances
|
|
# on random ports with unique data directories. This enables parallel test execution.
|
|
# Set to False to use default ports (4222, 2379) for sequential execution.
|
|
# Can be overridden by environment variable: ENABLE_ISOLATED_ETCD_AND_NATS=0 or =1
|
|
ENABLE_ISOLATED_ETCD_AND_NATS = (
|
|
os.environ.get("ENABLE_ISOLATED_ETCD_AND_NATS", "0") == "1"
|
|
)
|
|
|
|
# Check if pytest-forked is installed (only when using isolated NATS/ETCD)
|
|
# This is REQUIRED when ENABLE_ISOLATED_ETCD_AND_NATS=1 because each test gets
|
|
# fresh services and the DistributedRuntime singleton needs process isolation
|
|
if ENABLE_ISOLATED_ETCD_AND_NATS:
|
|
try:
|
|
import pytest_forked # noqa: F401
|
|
except ImportError:
|
|
pytest.exit(
|
|
"""
|
|
pytest-forked is required when ENABLE_ISOLATED_ETCD_AND_NATS=1.
|
|
Install it with: uv pip install pytest-forked
|
|
|
|
This is needed because DistributedRuntime is a process-level singleton
|
|
and tests must run in separate processes to avoid 'Worker already initialized' errors.
|
|
|
|
Alternatively, set ENABLE_ISOLATED_ETCD_AND_NATS=0 to use default ports (slower, sequential only).
|
|
""",
|
|
returncode=1,
|
|
)
|
|
|
|
# Timeout constants
|
|
SERVICE_STARTUP_TIMEOUT = 5
|
|
SERVICE_SHUTDOWN_TIMEOUT = 5
|
|
|
|
|
|
def get_free_port():
|
|
"""Find and return an available port."""
|
|
sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
|
|
sock.bind(("", 0))
|
|
port = sock.getsockname()[1]
|
|
sock.close()
|
|
return port
|
|
|
|
|
|
def wait_for_port(host, port, timeout: float = SERVICE_STARTUP_TIMEOUT):
|
|
"""Wait for a port to be available."""
|
|
start = time.time()
|
|
while time.time() - start < timeout:
|
|
try:
|
|
sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
|
|
sock.settimeout(1)
|
|
sock.connect((host, port))
|
|
sock.close()
|
|
return True
|
|
except (socket.error, ConnectionRefusedError):
|
|
time.sleep(0.1)
|
|
return False
|
|
|
|
|
|
def start_nats_and_etcd_default_ports():
|
|
"""
|
|
Start NATS and ETCD on default ports (4222, 2379).
|
|
|
|
Use this for sequential test execution or when running tests alone.
|
|
Faster startup if services are already running.
|
|
"""
|
|
# Use default ports
|
|
nats_port = 4222
|
|
etcd_client_port = 2379
|
|
|
|
# No data directories needed - use defaults
|
|
nats_data_dir = None
|
|
etcd_data_dir = None
|
|
|
|
# Check if ports are already in use (reuse them if so)
|
|
# TODO: In the future, error out to ensure proper test isolation
|
|
nats_already_running = wait_for_port("localhost", nats_port, timeout=0.1)
|
|
etcd_already_running = wait_for_port("localhost", etcd_client_port, timeout=0.1)
|
|
|
|
if nats_already_running and etcd_already_running:
|
|
print(
|
|
f"Reusing existing NATS on port {nats_port} and ETCD on port {etcd_client_port}"
|
|
)
|
|
# Set environment variables for the runtime to use
|
|
os.environ["NATS_SERVER"] = f"nats://localhost:{nats_port}"
|
|
os.environ["ETCD_ENDPOINTS"] = f"http://localhost:{etcd_client_port}"
|
|
# Return None for processes since we're reusing existing services
|
|
return None, None, nats_port, etcd_client_port, None, None
|
|
|
|
# Set environment variables for the runtime to use
|
|
os.environ["NATS_SERVER"] = f"nats://localhost:{nats_port}"
|
|
os.environ["ETCD_ENDPOINTS"] = f"http://localhost:{etcd_client_port}"
|
|
|
|
print(f"Using NATS on default port {nats_port}")
|
|
print(f"Using ETCD on default client port {etcd_client_port}")
|
|
|
|
# Start services with default ports
|
|
nats_server = subprocess.Popen(["nats-server", "-js", "--trace"])
|
|
etcd = subprocess.Popen(["etcd"])
|
|
|
|
return nats_server, etcd, nats_port, etcd_client_port, nats_data_dir, etcd_data_dir
|
|
|
|
|
|
def start_nats_and_etcd_random_ports():
|
|
"""
|
|
Start NATS and ETCD with random ports and unique data directories.
|
|
|
|
This ensures test isolation by giving each test module (or parallel worker)
|
|
its own NATS/ETCD instances on different ports with separate data directories.
|
|
This allows tests to run in parallel without port or filesystem conflicts.
|
|
|
|
Note: etcd uses port 0 (OS-assigned port) to eliminate race conditions.
|
|
NATS uses get_free_port() with retry logic since it doesn't support port 0.
|
|
Port collision probability per NATS attempt: ~1% (heavy parallel testing), ~0.05% (normal load).
|
|
With 5 retries, probability of all NATS attempts failing: ~1.5e-10 (essentially never).
|
|
"""
|
|
# Create unique temporary data directories
|
|
nats_data_dir = tempfile.mkdtemp(prefix="nats_data_")
|
|
etcd_data_dir = tempfile.mkdtemp(prefix="etcd_data_")
|
|
|
|
# Start etcd first with port 0 (no race condition, no retries needed)
|
|
print(f"Starting ETCD with port 0 (OS-assigned), data dir: {etcd_data_dir}")
|
|
etcd = subprocess.Popen(
|
|
[
|
|
"etcd",
|
|
"--logger",
|
|
"zap",
|
|
"--data-dir",
|
|
str(etcd_data_dir),
|
|
"--listen-client-urls",
|
|
"http://localhost:0",
|
|
"--advertise-client-urls",
|
|
"http://localhost:0",
|
|
"--listen-peer-urls",
|
|
"http://localhost:0",
|
|
"--initial-advertise-peer-urls",
|
|
"http://localhost:0",
|
|
"--initial-cluster",
|
|
"default=http://localhost:0",
|
|
],
|
|
stderr=subprocess.PIPE,
|
|
stdout=subprocess.PIPE,
|
|
text=True,
|
|
bufsize=1,
|
|
)
|
|
|
|
# Parse etcd's stderr to discover the actual client port it bound to
|
|
etcd_client_port = None
|
|
timeout_at = time.time() + 5.0
|
|
|
|
while time.time() < timeout_at:
|
|
if etcd.poll() is not None:
|
|
stderr = etcd.stderr.read() if etcd.stderr else ""
|
|
shutil.rmtree(nats_data_dir, ignore_errors=True)
|
|
shutil.rmtree(etcd_data_dir, ignore_errors=True)
|
|
raise RuntimeError(f"ETCD failed to start: {stderr}")
|
|
|
|
line = etcd.stderr.readline() if etcd.stderr else ""
|
|
if not line:
|
|
time.sleep(0.01)
|
|
continue
|
|
|
|
try:
|
|
log = json.loads(line)
|
|
msg = log.get("msg", "")
|
|
|
|
# Look for the client port
|
|
if (
|
|
"serving client traffic" in msg
|
|
or "serving client" in msg
|
|
or "serving insecure client" in msg
|
|
):
|
|
address = log.get("address", "")
|
|
match = re.search(r":(\d+)$", address)
|
|
if match:
|
|
etcd_client_port = int(match.group(1))
|
|
print(f"ETCD bound to client port: {etcd_client_port}")
|
|
break
|
|
except (json.JSONDecodeError, ValueError):
|
|
continue
|
|
|
|
if etcd_client_port is None:
|
|
etcd.terminate()
|
|
etcd.wait()
|
|
shutil.rmtree(nats_data_dir, ignore_errors=True)
|
|
shutil.rmtree(etcd_data_dir, ignore_errors=True)
|
|
raise RuntimeError("Failed to discover ETCD client port from logs")
|
|
|
|
# Now start NATS with retry logic (up to 5 attempts due to race condition)
|
|
max_nats_retries = 5
|
|
nats_server = None
|
|
nats_port = None
|
|
last_error = None
|
|
|
|
for attempt in range(max_nats_retries):
|
|
try:
|
|
nats_port = get_free_port()
|
|
print(
|
|
f"Attempt {attempt + 1}: Starting NATS on port {nats_port}, data dir: {nats_data_dir}"
|
|
)
|
|
|
|
nats_server = subprocess.Popen(
|
|
["nats-server", "-js", "-p", str(nats_port), "-sd", str(nats_data_dir)],
|
|
stderr=subprocess.PIPE,
|
|
)
|
|
|
|
# Give NATS a moment to bind to the port
|
|
time.sleep(0.1)
|
|
|
|
# Check if NATS failed to start
|
|
if nats_server.poll() is not None:
|
|
stderr = (
|
|
nats_server.stderr.read().decode() if nats_server.stderr else ""
|
|
)
|
|
if "address already in use" in stderr.lower():
|
|
print(f"NATS port {nats_port} already in use, retrying...")
|
|
time.sleep(0.1)
|
|
continue
|
|
etcd.terminate()
|
|
etcd.wait()
|
|
shutil.rmtree(nats_data_dir, ignore_errors=True)
|
|
shutil.rmtree(etcd_data_dir, ignore_errors=True)
|
|
raise RuntimeError(f"NATS failed to start: {stderr}")
|
|
|
|
# Success - NATS started
|
|
break
|
|
|
|
except Exception as e:
|
|
last_error = e
|
|
print(f"Attempt {attempt + 1} failed: {e}")
|
|
if attempt < max_nats_retries - 1:
|
|
time.sleep(0.2)
|
|
else:
|
|
etcd.terminate()
|
|
etcd.wait()
|
|
shutil.rmtree(nats_data_dir, ignore_errors=True)
|
|
shutil.rmtree(etcd_data_dir, ignore_errors=True)
|
|
raise RuntimeError(
|
|
f"Failed to start NATS after {max_nats_retries} attempts: {last_error}"
|
|
)
|
|
|
|
# Set environment variables for the runtime to use
|
|
os.environ["NATS_SERVER"] = f"nats://localhost:{nats_port}"
|
|
os.environ["ETCD_ENDPOINTS"] = f"http://localhost:{etcd_client_port}"
|
|
|
|
return nats_server, etcd, nats_port, etcd_client_port, nats_data_dir, etcd_data_dir
|
|
|
|
|
|
@pytest.fixture(scope="module", autouse=True)
|
|
def nats_and_etcd():
|
|
"""
|
|
Start NATS and ETCD for testing.
|
|
|
|
Scope is "module" which means each test module shares the same NATS/ETCD instance.
|
|
|
|
Behavior is controlled by ENABLE_ISOLATED_ETCD_AND_NATS constant:
|
|
- True (default): Random ports + unique data dirs for parallel execution
|
|
- False: Default ports (4222, 2379) for sequential execution
|
|
"""
|
|
if ENABLE_ISOLATED_ETCD_AND_NATS:
|
|
(
|
|
nats_server,
|
|
etcd,
|
|
nats_port,
|
|
etcd_client_port,
|
|
nats_data_dir,
|
|
etcd_data_dir,
|
|
) = start_nats_and_etcd_random_ports()
|
|
else:
|
|
(
|
|
nats_server,
|
|
etcd,
|
|
nats_port,
|
|
etcd_client_port,
|
|
nats_data_dir,
|
|
etcd_data_dir,
|
|
) = start_nats_and_etcd_default_ports()
|
|
|
|
try:
|
|
# Wait for services to be ready
|
|
if not wait_for_port("localhost", nats_port, timeout=SERVICE_STARTUP_TIMEOUT):
|
|
raise RuntimeError(f"NATS server failed to start on port {nats_port}")
|
|
if not wait_for_port(
|
|
"localhost", etcd_client_port, timeout=SERVICE_STARTUP_TIMEOUT
|
|
):
|
|
raise RuntimeError(f"ETCD failed to start on port {etcd_client_port}")
|
|
|
|
print(f"NATS ({nats_port}) and ETCD ({etcd_client_port}) services ready")
|
|
yield
|
|
finally:
|
|
# Teardown code - always runs even if setup fails or tests error
|
|
print("Tearing down resources")
|
|
|
|
# Only terminate services if we started them (not reusing existing)
|
|
if nats_server is None and etcd is None:
|
|
print("Reused existing services, not stopping them")
|
|
else:
|
|
# Terminate both processes first (parallel shutdown)
|
|
try:
|
|
if nats_server:
|
|
nats_server.terminate()
|
|
except Exception as e:
|
|
print(f"Error terminating NATS: {e}")
|
|
try:
|
|
if etcd:
|
|
etcd.terminate()
|
|
except Exception as e:
|
|
print(f"Error terminating ETCD: {e}")
|
|
|
|
# Wait for both processes to finish
|
|
try:
|
|
if nats_server:
|
|
nats_server.wait(timeout=SERVICE_SHUTDOWN_TIMEOUT)
|
|
except subprocess.TimeoutExpired:
|
|
print("NATS did not terminate gracefully, killing")
|
|
try:
|
|
nats_server.kill()
|
|
except Exception:
|
|
pass
|
|
except Exception as e:
|
|
print(f"Error waiting for NATS: {e}")
|
|
|
|
try:
|
|
if etcd:
|
|
etcd.wait(timeout=SERVICE_SHUTDOWN_TIMEOUT)
|
|
except subprocess.TimeoutExpired:
|
|
print("ETCD did not terminate gracefully, killing")
|
|
try:
|
|
etcd.kill()
|
|
except Exception:
|
|
pass
|
|
except Exception as e:
|
|
print(f"Error waiting for ETCD: {e}")
|
|
|
|
# Clean up temporary data directories (if created)
|
|
if nats_data_dir:
|
|
try:
|
|
shutil.rmtree(nats_data_dir, ignore_errors=True)
|
|
except Exception as e:
|
|
print(f"Error removing NATS data dir: {e}")
|
|
if etcd_data_dir:
|
|
try:
|
|
shutil.rmtree(etcd_data_dir, ignore_errors=True)
|
|
except Exception as e:
|
|
print(f"Error removing ETCD data dir: {e}")
|
|
|
|
|
|
@pytest.fixture(scope="function")
|
|
def temp_file_store():
|
|
"""
|
|
A temporary directory to use as the key-value store. Cleaned up on test exit.
|
|
Local to the unit test using it.
|
|
"""
|
|
with tempfile.TemporaryDirectory() as tmpdir:
|
|
os.environ["DYN_FILE_KV"] = tmpdir
|
|
yield tmpdir
|
|
|
|
|
|
@pytest.fixture
|
|
def discovery_backend(request):
|
|
"""
|
|
Discovery backend for runtime. Defaults to "file".
|
|
|
|
To iterate over multiple backends in a test:
|
|
@pytest.mark.parametrize("discovery_backend", ["file", "etcd"], indirect=True)
|
|
async def test_example(runtime):
|
|
...
|
|
"""
|
|
return getattr(request, "param", "file")
|
|
|
|
|
|
@pytest.fixture
|
|
def request_plane(request):
|
|
"""
|
|
Request plane for runtime. Defaults to "nats".
|
|
|
|
To iterate over multiple transports in a test:
|
|
@pytest.mark.parametrize("request_plane", ["tcp", "nats"], indirect=True)
|
|
async def test_example(runtime):
|
|
...
|
|
"""
|
|
return getattr(request, "param", "nats")
|
|
|
|
|
|
@pytest.fixture(scope="function", autouse=False)
|
|
async def runtime(request, discovery_backend, request_plane):
|
|
"""
|
|
Create a DistributedRuntime for testing.
|
|
|
|
IMPORTANT: DistributedRuntime is a process-level singleton. When using isolated
|
|
NATS/ETCD (ENABLE_ISOLATED_ETCD_AND_NATS=1), tests using this fixture MUST be
|
|
marked with `@pytest.mark.forked` to run in a separate process for isolation.
|
|
|
|
Without @pytest.mark.forked in isolated mode, you will get "Worker already initialized"
|
|
errors when multiple tests try to create runtimes in the same process.
|
|
|
|
The discovery_backend and request_plane can be customized by overriding their fixtures
|
|
or using @pytest.mark.parametrize with indirect=True:
|
|
|
|
@pytest.mark.forked
|
|
@pytest.mark.parametrize("discovery_backend", ["etcd"], indirect=True)
|
|
async def test_with_etcd(runtime):
|
|
...
|
|
"""
|
|
# Check if the test is marked with @pytest.mark.forked (only in isolated mode)
|
|
if ENABLE_ISOLATED_ETCD_AND_NATS:
|
|
forked_marker = request.node.get_closest_marker("forked")
|
|
if forked_marker is None:
|
|
pytest.fail(
|
|
f"""
|
|
Test '{request.node.name}' uses the 'runtime' fixture but is not marked with @pytest.mark.forked.
|
|
This is required when ENABLE_ISOLATED_ETCD_AND_NATS=1.
|
|
|
|
Add @pytest.mark.forked decorator to run this test in a separate process:
|
|
@pytest.mark.forked
|
|
async def test_my_test(runtime):
|
|
...
|
|
|
|
Or set ENABLE_ISOLATED_ETCD_AND_NATS=0 to use default ports (no forking needed).
|
|
|
|
This is required because DistributedRuntime is a process-level singleton.
|
|
"""
|
|
)
|
|
|
|
loop = asyncio.get_running_loop()
|
|
runtime = DistributedRuntime(loop, discovery_backend, request_plane)
|
|
yield runtime
|
|
runtime.shutdown()
|