adam/process_control.py

from __future__ import annotations

import subprocess
from typing import Any


def set_process_tree_paused(process: subprocess.Popen, paused: bool) -> bool:
    """Suspend/resume a trainer process tree when psutil is available."""
    try:
        import psutil

        parent = psutil.Process(process.pid)
        children = parent.children(recursive=True)
        targets = [*children, parent] if paused else [parent, *children]
        for target in targets:
            try:
                target.suspend() if paused else target.resume()
            except (psutil.NoSuchProcess, psutil.AccessDenied):
                continue
        return True
    except Exception:
        return False


def terminate_process_tree(process: Any, *, timeout: float = 3.0) -> None:
    """Terminate a process and any children it launched."""
    try:
        import psutil

        parent = psutil.Process(process.pid)
        children = parent.children(recursive=True)
        targets = [*children, parent]
        for target in targets:
            try:
                target.terminate()
            except (psutil.NoSuchProcess, psutil.AccessDenied):
                continue
        _gone, alive = psutil.wait_procs(targets, timeout=timeout)
        for target in alive:
            try:
                target.kill()
            except (psutil.NoSuchProcess, psutil.AccessDenied):
                continue
        return
    except Exception:
        pass
    try:
        process.terminate()
    except Exception:
        return
    try:
        process.wait(timeout=timeout)
    except Exception:
        try:
            process.kill()
        except Exception:
            pass