managed_process.py 33.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


Neelay Shah's avatar
Neelay Shah committed
97
98
@dataclass
class ManagedProcess:
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
124
    """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(
            command=["vllm", "serve", "--port", "8000"],
            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
125
126
127
128
    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)
129
    health_check_funcs: List[Any] = field(default_factory=list)
130
    delayed_start: int = 0
Neelay Shah's avatar
Neelay Shah committed
131
132
133
134
    timeout: int = 300
    working_dir: Optional[str] = None
    display_output: bool = False
    data_dir: Optional[str] = None
135
136
137
138
139
    # 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
140
    stragglers: List[str] = field(default_factory=list)
141
    straggler_commands: List[str] = field(default_factory=list)
Neelay Shah's avatar
Neelay Shah committed
142
    log_dir: str = os.getcwd()
143
    display_name: Optional[str] = None
Neelay Shah's avatar
Neelay Shah committed
144

145
146
147
148
    # Ensure attributes exist even if startup fails early
    proc: Optional[subprocess.Popen] = None
    _pgid: Optional[int] = None

Neelay Shah's avatar
Neelay Shah committed
149
150
151
152
153
154
    _logger = logging.getLogger()
    _command_name = None
    _log_path = None
    _tee_proc = None
    _sed_proc = None

155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
    @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
173
174
175
    def __enter__(self):
        try:
            self._logger = logging.getLogger(self.__class__.__name__)
176
177
178
179
180
            # self._command_name = self.command[0]
            if self.display_name:
                self._command_name = self.display_name
            else:
                self._command_name = self.command[0]
181
182
183
184

            # 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.
185
            self.log_dir = resolve_test_output_path(self.log_dir)
186

Neelay Shah's avatar
Neelay Shah committed
187
188
189
190
191
192
193
            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)

194
            self._terminate_all_matching_process_names()  # Name-based cleanup (NOT xdist safe)
Neelay Shah's avatar
Neelay Shah committed
195
            self._start_process()
196
            time.sleep(self.delayed_start)
Neelay Shah's avatar
Neelay Shah committed
197
198
            elapsed = self._check_ports(self.timeout)
            self._check_urls(self.timeout - elapsed)
199
            self._check_funcs(self.timeout - elapsed)
Neelay Shah's avatar
Neelay Shah committed
200
201
202

            return self

203
204
205
206
207
208
209
210
        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
211

212
213
214
215
216
217
218
219
220
    def _cleanup_stragglers(self):
        """Clean up straggler processes - called during exit and signal handling"""
        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,
                )
221

222
            for ps_process in psutil.process_iter(["name", "cmdline"]):
223
                try:
224
225
                    process_name = ps_process.name()
                    if process_name in self.stragglers:
226
                        self._logger.info(
227
                            "Terminating Straggler %s %s", process_name, ps_process.pid
228
229
                        )
                        terminate_process_tree(ps_process.pid, self._logger)
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258

                    # 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,
                            )
                            terminate_process_tree(ps_process.pid, self._logger)
                            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):
259
260
261
262
        """Cleanup: Terminate launched processes.

        Termination Strategy (Graceful → Escalate to Force):
        =====================================================
263
264
265
266
267
        1. Snapshot child processes (before killing parent, so we can find them)
        2. Send SIGTERM to process group immediately (no delay before SIGTERM)
        3. Wait up to 2s (poll every 0.1s) for processes to exit
        4. If still alive after 2s: Send SIGKILL (force kill)
        5. Terminate individual processes (self.proc, tee, sed):
268
269
270
           - Send SIGTERM immediately (no delay)
           - Wait up to 10s for exit
           - If still alive after 10s: Send SIGKILL (force kill)
271
272
273
        6. Kill any orphaned child processes that escaped the process group
           (e.g. TRT-LLM engine workers started via MPI in separate PGIDs)
        7. Clean up straggler processes (if configured)
274
275
276
277
278
279
280
281
282
283
284

        Signal Details:
        - SIGTERM (15): Graceful - allows cleanup handlers to run
        - SIGKILL (9): Force kill - immediate, cannot be caught

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

        This ALWAYS runs regardless of terminate_all_matching_process_names setting.
        """
285
        try:
286
287
288
289
290
291
292
293
294
295
296
297
298
            # Snapshot all child processes BEFORE killing the parent.
            # Some frameworks (e.g. TRT-LLM via MPI) spawn children in separate
            # process groups. After the parent dies these children become orphans
            # reparented to init, so psutil can no longer find them via the parent
            # PID. We must capture them now while the parent is still alive.
            orphan_candidates = []
            if self.proc:
                try:
                    parent = psutil.Process(self.proc.pid)
                    orphan_candidates = parent.children(recursive=True)
                except psutil.NoSuchProcess:
                    pass

299
300
301
302
303
304
305
306
307
308
309
310
311
312
            self._terminate_process_group()

            process_list = [self.proc, self._tee_proc, self._sed_proc]
            for process in process_list:
                if process:
                    try:
                        if process.stdout:
                            process.stdout.close()
                        if process.stdin:
                            process.stdin.close()
                        terminate_process_tree(process.pid, self._logger)
                        process.wait()
                    except Exception as e:
                        self._logger.warning("Error terminating process: %s", e)
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328

            # Kill any child processes that survived the process group kill.
            # This catches children in different PGIDs (e.g. MPI workers, engine
            # cores) that os.killpg() could not reach.
            for child in orphan_candidates:
                try:
                    if child.is_running():
                        self._logger.info(
                            "Killing orphaned child process: PID %s name=%s",
                            child.pid,
                            child.name(),
                        )
                        terminate_process_tree(child.pid, self._logger)
                except (psutil.NoSuchProcess, psutil.AccessDenied):
                    pass

329
330
331
332
333
            if self.data_dir:
                self._remove_directory(self.data_dir)
        finally:
            # Always run straggler cleanup, even if interrupted
            self._cleanup_stragglers()
Neelay Shah's avatar
Neelay Shah committed
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356

    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,
357
                start_new_session=True,  # Isolate process group to prevent kill 0 from affecting parent
Neelay Shah's avatar
Neelay Shah committed
358
            )
359
360
361
362
363
364
            # 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
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
            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,
384
                    start_new_session=True,  # Isolate process group to prevent kill 0 from affecting parent
Neelay Shah's avatar
Neelay Shah committed
385
                )
386
387
388
389
390
391
                # 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
392
393
394
395
396
397
398
399

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

400
    def _terminate_process_group(self, timeout: float = 8.0):
401
402
        """Terminate the entire process group/session started for the child.

403
404
405
        Kill Sequence:
        ==============
        1. Send SIGTERM to entire process group IMMEDIATELY (no delay)
406
        2. Wait up to `timeout` seconds (default 8s), polling every 0.1s
407
408
409
410
411
        3. If still alive after timeout: Send SIGKILL (force kill, immediate)

        Timeout Parameter:
        - Controls how long to WAIT AFTER SIGTERM before escalating to SIGKILL
        - NOT a delay before sending SIGTERM (SIGTERM is sent immediately)
412
413
414
        - 8s gives engines (TRT-LLM, vLLM, etc.) enough time to gracefully
          shut down MPI workers, release GPU memory, and drain pending requests
        - Polling at 0.1s intervals means fast exits are not penalized
415
416
417
418

        Process groups catch cases where the launcher shell exits and its
        children are reparented, leaving no parent PID to traverse, but they
        remain in the same process group.
419
420
421
422
423
        """
        if self._pgid is None:
            return
        try:
            self._logger.info("Terminating process group: %s", self._pgid)
424
            os.killpg(self._pgid, signal.SIGTERM)  # Step 1: Graceful SIGTERM
425
426
427
428
429
430
431
432
        except ProcessLookupError:
            return
        except Exception as e:
            self._logger.warning(
                "Error sending SIGTERM to process group %s: %s", self._pgid, e
            )
            return

433
        # Step 2: Poll for process exit instead of fixed sleep to minimize teardown time
434
435
436
437
438
        poll_interval = 0.1
        elapsed = 0.0
        while elapsed < timeout:
            try:
                # Check if any process in the group is still alive
439
                os.killpg(self._pgid, 0)  # Signal 0 = check existence (no kill)
440
441
442
443
444
445
446
447
448
            except ProcessLookupError:
                # Process group no longer exists - done
                return
            except Exception:
                # Other errors (e.g., permission) - assume done
                return
            time.sleep(poll_interval)
            elapsed += poll_interval

449
        # Step 3: Force kill if anything remains after timeout
450
        try:
451
            os.killpg(self._pgid, signal.SIGKILL)  # SIGKILL (kill -9) - immediate
452
453
454
455
456
457
458
        except ProcessLookupError:
            pass
        except Exception as e:
            self._logger.warning(
                "Error sending SIGKILL to process group %s: %s", self._pgid, e
            )

Neelay Shah's avatar
Neelay Shah committed
459
460
461
462
463
464
465
    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)

466
467
468
469
470
471
472
473
474
475
476
477
478
479
480
481
482
483
484
485
486
487
488
489
490
491
492
493
494
495
496
497
498
    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
499
500
501
502
503
504
505
506
507
508
509
510
    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:
511
512
513
            # 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
514
515
516
517
518
519
520
521
522
523
524
525
526
527
528
            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

529
    def _check_url(self, url, timeout=30, sleep=1, log_interval=20):
Neelay Shah's avatar
Neelay Shah committed
530
531
532
533
534
535
536
537
        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
538
539
540
        attempt = 0
        last_log_time = 0.0

Neelay Shah's avatar
Neelay Shah committed
541
        while elapsed < timeout:
542
543
544
545
546
547
            self._check_process_alive("while waiting for health check")

            attempt += 1
            check_failed = False
            failure_reason = None

Neelay Shah's avatar
Neelay Shah committed
548
549
550
551
            try:
                response = requests.get(url, timeout=timeout - elapsed)
                if response.status_code == 200:
                    if response_check is None or response_check(response):
552
553
554
555
556
557
558
559
560
561
562
563
564
565
566
567
568
569
570
571
572
573
574
                        # 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
575
                        return time.time() - start_time
576
577
578
579
580
581
                    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
582
            except requests.RequestException as e:
583
584
585
586
587
588
589
590
591
592
593
594
595
596
                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
597
598
599
            time.sleep(sleep)
            elapsed = time.time() - start_time

600
        self._logger.error(
601
602
603
604
605
606
607
608
609
            "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)
610
        )
Neelay Shah's avatar
Neelay Shah committed
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
639
640
641
642
643
644
645
646
647
648
649
650
651
652
653
654
655
656
657
658
659
660
661
662
663
664
665
666
667
668
669
670
671
672
673
674
675
    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")

676
677
678
679
680
681
682
683
684
685
686
687
688
    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
689
            for proc in psutil.process_iter(["name", "cmdline"]):
690
691
692
693
694
                try:
                    if (
                        proc.name() == self._command_name
                        or proc.name() in self.stragglers
                    ):
695
                        self._logger.info(
696
                            "Terminating Existing %s %s", proc.name(), proc.pid
697
                        )
698

699
                        terminate_process_tree(proc.pid, self._logger)
700
701
702
703
704
705
706
707
708
709
710
711
712
713
714
715
                    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
716

717
718
719
720
721
722
    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
        )

723
724
725
726
    def get_pid(self) -> int | None:
        """Get the PID of the managed process."""
        return self.proc.pid if self.proc else None

727
728
729
730
731
732
733
734
735
736
737
738
739
740
741
    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
742

743
744
745
746
747
class DynamoFrontendProcess(ManagedProcess):
    """Process manager for Dynamo frontend"""

    _logger = logging.getLogger()

748
749
750
751
752
753
    def __init__(
        self,
        request: Any,
        *,
        frontend_port: Optional[int] = None,
        router_mode: str = "round-robin",
754
        migration_limit: int = 0,
755
756
        extra_args: Optional[list[str]] = None,
        extra_env: Optional[dict[str, str]] = None,
757
        # Default to false so pytest-xdist workers don't kill each other's frontends.
758
        terminate_all_matching_process_names: bool = False,
759
        display_name: Optional[str] = None,
760
761
762
    ):
        # TODO: Refactor remaining duplicate "DynamoFrontendProcess" helpers in tests to
        # use this shared implementation (and delete the copies):
763
764
765
766
        # - 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
767
768
769
770
771
772
773
774
775
776
777
778
779
780
781
782
783
        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)])
784
785
        # Migration limit is configured at the frontend level
        command.extend(["--migration-limit", str(migration_limit)])
786
787
        if extra_args:
            command.extend(extra_args)
788

789
790
791
        # Unset DYN_SYSTEM_PORT - frontend doesn't use system metrics server
        env = os.environ.copy()
        env.pop("DYN_SYSTEM_PORT", None)
792
793
        if extra_env:
            env.update(extra_env)
794

795
796
797
798
799
800
801
802
803
804
805
806
        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,
807
            env=env,
808
            display_output=True,
809
            terminate_all_matching_process_names=terminate_all_matching_process_names,
810
            log_dir=log_dir,
811
            display_name=display_name,
812
813
        )

814
815
816
817
818
819
820
821
    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

822
823
824
825
826
    @property
    def frontend_port(self) -> int:
        """Back-compat alias for older tests that expect `frontend.frontend_port`."""
        return self.http_port

827

Neelay Shah's avatar
Neelay Shah committed
828
def main():
829
830
    # NOTE: This entrypoint is for manual testing/debugging of `ManagedProcess` only.
    # It is not used by the pytest suite.
831
832

    # Example: Safe for parallel tests with terminate_all_matching_process_names=False
Neelay Shah's avatar
Neelay Shah committed
833
    with ManagedProcess(
834
        command=["python", "-m", "dynamo.frontend", "--port", "8000"],
Neelay Shah's avatar
Neelay Shah committed
835
        display_output=True,
836
        terminate_all_matching_process_names=False,  # ✅ Safe - only kills what we launch
837
838
        health_check_ports=[8000],
        health_check_urls=["http://localhost:8000/v1/models"],
Neelay Shah's avatar
Neelay Shah committed
839
840
841
842
843
844
845
846
        timeout=10,
    ):
        time.sleep(60)
        pass


if __name__ == "__main__":
    main()