conftest.py 35 KB
Newer Older
1
# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
Neelay Shah's avatar
Neelay Shah committed
2
3
# SPDX-License-Identifier: Apache-2.0

4
import importlib.util
Neelay Shah's avatar
Neelay Shah committed
5
6
import logging
import os
7
import shutil
Neelay Shah's avatar
Neelay Shah committed
8
import tempfile
9
from pathlib import Path
10
from typing import Generator, Optional
Neelay Shah's avatar
Neelay Shah committed
11
12

import pytest
13
from filelock import FileLock
Neelay Shah's avatar
Neelay Shah committed
14

15
from tests.utils.constants import TEST_MODELS, DefaultPort
Neelay Shah's avatar
Neelay Shah committed
16
from tests.utils.managed_process import ManagedProcess
17
from tests.utils.port_utils import (
18
    ServicePorts,
19
20
21
22
23
    allocate_port,
    allocate_ports,
    deallocate_port,
    deallocate_ports,
)
24
from tests.utils.test_output import resolve_test_output_path
25
26

_logger = logging.getLogger(__name__)
Neelay Shah's avatar
Neelay Shah committed
27

Alec's avatar
Alec committed
28
29

def pytest_configure(config):
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
    # Defining markers to avoid `<marker> not found in 'markers' configuration option`
    # errors when pyproject.toml is not available in the container (e.g. some CI jobs).
    # IMPORTANT: Keep this marker list in sync with [tool.pytest.ini_options].markers
    # in pyproject.toml. If you add or remove markers there, mirror the change here.
    markers = [
        "pre_merge: marks tests to run before merging",
        "post_merge: marks tests to run after merge",
        "parallel: marks tests that can run in parallel with pytest-xdist",
        "nightly: marks tests to run nightly",
        "weekly: marks tests to run weekly",
        "gpu_0: marks tests that don't require GPU",
        "gpu_1: marks tests to run on GPU",
        "gpu_2: marks tests to run on 2GPUs",
        "gpu_4: marks tests to run on 4GPUs",
        "gpu_8: marks tests to run on 8GPUs",
45
        "max_vram_gib(N): peak VRAM in GiB (with 10% safety). Filter with --max-vram-gib=N",
46
47
48
49
50
51
52
53
        "e2e: marks tests as end-to-end tests",
        "integration: marks tests as integration tests",
        "unit: marks tests as unit tests",
        "stress: marks tests as stress tests",
        "performance: marks tests as performance tests",
        "vllm: marks tests as requiring vllm",
        "trtllm: marks tests as requiring trtllm",
        "sglang: marks tests as requiring sglang",
54
        "lmcache: mark tests as requiring lmcache",
55
56
57
        "multimodal: marks tests as multimodal (image/video) tests",
        "slow: marks tests as known to be slow",
        "h100: marks tests to run on H100",
58
        "aiconfigurator: marks e2e tests that cover aiconfigurator functionality",
59
60
61
62
        "router: marks tests for router component",
        "planner: marks tests for planner component",
        "kvbm: marks tests for KV behavior and model determinism",
        "kvbm_v2: marks tests using KVBM V2",
63
        "kvbm_concurrency: marks concurrency stress tests for KVBM (runs separately)",
64
65
66
67
        "model: model id used by a test or parameter",
        "custom_build: marks tests that require custom builds or special setup (e.g., MoE models)",
        "k8s: marks tests as requiring Kubernetes",
        "fault_tolerance: marks tests as fault tolerance tests",
68
        "deploy: marks tests as deployment tests",
69
70
        # Third-party plugin markers
        "timeout: test timeout in seconds (pytest-timeout plugin)",
71
72
73
    ]
    for marker in markers:
        config.addinivalue_line("markers", marker)
Alec's avatar
Alec committed
74
75


76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
def pytest_addoption(parser: pytest.Parser) -> None:
    """Add shared command-line options for all tests.

    Shared options that apply across multiple test suites are defined here.
    Suite-specific options (e.g., deploy, fault-tolerance) are defined in
    their respective subdirectory conftest.py files.
    """
    # -------------------------------------------------------------------------
    # Shared Deployment Options (used by multiple test suites)
    # -------------------------------------------------------------------------
    parser.addoption(
        "--image",
        type=str,
        default=None,
        help="Container image to use for deployment (overrides YAML default)",
    )
    parser.addoption(
        "--namespace",
        type=str,
        default=None,  # No default here - subdirectories provide their own
        help="Kubernetes namespace for deployment",
    )
    parser.addoption(
        "--skip-service-restart",
        action="store_true",
        default=None,  # None = use fixture's default behavior
        help="Skip restarting NATS and etcd services before deployment. "
        "Default: deploy tests skip (for speed), fault-tolerance tests restart (for clean state).",
    )
105
106
107
108
109
110
    parser.addoption(
        "--max-vram-gib",
        type=float,
        default=None,
        help="Skip tests whose @pytest.mark.max_vram_gib(N) exceeds this value (GiB).",
    )
111
112
113
114
115
116
    parser.addoption(
        "--dry-run",
        action="store_true",
        default=False,
        help="Show which tests would run vs skip based on --max-vram-gib, then exit.",
    )
117
118


Neelay Shah's avatar
Neelay Shah committed
119
LOG_FORMAT = "[TEST] %(asctime)s %(levelname)s %(name)s: %(message)s"
120
DATE_FORMAT = "%Y-%m-%dT%H:%M:%S"
Neelay Shah's avatar
Neelay Shah committed
121
122
123
124

logging.basicConfig(
    level=logging.INFO,
    format=LOG_FORMAT,
125
    datefmt=DATE_FORMAT,  # ISO 8601 UTC format
Neelay Shah's avatar
Neelay Shah committed
126
127
)

128

Alec's avatar
Alec committed
129
130
131
132
133
134
135
136
137
138
@pytest.fixture()
def set_ucx_tls_no_mm():
    """Set UCX env defaults for all tests."""
    mp = pytest.MonkeyPatch()
    # CI note:
    # - Affected test: tests/fault_tolerance/cancellation/test_vllm.py::test_request_cancellation_vllm_decode_cancel
    # - Symptom on L40 CI: UCX/NIXL mm transport assertion during worker init
    #   (uct_mem.c:482: mem.memh != UCT_MEM_HANDLE_NULL) when two workers
    #   start on the same node (maybe a shared-memory segment collision/limits).
    # - Mitigation: disable UCX "mm" shared-memory transport globally for tests
139
140
141
142
143
    #
    # Also exclude gdr_copy transport to prevent GDRCopy driver initialization
    # failures (driverInitFileInfo result=11) that can abort the process when
    # the gdrdrv kernel module is not loaded.
    mp.setenv("UCX_TLS", "^mm,gdr_copy")
Alec's avatar
Alec committed
144
145
146
147
    yield
    mp.undo()


148
def download_models(model_list=None, ignore_weights=False):
149
150
151
152
    """Download models - can be called directly or via fixture

    Args:
        model_list: List of model IDs to download. If None, downloads TEST_MODELS.
153
        ignore_weights: If True, skips downloading model weight files. Default is False.
154
155
156
157
158
    """
    if model_list is None:
        model_list = TEST_MODELS

    # Check for HF_TOKEN in environment
159
    hf_token = os.environ.get("HF_TOKEN", "").strip() or None
160
161
162
163
164
165
166
167
168
169
170
    if hf_token:
        logging.info("HF_TOKEN found in environment")
    else:
        logging.warning(
            "HF_TOKEN not found in environment. "
            "Some models may fail to download or you may encounter rate limits. "
            "Get a token from https://huggingface.co/settings/tokens"
        )

    try:
        from huggingface_hub import snapshot_download
171
172
173
174
175
176
177
178
179
180
    except ImportError as exc:
        raise RuntimeError(
            "huggingface_hub is required to pre-download models for tests"
        ) from exc

    failures = []
    for model_id in model_list:
        logging.info(
            f"Pre-downloading {'model (no weights)' if ignore_weights else 'model'}: {model_id}"
        )
181

182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
        try:
            if ignore_weights:
                # Weight file patterns to exclude (based on hub.rs implementation)
                weight_patterns = [
                    "*.bin",
                    "*.safetensors",
                    "*.h5",
                    "*.msgpack",
                    "*.ckpt.index",
                ]

                # Download everything except weight files
                snapshot_download(
                    repo_id=model_id,
                    token=hf_token,
                    ignore_patterns=weight_patterns,
                )
            else:
                # Download the full model snapshot (includes all files)
                snapshot_download(
                    repo_id=model_id,
                    token=hf_token,
                )
            logging.info(f"Successfully pre-downloaded: {model_id}")

        except Exception as exc:
            logging.error(f"Failed to pre-download {model_id}: {exc}")
            failures.append(f"{model_id}: {exc}")

    if failures:
        raise RuntimeError(
            "Failed to pre-download required Hugging Face models:\n"
            + "\n".join(failures)
215
216
217
218
        )


@pytest.fixture(scope="session")
Alec's avatar
Alec committed
219
220
221
222
223
224
225
226
227
228
229
230
def predownload_models(pytestconfig):
    """Fixture wrapper around download_models for models used in collected tests"""
    # Get models from pytest config if available, otherwise fall back to TEST_MODELS
    models = getattr(pytestconfig, "models_to_download", None)
    if models:
        logging.info(
            f"Downloading {len(models)} models needed for collected tests\nModels: {models}"
        )
        download_models(model_list=list(models))
    else:
        # Fallback to original behavior if extraction failed
        download_models()
