conftest.py 28.6 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
4
5
# SPDX-License-Identifier: Apache-2.0

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

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

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

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

Alec's avatar
Alec committed
26
27

def pytest_configure(config):
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
    # 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",
        "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",
        "multimodal: marks tests as multimodal (image/video) tests",
        "slow: marks tests as known to be slow",
        "h100: marks tests to run on H100",
        "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",
58
        "kvbm_concurrency: marks concurrency stress tests for KVBM (runs separately)",
59
60
61
62
        "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",
63
64
        # Third-party plugin markers
        "timeout: test timeout in seconds (pytest-timeout plugin)",
65
66
67
    ]
    for marker in markers:
        config.addinivalue_line("markers", marker)
Alec's avatar
Alec committed
68
69


Neelay Shah's avatar
Neelay Shah committed
70
LOG_FORMAT = "[TEST] %(asctime)s %(levelname)s %(name)s: %(message)s"
71
DATE_FORMAT = "%Y-%m-%dT%H:%M:%S"
Neelay Shah's avatar
Neelay Shah committed
72
73
74
75

logging.basicConfig(
    level=logging.INFO,
    format=LOG_FORMAT,
76
    datefmt=DATE_FORMAT,  # ISO 8601 UTC format
Neelay Shah's avatar
Neelay Shah committed
77
78
)

79

Alec's avatar
Alec committed
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
@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
    mp.setenv("UCX_TLS", "^mm")
    yield
    mp.undo()


95
def download_models(model_list=None, ignore_weights=False):
96
97
98
99
    """Download models - can be called directly or via fixture

    Args:
        model_list: List of model IDs to download. If None, downloads TEST_MODELS.
100
        ignore_weights: If True, skips downloading model weight files. Default is False.
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
    """
    if model_list is None:
        model_list = TEST_MODELS

    # Check for HF_TOKEN in environment
    hf_token = os.environ.get("HF_TOKEN")
    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

        for model_id in model_list:
120
121
122
            logging.info(
                f"Pre-downloading {'model (no weights)' if ignore_weights else 'model'}: {model_id}"
            )
123
124

            try:
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
                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,
                    )
147
148
149
150
151
152
153
154
155
156
157
158
159
160
                logging.info(f"Successfully pre-downloaded: {model_id}")

            except Exception as e:
                logging.error(f"Failed to pre-download {model_id}: {e}")
                # Don't fail the fixture - let individual tests handle missing models

    except ImportError:
        logging.warning(
            "huggingface_hub not installed. "
            "Models will be downloaded during test execution."
        )


@pytest.fixture(scope="session")
Alec's avatar
Alec committed
161
162
163
164
165
166
167
168
169
170
171
172
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()
173
174
    yield

Neelay Shah's avatar
Neelay Shah committed
175

176
@pytest.fixture(scope="session")
Alec's avatar
Alec committed
177
178
179
180
181
182
183
184
185
186
187
188
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)
189
190
191
    yield


192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
@pytest.fixture(autouse=True)
def logger(request):
    log_path = os.path.join(request.node.name, "test.log.txt")
    logger = logging.getLogger()
    shutil.rmtree(request.node.name, ignore_errors=True)
    os.makedirs(request.node.name, exist_ok=True)
    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)


Alec's avatar
Alec committed
207
@pytest.hookimpl(trylast=True)
208
209
210
211
def pytest_collection_modifyitems(config, items):
    """
    This function is called to modify the list of tests to run.
    """
Alec's avatar
Alec committed
212
213
    # Collect models via explicit pytest mark from final filtered items only
    models_to_download = set()
214
    for item in items:
Alec's avatar
Alec committed
215
216
217
218
219
220
221
222
223
224
225
226
        # 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
227

228

Neelay Shah's avatar
Neelay Shah committed
229
230
class EtcdServer(ManagedProcess):
    def __init__(self, request, port=2379, timeout=300):
231
232
233
234
235
236
237
238
239
240
241
242
        # 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
243
244
245
246
        port_string = str(port)
        etcd_env = os.environ.copy()
        etcd_env["ALLOW_NONE_AUTHENTICATION"] = "yes"
        data_dir = tempfile.mkdtemp(prefix="etcd_")
247

Neelay Shah's avatar
Neelay Shah committed
248
249
250
251
252
253
254
        command = [
            "etcd",
            "--listen-client-urls",
            f"http://0.0.0.0:{port_string}",
            "--advertise-client-urls",
            f"http://0.0.0.0:{port_string}",
        ]
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275

        # 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
276
277
278
279
280
        super().__init__(
            env=etcd_env,
            command=command,
            timeout=timeout,
            display_output=False,
281
            terminate_existing=not use_random_port,  # Disabled for parallel test execution with random ports
Neelay Shah's avatar
Neelay Shah committed
282
            health_check_ports=[port],
283
            data_dir=data_dir,
Neelay Shah's avatar
Neelay Shah committed
284
285
286
            log_dir=request.node.name,
        )

287
288
289
290
291
292
293
294
295
296
297
298
299
300
    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
301
302

