From e3fe60118950ecc734642eda5a328127ba593a7d Mon Sep 17 00:00:00 2001 From: gspark Date: Fri, 23 Oct 2026 09:54:41 +0900 Subject: [PATCH] fix: scope MPS daemon detection and cleanup to our pipe dir (#166) _mps_daemon_running fell back to a host-wide pgrep, so another job's MPS daemon (other pipe directory) looked like ours; the health probe on our pipe dir then failed and the stale cleanup ran pkill -u $UID on every MPS server and control process, killing the other job's daemon. Liveness is now our control socket only, and the stale cleanup sends quit to our pipe dir and removes our socket, without pkill. Adds tests/test_mps_manager.py with a fake pipe dir. Co-Authored-By: Claude Opus 5.5 --- opt_prime/opt_prime/mps_manager.py | 44 +++++------------------------ opt_prime/tests/test_mps_manager.py | 36 +++++++++++++++++++++++ 2 files changed, 43 insertions(+), 37 deletions(-) create mode 100644 opt_prime/tests/test_mps_manager.py diff --git a/opt_prime/opt_prime/mps_manager.py b/opt_prime/opt_prime/mps_manager.py index 71b5cfb7..63ba6fb8 100644 --- a/opt_prime/opt_prime/mps_manager.py +++ b/opt_prime/opt_prime/mps_manager.py @@ -502,34 +502,12 @@ def cleanup_barrier_files(pipe_dir: str) -> None: # --------------------------------------------------------------------------- def _mps_daemon_running(pipe_dir: str) -> bool: - """Best-effort check whether an MPS daemon is already running. + """Whether an MPS daemon serves this pipe directory (its control socket exists). - Looks for the pipe directory's control socket and probes pgrep. - Returns True if either signal is present. + A host-wide process check would mistake another job's daemon (other pipe + directory) for ours, and the stale cleanup would then hit it. """ - # 1) control socket presence (most reliable when MPS daemon started cleanly) - if os.path.exists(os.path.join(pipe_dir, "control")): - return True - # 2) Process check - try: - out = subprocess.run( - ["pgrep", "-x", "nvidia-cuda-mps-server"], - capture_output=True, text=True, timeout=2, - ) - if out.returncode == 0 and out.stdout.strip(): - return True - except (FileNotFoundError, subprocess.TimeoutExpired, OSError): - pass - try: - out = subprocess.run( - ["pgrep", "-x", "nvidia-cuda-mps-control"], - capture_output=True, text=True, timeout=2, - ) - if out.returncode == 0 and out.stdout.strip(): - return True - except (FileNotFoundError, subprocess.TimeoutExpired, OSError): - pass - return False + return os.path.exists(os.path.join(pipe_dir, "control")) # --------------------------------------------------------------------------- @@ -786,7 +764,7 @@ def _daemon_responsive(self, timeout_s: float = 2.0) -> bool: return False def _force_cleanup_stale_daemon(self) -> None: - """Best-effort cleanup of a stale daemon: try quit, then pkill (own user).""" + """Best-effort cleanup of a stale daemon on our pipe dir: quit, drop socket.""" try: subprocess.run( ["nvidia-cuda-mps-control"], @@ -798,16 +776,8 @@ def _force_cleanup_stale_daemon(self) -> None: ) except (subprocess.TimeoutExpired, OSError): pass - # Force kill stragglers (only this user's processes) - for proc_name in ("nvidia-cuda-mps-server", "nvidia-cuda-mps-control"): - try: - subprocess.run( - ["pkill", "-u", str(os.geteuid()), "-x", proc_name], - capture_output=True, timeout=5, - ) - except (FileNotFoundError, subprocess.TimeoutExpired, OSError): - pass - # Remove control socket if present + # No pkill: other jobs' daemons (other pipe directories) share the + # process names. Remove our control socket so a new daemon can bind. ctrl = os.path.join(self.pipe_dir, "control") try: os.unlink(ctrl) diff --git a/opt_prime/tests/test_mps_manager.py b/opt_prime/tests/test_mps_manager.py new file mode 100644 index 00000000..d72d7237 --- /dev/null +++ b/opt_prime/tests/test_mps_manager.py @@ -0,0 +1,36 @@ +# +# Copyright (c) 2026-present, ETRI, All rights reserved. +# +# opt_prime.mps_manager: daemon detection and stale cleanup stay scoped to this job's pipe directory. +import os +import subprocess +import sys + +sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) +from opt_prime import mps_manager # noqa: E402 + + +def _fake_run(calls): + # Every command "succeeds" as if another job's MPS daemon were running on this host. + def run(cmd, *a, **kw): + calls.append(cmd) + return subprocess.CompletedProcess(cmd, 0, stdout="4242\n", stderr="") + return run + + +def test_daemon_running_only_from_own_control_socket(tmp_path, monkeypatch): + calls = [] + monkeypatch.setattr(subprocess, "run", _fake_run(calls)) + assert not mps_manager._mps_daemon_running(str(tmp_path)) # another job's daemon is not ours + (tmp_path / "control").touch() + assert mps_manager._mps_daemon_running(str(tmp_path)) + + +def test_stale_cleanup_does_not_kill_other_daemons(tmp_path, monkeypatch): + calls = [] + monkeypatch.setattr(subprocess, "run", _fake_run(calls)) + monkeypatch.setattr(mps_manager.time, "sleep", lambda s: None) + (tmp_path / "control").touch() + mps_manager.MPSManager(pipe_dir=str(tmp_path), log_dir=str(tmp_path / "log"))._force_cleanup_stale_daemon() + assert not any(c[0] in ("pkill", "killall") for c in calls), calls + assert not (tmp_path / "control").exists()