diff --git a/unsloth_cli/commands/studio.py b/unsloth_cli/commands/studio.py index 57c1c3de1b..210ce3b25f 100644 --- a/unsloth_cli/commands/studio.py +++ b/unsloth_cli/commands/studio.py @@ -2453,6 +2453,17 @@ def _pid_file_entries() -> "list[tuple[int, list[Path]]]": return list(by_pid.items()) +def _pid_is_studio_server(pid: int) -> bool: + """Guard against PID reuse: a stale record must not get an unrelated process + killed. Unknowable without psutil, so trust the record there.""" + try: + import psutil + cmdline = " ".join(psutil.Process(pid).cmdline()).lower() + except Exception: + return True + return "run.py" in cmdline or "unsloth" in cmdline + + def _signal_stop(pid: int) -> "str | None": """SIGTERM (or taskkill) the pid. Returns an error string, or None on success.""" import signal as _signal @@ -2483,7 +2494,7 @@ def stop(): signalled, failed = [], [] for pid, paths in entries: - if not _pid_alive(pid): + if not _pid_alive(pid) or not _pid_is_studio_server(pid): for path in paths: path.unlink(missing_ok = True) continue diff --git a/unsloth_cli/tests/test_studio_stop.py b/unsloth_cli/tests/test_studio_stop.py index e7a6fe22db..515540c1b9 100644 --- a/unsloth_cli/tests/test_studio_stop.py +++ b/unsloth_cli/tests/test_studio_stop.py @@ -11,6 +11,7 @@ from __future__ import annotations import sys from pathlib import Path +from types import SimpleNamespace import pytest from typer.testing import CliRunner @@ -43,6 +44,7 @@ def _install( killed = killed if killed is not None else [] monkeypatch.setattr(studio_mod, "_pid_alive", lambda pid: pid in live) + monkeypatch.setattr(studio_mod, "_pid_is_studio_server", lambda pid: True) def fake_kill(pid, _sig): killed.append(pid) @@ -98,6 +100,7 @@ def test_stop_signals_each_server_once(monkeypatch, tmp_path): monkeypatch.setattr(studio_mod, "_PID_FILE", tmp_path / "studio.pid") monkeypatch.setattr(studio_mod.time, "sleep", lambda _s: None) monkeypatch.setattr(studio_mod, "_pid_alive", lambda pid: True) + monkeypatch.setattr(studio_mod, "_pid_is_studio_server", lambda pid: True) killed = [] monkeypatch.setattr(studio_mod.os, "kill", lambda pid, _sig: killed.append(pid)) monkeypatch.setattr(sys, "platform", "linux") @@ -123,6 +126,46 @@ def test_stop_removes_every_stale_file_for_one_pid(monkeypatch, tmp_path): assert not list(tmp_path.glob("*.pid")) +def test_stop_does_not_signal_a_reused_pid(monkeypatch, tmp_path): + # Crash leaves a per-port file behind, the OS hands that PID to something + # else: stop must drop the record, not SIGTERM an unrelated process. + studio_mod, _live, killed = _install(monkeypatch, tmp_path, alive = {8550}) + monkeypatch.setattr(studio_mod, "_pid_is_studio_server", lambda pid: False) + _write_pid(tmp_path, "studio-8901-8550.pid", 8550) + + result = _run_stop(studio_mod) + + assert result.exit_code == 0, result.output + assert killed == [] + assert not (tmp_path / "studio-8901-8550.pid").exists() + + +def test_pid_identity_check_trusts_the_record_without_psutil(monkeypatch): + # No psutil: fall back to trusting the record rather than never stopping. + studio_mod = _studio() + monkeypatch.setitem(sys.modules, "psutil", None) + + assert studio_mod._pid_is_studio_server(8550) is True + + +def test_pid_identity_check_matches_a_studio_command_line(monkeypatch): + studio_mod = _studio() + + class _FakeProcess: + def __init__(self, pid): + self.pid = pid + + def cmdline(self): + if self.pid == 8550: + return ["/venv/bin/python", "/pkg/studio/backend/run.py", "--port", "8901"] + return ["/usr/bin/postgres", "-D", "/var/lib/pg"] + + monkeypatch.setitem(sys.modules, "psutil", SimpleNamespace(Process = _FakeProcess)) + + assert studio_mod._pid_is_studio_server(8550) is True + assert studio_mod._pid_is_studio_server(9999) is False + + def test_stop_reads_the_legacy_single_pid_file(monkeypatch, tmp_path): studio_mod, _live, killed = _install(monkeypatch, tmp_path, alive = {4242}) _write_pid(tmp_path, "studio.pid", 4242) @@ -162,6 +205,7 @@ def test_stop_does_not_claim_a_stop_while_a_server_is_still_alive(monkeypatch, t monkeypatch.setattr(studio_mod, "_PID_FILE", tmp_path / "studio.pid") monkeypatch.setattr(studio_mod.time, "sleep", lambda _s: None) monkeypatch.setattr(studio_mod, "_pid_alive", lambda pid: True) + monkeypatch.setattr(studio_mod, "_pid_is_studio_server", lambda pid: True) monkeypatch.setattr(studio_mod.os, "kill", lambda pid, sig: None) monkeypatch.setattr(sys, "platform", "linux") _write_pid(tmp_path, "studio-8901-8550.pid", 8550) @@ -181,6 +225,7 @@ def test_stop_continues_after_one_server_fails_to_stop(monkeypatch, tmp_path): monkeypatch.setattr(studio_mod.time, "sleep", lambda _s: None) live = {8550, 8600} monkeypatch.setattr(studio_mod, "_pid_alive", lambda pid: pid in live) + monkeypatch.setattr(studio_mod, "_pid_is_studio_server", lambda pid: True) def fake_kill(pid, _sig): if pid == 8550: