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

Match the blocking listener address and stop trusting unverifiable PID records

Browse source
This commit is contained in:
Nilay Yadav 2026-07-29 03:35:39 +05:30
commit 24b53bc251
4 changed files with 146 additions and 32 deletions
Download patch file Download diff file Expand all files Collapse all files

65
studio/backend/run.py
View file

@ -689,27 +689,47 @@ 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."""
def _get_pids_on_port(port: int) -> "list[tuple[int, str, str]]":
"""(pid, name, bind address) for every listener on *port*. A port can have one
per address, and checking only the first makes the own-Studio call arbitrary."""
try:
import psutil
except ImportError:
return []
found: "dict[int, str]" = {}
found: "dict[tuple[int, str], tuple[int, str, 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:
key = (conn.pid, conn.laddr.ip)
if key in found:
continue
try:
found[conn.pid] = psutil.Process(conn.pid).name()
name = psutil.Process(conn.pid).name()
except (psutil.NoSuchProcess, psutil.AccessDenied):
found[conn.pid] = "<unknown>"
name = "<unknown>"
found[key] = (conn.pid, name, conn.laddr.ip)
except (psutil.AccessDenied, OSError) as e:
logger.debug("Failed to scan network connections for port %s: %s", port, e)
return list(found.items())
return list(found.values())
def _listener_blocks_host(listen_ip: str, host: str) -> bool:
"""Would a listener on *listen_ip* block a bind to *host*?
Wildcards on either side collide with everything; otherwise compare resolved
addresses. Unresolvable means assume a collision rather than start a duplicate.
"""
wildcards = ("0.0.0.0", "::", "")
if listen_ip in wildcards or host in wildcards:
return True
import socket
try:
host_ips = {info[4][0] for info in socket.getaddrinfo(host, None)}
except OSError:
return True
return listen_ip in host_ips
def _is_port_free(host: str, port: int) -> bool:
@ -768,7 +788,7 @@ def _find_free_port(
if _is_port_free(host, candidate):
return candidate
if avoid_own_studio:
own = _blocker_is_own_studio(_get_pids_on_port(candidate))
own = _blocker_is_own_studio(_get_pids_on_port(candidate), host)
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}")
@ -836,22 +856,35 @@ def _pid_is_studio_backend(pid: int, created_times: "Sequence[float | None]" = (
in-venv path runs run_server() in-process, so its argv is `unsloth studio ...`.
"""
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)
# Unknowable without psutil: trust the record rather than guess.
return actual is None or any(abs(actual - c) < 1.0 for c in known)
if actual is None:
return untimed
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:
return True
return untimed
return "studio" in cmdline and ("run.py" in cmdline or "unsloth" in cmdline)
def _blocker_is_own_studio(blockers: "list[tuple[int, str]]") -> "int | None":
"""PID of one of our recorded servers holding the port, if any."""
def _blocker_is_own_studio(
blockers: "list[tuple[int, str, str]]", host: "str | None" = None
) -> "int | None":
"""PID of one of our recorded servers actually blocking *host*, if any.
Address-matched: Jupyter on 127.0.0.1:8889 and our server on ::1:8889 is not a
conflict for `-H 127.0.0.1`, and must still fall through to the next port.
"""
recorded = _recorded_studio_records()
for pid, _name in blockers:
for pid, _name, listen_ip in blockers:
if host is not None and not _listener_blocks_host(listen_ip, host):
continue
if pid in recorded and _pid_is_studio_backend(pid, recorded[pid]):
return pid
return None
@ -1691,7 +1724,7 @@ def run_server(
original_port = port
blocker = _get_pid_on_port(port)
# Falling back past our own server is what creates the orphan.
own = _blocker_is_own_studio(_get_pids_on_port(port))
own = _blocker_is_own_studio(_get_pids_on_port(port), host)
if own is not None:
_abort_already_running(own, port)
port = _find_free_port(host, port + 1, avoid_own_studio = True)

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

@ -130,15 +130,15 @@ def test_recorded_records_prune_dead_entries(tmp_path, monkeypatch):
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")]) == 8550
assert run._blocker_is_own_studio([(8550, "python", "127.0.0.1")], "127.0.0.1") == 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 None
assert run._blocker_is_own_studio([]) is None
assert run._blocker_is_own_studio([(117, "jupyter-lab", "127.0.0.1")], "127.0.0.1") is None
assert run._blocker_is_own_studio([], "127.0.0.1") is None
def test_our_studio_is_found_behind_a_foreign_listener(tmp_path):
@ -146,7 +146,12 @@ def test_our_studio_is_found_behind_a_foreign_listener(tmp_path):
# 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
assert (
run._blocker_is_own_studio(
[(117, "jupyter-lab", "127.0.0.1"), (8550, "python", "127.0.0.1")], "127.0.0.1"
)
== 8550
)
def test_a_reused_pid_is_not_treated_as_our_studio(tmp_path, monkeypatch):
@ -154,7 +159,7 @@ def test_a_reused_pid_is_not_treated_as_our_studio(tmp_path, monkeypatch):
monkeypatch.setattr(run, "_pid_is_studio_backend", lambda pid, created_times = (): False)
(tmp_path / "studio-8901-8550.pid").write_text("8550", encoding = "utf-8")
assert run._blocker_is_own_studio([(8550, "postgres")]) is None
assert run._blocker_is_own_studio([(8550, "postgres", "127.0.0.1")], "127.0.0.1") is None
def test_start_time_mismatch_rejects_a_reused_pid(monkeypatch):
@ -218,6 +223,36 @@ def test_legacy_records_do_not_match_a_training_run(monkeypatch):
assert run._pid_is_studio_backend(9999) is False
def test_our_studio_on_another_bind_address_does_not_abort(tmp_path):
# Jupyter holds 127.0.0.1:8889, our server holds ::1:8889. Binding 127.0.0.1
# conflicts with Jupyter, not with us, so fall through to the next port.
(tmp_path / "studio-8889-8550.pid").write_text("8550", encoding = "utf-8")
blockers = [(117, "jupyter-lab", "127.0.0.1"), (8550, "python", "::1")]
assert run._blocker_is_own_studio(blockers, "127.0.0.1") is None
assert run._blocker_is_own_studio(blockers, "::1") == 8550
def test_a_wildcard_listener_blocks_any_bind(tmp_path):
(tmp_path / "studio-8889-8550.pid").write_text("8550", encoding = "utf-8")
assert run._blocker_is_own_studio([(8550, "python", "0.0.0.0")], "127.0.0.1") == 8550
def test_a_wildcard_bind_is_blocked_by_any_listener(tmp_path):
(tmp_path / "studio-8889-8550.pid").write_text("8550", encoding = "utf-8")
assert run._blocker_is_own_studio([(8550, "python", "127.0.0.1")], "0.0.0.0") == 8550
def test_listener_address_matching(monkeypatch):
assert run._listener_blocks_host("0.0.0.0", "127.0.0.1") is True
assert run._listener_blocks_host("127.0.0.1", "0.0.0.0") is True
assert run._listener_blocks_host("127.0.0.1", "127.0.0.1") is True
assert run._listener_blocks_host("::1", "127.0.0.1") is False
def test_fallback_aborts_on_our_own_server_further_up_the_range(tmp_path, monkeypatch):
# jupyter holds 8888, our server holds 8889: skipping to 8890 is the duplicate.
(tmp_path / "studio-8889-8550.pid").write_text("8550", encoding = "utf-8")
@ -225,7 +260,9 @@ def test_fallback_aborts_on_our_own_server_further_up_the_range(tmp_path, monkey
monkeypatch.setattr(
run,
"_get_pids_on_port",
lambda p: [(8550, "python")] if p == 8889 else [(117, "jupyter-lab")],
lambda p: (
[(8550, "python", "127.0.0.1")] if p == 8889 else [(117, "jupyter-lab", "127.0.0.1")]
),
)
with pytest.raises(SystemExit) as excinfo:
@ -236,6 +273,6 @@ def test_fallback_aborts_on_our_own_server_further_up_the_range(tmp_path, monkey
def test_fallback_still_skips_foreign_processes(monkeypatch):
monkeypatch.setattr(run, "_is_port_free", lambda host, p: p >= 8890)
monkeypatch.setattr(run, "_get_pids_on_port", lambda p: [(117, "jupyter-lab")])
monkeypatch.setattr(run, "_get_pids_on_port", lambda p: [(117, "jupyter-lab", "127.0.0.1")])
assert run._find_free_port("127.0.0.1", 8889, avoid_own_studio = True) == 8890

30
unsloth_cli/commands/studio.py
View file

@ -2472,22 +2472,34 @@ 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:
"""Guard against PID reuse: a stale record must not get an unrelated process
killed. Any recorded start time matching is enough. The cmdline check only
covers legacy records -- and the in-venv path runs run_server() in-process, so
its argv is `unsloth studio ...` with no run.py."""
"""Guard against PID reuse: a stale record must not get an unrelated process killed.
A start time only exists if psutil was available when the server started, so
without psutil now such a record cannot be trusted -- but untimed (legacy)
records must keep working, or `stop` stops nothing on a psutil-less install.
The cmdline check covers those: the in-venv path runs run_server() in-process,
so its argv is `unsloth studio ...` with no run.py.
"""
known = [c for c in created_times if c is not None]
untimed = not created_times or len(known) < len(created_times)
try:
import psutil
proc = psutil.Process(pid)
if known:
except Exception:
return untimed
if known:
try:
actual = proc.create_time()
return any(abs(actual - c) < 1.0 for c in known)
except Exception:
return untimed
if any(abs(actual - c) < 1.0 for c in known):
return True
if not untimed:
return False
try:
cmdline = " ".join(proc.cmdline()).lower()
except Exception:
# Unknowable without psutil: trust the record rather than never stopping.
return True
return untimed
return "studio" in cmdline and ("run.py" in cmdline or "unsloth" in cmdline)

32
unsloth_cli/tests/test_studio_stop.py
View file

@ -184,6 +184,38 @@ def test_pid_identity_check_accepts_an_in_process_studio(monkeypatch):
assert studio_mod._pid_is_studio_server(8550) is True
def test_a_timestamped_record_is_not_trusted_without_psutil(monkeypatch):
# psutil is not a base CLI dependency. A start time only exists if psutil was
# present when the server started, so without it now the record is unverifiable
# and must not get a reused PID killed.
studio_mod = _studio()
monkeypatch.setitem(sys.modules, "psutil", None)
assert studio_mod._pid_is_studio_server(8550, [111.5]) is False
assert studio_mod._pid_is_studio_server(8550, [None]) is True
def test_a_legacy_record_survives_a_stale_timestamp_for_the_same_pid(monkeypatch):
# Stale per-port file + live legacy studio.pid sharing a reused PID: judging
# only by the stale timestamp would drop the live server.
studio_mod = _studio()
class _FakeProcess:
def __init__(self, pid):
self.pid = pid
def create_time(self):
return 999.0
def cmdline(self):
return ["/venv/bin/unsloth", "studio", "-p", "8901"]
monkeypatch.setitem(sys.modules, "psutil", SimpleNamespace(Process = _FakeProcess))
assert studio_mod._pid_is_studio_server(8550, [111.5, None]) is True
assert studio_mod._pid_is_studio_server(8550, [111.5]) is False
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()