231
232

    os.environ["HF_HUB_OFFLINE"] = "1"
233
    yield
234
    os.environ.pop("HF_HUB_OFFLINE", None)
235

Neelay Shah's avatar
Neelay Shah committed
236

237
@pytest.fixture(scope="session")
Alec's avatar
Alec committed
238
239
240
241
242
243
244
245
246
247
248
249
def predownload_tokenizers(pytestconfig):
    """Fixture wrapper around download_models for tokenizers used in collected tests"""
    # Get models from pytest config if available, otherwise fall back to TEST_MODELS
    models = getattr(pytestconfig, "models_to_download", None)
    if models:
        logging.info(
            f"Downloading tokenizers for {len(models)} models needed for collected tests\nModels: {models}"
        )
        download_models(model_list=list(models), ignore_weights=True)
    else:
        # Fallback to original behavior if extraction failed
        download_models(ignore_weights=True)
250
251
252
253
254

    # Skip redundant HuggingFace API calls in worker subprocesses since
    # tokenizers are already cached. This avoids flaky timeouts from slow
    # HF API responses (the RepoInfo fetch still happens even for cached models).
    os.environ["HF_HUB_OFFLINE"] = "1"
255
    yield
256
    os.environ.pop("HF_HUB_OFFLINE", None)
257
258


259
260
@pytest.fixture(autouse=True)
def logger(request):
261
262
    log_dir = resolve_test_output_path(request.node.name)
    log_path = os.path.join(log_dir, "test.log.txt")
263
    logger = logging.getLogger()
264
265
    shutil.rmtree(log_dir, ignore_errors=True)
    os.makedirs(log_dir, exist_ok=True)
266
267
268
269
270
271
272
273
274
    handler = logging.FileHandler(log_path, mode="w")
    formatter = logging.Formatter(LOG_FORMAT, datefmt=DATE_FORMAT)
    handler.setFormatter(formatter)
    logger.addHandler(handler)
    yield
    handler.close()
    logger.removeHandler(handler)


275
276
277
278
279
280
281
282
283
284
285
286
287
288
def _item_has_marker(item, marker_name):
    """Check if a test item has a marker, including module-level pytestmark."""
    if item.get_closest_marker(marker_name):
        return True
    module = getattr(item, "module", None)
    if module is not None:
        marks = getattr(module, "pytestmark", [])
        if not isinstance(marks, list):
            marks = [marks]
        if any(getattr(m, "name", "") == marker_name for m in marks):
            return True
    return False


Alec's avatar
Alec committed
289
@pytest.hookimpl(trylast=True)
290
291
292
293
def pytest_collection_modifyitems(config, items):
    """
    This function is called to modify the list of tests to run.
    """
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
    # Auto-skip tests marked with a framework marker when the framework is not installed
    framework_markers = {
        "trtllm": "tensorrt_llm",
        "vllm": "vllm",
        "sglang": "sglang",
        "kvbm": "kvbm",
        "lmcache": "lmcache",
    }
    for marker_name, module_name in framework_markers.items():
        if importlib.util.find_spec(module_name) is None:
            skip = pytest.mark.skip(reason=f"{module_name} is not installed")
            for item in items:
                if _item_has_marker(item, marker_name):
                    item.add_marker(skip)

309
310
311
312
313
314
315
316
317
318
319
    # Skip tests that exceed --max-vram-gib
    vram_limit = config.getoption("--max-vram-gib", default=None)
    if vram_limit is not None:
        skip_vram = pytest.mark.skip(
            reason=f"requires more than {vram_limit} GiB VRAM (--max-vram-gib={vram_limit})"
        )
        for item in items:
            vram_mark = item.get_closest_marker("max_vram_gib")
            if vram_mark and vram_mark.args and vram_mark.args[0] > vram_limit:
                item.add_marker(skip_vram)