class NatsServer(ManagedProcess):
303
    def __init__(self, request, port=4222, timeout=300, disable_jetstream=False):
304
305
306
307
308
309
310
311
        # 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
312
313
        self._request = request  # Store for restart
        self._timeout = timeout
314
315
        self._disable_jetstream = disable_jetstream
        data_dir = tempfile.mkdtemp(prefix="nats_") if not disable_jetstream else None
316
317
318
319
320
321
        command = [
            "nats-server",
            "--trace",
            "-p",
            str(port),
        ]
322
323
        if not disable_jetstream and data_dir:
            command.extend(["-js", "--store_dir", data_dir])
Neelay Shah's avatar
Neelay Shah committed
324
325
326
327
        super().__init__(
            command=command,
            timeout=timeout,
            display_output=False,
328
            terminate_existing=not use_random_port,  # Disabled for parallel test execution with random ports
Neelay Shah's avatar
Neelay Shah committed
329
330
            data_dir=data_dir,
            health_check_ports=[port],
331
            health_check_funcs=[self._nats_ready],
Neelay Shah's avatar
Neelay Shah committed
332
333
334
            log_dir=request.node.name,
        )

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
    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())

370
371
372
373
374
375
376
377
378
379
380
    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)

381
382
383
384
    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}")
        self._terminate_process_group()
385
386
        proc = self.proc  # type: ignore[has-type]
        if proc is not None:
387
            try:
388
                proc.wait(timeout=10)
389
390
391
392
393
394
395
            except Exception as e:
                _logger.warning(f"Error waiting for NATS process to stop: {e}")
            self.proc = None

    def start(self):
        """Restart a stopped NATS server with fresh state."""
        _logger.info(f"Starting NATS server on port {self.port} with fresh state")
396
397
398
399
400
401
402
403
        # 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
404
405
406
407
408
409
        self.command = [
            "nats-server",
            "--trace",
            "-p",
            str(self.port),
        ]
410
411
        if not self._disable_jetstream and self.data_dir:
            self.command.extend(["-js", "--store_dir", self.data_dir])
412
413

        self._start_process()
414
415
        elapsed = self._check_ports(self._timeout)
        self._check_funcs(self._timeout - elapsed)
416

Neelay Shah's avatar
Neelay Shah committed
417

418
class SharedManagedProcess:
419
420
421
422
423
424
    """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.
    """
425
426
427
428
429
430

    def __init__(
        self,
        request,
        tmp_path_factory,
        resource_name: str,
431
        start_port: int,
432
433
434
        timeout: int = 300,
    ):
        self.request = request
435
436
        self.start_port = start_port
        self.port: Optional[int] = None  # Set when entering context
437
438
439
440
        self.timeout = timeout
        self.resource_name = resource_name
        self._server: Optional[ManagedProcess] = None

441
        root_tmp = Path(tempfile.gettempdir()) / "pytest_shared_services"
442
443
        root_tmp.mkdir(parents=True, exist_ok=True)

444
445
        self.port_file = root_tmp / f"{resource_name}_port"
        self.lock_file = str(self.port_file) + ".lock"
446

447
    def _create_server(self, port: int) -> ManagedProcess:
448
449
450
        """Create the underlying server instance. Must be implemented by subclasses."""
        raise NotImplementedError

451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
466
    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():
467
            try:
468
                return int(self.port_file.read_text().strip())
469
            except (ValueError, IOError):
470
471
472
473
474
475
                return None
        return None

    def _write_port(self, port: int):
        """Write port to file."""
        self.port_file.write_text(str(port))
476
477
478

    def __enter__(self):
        with FileLock(self.lock_file):
479
480
481
482
483
484
485
486
487
            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}"
                )
488
            else:
489
490
491
492
493
494
495
496
497
                # 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__()
498
                logging.info(
499
                    f"[{self.resource_name}] Started process on port {self.port}"
500
501
502
503
                )
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
504
505
506
        # Never tear down - let the process live until the container dies.
        # This avoids race conditions and simplifies the logic.
        pass
507
508
509
510
511


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

512
513
    def __init__(self, request, tmp_path_factory, start_port=2380, timeout=300):
        super().__init__(request, tmp_path_factory, "etcd", start_port, timeout)
514
515
516
        # Create a log directory for session-scoped servers
        self._log_dir = tempfile.mkdtemp(prefix=f"pytest_{self.resource_name}_logs_")

517
    def _create_server(self, port: int) -> ManagedProcess:
518
        """Create EtcdServer instance."""
519
        server = EtcdServer(self.request, port=port, timeout=self.timeout)
520
521
522
523
524
525
526
527
        # 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."""

528
529
530
531
532
533
534
535
536
    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)
537
538
        # Create a log directory for session-scoped servers
        self._log_dir = tempfile.mkdtemp(prefix=f"pytest_{self.resource_name}_logs_")
539
        self._disable_jetstream = disable_jetstream
540

541
    def _create_server(self, port: int) -> ManagedProcess:
542
        """Create NatsServer instance."""
543
544
545
546
547
548
        server = NatsServer(
            self.request,
            port=port,
            timeout=self.timeout,
            disable_jetstream=self._disable_jetstream,
        )
