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