320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
    # --dry-run: print run/skip breakdown and exit without executing tests
    if config.getoption("--dry-run", default=False):
        would_run = []
        would_skip = []
        unmarked = []
        for item in items:
            vram_mark = item.get_closest_marker("max_vram_gib")
            vram_val = vram_mark.args[0] if vram_mark and vram_mark.args else None
            name = item.nodeid.split("::", 1)[1] if "::" in item.nodeid else item.nodeid

            skip_reasons = []
            for marker in item.iter_markers("skip"):
                reason = marker.kwargs.get("reason", "")
                if not reason and marker.args:
                    reason = marker.args[0]
                skip_reasons.append(reason or "no reason given")

            vram_skipped = (
                vram_limit is not None
                and vram_val is not None
                and vram_val > vram_limit
            )
            if vram_skipped:
                skip_reasons.insert(0, f"{vram_val} GiB > {vram_limit} GiB VRAM limit")

            if skip_reasons:
                would_skip.append((name, vram_val, skip_reasons))
            elif vram_val is not None:
                would_run.append((name, vram_val))
            else:
                unmarked.append(name)

        print(f"\n{'=' * 60}")
        print(
            f"--max-vram-gib={vram_limit or 'not set'}  |  {len(items)} tests selected"
        )
        print(f"{'=' * 60}")
        if would_run:
            print(f"\nWould RUN ({len(would_run)}):")
            for name, gib in would_run:
                print(f"  {name}  ({gib} GiB)")
        if would_skip:
            print(f"\nWould SKIP ({len(would_skip)}):")
            for name, vram_val, reasons in would_skip:
                vram_str = f"  ({vram_val} GiB)" if vram_val is not None else ""
                print(f"  {name}{vram_str}  -- {'; '.join(reasons)}")
        if unmarked:
            print(f"\nNo VRAM marker — always run ({len(unmarked)}):")
            for name in unmarked:
                print(f"  {name}")
        print()
        items.clear()
        return

Alec's avatar
Alec committed
374
375
    # Collect models via explicit pytest mark from final filtered items only
    models_to_download = set()
376
    for item in items:
Alec's avatar
Alec committed
377
378
379
380
381
382
383
384
385
386
387
388
        # Only collect from items that are not skipped
        if any(
            getattr(m, "name", "") == "skip" for m in getattr(item, "own_markers", [])
        ):
            continue
        model_mark = item.get_closest_marker("model")
        if model_mark and model_mark.args:
            models_to_download.add(model_mark.args[0])

    # Store models to download in pytest config for fixtures to access
    if models_to_download:
        config.models_to_download = models_to_download
389

390

Neelay Shah's avatar
Neelay Shah committed
391
392
class EtcdServer(ManagedProcess):
    def __init__(self, request, port=2379, timeout=300):
393
394
395
396
397
398
399
400
401
402
403
404
        # Allocate free ports if port is 0
        use_random_port = port == 0
        if use_random_port:
            # Need two ports: client port and peer port for parallel execution
            # Start from 2380 (etcd default 2379 + 1)
            port, peer_port = allocate_ports(2, 2380)
        else:
            peer_port = None

        self.port = port
        self.peer_port = peer_port  # Store for cleanup
        self.use_random_port = use_random_port  # Track if we allocated the port
Neelay Shah's avatar
Neelay Shah committed
405
406
407
408
        port_string = str(port)
        etcd_env = os.environ.copy()
        etcd_env["ALLOW_NONE_AUTHENTICATION"] = "yes"
        data_dir = tempfile.mkdtemp(prefix="etcd_")
409

Neelay Shah's avatar
Neelay Shah committed
410
411
412
413
414
415
416
        command = [
            "etcd",
            "--listen-client-urls",
            f"http://0.0.0.0:{port_string}",
            "--advertise-client-urls",
            f"http://0.0.0.0:{port_string}",
        ]
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437

        # Add peer port configuration only for random ports (parallel execution)
        if peer_port is not None:
            peer_port_string = str(peer_port)
            command.extend(
                [
                    "--listen-peer-urls",
                    f"http://0.0.0.0:{peer_port_string}",
                    "--initial-advertise-peer-urls",
                    f"http://localhost:{peer_port_string}",
                    "--initial-cluster",
                    f"default=http://localhost:{peer_port_string}",
                ]
            )

        command.extend(
            [
                "--data-dir",
                data_dir,
            ]
        )
Neelay Shah's avatar
Neelay Shah committed
438
439
440
441
442
        super().__init__(
            env=etcd_env,
            command=command,
            timeout=timeout,
            display_output=False,
443
            terminate_all_matching_process_names=not use_random_port,  # For distributed tests, do not terminate all matching processes
Neelay Shah's avatar
Neelay Shah committed
444
            health_check_ports=[port],
445
            data_dir=data_dir,
Neelay Shah's avatar
Neelay Shah committed
446
447
448
            log_dir=request.node.name,
        )

449
450
451
452
453
454
455
456
457
458
459
460
461
462
    def __exit__(self, exc_type, exc_val, exc_tb):
        """Release allocated ports when server exits."""
        try:
            # Only deallocate ports that were dynamically allocated (not default ports)
            if self.use_random_port:
                ports_to_release = [self.port]
                if self.peer_port is not None:
                    ports_to_release.append(self.peer_port)
                deallocate_ports(ports_to_release)
        except Exception as e:
            logging.warning(f"Failed to release EtcdServer port: {e}")

        return super().__exit__(exc_type, exc_val, exc_tb)

