aides/_arret.py

# Arrete le chien et le moteur en lisant /proc, jamais par motif.
#
# L'agent du pod porte le TEXTE des scripts dans sa ligne de commande : un
# `pkill -f "vllm serve"` matche donc l'agent lui-meme et se suicide (constate,
# `[FIN] code=-15`). On exclut toute notre ascendance a la place.
#
# `bash /run.sh` est intouchable : c'est le processus dont depend le conteneur,
# quel que soit son PID.
import os
import signal
import time


def ascendance(pid):
    vus = set()
    while pid and pid > 1:
        vus.add(pid)
        try:
            pid = int(open("/proc/%d/stat" % pid).read().rsplit(") ", 1)[1].split()[1])
        except Exception:
            break
    return vus


sur = ascendance(os.getpid())


def cibles():
    trouves = []
    for entree in os.listdir("/proc"):
        if not entree.isdigit():
            continue
        pid = int(entree)
        if pid in sur:
            continue
        try:
            argv = [a for a in open("/proc/%d/cmdline" % pid, "rb").read().split(b"\0") if a]
            comm = open("/proc/%d/comm" % pid).read().strip()
        except Exception:
            continue
        if not argv:
            continue
        ligne = b" ".join(argv)
        if b"/run.sh" in ligne:
            continue
        chien = b"chien.py" in ligne
        # EngineCore se renomme et survit a un kill par nom de commande en
        # retenant toute la VRAM : le viser aussi par sa ligne complete.
        moteur = ((b"vllm" in argv[0] and b"serve" in ligne)
                  or argv[0].endswith(b"vllm")
                  or comm in ("VLLM::EngineCore", "EngineCore", "PleOffloadWorker")
                  or b"VLLM::" in ligne)
        if chien or moteur:
            trouves.append((pid, comm))
    return trouves


c = cibles()
print("   arret de :", [(p, n) for p, n in c] or "rien")
for pid, _ in c:
    try:
        os.kill(pid, signal.SIGTERM)
    except Exception:
        pass
time.sleep(8)
for pid, _ in cibles():
    try:
        os.kill(pid, signal.SIGKILL)
    except Exception:
        pass
print("   restants :", [p for p, _ in cibles()] or "aucun")