managed_process.py 38.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
# SPDX-License-Identifier: Apache-2.0

4
import json
Neelay Shah's avatar
Neelay Shah committed
5
6
7
import logging
import os
import shutil
8
import signal
Neelay Shah's avatar
Neelay Shah committed
9
10
11
12
13
14
15
16
17
import socket
import subprocess
import time
from dataclasses import dataclass, field
from typing import Any, List, Optional

import psutil
import requests

18
19
from tests.utils.constants import DefaultPort
from tests.utils.port_utils import allocate_port, deallocate_port
20
from tests.utils.test_output import resolve_test_output_path
21

Neelay Shah's avatar
Neelay Shah committed
22

23
def terminate_process(process, logger=logging.getLogger(), immediate_kill=False):
24
25
26
27
28
29
    """Terminate a single process.

    Kill Strategy:
    - immediate_kill=False: Send SIGTERM (graceful, signal 15)
    - immediate_kill=True:  Send SIGKILL (force, signal 9, aka kill -9)
    """
30
31
32
    try:
        logger.info("Terminating PID: %s name: %s", process.pid, process.name())
        if immediate_kill:
33
34
            logger.info("Sending SIGKILL (kill -9): %s %s", process.pid, process.name())
            process.kill()  # SIGKILL (9)
35
        else:
36
37
            logger.info("Sending SIGTERM (kill): %s %s", process.pid, process.name())
            process.terminate()  # SIGTERM (15)
38
39
40
41
42
43
44
    except psutil.AccessDenied:
        logger.warning("Access denied for PID %s", process.pid)
    except psutil.NoSuchProcess:
        logger.warning("PID %s no longer exists", process.pid)


def terminate_process_tree(
45
    pid, logger=logging.getLogger(), immediate_kill=False, timeout=10
46
):
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
    """Terminate a process and all its children.

    Kill Sequence:
    ==============
    1. Snapshot all children (recursive)
    2. Send SIGTERM to parent IMMEDIATELY (no delay)
    3. Wait up to `timeout` seconds (default 10s) for parent to exit
    4. If parent still alive after timeout: Send SIGKILL to parent (force)
    5. Send SIGTERM to all children IMMEDIATELY
    6. Wait up to `timeout` seconds for all processes to exit
    7. If any still alive after timeout: Send SIGKILL to remaining (force)

    Timeout Parameter:
    - Controls how long to WAIT AFTER SIGTERM before escalating to SIGKILL
    - NOT a delay before sending SIGTERM (SIGTERM is sent immediately)

    Summary: SIGTERM (immediate) → wait → SIGKILL if still alive
    """
65
66
    try:
        parent = psutil.Process(pid)
67
68
    except psutil.NoSuchProcess:
        return
69

70
71
    # 1. Snapshot children before signaling parent
    children = parent.children(recursive=True)
72

73
74
    # 2. Terminate parent first (graceful)
    terminate_process(parent, logger, immediate_kill=immediate_kill)
75

76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
    # 3. Wait for parent to exit
    try:
        parent.wait(timeout=timeout)
    except psutil.TimeoutExpired:
        logger.warning("Parent process did not exit within timeout")
        terminate_process(parent, logger, immediate_kill=True)

    # 4. Terminate children if still alive
    for child in children:
        terminate_process(child, logger, immediate_kill=immediate_kill)

    # 5. Wait for all processes to exit
    all_procs = [parent] + children
    gone, alive = psutil.wait_procs(all_procs, timeout=timeout)

    # 6. Escalate remaining alive processes if needed
    for p in alive:
        terminate_process(p, logger, immediate_kill=True)
    psutil.wait_procs(alive, timeout=timeout)
95
96


97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
class ManagedProcessStopError(RuntimeError):
    """Raised when a ManagedProcess subprocess cannot be stopped cleanly."""

    def __init__(self, attr_name: str, pid: int, detail: str):
        super().__init__(
            f"Failed to stop managed process '{attr_name}' (pid={pid}): {detail}"
        )
        self.attr_name = attr_name
        self.pid = pid
        self.detail = detail


class ManagedProcessStopTimeoutError(ManagedProcessStopError):
    """Raised when a subprocess remains alive after forceful shutdown."""

    def __init__(self, attr_name: str, pid: int, wait_timeout: float):
        super().__init__(
            attr_name,
            pid,
            (
                "process did not exit after forceful shutdown "
                f"within {wait_timeout:.1f}s"
            ),
        )
        self.wait_timeout = wait_timeout


Neelay Shah's avatar
Neelay Shah committed
124
125
@dataclass
class ManagedProcess:
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
    """Manages a subprocess with health checks and automatic cleanup.

    Parallel Execution Safety:
    -------------------------
    For pytest-xdist or parallel test execution, use terminate_all_matching_process_names=False:

        # ❌ NOT SAFE for parallel tests (kills ALL vllm processes system-wide):
        ManagedProcess(
            command=["vllm", "serve", "--port", "8000"],
            terminate_all_matching_process_names=True,  # Kills all vllm, including other tests!
            stragglers=["vllm"],
        )

        # ✅ SAFE for parallel tests (only kills what we launch):
        ManagedProcess(
141
            command=["vllm", "serve", "--port", str(dynamically_allocated_port)],
142
143
144
145
146
147
148
149
150
151
            terminate_all_matching_process_names=False,  # Don't kill other processes
        )

    With terminate_all_matching_process_names=False and dynamic/randomized ports:
    - Each test gets unique ports, so zombies from previous runs won't conflict
    - ManagedProcess only terminates the process it launched (via self.proc.pid)
    - No risk of killing other tests' processes or developer's manual processes
    - Pytest will clean up any remaining processes when the test process exits
    """

Neelay Shah's avatar
Neelay Shah committed
152
153
154
155
    command: List[str]
    env: Optional[dict] = None
    health_check_ports: List[int] = field(default_factory=list)
    health_check_urls: List[Any] = field(default_factory=list)
156
    health_check_funcs: List[Any] = field(default_factory=list)
157
    delayed_start: int = 0
Neelay Shah's avatar
Neelay Shah committed
158
159
160
161
    timeout: int = 300
    working_dir: Optional[str] = None
    display_output: bool = False
    data_dir: Optional[str] = None
162
163
164
165
166
    # WARNING: terminate_all_matching_process_names=True is NOT safe for pytest-xdist
    # Kills ALL processes system-wide with matching name (other tests, dev processes, etc.)
    # Use False for parallel tests with dynamic ports to only kill launched process.
    # TODO: Change default to False once all tests use dynamic ports
    terminate_all_matching_process_names: bool = True
Neelay Shah's avatar
Neelay Shah committed
167
    stragglers: List[str] = field(default_factory=list)
168
    straggler_commands: List[str] = field(default_factory=list)
Neelay Shah's avatar
Neelay Shah committed
169
    log_dir: str = os.getcwd()
170
    display_name: Optional[str] = None
Neelay Shah's avatar
Neelay Shah committed
171

172
173
174
175
    # Ensure attributes exist even if startup fails early
    proc: Optional[subprocess.Popen] = None
    _pgid: Optional[int] = None

Neelay Shah's avatar
Neelay Shah committed
176
177
178
179
180
181
    _logger = logging.getLogger()
    _command_name = None
    _log_path = None
    _tee_proc = None
    _sed_proc = None

182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
    @property
    def log_path(self):
        """Return the absolute path to the process log file if available."""
        return self._log_path

    def read_logs(self) -> str:
        """Read and return the entire contents of the process log file.

        Returns an empty string if the log file is not yet available.
        """
        try:
            if self._log_path and os.path.exists(self._log_path):
                with open(self._log_path, "r", encoding="utf-8", errors="ignore") as f:
                    return f.read()
        except Exception as e:
            self._logger.warning("Could not read log file %s: %s", self._log_path, e)
        return ""

Neelay Shah's avatar
Neelay Shah committed
200
201
202
    def __enter__(self):
        try:
            self._logger = logging.getLogger(self.__class__.__name__)
203
204
205
206
207
            # self._command_name = self.command[0]
            if self.display_name:
                self._command_name = self.display_name
            else:
                self._command_name = self.command[0]
208
209
210
211

            # Keep test logs out of the git working tree: many tests pass a relative
            # `log_dir` derived from `request.node.name`, which otherwise creates a large
            # number of untracked directories under the repo root during pytest runs.
212
            self.log_dir = resolve_test_output_path(self.log_dir)
213

Neelay Shah's avatar
Neelay Shah committed
214
215
216
217
218
219
220
            os.makedirs(self.log_dir, exist_ok=True)
            log_name = f"{self._command_name}.log.txt"
            self._log_path = os.path.join(self.log_dir, log_name)

            if self.data_dir:
                self._remove_directory(self.data_dir)

221
            self._terminate_all_matching_process_names()  # Name-based cleanup (NOT xdist safe)
Neelay Shah's avatar
Neelay Shah committed
222
            self._start_process()
223
            time.sleep(self.delayed_start)
Neelay Shah's avatar
Neelay Shah committed
224
225
            elapsed = self._check_ports(self.timeout)
            self._check_urls(self.timeout - elapsed)
226
            self._check_funcs(self.timeout - elapsed)
Neelay Shah's avatar
Neelay Shah committed
227
228
229

            return self

230
231
232
233
234
235
236
237
        except Exception:
            try:
                self.__exit__(None, None, None)
            except Exception as cleanup_err:
                self._logger.warning(
                    "Error during cleanup in __enter__: %s", cleanup_err
                )
            raise
Neelay Shah's avatar
Neelay Shah committed
238