549
550
551
552
553
        # Override log_dir since request.node.name is empty in session scope
        server.log_dir = self._log_dir
        return server


554
555
556
557
558
559
560
561
562
563
564
565
566
567
568
569
570
571
572
573
574
575
576
577
578
579
@pytest.fixture
def store_kv(request):
    """
    KV store for runtime. Defaults to "etcd".

    To iterate over multiple stores in a test:
        @pytest.mark.parametrize("store_kv", ["file", "etcd"], indirect=True)
        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")


580
581
582
583
584
585
586
587
588
589
590
591
592
593
594
595
596
597
598
599
600
@pytest.fixture
def use_nats_core(request):
    """
    Whether to use NATS Core mode (local indexer) instead of JetStream. Defaults to False.

    When True:
    - NATS server starts without JetStream (-js flag omitted) for faster startup
    - Tests should use enable_local_indexer=True in mocker_args

    When False (default):
    - NATS server starts with JetStream for KV event distribution
    - Tests use JetStream-based indexer synchronization

    To use NATS Core mode:
        @pytest.mark.parametrize("use_nats_core", [True], indirect=True)
        def test_example(runtime_services_dynamic_ports):
            ...
    """
    return getattr(request, "param", False)


Neelay Shah's avatar
Neelay Shah committed
601
@pytest.fixture()
602
603
604
605
606
607
def runtime_services(request, store_kv, request_plane):
    """
    Start runtime services (NATS and/or etcd) based on store_kv and request_plane.

    - If store_kv != "etcd", etcd is not started (returns None)
    - If request_plane != "nats", NATS is not started (returns None)
608
609

    Returns a tuple of (nats_process, etcd_process) where each has a .port attribute.
610
    """
611
    # Port cleanup is now handled in NatsServer and EtcdServer __exit__ methods
612
613
614
615
616
617
618
619
    if request_plane == "nats" and store_kv == "etcd":
        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
    elif store_kv == "etcd":
Neelay Shah's avatar
Neelay Shah committed
620
        with EtcdServer(request) as etcd_process:
621
622
623
            yield None, etcd_process
    else:
        yield None, None
624
625


626
@pytest.fixture()
627
def runtime_services_dynamic_ports(request, store_kv, request_plane, use_nats_core):
628
629
630
631
632
633
    """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.

634
635
636
637
638
    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.

639
    - If store_kv != "etcd", etcd is not started (returns None)
640
641
    - NATS is always started when etcd is used, because KV events require NATS
      regardless of the request_plane (tcp/nats only affects request transport)
642
    - JetStream is enabled by default; disabled when use_nats_core=True for faster startup
643
644
645
646
647
648

    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
649
    # Always start NATS when etcd is used - KV events require NATS regardless of request_plane
650
    # When use_nats_core=True, disable JetStream for faster startup
651
    if store_kv == "etcd":
652
653
654
        with NatsServer(
            request, port=0, disable_jetstream=use_nats_core
        ) as nats_process:
655
            with EtcdServer(request, port=0) as etcd_process:
656
657
658
659
660
                # 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
661
662
663
664
665
                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

666
667
668
669
670
671
672
673
674
                # 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)
675
    elif request_plane == "nats":
676
677
678
679
        with NatsServer(
            request, port=0, disable_jetstream=use_nats_core
        ) as nats_process:
            orig_nats = os.environ.get("NATS_SERVER")
680
681
            os.environ["NATS_SERVER"] = f"nats://localhost:{nats_process.port}"
            yield nats_process, None
682
683
684
685
            if orig_nats is not None:
                os.environ["NATS_SERVER"] = orig_nats
            else:
                os.environ.pop("NATS_SERVER", None)
686
687
688
689
    else:
        yield None, None


690
691
692
693
@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.

694
695
696
    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.
697

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

701
702
    For tests that need to restart NATS (e.g. indexer sync), use `runtime_services_dynamic_ports`
    which provides per-test isolated instances.
703
704
705
    """
    with SharedNatsServer(request, tmp_path_factory) as nats:
        with SharedEtcdServer(request, tmp_path_factory) as etcd:
706
707
708
709
            # 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}"

710
711
            yield nats, etcd

712
713
714
715
            # Clean up environment variables
            os.environ.pop("NATS_SERVER", None)
            os.environ.pop("ETCD_ENDPOINTS", None)

716

717
718
719
720
721
722
723
724
725
726
727
728
729
730
731
732
733
@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)
734
735
736
737
738
739
740
741
742
743
744
745
746
747
748
749
750
751
752
753
754
755
756
757
758
759
760
761
762
763
764
765
766
767


########################################################
# 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)
    """
    frontend_port = allocate_port(DefaultPort.FRONTEND.value)
    system_port_list = allocate_ports(num_system_ports, DefaultPort.SYSTEM1.value)
    all_ports = [frontend_port, *system_port_list]
    try:
        yield ServicePorts(frontend_port=frontend_port, system_ports=system_port_list)
    finally:
        deallocate_ports(all_ports)