Neelay Shah's avatar
Neelay Shah committed
463
464

class NatsServer(ManagedProcess):
465
    def __init__(self, request, port=4222, timeout=300, disable_jetstream=False):
466
467
468
469
470
471
472
473
        # Allocate a free port if port is 0
        use_random_port = port == 0
        if use_random_port:
            # Start from 4223 (nats-server default 4222 + 1)
            port = allocate_port(4223)

        self.port = port
        self.use_random_port = use_random_port  # Track if we allocated the port
474
475
        self._request = request  # Store for restart
        self._timeout = timeout
476
477
        self._disable_jetstream = disable_jetstream
        data_dir = tempfile.mkdtemp(prefix="nats_") if not disable_jetstream else None
478
479
480
481
482
483
        command = [
            "nats-server",
            "--trace",
            "-p",
            str(port),
        ]
484
485
        if not disable_jetstream and data_dir:
            command.extend(["-js", "--store_dir", data_dir])
Neelay Shah's avatar
Neelay Shah committed
486
487
488
489
        super().__init__(
            command=command,
            timeout=timeout,
            display_output=False,
490
            terminate_all_matching_process_names=not use_random_port,  # For distributed tests, do not terminate all matching processes
Neelay Shah's avatar
Neelay Shah committed
491
492
            data_dir=data_dir,
            health_check_ports=[port],
493
            health_check_funcs=[self._nats_ready],
Neelay Shah's avatar
Neelay Shah committed
494
495
496
            log_dir=request.node.name,
        )

497
498
499
500
501
502
503
504
505
506
507
508
509
510
511
512
513
514
515
516
517
518
519
520
521
522
523
524
525
526
527
528
529
530
531
    def _nats_ready(self, timeout: float = 5) -> bool:
        """Verify NATS server is ready by connecting and optionally checking JetStream."""
        import asyncio

        import nats

        async def check():
            try:
                nc = await nats.connect(
                    f"nats://localhost:{self.port}",
                    connect_timeout=min(timeout, 2),
                )
                try:
                    if not self._disable_jetstream:
                        # Verify JetStream is initialized
                        js = nc.jetstream()
                        await js.account_info()
                    return True
                finally:
                    await nc.close()
            except Exception:
                return False

        # Handle both sync and async contexts
        try:
            asyncio.get_running_loop()  # Check if we're in async context
            # Already in async context - run in a thread to avoid blocking
            import concurrent.futures

            with concurrent.futures.ThreadPoolExecutor() as pool:
                return pool.submit(asyncio.run, check()).result(timeout=timeout)
        except RuntimeError:
            # No running loop - safe to use asyncio.run()
            return asyncio.run(check())

532
533
534
535
536
537
538
539
540
541
542
    def __exit__(self, exc_type, exc_val, exc_tb):
        """Release allocated port when server exits."""
        try:
            # Only deallocate ports that were dynamically allocated (not default ports)
            if self.use_random_port:
                deallocate_port(self.port)
        except Exception as e:
            logging.warning(f"Failed to release NatsServer port: {e}")

        return super().__exit__(exc_type, exc_val, exc_tb)

543
544
545
    def stop(self):
        """Stop the NATS server for restart. Does not release port or clean up fully."""
        _logger.info(f"Stopping NATS server on port {self.port}")
546
        self._stop_started_processes()
547
548
549
550

    def start(self):
        """Restart a stopped NATS server with fresh state."""
        _logger.info(f"Starting NATS server on port {self.port} with fresh state")
551
552
553
554
555
556
557
558
        # Clean up old data directory and create fresh one (only if JetStream enabled)
        if not self._disable_jetstream:
            old_data_dir = self.data_dir  # type: ignore[has-type]
            if old_data_dir is not None:
                shutil.rmtree(old_data_dir, ignore_errors=True)
            self.data_dir = tempfile.mkdtemp(prefix="nats_")

        # Rebuild command
559
560
561
562
563
564
        self.command = [
            "nats-server",
            "--trace",
            "-p",
            str(self.port),
        ]
565
566
        if not self._disable_jetstream and self.data_dir:
            self.command.extend(["-js", "--store_dir", self.data_dir])
567
568

        self._start_process()
569
570
        elapsed = self._check_ports(self._timeout)
        self._check_funcs(self._timeout - elapsed)
571

Neelay Shah's avatar
Neelay Shah committed
572

573
class SharedManagedProcess:
574
575
576
577
578
579
    """Base class for persistent shared processes across pytest-xdist workers.

    Simplified design: first worker starts the process on a dynamic port, it lives forever
    (until the container dies). No ref counting, no teardown. Subsequent workers just
    reuse via port check. This eliminates race conditions and simplifies the logic.
    """
580
581
582
583
584
585

    def __init__(
        self,
        request,
        tmp_path_factory,
        resource_name: str,
586
        start_port: int,
587
588
589
        timeout: int = 300,
    ):
        self.request = request
590
591
        self.start_port = start_port
        self.port: Optional[int] = None  # Set when entering context
592
593
594
595
        self.timeout = timeout
        self.resource_name = resource_name
        self._server: Optional[ManagedProcess] = None

596
        root_tmp = Path(tempfile.gettempdir()) / "pytest_shared_services"
597
598
        root_tmp.mkdir(parents=True, exist_ok=True)

599
600
        self.port_file = root_tmp / f"{resource_name}_port"
        self.lock_file = str(self.port_file) + ".lock"
601

602
    def _create_server(self, port: int) -> ManagedProcess:
603
604
605
        """Create the underlying server instance. Must be implemented by subclasses."""
        raise NotImplementedError

606
607
608
609
610
611
612
613
614
615
616
617
618
619
620
621
    def _is_port_in_use(self, port: int) -> bool:
        """Check if a port is in use (i.e., a process is listening on it)."""
        import socket

        try:
            sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
            sock.settimeout(1)
            result = sock.connect_ex(("localhost", port))
            sock.close()
            return result == 0  # 0 means connection succeeded (port in use)
        except Exception:
            return False

    def _read_port(self) -> Optional[int]:
        """Read stored port from file."""
        if self.port_file.exists():
622
            try:
623
                return int(self.port_file.read_text().strip())
624
            except (ValueError, IOError):
625
626
627
628
629
630
                return None
        return None

    def _write_port(self, port: int):
        """Write port to file."""
        self.port_file.write_text(str(port))
631
632
633

    def __enter__(self):
        with FileLock(self.lock_file):
634
635
636
637
638
639
640
641
642
            stored_port = self._read_port()

            # Check if a process is already running on the stored port
            if stored_port is not None and self._is_port_in_use(stored_port):
                # Reuse existing process
                self.port = stored_port
                logging.info(
                    f"[{self.resource_name}] Reusing existing process on port {self.port}"
                )
643
            else:
644
645
646
647
648
649
650
651
652
                # Start new process
                if stored_port is not None:
                    logging.warning(
                        f"[{self.resource_name}] Stale port file: port {stored_port} not in use, starting fresh"
                    )
                self.port = allocate_port(self.start_port)
                self._write_port(self.port)
                self._server = self._create_server(self.port)
                self._server.__enter__()
653
                logging.info(
654
                    f"[{self.resource_name}] Started process on port {self.port}"
655
656
657
658
                )
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
659
660
661
        # Never tear down - let the process live until the container dies.
        # This avoids race conditions and simplifies the logic.
        pass
662
663
664
665
666


class SharedEtcdServer(SharedManagedProcess):
    """EtcdServer with file-based reference counting for multi-process sharing."""

667
668
    def __init__(self, request, tmp_path_factory, start_port=2380, timeout=300):
        super().__init__(request, tmp_path_factory, "etcd", start_port, timeout)
669
670
671
        # Create a log directory for session-scoped servers
        self._log_dir = tempfile.mkdtemp(prefix=f"pytest_{self.resource_name}_logs_")

672
    def _create_server(self, port: int) -> ManagedProcess:
673
        """Create EtcdServer instance."""
674
        server = EtcdServer(self.request, port=port, timeout=self.timeout)
675
676
677
678
679
680
681
682
        # Override log_dir since request.node.name is empty in session scope
        server.log_dir = self._log_dir
        return server


class SharedNatsServer(SharedManagedProcess):
    """NatsServer with file-based reference counting for multi-process sharing."""

683
684
685
686
687
688
689
690
691
    def __init__(
        self,
        request,
        tmp_path_factory,
        start_port=4223,
        timeout=300,
        disable_jetstream=False,
    ):
        super().__init__(request, tmp_path_factory, "nats", start_port, timeout)
692
693
        # Create a log directory for session-scoped servers
        self._log_dir = tempfile.mkdtemp(prefix=f"pytest_{self.resource_name}_logs_")
694
        self._disable_jetstream = disable_jetstream
695

696
    def _create_server(self, port: int) -> ManagedProcess:
697
        """Create NatsServer instance."""
698
699
700
701
702
703
        server = NatsServer(
            self.request,
            port=port,
            timeout=self.timeout,
            disable_jetstream=self._disable_jetstream,
        )
704
705
706
707
708
        # Override log_dir since request.node.name is empty in session scope
        server.log_dir = self._log_dir
        return server


709
@pytest.fixture
710
def discovery_backend(request):
711
    """
712
    Discovery backend for runtime. Defaults to "etcd".
713

714
715
    To iterate over multiple backends in a test:
        @pytest.mark.parametrize("discovery_backend", ["file", "etcd"], indirect=True)
716
717
718
719
720
721
722
723
724
725
726
727
728
729
730
731
732
733
734
        def test_example(runtime_services):
            ...
    """
    return getattr(request, "param", "etcd")


@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", ["nats", "tcp"], indirect=True)
        def test_example(runtime_services):
            ...
    """
    return getattr(request, "param", "nats")


735
@pytest.fixture
736
def durable_kv_events(request):
737
    """
738
    Whether to use durable KV events via JetStream. Defaults to False (NATS Core mode).
739

740
    When False (default):
741
    - NATS server starts without JetStream (-js flag omitted) for faster startup
742
    - Workers use local indexer mode (NATS Core / fire-and-forget events)
743

744
745
746
    When True:
    - NATS server starts with JetStream for durable KV event distribution
    - Workers use --durable-kv-events flag to publish to JetStream
747

748
749
    To use JetStream mode:
        @pytest.mark.parametrize("durable_kv_events", [True], indirect=True)
750
751
752
753
754
755
        def test_example(runtime_services_dynamic_ports):
            ...
    """
    return getattr(request, "param", False)


Neelay Shah's avatar
Neelay Shah committed
756
@pytest.fixture()
757
def runtime_services(request, discovery_backend, request_plane):
758
    """
759
    Start runtime services (NATS and/or etcd) based on discovery_backend and request_plane.
760

761
    - If discovery_backend != "etcd", etcd is not started (returns None)
762
    - If request_plane != "nats", NATS is not started (returns None)
763
764

    Returns a tuple of (nats_process, etcd_process) where each has a .port attribute.
765
    """
766
    # Port cleanup is now handled in NatsServer and EtcdServer __exit__ methods
767
    if request_plane == "nats" and discovery_backend == "etcd":
768
769
770
771
772
773
        with NatsServer(request) as nats_process:
            with EtcdServer(request) as etcd_process:
                yield nats_process, etcd_process
    elif request_plane == "nats":
        with NatsServer(request) as nats_process:
            yield nats_process, None
774
    elif discovery_backend == "etcd":
Neelay Shah's avatar
Neelay Shah committed
775
        with EtcdServer(request) as etcd_process:
776
777
778
            yield None, etcd_process
    else:
        yield None, None
779
780


781
@pytest.fixture()
782
783
784
def runtime_services_dynamic_ports(
    request, discovery_backend, request_plane, durable_kv_events
):
785
786
787
788
789
790
    """Provide NATS and Etcd servers with truly dynamic ports per test.

    This fixture actually allocates dynamic ports by passing port=0 to the servers.
    It also sets the NATS_SERVER and ETCD_ENDPOINTS environment variables so that
    Dynamo processes can find the services on the dynamic ports.

791
792
793
794
795
    xdist/parallel safety:
    - Function-scoped: each test gets its own NATS/etcd instances and ports.
    - Each pytest-xdist worker runs tests in a separate process, so env vars do not
      leak across workers.

796
    - If discovery_backend != "etcd", etcd is not started (returns None)
797
798
    - NATS is always started when etcd is used, because KV events require NATS
      regardless of the request_plane (tcp/nats only affects request transport)
799
    - NATS Core mode (no JetStream) is the default; JetStream is enabled when durable_kv_events=True
800
801
802
803
804
805

    Returns a tuple of (nats_process, etcd_process) where each has a .port attribute.
    """
    import os

    # Port cleanup is now handled in NatsServer and EtcdServer __exit__ methods
806
    # Always start NATS when etcd is used - KV events require NATS regardless of request_plane
807
    # When durable_kv_events=False (default), disable JetStream for faster startup
808
    if discovery_backend == "etcd":
809
        with NatsServer(
810
            request, port=0, disable_jetstream=not durable_kv_events
811
        ) as nats_process:
812
            with EtcdServer(request, port=0) as etcd_process:
813
814
815
816
817
                # Save original env vars (may be set by session-scoped fixture)
                orig_nats = os.environ.get("NATS_SERVER")
                orig_etcd = os.environ.get("ETCD_ENDPOINTS")

                # Set environment variables for this test's dynamic ports
818
819
820
821
822
                os.environ["NATS_SERVER"] = f"nats://localhost:{nats_process.port}"
                os.environ["ETCD_ENDPOINTS"] = f"http://localhost:{etcd_process.port}"

                yield nats_process, etcd_process

823
824
825
826
827
828
829
830
831
                # Restore original env vars (or remove if they weren't set)
                if orig_nats is not None:
                    os.environ["NATS_SERVER"] = orig_nats
                else:
                    os.environ.pop("NATS_SERVER", None)
                if orig_etcd is not None:
                    os.environ["ETCD_ENDPOINTS"] = orig_etcd
                else:
                    os.environ.pop("ETCD_ENDPOINTS", None)
832
    elif request_plane == "nats":
833
        with NatsServer(
834
            request, port=0, disable_jetstream=not durable_kv_events
835
836
        ) as nats_process:
            orig_nats = os.environ.get("NATS_SERVER")
837
838
            os.environ["NATS_SERVER"] = f"nats://localhost:{nats_process.port}"
            yield nats_process, None
839
840
841
842
            if orig_nats is not None:
                os.environ["NATS_SERVER"] = orig_nats
            else:
                os.environ.pop("NATS_SERVER", None)
843
844
845
846
    else:
        yield None, None


847
848
849
850
@pytest.fixture(scope="session")
def runtime_services_session(request, tmp_path_factory):
    """Session-scoped fixture that provides shared NATS and etcd instances for all tests.

