ceos/unsloth
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Drop the command-line guess, fix Windows liveness, and free the PID record last

Browse source
This commit is contained in:
Nilay Yadav 2026-07-29 05:37:01 +05:30
commit 4d472ed392
4 changed files with 127 additions and 87 deletions
Download patch file Download diff file Expand all files Collapse all files

55
studio/backend/run.py
View file

@ -797,6 +797,20 @@ def _pid_alive(pid: int) -> bool:
return psutil.pid_exists(pid)
except ImportError:
pass
if sys.platform == "win32":
# os.kill(pid, 0) raises OSError for every pid on Windows, so a stale record
# would look alive forever and block this port. Unconfirmed means prune.
import subprocess
try:
out = subprocess.run(
["tasklist", "/FI", f"PID eq {int(pid)}", "/NH", "/FO", "CSV"],
capture_output = True,
text = True,
timeout = 10,
).stdout
except Exception:
return False
return f'"{int(pid)}"' in out
try:
os.kill(pid, 0)
except ProcessLookupError:
@ -836,29 +850,21 @@ def _read_pid_record(path: Path) -> "tuple[int, float | None, str | None] | None
return pid, created, address
def _pid_is_studio_backend(pid: int, created_times: "Sequence[float | None]" = ()) -> "bool | None":
"""Guard against PID reuse. True = ours, False = not ours, None = can't tell.
def _pid_is_studio_backend(pid: int, created_times: "Sequence[float | None]" = ()) -> bool:
"""False only when a recorded start time proves this PID is a different process.
Any recorded start time matching is enough: a stale record must not veto a live
server that reused the PID. The cmdline check covers untimed legacy records --
the in-venv path runs run_server() in-process, so argv is `unsloth studio ...`.
Any recorded time matching is enough -- a stale record must not veto a live
server that reused the PID. Untimed records cannot be checked at all, so they
are trusted: a legacy `python run.py` has no telltale argv, and guessing from
the command line rejected real servers.
"""
known = [c for c in created_times if c is not None]
untimed = not created_times or len(known) < len(created_times)
if known:
actual = _process_create_time(pid)
if actual is None:
return True if untimed else None
if any(abs(actual - c) < 1.0 for c in known):
return True
if not untimed:
return False
try:
import psutil
cmdline = " ".join(psutil.Process(pid).cmdline()).lower()
except Exception:
if not known or len(known) < len(created_times):
return True
return "studio" in cmdline and ("run.py" in cmdline or "unsloth" in cmdline)
actual = _process_create_time(pid)
if actual is None:
return True
return any(abs(actual - c) < 1.0 for c in known)
def _own_studio_on_port(port: int, host: str) -> "int | None":
@ -881,8 +887,7 @@ def _own_studio_on_port(port: int, host: str) -> "int | None":
continue
if not _addresses_collide(address, host, port):
continue
# None (unverifiable) counts as ours: refusing beats a silent duplicate.
if _pid_is_studio_backend(pid, [created]) is not False:
if _pid_is_studio_backend(pid, [created]):
return pid
return _legacy_studio_on_port(port)
@ -903,12 +908,12 @@ def _legacy_studio_on_port(port: int) -> "int | None":
# and this port's records were just checked. Only count a record that still
# matches the live process: a stale one may just share a reused PID.
for other in _per_port_records():
if other and other[0] == pid and _pid_is_studio_backend(pid, [other[1]]) is not False:
if other and other[0] == pid and _pid_is_studio_backend(pid, [other[1]]):
return None
blocker = _get_pid_on_port(port)
if blocker is not None and blocker[0] != pid:
return None
if _pid_is_studio_backend(pid, [created]) is False:
if not _pid_is_studio_backend(pid, [created]):
return None
return pid
@ -995,7 +1000,6 @@ def _graceful_shutdown(server = None):
Called from signal handlers to clean up children before exit. Critical on
Windows where atexit handlers are unreliable after Ctrl+C.
"""
_remove_pid_file()
logger.info("Graceful shutdown initiated -- cleaning up subprocesses...")
# 1. Shut down uvicorn (releases the listening socket).
@ -1048,6 +1052,9 @@ def _graceful_shutdown(server = None):
except Exception as e:
logger.warning("Error in process-lifetime sweep: %s", e)
# Last: while cleanup runs the server is still alive, and dropping the record
# early leaves a retried `stop` or a new launch unable to find it.
_remove_pid_file()
logger.info("All subprocesses cleaned up")

89
studio/backend/tests/test_studio_pid_files.py
View file