239
    def _cleanup_stragglers(self):
240
241
242
243
244
245
246
247
248
249
        """Clean up straggler processes - called during exit and signal handling.

        WARNING: NOT pytest-xdist safe! This does a system-wide sweep matching by
        process name (self.stragglers) and command-line substring (self.straggler_commands),
        similar to _terminate_all_matching_process_names. Skipped when
        terminate_all_matching_process_names=False (i.e. xdist-safe mode) to avoid
        killing other workers' processes.
        """
        if not self.terminate_all_matching_process_names:
            return
250
251
252
253
254
255
256
        try:
            if self.stragglers or self.straggler_commands:
                self._logger.info(
                    "Checking for straggler processes: stragglers=%s, straggler_commands=%s",
                    self.stragglers,
                    self.straggler_commands,
                )
257

258
            for ps_process in psutil.process_iter(["name", "cmdline"]):
259
                try:
260
261
                    process_name = ps_process.name()
                    if process_name in self.stragglers:
262
                        self._logger.info(
263
                            "Terminating Straggler %s %s", process_name, ps_process.pid
264
                        )
265
266
267
                        terminate_process_tree(
                            ps_process.pid, self._logger, immediate_kill=True
                        )
268
269
270
271
272
273
274
275
276
277
278
279

                    # Check command line arguments
                    cmdline = ps_process.cmdline()
                    cmdline_str = " ".join(cmdline) if cmdline else ""
                    for straggler_cmd in self.straggler_commands:
                        if straggler_cmd in cmdline_str:
                            self._logger.info(
                                "Terminating Straggler Cmdline %s %s %s",
                                process_name,
                                ps_process.pid,
                                straggler_cmd,
                            )
280
281
282
                            terminate_process_tree(
                                ps_process.pid, self._logger, immediate_kill=True
                            )
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
                            break  # Avoid terminating the same process multiple times
                except (
                    psutil.NoSuchProcess,
                    psutil.AccessDenied,
                    psutil.ZombieProcess,
                ):
                    # Process may have terminated or become inaccessible during iteration
                    pass
                except Exception as e:
                    # Catch any other unexpected errors to ensure cleanup continues
                    self._logger.warning("Error checking process: %s", e)
        except Exception as e:
            # Ensure that any error in straggler cleanup doesn't prevent other cleanup
            self._logger.error("Error during straggler cleanup: %s", e)

    def __exit__(self, exc_type, exc_val, exc_tb):
299
300
        """Cleanup: Terminate launched processes.

301
302
303
304
305
306
307
        Termination Strategy:
        =====================
        1. _stop_started_processes: calls _terminate_process_group (SIGTERM →
           poll → SIGKILL all snapshotted pgids), closes pipes, waits/escalates
           per-process.
        2. Clean up data directory if configured.
        3. Run name-based straggler cleanup (only when not in xdist-safe mode).
308
        """
309
        try:
310
            self._stop_started_processes()
311

312
313
314
315
            if self.data_dir:
                self._remove_directory(self.data_dir)
        finally:
            self._cleanup_stragglers()
Neelay Shah's avatar
Neelay Shah committed
316

317
318
319
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
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
    def _stop_started_processes(self, wait_timeout: float = 10.0):
        """Terminate launched subprocesses and close any open pipe handles.

        This is used both during normal teardown and when a managed service
        needs to stop and restart in-place without releasing higher-level
        resources such as ports.
        """
        self._terminate_process_group()

        process_entries = [
            ("proc", self.proc),
            ("_tee_proc", self._tee_proc),
            ("_sed_proc", self._sed_proc),
        ]
        for attr_name, process in process_entries:
            if process is None:
                continue

            try:
                for stream_name in ("stdout", "stdin", "stderr"):
                    stream = getattr(process, stream_name, None)
                    if stream is not None:
                        stream.close()
            except OSError as e:
                self._logger.warning("Error closing process streams: %s", e)

            try:
                terminate_process_tree(process.pid, self._logger)
                process.wait(timeout=wait_timeout)
            except (psutil.TimeoutExpired, subprocess.TimeoutExpired) as e:
                self._logger.warning(
                    "Process '%s' (pid=%s) timed out during shutdown: %s. "
                    "Retrying with forceful termination.",
                    attr_name,
                    process.pid,
                    e,
                )
                self._force_stop_process(process, attr_name, wait_timeout)
            except OSError as e:
                if process.poll() is None:
                    self._logger.warning(
                        "Process '%s' (pid=%s) hit an OS error during shutdown: %s. "
                        "Retrying with forceful termination.",
                        attr_name,
                        process.pid,
                        e,
                    )
                    self._force_stop_process(process, attr_name, wait_timeout)

            if process.poll() is None:
                raise ManagedProcessStopTimeoutError(
                    attr_name, process.pid, wait_timeout
                )

            setattr(self, attr_name, None)

        self._pgid = None

    def _force_stop_process(
        self,
        process: subprocess.Popen,
        attr_name: str,
        wait_timeout: float,
    ) -> None:
        """Forcefully stop a process tree and confirm the launched child exits."""
        try:
            terminate_process_tree(
                process.pid,
                self._logger,
                immediate_kill=True,
                timeout=0,
            )
        except (psutil.TimeoutExpired, OSError) as e:
            self._logger.warning(
                "Forceful tree termination failed for process '%s' (pid=%s): %s. "
                "Falling back to process.kill().",
                attr_name,
                process.pid,
                e,
            )
            try:
                process.kill()
            except OSError as kill_err:
                if process.poll() is None:
                    raise ManagedProcessStopError(
                        attr_name,
                        process.pid,
                        f"forceful shutdown failed: {kill_err}",
                    ) from kill_err

        try:
            process.wait(timeout=wait_timeout)
        except subprocess.TimeoutExpired as e:
            if process.poll() is None:
                raise ManagedProcessStopTimeoutError(
                    attr_name, process.pid, wait_timeout
                ) from e
        except OSError as e:
            if process.poll() is None:
                raise ManagedProcessStopError(
                    attr_name,
                    process.pid,
                    f"error waiting for exit after forceful shutdown: {e}",
                ) from e

        if process.poll() is None:
            raise ManagedProcessStopTimeoutError(attr_name, process.pid, wait_timeout)

Neelay Shah's avatar
Neelay Shah committed
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
    def _start_process(self):
        assert self._command_name
        assert self._log_path

        self._logger.info(
            "Running command: %s in %s",
            " ".join(self.command),
            self.working_dir or os.getcwd(),
        )

        stdin = subprocess.DEVNULL
        stdout = subprocess.PIPE
        stderr = subprocess.STDOUT

        if self.display_output:
            self.proc = subprocess.Popen(
                self.command,
                env=self.env or os.environ.copy(),
                cwd=self.working_dir,
                stdin=stdin,
                stdout=stdout,
                stderr=stderr,
447
                start_new_session=True,  # Isolate process group to prevent kill 0 from affecting parent
Neelay Shah's avatar
Neelay Shah committed
448
            )
449
450
451
452
453
454
            # Capture the child's process group id for robust cleanup even if parent shell exits
            try:
                self._pgid = os.getpgid(self.proc.pid)
            except Exception as e:
                self._logger.warning("Could not get process group id: %s", e)
                self._pgid = None
Neelay Shah's avatar
Neelay Shah committed
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
473
            self._sed_proc = subprocess.Popen(
                ["sed", "-u", f"s/^/[{self._command_name.upper()}] /"],
                stdin=self.proc.stdout,
                stdout=subprocess.PIPE,
            )

            self._tee_proc = subprocess.Popen(
                ["tee", self._log_path], stdin=self._sed_proc.stdout
            )

        else:
            with open(self._log_path, "w", encoding="utf-8") as f:
                self.proc = subprocess.Popen(
                    self.command,
                    env=self.env or os.environ.copy(),
                    cwd=self.working_dir,
                    stdin=stdin,
                    stdout=stdout,
                    stderr=stderr,
474
                    start_new_session=True,  # Isolate process group to prevent kill 0 from affecting parent
Neelay Shah's avatar
Neelay Shah committed
475
                )
476
477
478
479
480
481
                # Capture the child's process group id for robust cleanup even if parent shell exits
                try:
                    self._pgid = os.getpgid(self.proc.pid)
                except Exception as e:
                    self._logger.warning("Could not get process group id: %s", e)
                    self._pgid = None
Neelay Shah's avatar
Neelay Shah committed
482
483
484
485
486
487
488
489

                self._sed_proc = subprocess.Popen(
                    ["sed", "-u", f"s/^/[{self._command_name.upper()}] /"],
                    stdin=self.proc.stdout,
                    stdout=f,
                )
            self._tee_proc = None

490
    def _terminate_process_group(self, timeout: float = 8.0):
491
492
        """Terminate the entire process group/session started for the child.

493
494
        Kill Sequence:
        ==============
495
496
497
498
499
500
501
502
503
504
505
506
507
508
509
510
511
512
513
514
515
516
517
518
519
        1. Snapshot all child process groups (the main pgid plus any distinct
           child pgids) before sending any signals.
        2. Send SIGTERM to ALL snapshotted process groups (no delay).
           os.killpg delivers the signal atomically to every group member.
        3. Wait up to `timeout` seconds (default 8s), polling every 0.1s.
           8s usually gives engines (e.g. vLLM EngineCore) enough time to flush
           in-flight work and release resources before we escalate.
        4. SIGKILL ALL snapshotted process groups (those that ignored SIGTERM
           or timed out).

        Limitation: if a child calls setsid()/setpgid() AFTER our snapshot,
        its new group won't be in our set. This is rare -- Python
        multiprocessing and vLLM engine core inherit the parent's group.

        Post-SIGKILL resource cleanup notes:
        - GPU memory: vLLM Engine releases VRAM on process death (driver-level
          reclaim). No manual GPU cleanup needed after SIGKILL.
        - Distributed state: NATS subscriptions, etcd leases, and KV cache
          index entries are NOT cleaned up by SIGKILL. These persist until
          TTL expiry or explicit purge.
        - Shared memory: POSIX shm segments (/dev/shm/) survive process
          death and must be unlinked separately if used.
        - Network ports: TCP sockets enter TIME_WAIT (~60s); rebinding the
          same port immediately may fail. Prefer allocate_port() from
          tests.utils.port_utils to avoid collisions.
520
521
522
523
        """
        if self._pgid is None:
            return

524
525
526
527
528
529
530
531
532
533
534
535
536
537
538
539
540
541
542
543
544
545
546
547
548
549
550
551
552
553
554
555
556
557
558
559
560
        # Step 1: Snapshot all process groups before signaling.  Only self.proc runs
        # in self._pgid (start_new_session=True); _tee_proc/_sed_proc are pipe
        # helpers in pytest's group with no children.
        all_pgids: set[int] = {self._pgid}
        if self.proc and self.proc.poll() is None:
            try:
                for child in psutil.Process(self.proc.pid).children(recursive=True):
                    try:
                        all_pgids.add(os.getpgid(child.pid))
                    except (ProcessLookupError, OSError):
                        pass
            except psutil.NoSuchProcess:
                pass

        # Step 2: SIGTERM all snapshotted process groups (graceful shutdown).
        # Delivers the signal to every group so children that called
        # setpgid()/setsid() also get a chance to shut down gracefully.
        sigtermed_pgids: set[int] = set()
        for pgid in all_pgids:
            try:
                self._logger.info("Sending SIGTERM to process group: %s", pgid)
                os.killpg(pgid, signal.SIGTERM)
                sigtermed_pgids.add(pgid)
            except ProcessLookupError:
                pass  # Already gone
            except Exception as e:
                self._logger.warning(
                    "Error sending SIGTERM to process group %s: %s", pgid, e
                )

        if not sigtermed_pgids:
            return  # All groups already gone

        # Step 3: poll for all process groups to exit.
        # self.proc.poll() reaps the zombie so os.killpg(pgid, 0) can
        # detect an empty group; without this the zombie keeps the group
        # "alive" and the loop burns the full timeout.
561
562
563
        poll_interval = 0.1
        elapsed = 0.0
        while elapsed < timeout:
564
565
            if self.proc is not None:
                self.proc.poll()
566
567
568
569
570
571
572
573
574
575
576
            # Check if any signaled group is still alive
            still_alive = False
            for pgid in sigtermed_pgids:
                try:
                    os.killpg(pgid, 0)  # signal 0 = check existence
                    still_alive = True
                    break  # At least one alive, keep waiting
                except (ProcessLookupError, OSError):
                    pass
            if not still_alive:
                break
577
578
579
            time.sleep(poll_interval)
            elapsed += poll_interval

580
581
582
583
584
585
586
587
588
589
590
591
592
593
594
595
596
597
        # Step 4: SIGKILL all snapshotted process groups to ensure nothing
        # survives (e.g. processes that ignored SIGTERM or timed out).
        # _stop_started_processes handles per-process wait/escalation afterward.
        for pgid in all_pgids:
            try:
                os.killpg(pgid, signal.SIGKILL)
            except ProcessLookupError:
                pass
            except PermissionError:
                self._logger.warning(
                    "Permission denied sending SIGKILL to process group %s "
                    "(should not happen for our own children)",
                    pgid,
                )
            except Exception as e:
                self._logger.warning(
                    "Error sending SIGKILL to process group %s: %s", pgid, e
                )
598

Neelay Shah's avatar
Neelay Shah committed
599
600
601
602
603
604
605
    def _remove_directory(self, path: str) -> None:
        """Remove a directory."""
        try:
            shutil.rmtree(path, ignore_errors=True)
        except (OSError, IOError) as e:
            self._logger.warning("Warning: Failed to remove directory %s: %s", path, e)

606
607
608
609
610
611
612
613
614
615
616
617
618
619
620
621
622
623
624
625
626
627
628
629
630
631
632
633
634
635
636
637
638
    def _log_tail_on_error(self, lines=20):
        """Print the last few lines of the log file when process dies."""
        if self._log_path and os.path.exists(self._log_path):
            try:
                with open(self._log_path, "r") as f:
                    log_lines = f.readlines()
                    if log_lines:
                        self._logger.error(
                            "=== Last %d lines from %s ===",
                            min(lines, len(log_lines)),
                            self._log_path,
                        )
                        for line in log_lines[-lines:]:
                            self._logger.error(line.rstrip())
                        self._logger.error("=== End of log tail ===")
            except Exception as e:
                self._logger.warning("Could not read log file: %s", e)

    def _check_process_alive(self, context=""):
        """Check if the main process is still alive. Raises RuntimeError if dead."""
        if self.proc and self.proc.poll() is not None:
            returncode = self.proc.returncode
            self._logger.error(
                "Main server process died with exit code %d%s",
                returncode,
                f" {context}" if context else "",
            )
            # Try to get last few lines from log for debugging
            self._log_tail_on_error()
            raise RuntimeError(
                f"Main server process exited with code {returncode}{f' {context}' if context else ''}"
            )

Neelay Shah's avatar
Neelay Shah committed
639
640
641
642
643
644
645
646
647
648
649
650
    def _check_ports(self, timeout):
        elapsed = 0.0
        for port in self.health_check_ports:
            elapsed += self._check_port(port, timeout - elapsed)
        return elapsed

    def _check_port(self, port, timeout=30, sleep=0.1):
        """Check if a port is open on localhost."""
        start_time = time.time()
        self._logger.info("Checking Port: %s", port)
        elapsed = 0.0
        while elapsed < timeout:
651
652
653
            # Check if the main process is still alive
            self._check_process_alive(f"while waiting for port {port}")

Neelay Shah's avatar
Neelay Shah committed
654
655
656
657
658
659
660
661
662
663
664
665
666
667
668
            with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
                if s.connect_ex(("localhost", port)) == 0:
                    self._logger.info("SUCCESS: Check Port: %s", port)
                    return time.time() - start_time
            time.sleep(sleep)
            elapsed = time.time() - start_time
        self._logger.error("FAILED: Check Port: %s", port)
        raise RuntimeError("FAILED: Check Port: %s" % port)

    def _check_urls(self, timeout):
        elapsed = 0.0
        for url in self.health_check_urls:
            elapsed += self._check_url(url, timeout - elapsed)
        return elapsed

669
    def _check_url(self, url, timeout=30, sleep=1, log_interval=20):
Neelay Shah's avatar
Neelay Shah committed
670
671
672
673
674
675
676
677
        if isinstance(url, tuple):
            response_check = url[1]
            url = url[0]
        else:
            response_check = None
        start_time = time.time()
        self._logger.info("Checking URL %s", url)
        elapsed = 0.0
678
679
680
        attempt = 0
        last_log_time = 0.0

Neelay Shah's avatar
Neelay Shah committed
681
        while elapsed < timeout:
682
683
684
685
686
687
            self._check_process_alive("while waiting for health check")

            attempt += 1
            check_failed = False
            failure_reason = None

Neelay Shah's avatar
Neelay Shah committed
688
689
690
691
            try:
                response = requests.get(url, timeout=timeout - elapsed)
                if response.status_code == 200:
                    if response_check is None or response_check(response):
692
693
694
695
696
697
698
699
700
701
702
703
704
705
706
707
708
709
710
711
712
713
714
                        # Try to format JSON response nicely, otherwise show raw text
                        try:
                            response_data = response.json()
                            response_str = json.dumps(response_data, indent=2)
                            self._logger.info(
                                "SUCCESS: Check URL: %s (attempt=%d, elapsed=%.1fs)\nResponse:\n%s",
                                url,
                                attempt,
                                elapsed,
                                response_str,
                            )
                        except (json.JSONDecodeError, Exception):
                            # If not JSON or any error, show raw text (truncated if too long)
                            response_text = response.text
                            if len(response_text) > 500:
                                response_text = response_text[:500] + "... (truncated)"
                            self._logger.info(
                                "SUCCESS: Check URL: %s (attempt=%d, elapsed=%.1fs)\nResponse: %s",
                                url,
                                attempt,
                                elapsed,
                                response_text,
                            )
Neelay Shah's avatar
Neelay Shah committed
715
                        return time.time() - start_time
716
717
718
719
720
721
                    else:
                        check_failed = True
                        failure_reason = "custom check failed"
                else:
                    check_failed = True
                    failure_reason = f"status code {response.status_code}"
Neelay Shah's avatar
Neelay Shah committed
722
            except requests.RequestException as e:
723
724
725
726
727
728
729
730
731
732
733
734
735
736
                check_failed = True
                failure_reason = f"request exception: {e}"

            # Log progress every log_interval seconds for any failure
            if check_failed and elapsed - last_log_time >= log_interval:
                self._logger.info(
                    "Still waiting for URL %s (%s) (attempt=%d, elapsed=%.1fs)",
                    url,
                    failure_reason,
                    attempt,
                    elapsed,
                )
                last_log_time = elapsed

Neelay Shah's avatar
Neelay Shah committed
737
738
739
            time.sleep(sleep)
            elapsed = time.time() - start_time

740
        self._logger.error(
741
742
743
744
745
746
747
748
749
            "TIMEOUT: Check URL: %s failed after %.1fs (attempts=%d, timeout=%.1fs)",
            url,
            elapsed,
            attempt,
            timeout,
        )
        raise RuntimeError(
            "TIMEOUT: Check URL: %s failed after %.1fs (timeout=%.1fs)"
            % (url, elapsed, timeout)
750
        )
Neelay Shah's avatar
Neelay Shah committed
751

752
753
754
755
756
757
758
759
760
761
762
763
764
765
766
767
768
769
770
771
772
773
774
775
776
777
778
779
780
781
782
783
784
785
786
787
788
789
790
791
792
793
794
795
796
797
798
799
800
801
802
803
804
805
806
807
808
809
810
811
812
813
814
815
    def _check_funcs(self, timeout):
        elapsed = 0.0
        for func in self.health_check_funcs:
            elapsed += self._check_func(func, timeout - elapsed)
        return elapsed

    def _check_func(self, func, timeout=30, sleep=1, log_interval=20):
        start_time = time.time()
        func_name = getattr(func, "__name__", str(func))
        self._logger.info("Running custom health check '%s'", func_name)
        elapsed = 0.0
        attempt = 0
        last_log_time = 0.0

        while elapsed < timeout:
            self._check_process_alive("while waiting for health check")

            attempt += 1
            check_failed = False
            failure_reason = None

            try:
                # Prefer functions that accept remaining timeout; fall back to no-arg call
                try:
                    result = func(timeout - elapsed)
                except TypeError:
                    result = func()

                if bool(result):
                    self._logger.info(
                        "SUCCESS: Custom health check '%s' passed (attempt=%d, elapsed=%.1fs)",
                        func_name,
                        attempt,
                        elapsed,
                    )
                    return time.time() - start_time
                else:
                    check_failed = True
                    failure_reason = "returned False"
            except Exception as e:
                check_failed = True
                failure_reason = f"exception: {e}"

            if check_failed and elapsed - last_log_time >= log_interval:
                self._logger.info(
                    "Still waiting on custom health check '%s' (%s) (attempt=%d, elapsed=%.1fs)",
                    func_name,
                    failure_reason,
                    attempt,
                    elapsed,
                )
                last_log_time = elapsed

            time.sleep(sleep)
            elapsed = time.time() - start_time

        self._logger.error(
            "FAILED: Custom health check '%s' (attempts=%d, elapsed=%.1fs)",
            func_name,
            attempt,
            elapsed,
        )
        raise RuntimeError("FAILED: Custom health check")

816
817
818
819
820
821
822
823
824
825
826
827
828
    def _terminate_all_matching_process_names(self):
        # WARNING: This method is NOT pytest-xdist safe! Kills ALL matching processes system-wide.
        # DANGER: When terminate_all_matching_process_names=True, this kills:
        #   - Processes with matching name (self._command_name)
        #   - Processes in stragglers list
        #   - Processes with matching command substring (straggler_commands)
        # It does NOT check:
        #   - Port numbers (kills processes on ALL ports, not just ours)
        #   - Process ownership (kills other tests' processes)
        #   - Working directory (kills unrelated processes)
        # RECOMMENDED: Use terminate_all_matching_process_names=False for parallel tests.
        # ManagedProcess already only kills what it launched (self.proc.pid) on exit.
        if self.terminate_all_matching_process_names:
Neelay Shah's avatar
Neelay Shah committed
829
            for proc in psutil.process_iter(["name", "cmdline"]):
830
831
832
833
834
                try:
                    if (
                        proc.name() == self._command_name
                        or proc.name() in self.stragglers
                    ):
835
                        self._logger.info(
836
                            "Terminating Existing %s %s", proc.name(), proc.pid
837
                        )
838

839
                        terminate_process_tree(proc.pid, self._logger)
840
841
842
843
844
845
846
847
848
849
850
851
852
853
854
855
                    for cmdline in self.straggler_commands:
                        if cmdline in " ".join(proc.cmdline()):
                            self._logger.info(
                                "Terminating Existing CmdLine %s %s %s",
                                proc.name(),
                                proc.pid,
                                proc.cmdline(),
                            )
                            terminate_process_tree(proc.pid, self._logger)
                except (
                    psutil.NoSuchProcess,
                    psutil.AccessDenied,
                    psutil.ZombieProcess,
                ):
                    # Process may have terminated or become inaccessible during iteration
                    pass
Neelay Shah's avatar
Neelay Shah committed
856

857
858
859
860
861
862
    def is_running(self) -> bool:
        """Check if the process is still running"""
        return (
            hasattr(self, "proc") and self.proc is not None and self.proc.poll() is None
        )

863
864
865
866
    def get_pid(self) -> int | None:
        """Get the PID of the managed process."""
        return self.proc.pid if self.proc else None

867
868
869
870
871
872
873
874
875
876
877
878
879
880
881
    def subprocesses(self) -> list[psutil.Process]:
        """Find child processes of the current process."""
        if (
            not hasattr(self, "proc")
            or self.proc is None
            or self.proc.poll() is not None
        ):
            return []

        try:
            parent = psutil.Process(self.proc.pid)
            return parent.children(recursive=True)
        except psutil.NoSuchProcess:
            return []

Neelay Shah's avatar
Neelay Shah committed
882

883
884
885
886
887
class DynamoFrontendProcess(ManagedProcess):
    """Process manager for Dynamo frontend"""

    _logger = logging.getLogger()

888
889
890
891
892
893
    def __init__(
        self,
        request: Any,
        *,
        frontend_port: Optional[int] = None,
        router_mode: str = "round-robin",
894
        migration_limit: int = 0,
895
896
        extra_args: Optional[list[str]] = None,
        extra_env: Optional[dict[str, str]] = None,
897
        # Default to false so pytest-xdist workers don't kill each other's frontends.
898
        terminate_all_matching_process_names: bool = False,
899
        display_name: Optional[str] = None,
900
901
902
    ):
        # TODO: Refactor remaining duplicate "DynamoFrontendProcess" helpers in tests to
        # use this shared implementation (and delete the copies):
903
904
905
906
        # - tests/fault_tolerance/cancellation/utils.py
        # - tests/fault_tolerance/migration/utils.py
        # - tests/fault_tolerance/etcd_ha/utils.py
        # - tests/fault_tolerance/test_vllm_health_check.py
907
908
909
910
911
912
913
914
915
916
917
918
919
920
921
922
923
        self._allocated_http_port: Optional[int] = None
        if frontend_port == 0:
            # Treat `0` as "allocate a random free port" for xdist-safe tests.
            # We allocate within the i16-safe range required by the Rust side.
            frontend_port = allocate_port(DefaultPort.FRONTEND.value)
            self._allocated_http_port = frontend_port

        # If frontend_port is unset, dynamo.frontend defaults to DefaultPort.FRONTEND.
        self.http_port = (
            DefaultPort.FRONTEND.value if frontend_port is None else int(frontend_port)
        )

        command = ["python", "-m", "dynamo.frontend", "--router-mode", router_mode]

        # dynamo.frontend defaults to 8000 when neither env nor flag is provided.
        if frontend_port is not None:
            command.extend(["--http-port", str(frontend_port)])
924
925
        # Migration limit is configured at the frontend level
        command.extend(["--migration-limit", str(migration_limit)])
926
927
        if extra_args:
            command.extend(extra_args)
928

929
930
931
        # Unset DYN_SYSTEM_PORT - frontend doesn't use system metrics server
        env = os.environ.copy()
        env.pop("DYN_SYSTEM_PORT", None)
932
933
        if extra_env:
            env.update(extra_env)
934

935
936
937
938
939
940
941
942
943
944
945
946
        log_dir = f"{request.node.name}_frontend"

        # Clean up any existing log directory from previous runs
        try:
            shutil.rmtree(log_dir)
            self._logger.info(f"Cleaned up existing log directory: {log_dir}")
        except FileNotFoundError:
            # Directory doesn't exist, which is fine
            pass

        super().__init__(
            command=command,
947
            env=env,
948
            display_output=True,
949
            terminate_all_matching_process_names=terminate_all_matching_process_names,
950
            log_dir=log_dir,
951
            display_name=display_name,
952
953
        )

954
955
956
957
958
959
960
961
    def __exit__(self, exc_type, exc_val, exc_tb):
        try:
            return super().__exit__(exc_type, exc_val, exc_tb)
        finally:
            if self._allocated_http_port is not None:
                deallocate_port(self._allocated_http_port)
                self._allocated_http_port = None

962
963
964
965
966
    @property
    def frontend_port(self) -> int:
        """Back-compat alias for older tests that expect `frontend.frontend_port`."""
        return self.http_port

967

Neelay Shah's avatar
Neelay Shah committed
968
def main():
969
970
    # NOTE: This entrypoint is for manual testing/debugging of `ManagedProcess` only.
    # It is not used by the pytest suite.
971
972

    # Example: Safe for parallel tests with terminate_all_matching_process_names=False
Neelay Shah's avatar
Neelay Shah committed
973
    with ManagedProcess(
974
        command=["python", "-m", "dynamo.frontend", "--port", "8000"],
Neelay Shah's avatar
Neelay Shah committed
975
        display_output=True,
976
        terminate_all_matching_process_names=False,  # ✅ Safe - only kills what we launch
977
978
        health_check_ports=[8000],
        health_check_urls=["http://localhost:8000/v1/models"],
Neelay Shah's avatar
Neelay Shah committed
979
980
981
982
983
984
985
986
        timeout=10,
    ):
        time.sleep(60)
        pass


if __name__ == "__main__":
    main()