managed_process.py

# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: Apache-2.0

import json
import logging
import os
import shutil
import signal
import socket
import subprocess
import time
from dataclasses import dataclass, field
from typing import Any, List, Optional

import psutil
import requests

from tests.utils.constants import DefaultPort
from tests.utils.port_utils import allocate_port, deallocate_port
from tests.utils.test_output import resolve_test_output_path


def terminate_process(process, logger=logging.getLogger(), immediate_kill=False):
    """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)
    """
    try:
        logger.info("Terminating PID: %s name: %s", process.pid, process.name())
        if immediate_kill:
            logger.info("Sending SIGKILL (kill -9): %s %s", process.pid, process.name())
            process.kill()  # SIGKILL (9)
        else:
            logger.info("Sending SIGTERM (kill): %s %s", process.pid, process.name())
            process.terminate()  # SIGTERM (15)
    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(
    pid, logger=logging.getLogger(), immediate_kill=False, timeout=10
):
    """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
    """
    try:
        parent = psutil.Process(pid)
    except psutil.NoSuchProcess:
        return

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

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

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


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


@dataclass
class ManagedProcess:
    """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
    """

    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)
    health_check_funcs: List[Any] = field(default_factory=list)
    delayed_start: int = 0
    timeout: int = 300
    working_dir: Optional[str] = None
    display_output: bool = False
    data_dir: Optional[str] = None
    # 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
    stragglers: List[str] = field(default_factory=list)
    straggler_commands: List[str] = field(default_factory=list)
    log_dir: str = os.getcwd()
    display_name: Optional[str] = None

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

    _logger = logging.getLogger()
    _command_name = None
    _log_path = None
    _tee_proc = None
    _sed_proc = None

    @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 ""

    def __enter__(self):
        try:
            self._logger = logging.getLogger(self.__class__.__name__)
            # self._command_name = self.command[0]
            if self.display_name:
                self._command_name = self.display_name
            else:
                self._command_name = self.command[0]

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

            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)

            self._terminate_all_matching_process_names()  # Name-based cleanup (NOT xdist safe)
            self._start_process()
            time.sleep(self.delayed_start)
            elapsed = self._check_ports(self.timeout)
            self._check_urls(self.timeout - elapsed)
            self._check_funcs(self.timeout - elapsed)

            return self

        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

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

            for ps_process in psutil.process_iter(["name", "cmdline"]):
                try:
                    process_name = ps_process.name()
                    if process_name in self.stragglers:
                        self._logger.info(
                            "Terminating Straggler %s %s", process_name, ps_process.pid
                        )
                        terminate_process_tree(ps_process.pid, self._logger)

                    # 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):
        """Cleanup: Terminate launched processes.

        Termination Strategy (Graceful → Escalate to Force):
        =====================================================
        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):
           - Send SIGTERM immediately (no delay)
           - Wait up to 10s for exit
           - If still alive after 10s: Send SIGKILL (force kill)
        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)

        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.
        """
        try:
            # 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

            self._stop_started_processes()

            # 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

            if self.data_dir:
                self._remove_directory(self.data_dir)
        finally:
            # Always run straggler cleanup, even if interrupted
            self._cleanup_stragglers()

    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)

    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,
                start_new_session=True,  # Isolate process group to prevent kill 0 from affecting parent
            )
            # 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
            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,
                    start_new_session=True,  # Isolate process group to prevent kill 0 from affecting parent
                )
                # 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

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

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

        Kill Sequence:
        ==============
        1. Send SIGTERM to entire process group IMMEDIATELY (no delay)
        2. Wait up to `timeout` seconds (default 8s), 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)
        - 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

        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.
        """
        if self._pgid is None:
            return
        try:
            self._logger.info("Terminating process group: %s", self._pgid)
            os.killpg(self._pgid, signal.SIGTERM)  # Step 1: Graceful SIGTERM
        except ProcessLookupError:
            return
        except Exception as e:
            self._logger.warning(
                "Error sending SIGTERM to process group %s: %s", self._pgid, e
            )
            return

        # Step 2: Poll for process exit instead of fixed sleep to minimize teardown time
        poll_interval = 0.1
        elapsed = 0.0
        while elapsed < timeout:
            # Reap the launched child if it has already exited. Without this,
            # the child can remain as a zombie and keep killpg(..., 0) reporting
            # the process group as alive until the timeout expires.
            if self.proc is not None:
                self.proc.poll()
            try:
                # Check if any process in the group is still alive
                os.killpg(self._pgid, 0)  # Signal 0 = check existence (no kill)
            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

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

    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)

    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 ''}"
            )

    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:
            # Check if the main process is still alive
            self._check_process_alive(f"while waiting for port {port}")

            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

    def _check_url(self, url, timeout=30, sleep=1, log_interval=20):
        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
        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:
                response = requests.get(url, timeout=timeout - elapsed)
                if response.status_code == 200:
                    if response_check is None or response_check(response):
                        # 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,
                            )
                        return time.time() - start_time
                    else:
                        check_failed = True
                        failure_reason = "custom check failed"
                else:
                    check_failed = True
                    failure_reason = f"status code {response.status_code}"
            except requests.RequestException as e:
                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

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

        self._logger.error(
            "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)
        )

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

    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:
            for proc in psutil.process_iter(["name", "cmdline"]):
                try:
                    if (
                        proc.name() == self._command_name
                        or proc.name() in self.stragglers
                    ):
                        self._logger.info(
                            "Terminating Existing %s %s", proc.name(), proc.pid
                        )

                        terminate_process_tree(proc.pid, self._logger)
                    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

    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
        )

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

    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 []


class DynamoFrontendProcess(ManagedProcess):
    """Process manager for Dynamo frontend"""

    _logger = logging.getLogger()

    def __init__(
        self,
        request: Any,
        *,
        frontend_port: Optional[int] = None,
        router_mode: str = "round-robin",
        migration_limit: int = 0,
        extra_args: Optional[list[str]] = None,
        extra_env: Optional[dict[str, str]] = None,
        # Default to false so pytest-xdist workers don't kill each other's frontends.
        terminate_all_matching_process_names: bool = False,
        display_name: Optional[str] = None,
    ):
        # TODO: Refactor remaining duplicate "DynamoFrontendProcess" helpers in tests to
        # use this shared implementation (and delete the copies):
        # - 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
        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)])
        # Migration limit is configured at the frontend level
        command.extend(["--migration-limit", str(migration_limit)])
        if extra_args:
            command.extend(extra_args)

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

        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,
            env=env,
            display_output=True,
            terminate_all_matching_process_names=terminate_all_matching_process_names,
            log_dir=log_dir,
            display_name=display_name,
        )

    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

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


def main():
    # NOTE: This entrypoint is for manual testing/debugging of `ManagedProcess` only.
    # It is not used by the pytest suite.

    # Example: Safe for parallel tests with terminate_all_matching_process_names=False
    with ManagedProcess(
        command=["python", "-m", "dynamo.frontend", "--port", "8000"],
        display_output=True,
        terminate_all_matching_process_names=False,  # ✅ Safe - only kills what we launch
        health_check_ports=[8000],
        health_check_urls=["http://localhost:8000/v1/models"],
        timeout=10,
    ):
        time.sleep(60)
        pass


if __name__ == "__main__":
    main()