@ -21,8 +21,9 @@ if str(_BACKEND) not in sys.path:
import run # noqa: E402
# Captured before the autouse fixture stubs it, for the tests that exercise it.
# Captured before the autouse fixture stubs them, for the tests that exercise them.
_REAL_IS_STUDIO_BACKEND = run._pid_is_studio_backend
_REAL_PID_ALIVE = run._pid_alive
@pytest.fixture(autouse = True)
@ -105,6 +106,36 @@ def test_remove_pid_file_leaves_a_reused_entry_alone(tmp_path):
assert own.read_text(encoding = "utf-8") == "999999"
def test_windows_liveness_does_not_call_every_pid_alive(monkeypatch):
# os.kill(pid, 0) raises OSError for every pid on Windows, so without the
# tasklist fallback a stale record would block its port forever.
import subprocess
monkeypatch.setattr(run, "_pid_alive", _REAL_PID_ALIVE)
monkeypatch.setitem(sys.modules, "psutil", None)
monkeypatch.setattr(sys, "platform", "win32")
monkeypatch.setattr(
subprocess, "run", lambda *a, **k: SimpleNamespace(stdout = '"python.exe","8550",...')
)
assert run._pid_alive(8550) is True
assert run._pid_alive(9999) is False
def test_windows_liveness_prunes_when_tasklist_fails(monkeypatch):
import subprocess
def _boom(*a, **k):
raise OSError("tasklist missing")
monkeypatch.setattr(run, "_pid_alive", _REAL_PID_ALIVE)
monkeypatch.setitem(sys.modules, "psutil", None)
monkeypatch.setattr(sys, "platform", "win32")
monkeypatch.setattr(subprocess, "run", _boom)
assert run._pid_alive(8550) is False
def test_read_pid_record_parses_pid_time_and_address(tmp_path):
(tmp_path / "r.pid").write_text("8550\n111.5\n127.0.0.1", encoding = "utf-8")
@ -132,6 +163,21 @@ def test_read_pid_record_rejects_a_corrupt_file(tmp_path):
assert run._read_pid_record(tmp_path / "r.pid") is None
def test_graceful_shutdown_drops_the_record_last(monkeypatch):
# Cleanup can take seconds while the server is still alive. Dropping the record
# first leaves a retried `stop` or a new launch unable to find it.
order = []
monkeypatch.setattr(run, "_remove_pid_file", lambda: order.append("remove_record"))
class _Server:
def __setattr__(self, name, value):
order.append("release_socket")
run._graceful_shutdown(_Server())
assert order == ["release_socket", "remove_record"]
def test_own_studio_on_port_is_found_without_psutil(tmp_path, monkeypatch):
# psutil is optional; a listener scan finds nothing without it, so detection
# must come from our own records or we silently start a duplicate.
@ -166,7 +212,7 @@ def test_a_reused_pid_is_not_treated_as_our_studio(tmp_path, monkeypatch):
def test_an_unverifiable_record_still_blocks_a_duplicate(tmp_path, monkeypatch):
# Can't tell: refusing with a clear message beats a silent second instance.
monkeypatch.setattr(run, "_pid_is_studio_backend", lambda pid, created_times = (): None)
monkeypatch.setattr(run, "_pid_is_studio_backend", lambda pid, created_times = (): True)
(tmp_path / "studio-8901-8550.pid").write_text("8550", encoding = "utf-8")
assert run._own_studio_on_port(8901, "127.0.0.1") == 8550
@ -201,9 +247,17 @@ def test_a_stale_record_on_another_port_does_not_hide_a_live_server(tmp_path, mo
assert run._own_studio_on_port(9000, "127.0.0.1") == 1234
def test_legacy_records_match_an_in_process_studio(monkeypatch):
# The in-venv path calls run_server() in-process, so argv is `unsloth studio`
# with no run.py. Rejecting it would strand a running server.
def test_a_start_time_is_the_only_thing_that_disproves_a_record(monkeypatch):
monkeypatch.setattr(run, "_pid_is_studio_backend", _REAL_IS_STUDIO_BACKEND)
monkeypatch.setattr(run, "_process_create_time", lambda pid: 999.0)
assert run._pid_is_studio_backend(8550, [999.0]) is True
assert run._pid_is_studio_backend(8550, [111.5]) is False
def test_a_bare_run_py_command_line_is_not_rejected(monkeypatch):
# `cd studio/backend && python run.py --port 8901` has no "studio" or "unsloth"
# in argv. Guessing from the command line called that "not ours".
monkeypatch.setattr(run, "_pid_is_studio_backend", _REAL_IS_STUDIO_BACKEND)
class _FakeProcess:
@ -211,30 +265,25 @@ def test_legacy_records_match_an_in_process_studio(monkeypatch):
self.pid = pid
def cmdline(self):
return ["/root/.unsloth/studio/unsloth_studio/bin/unsloth", "studio", "-p", "8901"]
return ["python", "run.py", "--port", "8901"]
def create_time(self):
return 111.5
monkeypatch.setitem(sys.modules, "psutil", SimpleNamespace(Process = _FakeProcess))
assert run._pid_is_studio_backend(8550) is True
def test_legacy_records_do_not_match_a_training_run(monkeypatch):
# No start time recorded: the cmdline check must not accept `unsloth train`.
def test_an_untimed_legacy_record_is_trusted(monkeypatch):
# `python run.py --port 8901` has no telltale argv, so guessing from the
# command line rejected real servers. Only a start time can disprove one.
monkeypatch.setattr(run, "_pid_is_studio_backend", _REAL_IS_STUDIO_BACKEND)
class _FakeProcess:
def __init__(self, pid):
self.pid = pid
def cmdline(self):
if self.pid == 8550:
return ["python", "/pkg/studio/backend/run.py", "--port", "8901"]
return ["python", "-m", "unsloth", "train", "run.py"]
monkeypatch.setitem(sys.modules, "psutil", SimpleNamespace(Process = _FakeProcess))
monkeypatch.setattr(run, "_process_create_time", lambda pid: 999.0)
assert run._pid_is_studio_backend(8550) is True
assert run._pid_is_studio_backend(9999) is False
assert run._pid_is_studio_backend(8550, [None]) is True
assert run._pid_is_studio_backend(8550, [111.5, None]) is True
def test_a_legacy_server_on_the_port_is_recognised(tmp_path, monkeypatch):

27
unsloth_cli/commands/studio.py
View file

@ -2476,30 +2476,23 @@ def _pid_file_entries() -> "list[tuple[int, list[float | None], list[Path]]]":
def _pid_is_studio_server(pid: int, created_times: "Sequence[float | None]" = ()) -> bool:
"""Best-effort PID-reuse guard: False only when we can prove it isn't ours.
"""False only when a recorded start time proves this PID is a different process.
psutil is not a base CLI dependency, so the CLI meets records it cannot check.
Those are stopped anyway -- the old `stop` signalled its PID with no checks at
all, and skipping a live server is the orphan bug this exists to fix.
The cmdline check covers untimed legacy records -- the in-venv path runs
run_server() in-process, so its argv is `unsloth studio ...` with no run.py.
Any recorded time matching is enough -- a stale record must not veto a live
server that reused the PID. Untimed records (legacy studio.pid, or a server
started without psutil) cannot be checked, so they are trusted: the old `stop`
signalled with no checks at all, and skipping a live server is the orphan bug
this exists to fix.
"""
known = [c for c in created_times if c is not None]
untimed = not created_times or len(known) < len(created_times)
if not known or len(known) < len(created_times):
return True
try:
import psutil
proc = psutil.Process(pid)
if known:
actual = proc.create_time()
if any(abs(actual - c) < 1.0 for c in known):
return True
if not untimed:
return False
cmdline = " ".join(proc.cmdline()).lower()
actual = psutil.Process(pid).create_time()
except Exception:
return True
return "studio" in cmdline and ("run.py" in cmdline or "unsloth" in cmdline)
return any(abs(actual - c) < 1.0 for c in known)
def _signal_stop(pid: int) -> "str | None":

43
unsloth_cli/tests/test_studio_stop.py
View file

@ -168,8 +168,9 @@ def test_stop_signals_a_live_server_whose_pid_has_a_stale_record(monkeypatch, tm
assert not list(tmp_path.glob("studio-*.pid"))
def test_pid_identity_check_accepts_an_in_process_studio(monkeypatch):
# The in-venv path calls run_server() in-process: argv has no run.py.
def test_a_bare_run_py_command_line_is_not_rejected(monkeypatch):
# `cd studio/backend && python run.py --port 8901` has no "studio" or "unsloth"
# in argv. Guessing from the command line deleted its record without stopping it.
studio_mod = _studio()
class _FakeProcess:
@ -177,13 +178,26 @@ def test_pid_identity_check_accepts_an_in_process_studio(monkeypatch):
self.pid = pid
def cmdline(self):
return ["/root/.unsloth/studio/unsloth_studio/bin/unsloth", "studio", "-p", "8901"]
return ["python", "run.py", "--port", "8901"]
def create_time(self):
return 111.5
monkeypatch.setitem(sys.modules, "psutil", SimpleNamespace(Process = _FakeProcess))
assert studio_mod._pid_is_studio_server(8550) is True
def test_an_untimed_record_is_trusted(monkeypatch):
# A legacy `python run.py --port 8901` has no telltale argv, and the in-venv
# path runs in-process. Guessing from the command line rejected real servers.
studio_mod = _studio()
assert studio_mod._pid_is_studio_server(8550) is True
assert studio_mod._pid_is_studio_server(8550, [None]) is True
assert studio_mod._pid_is_studio_server(8550, [111.5, None]) is True
def test_an_unverifiable_record_is_still_stopped(monkeypatch):
# psutil is not a base CLI dependency, so the CLI meets timestamped records it
# cannot check. The old `stop` signalled with no checks at all -- skipping one
@ -239,29 +253,6 @@ def test_pid_identity_check_trusts_the_record_without_psutil(monkeypatch):
assert studio_mod._pid_is_studio_server(8550) is True
def test_pid_identity_check_matches_a_studio_command_line(monkeypatch):
# Legacy records carry no start time, so fall back to the command line -- but
# `unsloth train` and a stray run.py must not match.
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 ["/venv/bin/python", "-m", "unsloth", "train", "run.py"]
def create_time(self):
return 111.5
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_pid_identity_check_uses_the_recorded_start_time(monkeypatch):
studio_mod = _studio()