851
852
853
    Uses file locking to coordinate between pytest-xdist worker processes.
    First worker starts services on dynamic ports, subsequent workers reuse them.
    Services are never torn down (live until container dies) to avoid race conditions.
854

855
856
    This fixture is xdist-safe when tests use unique namespaces (e.g. random suffixes)
    and do not assume exclusive access to global streams/keys.
857

858
859
    For tests that need to restart NATS (e.g. indexer sync), use `runtime_services_dynamic_ports`
    which provides per-test isolated instances.
860
861
862
    """
    with SharedNatsServer(request, tmp_path_factory) as nats:
        with SharedEtcdServer(request, tmp_path_factory) as etcd:
863
864
865
866
            # Set environment variables for Rust/Python runtime to use
            os.environ["NATS_SERVER"] = f"nats://localhost:{nats.port}"
            os.environ["ETCD_ENDPOINTS"] = f"http://localhost:{etcd.port}"

867
868
            yield nats, etcd

869
870
871
872
            # Clean up environment variables
            os.environ.pop("NATS_SERVER", None)
            os.environ.pop("ETCD_ENDPOINTS", None)

873

874
875
876
877
878
879
880
881
882
883
884
885
886
887
888
889
890
@pytest.fixture
def file_storage_backend():
    """Fixture that sets up and tears down file storage backend.

    Creates a temporary directory for file-based KV storage and sets
    the DYN_FILE_KV environment variable. Cleans up after the test.
    """
    with tempfile.TemporaryDirectory() as tmpdir:
        old_env = os.environ.get("DYN_FILE_KV")
        os.environ["DYN_FILE_KV"] = tmpdir
        logging.info(f"Set up file storage backend in: {tmpdir}")
        yield tmpdir
        # Cleanup
        if old_env is not None:
            os.environ["DYN_FILE_KV"] = old_env
        else:
            os.environ.pop("DYN_FILE_KV", None)
891
892
893
894
895
896
897
898
899
900
901
902
903
904
905
906
907
908
909
910
911
912
913
914
915
916


########################################################
# Shared Port Allocation (Dynamo deployments)
########################################################


@pytest.fixture(scope="function")
def num_system_ports(request) -> int:
    """Number of system ports to allocate for this test.

    Default: 1 port.

    Tests that need multiple system ports (e.g. SYSTEM_PORT1 + SYSTEM_PORT2) must
    explicitly request them via indirect parametrization:
      @pytest.mark.parametrize("num_system_ports", [2], indirect=True)
    """
    return getattr(request, "param", 1)


@pytest.fixture(scope="function")
def dynamo_dynamic_ports(num_system_ports) -> Generator[ServicePorts, None, None]:
    """Allocate per-test ports for Dynamo deployments.

    - frontend_port: OpenAI-compatible HTTP/gRPC ingress (dynamo.frontend)
    - system_ports: List of worker metrics/system ports (configurable count via num_system_ports)
917
    - kv_event_port: ZMQ port for vLLM KV event publishing (avoids collisions under xdist)
918
919
920
    """
    frontend_port = allocate_port(DefaultPort.FRONTEND.value)
    system_port_list = allocate_ports(num_system_ports, DefaultPort.SYSTEM1.value)
921
922
    kv_event_port = allocate_port(DefaultPort.SYSTEM1.value)
    all_ports = [frontend_port, *system_port_list, kv_event_port]
923
    try:
924
925
926
927
928
        yield ServicePorts(
            frontend_port=frontend_port,
            system_ports=system_port_list,
            kv_event_port=kv_event_port,
        )
929
930
    finally:
        deallocate_ports(all_ports)