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

Pin PID records to process start time and check every listener on a port

Browse source
This commit is contained in:
Nilay Yadav 2026-07-29 02:59:31 +05:30
commit 3fbd60fa5f
4 changed files with 264 additions and 75 deletions
Download patch file Download diff file Expand all files Collapse all files

123
studio/backend/run.py
View file

@ -689,6 +689,29 @@ def _get_pid_on_port(port: int) -> "tuple[int, str] | None":
return None
def _get_pids_on_port(port: int) -> "list[tuple[int, str]]":
"""Every listener on *port*. A port can have one per bind address, and
checking only the first makes the own-Studio decision arbitrary."""
try:
import psutil
except ImportError:
return []
found: "dict[int, str]" = {}
try:
for conn in psutil.net_connections(kind = "tcp"):
if conn.status != "LISTEN" or conn.laddr.port != port or conn.pid is None:
continue
if conn.pid in found:
continue
try:
found[conn.pid] = psutil.Process(conn.pid).name()
except (psutil.NoSuchProcess, psutil.AccessDenied):
found[conn.pid] = "<unknown>"
except (psutil.AccessDenied, OSError) as e:
logger.debug("Failed to scan network connections for port %s: %s", port, e)
return list(found.items())
def _is_port_free(host: str, port: int) -> bool:
"""Check if a port is available for binding.
@ -745,9 +768,9 @@ def _find_free_port(
if _is_port_free(host, candidate):
return candidate
if avoid_own_studio:
blocker = _get_pid_on_port(candidate)
if _blocker_is_own_studio(blocker):
_abort_already_running(blocker[0], candidate)
own = _blocker_is_own_studio(_get_pids_on_port(candidate))
if own is not None:
_abort_already_running(own, candidate)
raise RuntimeError(f"Could not find a free port in range {start}-{start + max_attempts - 1}")
@ -779,43 +802,75 @@ def _pid_alive(pid: int) -> bool:
return True
def _pid_is_studio_backend(pid: int) -> bool:
"""Guard against PID reuse: a stale record must not block an unrelated process."""
def _process_create_time(pid: int) -> "float | None":
try:
import psutil
return psutil.Process(pid).create_time()
except Exception:
return None
def _read_pid_record(path: Path) -> "tuple[int, float | None] | None":
"""Parse ``pid`` / optional ``create_time`` from a PID file."""
try:
lines = path.read_text(encoding = "utf-8").splitlines()
except (OSError, UnicodeDecodeError):
return None
if not lines or not lines[0].strip().isdigit():
return None
created = None
if len(lines) > 1:
try:
created = float(lines[1].strip())
except ValueError:
created = None
return int(lines[0].strip()), created
def _pid_is_studio_backend(pid: int, created: "float | None" = None) -> bool:
"""Guard against PID reuse.
Start time is the reliable signal; the cmdline check only covers legacy
records that predate it, and must not match `unsloth train` or a stray run.py.
"""
if created is not None:
actual = _process_create_time(pid)
# Unknowable without psutil: trust the record rather than guess.
return actual is None or abs(actual - created) < 1.0
try:
import psutil
cmdline = " ".join(psutil.Process(pid).cmdline()).lower()
except Exception:
return True
return "run.py" in cmdline or "unsloth" in cmdline
return "run.py" in cmdline and "studio" in cmdline
def _blocker_is_own_studio(blocker: "tuple[int, str] | None") -> bool:
"""True when the process holding the port is a server we recorded."""
if not blocker:
return False
return blocker[0] in _recorded_studio_pids() and _pid_is_studio_backend(blocker[0])
def _blocker_is_own_studio(blockers: "list[tuple[int, str]]") -> "int | None":
"""PID of one of our recorded servers holding the port, if any."""
recorded = _recorded_studio_records()
for pid, _name in blockers:
if pid in recorded and _pid_is_studio_backend(pid, recorded[pid]):
return pid
return None
def _recorded_studio_pids() -> "set[int]":
"""Live PIDs recorded under this STUDIO_HOME; prunes dead records."""
pids: "set[int]" = set()
def _recorded_studio_records() -> "dict[int, float | None]":
"""Live {pid: create_time} recorded under this STUDIO_HOME; prunes dead records."""
records: "dict[int, float | None]" = {}
try:
paths = list(_studio_root().glob(PID_FILE_GLOB)) + [_PID_FILE]
except OSError:
return pids
return records
for path in paths:
try:
text = path.read_text(encoding = "utf-8").strip()
except (OSError, UnicodeDecodeError):
record = _read_pid_record(path)
if record is None:
continue
if not text.isdigit():
continue
pid = int(text)
pid, created = record
if _pid_alive(pid):
pids.add(pid)
records.setdefault(pid, created)
elif path != _PID_FILE:
path.unlink(missing_ok = True)
return pids
return records
def _abort_already_running(pid: int, port: int) -> "NoReturn":
@ -861,8 +916,12 @@ def _write_pid_file(port: int):
path = _pid_file_for_port(port)
try:
path.parent.mkdir(parents = True, exist_ok = True)
path.write_text(str(os.getpid()), encoding = "utf-8")
# An older CLI's `stop` only reads this one.
# Start time pins the record to this process, so a reused PID is not
# mistaken for it later.
created = _process_create_time(os.getpid())
body = str(os.getpid()) if created is None else f"{os.getpid()}\n{created!r}"
path.write_text(body, encoding = "utf-8")
# An older CLI's `stop` only reads this one, and expects a bare PID.
_PID_FILE.write_text(str(os.getpid()), encoding = "utf-8")
except OSError:
return
@ -874,11 +933,12 @@ def _remove_pid_file():
if _OWN_PID_FILE is None:
return
for path in (_OWN_PID_FILE, _PID_FILE):
try:
if path.is_file() and path.read_text(encoding = "utf-8").strip() == str(os.getpid()):
record = _read_pid_record(path) if path.is_file() else None
if record is not None and record[0] == os.getpid():
try:
path.unlink(missing_ok = True)
except (OSError, UnicodeDecodeError):
pass
except OSError:
pass
def _graceful_shutdown(server = None):
@ -1625,8 +1685,9 @@ def run_server(
original_port = port
blocker = _get_pid_on_port(port)
# Falling back past our own server is what creates the orphan.
if _blocker_is_own_studio(blocker):
_abort_already_running(blocker[0], port)
own = _blocker_is_own_studio(_get_pids_on_port(port))
if own is not None:
_abort_already_running(own, port)
port = _find_free_port(host, port + 1, avoid_own_studio = True)
if not silent:
print("")

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

@ -11,6 +11,7 @@ from __future__ import annotations
import os
import sys
from pathlib import Path
from types import SimpleNamespace
import pytest
@ -20,6 +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.
_REAL_IS_STUDIO_BACKEND = run._pid_is_studio_backend
@pytest.fixture(autouse = True)
def isolated_root(tmp_path, monkeypatch):
@ -27,7 +31,7 @@ def isolated_root(tmp_path, monkeypatch):
monkeypatch.setattr(run, "_PID_FILE", tmp_path / "studio.pid")
monkeypatch.setattr(run, "_OWN_PID_FILE", None)
monkeypatch.setattr(run, "_pid_alive", lambda pid: True)
monkeypatch.setattr(run, "_pid_is_studio_backend", lambda pid: True)
monkeypatch.setattr(run, "_pid_is_studio_backend", lambda pid, created = None: True)
yield
@ -35,15 +39,29 @@ def _files(tmp_path):
return sorted(p.name for p in tmp_path.glob("studio-*.pid"))
def _pid_of(path):
return path.read_text(encoding = "utf-8").splitlines()[0]
def test_write_pid_file_records_port_and_pid(tmp_path):
run._write_pid_file(8901)
assert _files(tmp_path) == [f"studio-8901-{os.getpid()}.pid"]
assert (tmp_path / f"studio-8901-{os.getpid()}.pid").read_text() == str(os.getpid())
assert _pid_of(tmp_path / f"studio-8901-{os.getpid()}.pid") == str(os.getpid())
def test_write_pid_file_also_updates_the_legacy_file(tmp_path):
# An older CLI's `stop` only reads studio.pid.
def test_write_pid_file_records_the_start_time(tmp_path):
# Pins the record to this process, so a reused PID isn't mistaken for it.
run._write_pid_file(8901)
record = run._read_pid_record(tmp_path / f"studio-8901-{os.getpid()}.pid")
assert record[0] == os.getpid()
assert record[1] == pytest.approx(run._process_create_time(os.getpid()))
def test_write_pid_file_keeps_the_legacy_file_a_bare_pid(tmp_path):
# An older CLI's `stop` reads studio.pid and expects only digits.
run._write_pid_file(8901)
assert (tmp_path / "studio.pid").read_text(encoding = "utf-8") == str(os.getpid())
@ -54,7 +72,7 @@ def test_second_port_does_not_clobber_the_first(tmp_path):
run._write_pid_file(8902)
assert (tmp_path / "studio-8901-8550.pid").read_text(encoding = "utf-8") == "8550"
assert _pid_of(tmp_path / "studio-8901-8550.pid") == "8550"
assert (tmp_path / f"studio-8902-{os.getpid()}.pid").exists()
@ -87,48 +105,83 @@ def test_remove_pid_file_leaves_a_reused_entry_alone(tmp_path):
assert own.read_text(encoding = "utf-8") == "999999"
def test_recorded_studio_pids_reads_per_port_and_legacy_files(tmp_path):
(tmp_path / "studio-8901-8550.pid").write_text("8550", encoding = "utf-8")
def test_recorded_records_read_per_port_and_legacy_files(tmp_path):
(tmp_path / "studio-8901-8550.pid").write_text("8550\n111.5", encoding = "utf-8")
(tmp_path / "studio-8902-8600.pid").write_text("8600", encoding = "utf-8")
(tmp_path / "studio.pid").write_text("4242", encoding = "utf-8")
assert run._recorded_studio_pids() == {8550, 8600, 4242}
assert run._recorded_studio_records() == {8550: 111.5, 8600: None, 4242: None}
def test_recorded_studio_pids_ignores_corrupt_files(tmp_path):
def test_recorded_records_ignore_corrupt_files(tmp_path):
(tmp_path / "studio-8901-x.pid").write_text("not-a-pid", encoding = "utf-8")
assert run._recorded_studio_pids() == set()
assert run._recorded_studio_records() == {}
def test_recorded_studio_pids_prunes_dead_records(tmp_path, monkeypatch):
def test_recorded_records_prune_dead_entries(tmp_path, monkeypatch):
monkeypatch.setattr(run, "_pid_alive", lambda pid: False)
(tmp_path / "studio-8901-8550.pid").write_text("8550", encoding = "utf-8")
assert run._recorded_studio_pids() == set()
assert run._recorded_studio_records() == {}
assert not (tmp_path / "studio-8901-8550.pid").exists()
def test_own_studio_blocking_the_port_is_recognised(tmp_path):
(tmp_path / "studio-8901-8550.pid").write_text("8550", encoding = "utf-8")
assert run._blocker_is_own_studio((8550, "python")) is True
assert run._blocker_is_own_studio([(8550, "python")]) == 8550
def test_a_foreign_blocker_still_falls_back(tmp_path):
# jupyter-lab on 8888 must keep the fallback, not abort the launch.
(tmp_path / "studio-8901-8550.pid").write_text("8550", encoding = "utf-8")
assert run._blocker_is_own_studio((117, "jupyter-lab")) is False
assert run._blocker_is_own_studio(None) is False
assert run._blocker_is_own_studio([(117, "jupyter-lab")]) is None
assert run._blocker_is_own_studio([]) is None
def test_our_studio_is_found_behind_a_foreign_listener(tmp_path):
# One port, two bind addresses: checking only the first listener made the
# decision depend on psutil's ordering.
(tmp_path / "studio-8889-8550.pid").write_text("8550", encoding = "utf-8")
assert run._blocker_is_own_studio([(117, "jupyter-lab"), (8550, "python")]) == 8550
def test_a_reused_pid_is_not_treated_as_our_studio(tmp_path, monkeypatch):
# Stale record + the OS handing that PID to something else must not abort.
monkeypatch.setattr(run, "_pid_is_studio_backend", lambda pid: False)
monkeypatch.setattr(run, "_pid_is_studio_backend", lambda pid, created = None: False)
(tmp_path / "studio-8901-8550.pid").write_text("8550", encoding = "utf-8")
assert run._blocker_is_own_studio((8550, "postgres")) is False
assert run._blocker_is_own_studio([(8550, "postgres")]) is None
def test_start_time_mismatch_rejects_a_reused_pid(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, created = 111.5) is False
assert run._pid_is_studio_backend(8550, created = 999.0) 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`.
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))
assert run._pid_is_studio_backend(8550) is True
assert run._pid_is_studio_backend(9999) is False
def test_fallback_aborts_on_our_own_server_further_up_the_range(tmp_path, monkeypatch):
@ -136,7 +189,9 @@ def test_fallback_aborts_on_our_own_server_further_up_the_range(tmp_path, monkey
(tmp_path / "studio-8889-8550.pid").write_text("8550", encoding = "utf-8")
monkeypatch.setattr(run, "_is_port_free", lambda host, p: p >= 8890)
monkeypatch.setattr(
run, "_get_pid_on_port", lambda p: (8550, "python") if p == 8889 else (117, "jupyter-lab")
run,
"_get_pids_on_port",
lambda p: [(8550, "python")] if p == 8889 else [(117, "jupyter-lab")],
)
with pytest.raises(SystemExit) as excinfo:
@ -145,8 +200,8 @@ def test_fallback_aborts_on_our_own_server_further_up_the_range(tmp_path, monkey
assert excinfo.value.code == 1
def test_fallback_still_skips_foreign_processes(tmp_path, monkeypatch):
def test_fallback_still_skips_foreign_processes(monkeypatch):
monkeypatch.setattr(run, "_is_port_free", lambda host, p: p >= 8890)
monkeypatch.setattr(run, "_get_pid_on_port", lambda p: (117, "jupyter-lab"))
monkeypatch.setattr(run, "_get_pids_on_port", lambda p: [(117, "jupyter-lab")])
assert run._find_free_port("127.0.0.1", 8889, avoid_own_studio = True) == 8890

60
unsloth_cli/commands/studio.py
View file

@ -2424,14 +2424,31 @@ def _pid_alive(pid: int) -> bool:
return True
def _pid_file_entries() -> "list[tuple[int, list[Path]]]":
"""(pid, files) per recorded server, including the legacy studio.pid.
def _read_pid_record(path: Path) -> "tuple[int, float | None] | None":
"""Parse ``pid`` / optional ``create_time`` from a PID file."""
try:
lines = path.read_text(encoding = "utf-8").splitlines()
except (OSError, UnicodeDecodeError):
return None
if not lines or not lines[0].strip().isdigit():
return None
created = None
if len(lines) > 1:
try:
created = float(lines[1].strip())
except ValueError:
created = None
return int(lines[0].strip()), created
def _pid_file_entries() -> "list[tuple[int, float | None, list[Path]]]":
"""(pid, create_time, files) per recorded server, including the legacy studio.pid.
Grouped by PID: a server writes both its per-port file and studio.pid, and
signalling twice would hit the SIG_DFL the first SIGTERM installs, hard-killing
it mid-shutdown.
"""
by_pid: "dict[int, list[Path]]" = {}
by_pid: "dict[int, tuple[float | None, list[Path]]]" = {}
try:
paths = sorted(STUDIO_HOME.glob(PID_FILE_GLOB)) + [_PID_FILE]
except OSError:
@ -2441,27 +2458,34 @@ def _pid_file_entries() -> "list[tuple[int, list[Path]]]":
if path in seen or not path.is_file():
continue
seen.add(path)
try:
text = path.read_text(encoding = "utf-8").strip()
except (OSError, UnicodeDecodeError):
continue
if text.isdigit():
by_pid.setdefault(int(text), []).append(path)
else:
typer.echo(f"Ignoring invalid PID file {path.name}: {text}")
record = _read_pid_record(path)
if record is None:
typer.echo(f"Ignoring invalid PID file {path.name}")
path.unlink(missing_ok = True)
return list(by_pid.items())
continue
pid, created = record
known_created, files = by_pid.setdefault(pid, (created, []))
files.append(path)
if known_created is None and created is not None:
by_pid[pid] = (created, files)
return [(pid, created, files) for pid, (created, files) in by_pid.items()]
def _pid_is_studio_server(pid: int) -> bool:
def _pid_is_studio_server(pid: int, created: "float | None" = None) -> bool:
"""Guard against PID reuse: a stale record must not get an unrelated process
killed. Unknowable without psutil, so trust the record there."""
killed. Start time pins it exactly; the cmdline check only covers legacy
records, and must not match `unsloth train` or a stray run.py."""
try:
import psutil
cmdline = " ".join(psutil.Process(pid).cmdline()).lower()
proc = psutil.Process(pid)
if created is not None:
return abs(proc.create_time() - created) < 1.0
cmdline = " ".join(proc.cmdline()).lower()
except Exception:
# Unknowable without psutil: trust the record rather than never stopping.
return True
return "run.py" in cmdline or "unsloth" in cmdline
return "run.py" in cmdline and "studio" in cmdline
def _signal_stop(pid: int) -> "str | None":
@ -2493,8 +2517,8 @@ def stop():
raise typer.Exit(0)
signalled, failed = [], []
for pid, paths in entries:
if not _pid_alive(pid) or not _pid_is_studio_server(pid):
for pid, created, paths in entries:
if not _pid_alive(pid) or not _pid_is_studio_server(pid, created):
for path in paths:
path.unlink(missing_ok = True)
continue

61
unsloth_cli/tests/test_studio_stop.py
View file

@ -27,6 +27,10 @@ def _studio():
return _studio_mod
# Captured before _install stubs it, for the tests that exercise it.
_REAL_IS_STUDIO_SERVER = _studio()._pid_is_studio_server
def _install(
monkeypatch,
tmp_path,
@ -44,7 +48,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)
monkeypatch.setattr(studio_mod, "_pid_is_studio_server", lambda pid, created = None: True)
def fake_kill(pid, _sig):
killed.append(pid)
@ -100,7 +104,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)
monkeypatch.setattr(studio_mod, "_pid_is_studio_server", lambda pid, created = None: True)
killed = []
monkeypatch.setattr(studio_mod.os, "kill", lambda pid, _sig: killed.append(pid))
monkeypatch.setattr(sys, "platform", "linux")
@ -130,7 +134,7 @@ 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)
monkeypatch.setattr(studio_mod, "_pid_is_studio_server", lambda pid, created = None: False)
_write_pid(tmp_path, "studio-8901-8550.pid", 8550)
result = _run_stop(studio_mod)
@ -149,6 +153,8 @@ def test_pid_identity_check_trusts_the_record_without_psutil(monkeypatch):
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:
@ -158,7 +164,10 @@ def test_pid_identity_check_matches_a_studio_command_line(monkeypatch):
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"]
return ["/venv/bin/python", "-m", "unsloth", "train", "run.py"]
def create_time(self):
return 111.5
monkeypatch.setitem(sys.modules, "psutil", SimpleNamespace(Process = _FakeProcess))
@ -166,6 +175,46 @@ def test_pid_identity_check_matches_a_studio_command_line(monkeypatch):
assert studio_mod._pid_is_studio_server(9999) is False
def test_pid_identity_check_uses_the_recorded_start_time(monkeypatch):
studio_mod = _studio()
class _FakeProcess:
def __init__(self, pid):
self.pid = pid
def create_time(self):
return 111.5
monkeypatch.setitem(sys.modules, "psutil", SimpleNamespace(Process = _FakeProcess))
assert studio_mod._pid_is_studio_server(8550, created = 111.5) is True
assert studio_mod._pid_is_studio_server(8550, created = 999.0) is False
def test_stop_drops_a_record_whose_start_time_no_longer_matches(monkeypatch, tmp_path):
# The PID was reused: same number, different process.
studio_mod, _live, killed = _install(monkeypatch, tmp_path, alive = {8550})
monkeypatch.setattr(studio_mod, "_pid_is_studio_server", _REAL_IS_STUDIO_SERVER)
class _FakeProcess:
def __init__(self, pid):
self.pid = pid
def create_time(self):
return 999.0
monkeypatch.setitem(sys.modules, "psutil", SimpleNamespace(Process = _FakeProcess))
(tmp_path / "studio-8901-8550.pid").write_text("8550\n111.5", encoding = "utf-8")
result = _run_stop(studio_mod)
assert result.exit_code == 0, result.output
assert killed == []
assert not (tmp_path / "studio-8901-8550.pid").exists()
# Dropped for the start-time mismatch, not because the record looked corrupt.
assert "invalid pid file" not in result.output.lower()
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)
@ -205,7 +254,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, "_pid_is_studio_server", lambda pid, created = None: 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)
@ -225,7 +274,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)
monkeypatch.setattr(studio_mod, "_pid_is_studio_server", lambda pid, created = None: True)
def fake_kill(pid, _sig):
if pid == 8550: