managed_process.py 32 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
import socket
import subprocess
11
import tempfile
Neelay Shah's avatar
Neelay Shah committed
12
13
14
15
16
17
18
import time
from dataclasses import dataclass, field
from typing import Any, List, Optional

import psutil
import requests

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

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
185
186
187
188
189
190
191

            # 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.
            if not os.path.isabs(self.log_dir):
                log_root = os.environ.get(
                    "DYN_TEST_OUTPUT_PATH",
                    os.path.join(tempfile.gettempdir(), "dynamo_tests"),
                )
                self.log_dir = os.path.join(log_root, self.log_dir)

Neelay Shah's avatar
Neelay Shah committed
192
193
194
195
196
197
198
            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)

199
            self._terminate_all_matching_process_names()  # Name-based cleanup (NOT xdist safe)
Neelay Shah's avatar
Neelay Shah committed
200
            self._start_process()
201
            time.sleep(self.delayed_start)
Neelay Shah's avatar
Neelay Shah committed
202
203
            elapsed = self._check_ports(self.timeout)
            self._check_urls(self.timeout - elapsed)
204
            self._check_funcs(self.timeout - elapsed)
Neelay Shah's avatar
Neelay Shah committed
205
206
207

            return self

208
209
210
211
212
213
214
215
        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
216

217
218
219
220
221
222
223
224
225
    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,
                )
226

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

                    # 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):
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
        """Cleanup: Terminate launched processes.

        Termination Strategy (Graceful → Escalate to Force):
        =====================================================
        1. Send SIGTERM to process group immediately (no delay before SIGTERM)
        2. Wait up to 2s (poll every 0.1s) for processes to exit
        3. If still alive after 2s: Send SIGKILL (force kill)
        4. Terminate individual processes (self.proc, tee, sed):
           - Send SIGTERM immediately (no delay)
           - Wait up to 10s for exit
           - If still alive after 10s: Send SIGKILL (force kill)
        5. Clean up straggler processes (if configured)

        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.
        """
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
        try:
            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)
            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
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329

    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,
330
                start_new_session=True,  # Isolate process group to prevent kill 0 from affecting parent
Neelay Shah's avatar
Neelay Shah committed
331
            )
332
333
334
335
336
337
            # 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
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
            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,
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

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

373
    def _terminate_process_group(self, timeout: float = 2.0):
374
375
        """Terminate the entire process group/session started for the child.

376
377
378
379
380
381
382
383
384
385
386
387
388
        Kill Sequence:
        ==============
        1. Send SIGTERM to entire process group IMMEDIATELY (no delay)
        2. Wait up to `timeout` seconds (default 2s), polling every 0.1s
        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)

        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.
389
390
391
392
393
        """
        if self._pgid is None:
            return
        try:
            self._logger.info("Terminating process group: %s", self._pgid)
394
            os.killpg(self._pgid, signal.SIGTERM)  # Step 1: Graceful SIGTERM
395
396
397
398
399
400
401
402
        except ProcessLookupError:
            return
        except Exception as e:
            self._logger.warning(
                "Error sending SIGTERM to process group %s: %s", self._pgid, e
            )
            return

403
        # Step 2: Poll for process exit instead of fixed sleep to minimize teardown time
404
405
406
407
408
        poll_interval = 0.1
        elapsed = 0.0
        while elapsed < timeout:
            try:
                # Check if any process in the group is still alive
409
                os.killpg(self._pgid, 0)  # Signal 0 = check existence (no kill)
410
411
412
413
414
415
416
417
418
            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

419
        # Step 3: Force kill if anything remains after timeout
420
        try:
421
            os.killpg(self._pgid, signal.SIGKILL)  # SIGKILL (kill -9) - immediate
422
423
424
425
426
427
428
        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
429
430
431
432
433
434
435
    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)

436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
    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
469
470
471
472
473
474
475
476
477
478
479
480
    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:
481
482
483
            # 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
484
485
486
487
488
489
490
491
492
493
494
495
496
497
498
            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

499
    def _check_url(self, url, timeout=30, sleep=1, log_interval=20):
Neelay Shah's avatar
Neelay Shah committed
500
501
502
503
504
505
506
507
        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
508
509
510
        attempt = 0
        last_log_time = 0.0

Neelay Shah's avatar
Neelay Shah committed
511
        while elapsed < timeout:
512
513
514
515
516
517
            self._check_process_alive("while waiting for health check")

            attempt += 1
            check_failed = False
            failure_reason = None

Neelay Shah's avatar
Neelay Shah committed
518
519
520
521
            try:
                response = requests.get(url, timeout=timeout - elapsed)
                if response.status_code == 200:
                    if response_check is None or response_check(response):
522
523
524
525
526
527
528
529
530
531
532
533
534
535
536
537
538
539
540
541
542
543
544
                        # 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
545
                        return time.time() - start_time
546
547
548
549
550
551
                    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
552
            except requests.RequestException as e:
553
554
555
556
557
558
559
560
561
562
563
564
565
566
                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
567
568
569
            time.sleep(sleep)
            elapsed = time.time() - start_time

570
        self._logger.error(
571
572
573
574
575
576
577
578
579
            "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)
580
        )
Neelay Shah's avatar
Neelay Shah committed
581

582
583
584
585
586
587
588
589
590
591
592
593
594
595
596
597
598
599
600
601
602
603
604
605
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
639
640
641
642
643
644
645
    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")

646
647
648
649
650
651
652
653
654
655
656
657
658
    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
659
            for proc in psutil.process_iter(["name", "cmdline"]):
660
661
662
663
664
                try:
                    if (
                        proc.name() == self._command_name
                        or proc.name() in self.stragglers
                    ):
665
                        self._logger.info(
666
                            "Terminating Existing %s %s", proc.name(), proc.pid
667
                        )
668

669
                        terminate_process_tree(proc.pid, self._logger)
670
671
672
673
674
675
676
677
678
679
680
681
682
683
684
685
                    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
686

687
688
689
690
691
692
    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
        )

693
694
695
696
    def get_pid(self) -> int | None:
        """Get the PID of the managed process."""
        return self.proc.pid if self.proc else None

697
698
699
700
701
702
703
704
705
706
707
708
709
710
711
    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
712

713
714
715
716
717
class DynamoFrontendProcess(ManagedProcess):
    """Process manager for Dynamo frontend"""

    _logger = logging.getLogger()

718
719
720
721
722
723
    def __init__(
        self,
        request: Any,
        *,
        frontend_port: Optional[int] = None,
        router_mode: str = "round-robin",
724
        migration_limit: int = 0,
725
726
        extra_args: Optional[list[str]] = None,
        extra_env: Optional[dict[str, str]] = None,
727
        # Default to false so pytest-xdist workers don't kill each other's frontends.
728
        terminate_all_matching_process_names: bool = False,
729
        display_name: Optional[str] = None,
730
731
732
    ):
        # TODO: Refactor remaining duplicate "DynamoFrontendProcess" helpers in tests to
        # use this shared implementation (and delete the copies):
733
734
735
736
        # - 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
737
738
739
740
741
742
743
744
745
746
747
748
749
750
751
752
753
        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)])
754
755
        # Migration limit is configured at the frontend level
        command.extend(["--migration-limit", str(migration_limit)])
756
757
        if extra_args:
            command.extend(extra_args)
758

759
760
761
        # Unset DYN_SYSTEM_PORT - frontend doesn't use system metrics server
        env = os.environ.copy()
        env.pop("DYN_SYSTEM_PORT", None)
762
763
        if extra_env:
            env.update(extra_env)
764

765
766
767
768
769
770
771
772
773
774
775
776
        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,
777
            env=env,
778
            display_output=True,
779
            terminate_all_matching_process_names=terminate_all_matching_process_names,
780
            log_dir=log_dir,
781
            display_name=display_name,
782
783
        )

784
785
786
787
788
789
790
791
    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

792
793
794
795
796
    @property
    def frontend_port(self) -> int:
        """Back-compat alias for older tests that expect `frontend.frontend_port`."""
        return self.http_port

797

Neelay Shah's avatar
Neelay Shah committed
798
def main():
799
800
    # NOTE: This entrypoint is for manual testing/debugging of `ManagedProcess` only.
    # It is not used by the pytest suite.
801
802

    # Example: Safe for parallel tests with terminate_all_matching_process_names=False
Neelay Shah's avatar
Neelay Shah committed
803
    with ManagedProcess(
804
        command=["python", "-m", "dynamo.frontend", "--port", "8000"],
Neelay Shah's avatar
Neelay Shah committed
805
        display_output=True,
806
        terminate_all_matching_process_names=False,  # ✅ Safe - only kills what we launch
807
808
        health_check_ports=[8000],
        health_check_urls=["http://localhost:8000/v1/models"],
Neelay Shah's avatar
Neelay Shah committed
809
810
811
812
813
814
815
816
        timeout=10,
    ):
        time.sleep(60)
        pass


if __name__ == "__main__":
    main()