robertvoy-ComfyUI-Distributed/utils/process.py

"""
Process management utilities for ComfyUI-Distributed.
"""
import asyncio
import os
import platform
import signal
import subprocess
import threading
from collections.abc import Sequence
from typing import Any


def _run_coroutine_blocking(coro):
    """Run a coroutine in blocking contexts, even if an event loop is active."""
    try:
        asyncio.get_running_loop()
    except RuntimeError:
        return asyncio.run(coro)

    outcome = {}

    def _runner():
        try:
            outcome["value"] = asyncio.run(coro)
        except Exception as exc:  # pragma: no cover - thin bridge
            outcome["error"] = exc

    thread = threading.Thread(target=_runner, daemon=True)
    thread.start()
    thread.join()

    error = outcome.get("error")
    if error is not None:
        raise error
    return outcome.get("value")


class AsyncProcessAdapter:
    """Sync-friendly adapter around asyncio subprocess.Process."""

    def __init__(self, process):
        self._process = process
        self.pid = process.pid

    @property
    def returncode(self):
        return self._process.returncode

    @property
    def stdout(self):
        return self._process.stdout

    def poll(self):
        return self._process.returncode

    def terminate(self):
        self._process.terminate()

    def kill(self):
        self._process.kill()

    def wait(self, timeout=None):
        async def _wait():
            if timeout is None:
                return await self._process.wait()
            return await asyncio.wait_for(self._process.wait(), timeout=float(timeout))

        return _run_coroutine_blocking(_wait())


def _normalize_stdio_value(value):
    if value is subprocess.PIPE:
        return asyncio.subprocess.PIPE
    if value is subprocess.STDOUT:
        return asyncio.subprocess.STDOUT
    if value is subprocess.DEVNULL:
        return asyncio.subprocess.DEVNULL
    return value


def launch_process_with_timeout(
    command: Sequence[str | bytes | os.PathLike[str]],
    timeout_seconds: float = 10.0,
    **create_kwargs: Any,
) -> AsyncProcessAdapter:
    """Launch a child process with startup timeout and sync-friendly lifecycle methods."""
    if isinstance(command, (str, bytes)):
        raise TypeError("Command must be a sequence, not a string")

    normalized_kwargs = dict(create_kwargs)
    normalized_kwargs["stdout"] = _normalize_stdio_value(normalized_kwargs.get("stdout"))
    normalized_kwargs["stderr"] = _normalize_stdio_value(normalized_kwargs.get("stderr"))

    async def _launch():
        return await asyncio.wait_for(
            asyncio.create_subprocess_exec(*command, **normalized_kwargs),
            timeout=float(timeout_seconds),
        )

    try:
        process = _run_coroutine_blocking(_launch())
    except asyncio.TimeoutError as exc:
        raise TimeoutError(f"Timed out launching process after {timeout_seconds} seconds") from exc

    return AsyncProcessAdapter(process)

def is_process_alive(pid: int) -> bool:
    """Check if a process with given PID is still alive."""
    try:
        if platform.system() == "Windows":
            # Windows: use tasklist
            result = subprocess.run(
                ['tasklist', '/FI', f'PID eq {pid}'],
                capture_output=True,
                text=True,
                timeout=5.0,
            )
            return str(pid) in result.stdout
        else:
            # Unix: send signal 0
            os.kill(pid, 0)
            return True
    except (OSError, subprocess.SubprocessError):
        return False

def terminate_process(process, timeout=5):
    """Gracefully terminate a process with timeout."""
    wait_timeout = max(float(timeout), 0.1)
    if process.poll() is not None:
        return True

    process.terminate()
    try:
        process.wait(timeout=wait_timeout)
        return True
    except (subprocess.TimeoutExpired, TimeoutError):
        process.kill()
        try:
            process.wait(timeout=wait_timeout)
            return True
        except (subprocess.TimeoutExpired, TimeoutError) as exc:
            raise TimeoutError(
                f"Process did not terminate within {wait_timeout} seconds after kill."
            ) from exc

def get_python_executable():
    """Get the Python executable path."""
    import sys
    return sys.executable