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

Detect our own server from our own records instead of a psutil listener scan

Browse source
This commit is contained in:
Nilay Yadav 2026-07-29 04:03:10 +05:30
commit 279afd1024
2 changed files with 106 additions and 145 deletions
Download patch file Download diff file Expand all files Collapse all files

127
studio/backend/run.py
View file

@ -689,47 +689,26 @@ def _get_pid_on_port(port: int) -> "tuple[int, str] | None":
return None
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."""
def _bind_address(host: str, port: int) -> str:
"""The address a bind to *host* actually uses -- resolved exactly as
_is_port_free does, so a recorded address and a requested one normalize alike."""
import socket
try:
import psutil
except ImportError:
return []
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
key = (conn.pid, conn.laddr.ip)
if key in found:
continue
try:
name = psutil.Process(conn.pid).name()
except (psutil.NoSuchProcess, psutil.AccessDenied):
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.values())
return socket.getaddrinfo(host, port, socket.AF_UNSPEC, socket.SOCK_STREAM)[0][4][0]
except (OSError, IndexError):
return host
def _listener_blocks_host(listen_ip: str, host: str) -> bool:
"""Would a listener on *listen_ip* block a bind to *host*?
def _addresses_collide(recorded: "str | None", host: str, port: int) -> bool:
"""Would a server bound to *recorded* 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.
Unknown or wildcard on either side collides: refusing with a clear message
beats silently starting a duplicate.
"""
wildcards = ("0.0.0.0", "::", "")
if listen_ip in wildcards or host in wildcards:
if not recorded or recorded 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
return recorded == _bind_address(host, port)
def _is_port_free(host: str, port: int) -> bool:
@ -788,7 +767,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), host)
own = _own_studio_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}")
@ -830,8 +809,8 @@ def _process_create_time(pid: int) -> "float | None":
return None
def _read_pid_record(path: Path) -> "tuple[int, float | None] | None":
"""Parse ``pid`` / optional ``create_time`` from a PID file."""
def _read_pid_record(path: Path) -> "tuple[int, float | None, str | None] | None":
"""Parse ``pid`` / optional ``create_time`` / optional bind address."""
try:
lines = path.read_text(encoding = "utf-8").splitlines()
except (OSError, UnicodeDecodeError):
@ -844,23 +823,23 @@ def _read_pid_record(path: Path) -> "tuple[int, float | None] | None":
created = float(lines[1].strip())
except ValueError:
created = None
return int(lines[0].strip()), created
address = lines[2].strip() if len(lines) > 2 and lines[2].strip() else None
return int(lines[0].strip()), created, address
def _pid_is_studio_backend(pid: int, created_times: "Sequence[float | None]" = ()) -> bool:
"""Guard against PID reuse.
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.
Start time is the reliable signal; any recorded time matching is enough, since
a stale record must not veto a live server that reused the PID. The cmdline
check only covers legacy records that predate the recorded time -- and the
in-venv path runs run_server() in-process, so its argv is `unsloth studio ...`.
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 ...`.
"""
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 untimed
return True if untimed else None
if any(abs(actual - c) < 1.0 for c in known):
return True
if not untimed:
@ -869,47 +848,34 @@ def _pid_is_studio_backend(pid: int, created_times: "Sequence[float | None]" = (
import psutil
cmdline = " ".join(psutil.Process(pid).cmdline()).lower()
except Exception:
return untimed
return True
return "studio" in cmdline and ("run.py" in cmdline or "unsloth" in cmdline)
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.
def _own_studio_on_port(port: int, host: str) -> "int | None":
"""PID of one of our own servers already bound to *port* for *host*.
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.
Reads our own records rather than enumerating listeners: psutil is optional,
and without it a listener scan finds nothing and we silently start a duplicate.
"""
recorded = _recorded_studio_records()
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
def _recorded_studio_records() -> "dict[int, list[float | None]]":
"""Live {pid: [create_time, ...]} under this STUDIO_HOME; prunes dead records.
Every recorded time is kept: a stale file and a live server can share a PID.
"""
records: "dict[int, list[float | None]]" = {}
try:
paths = list(_studio_root().glob(PID_FILE_GLOB)) + [_PID_FILE]
paths = list(_studio_root().glob(f"studio-{port}-*.pid"))
except OSError:
return records
return None
for path in paths:
record = _read_pid_record(path)
if record is None:
continue
pid, created = record
if _pid_alive(pid):
records.setdefault(pid, []).append(created)
elif path != _PID_FILE:
pid, created, address = record
if not _pid_alive(pid):
path.unlink(missing_ok = True)
return records
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:
return pid
return None
def _abort_already_running(pid: int, port: int) -> "NoReturn":
@ -949,16 +915,17 @@ os.environ.setdefault("UNSLOTH_IS_PRESENT", "1")
_OWN_PID_FILE: "Path | None" = None
def _write_pid_file(port: int):
def _write_pid_file(port: int, host: str = ""):
"""Record this PID under its own port so `stop` can find every server."""
global _OWN_PID_FILE
path = _pid_file_for_port(port)
try:
path.parent.mkdir(parents = True, exist_ok = True)
# Start time pins the record to this process, so a reused PID is not
# mistaken for it later.
# Start time pins the record to this process; the bind address tells a
# later launch whether this server would actually block it.
created = _process_create_time(os.getpid())
body = str(os.getpid()) if created is None else f"{os.getpid()}\n{created!r}"
address = _bind_address(host, port) if host else ""
body = f"{os.getpid()}\n{'' if created is None else repr(created)}\n{address}"
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")
@ -1724,7 +1691,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), host)
own = _own_studio_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)
@ -1925,7 +1892,7 @@ def run_server(
(time.perf_counter() - boot_started) * 1000,
)
_write_pid_file(port)
_write_pid_file(port, host)
import atexit
atexit.register(_remove_pid_file)

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

@ -105,53 +105,46 @@ def test_remove_pid_file_leaves_a_reused_entry_alone(tmp_path):
assert own.read_text(encoding = "utf-8") == "999999"
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")
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")
assert run._recorded_studio_records() == {8550: [111.5], 8600: [None], 4242: [None]}
assert run._read_pid_record(tmp_path / "r.pid") == (8550, 111.5, "127.0.0.1")
def test_recorded_records_ignore_corrupt_files(tmp_path):
(tmp_path / "studio-8901-x.pid").write_text("not-a-pid", encoding = "utf-8")
def test_read_pid_record_tolerates_a_bare_pid(tmp_path):
(tmp_path / "r.pid").write_text("8550", encoding = "utf-8")
assert run._recorded_studio_records() == {}
assert run._read_pid_record(tmp_path / "r.pid") == (8550, None, None)
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")
def test_read_pid_record_rejects_a_corrupt_file(tmp_path):
(tmp_path / "r.pid").write_text("not-a-pid", encoding = "utf-8")
assert run._recorded_studio_records() == {}
assert not (tmp_path / "studio-8901-8550.pid").exists()
assert run._read_pid_record(tmp_path / "r.pid") is None
def test_own_studio_blocking_the_port_is_recognised(tmp_path):
(tmp_path / "studio-8901-8550.pid").write_text("8550", encoding = "utf-8")
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.
monkeypatch.setitem(sys.modules, "psutil", None)
(tmp_path / "studio-8901-8550.pid").write_text("8550\n\n127.0.0.1", encoding = "utf-8")
assert run._blocker_is_own_studio([(8550, "python", "127.0.0.1")], "127.0.0.1") == 8550
assert run._own_studio_on_port(8901, "127.0.0.1") == 8550
def test_a_foreign_blocker_still_falls_back(tmp_path):
def test_no_record_for_the_port_means_no_own_studio(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", "127.0.0.1")], "127.0.0.1") is None
assert run._blocker_is_own_studio([], "127.0.0.1") is None
assert run._own_studio_on_port(8888, "127.0.0.1") 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")
def test_own_studio_on_port_prunes_a_dead_record(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._blocker_is_own_studio(
[(117, "jupyter-lab", "127.0.0.1"), (8550, "python", "127.0.0.1")], "127.0.0.1"
)
== 8550
)
assert run._own_studio_on_port(8901, "127.0.0.1") is None
assert not (tmp_path / "studio-8901-8550.pid").exists()
def test_a_reused_pid_is_not_treated_as_our_studio(tmp_path, monkeypatch):
@ -159,7 +152,15 @@ 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", "127.0.0.1")], "127.0.0.1") is None
assert run._own_studio_on_port(8901, "127.0.0.1") is None
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)
(tmp_path / "studio-8901-8550.pid").write_text("8550", encoding = "utf-8")
assert run._own_studio_on_port(8901, "127.0.0.1") == 8550
def test_start_time_mismatch_rejects_a_reused_pid(monkeypatch):
@ -180,11 +181,15 @@ def test_a_stale_record_does_not_veto_a_live_server_sharing_the_pid(monkeypatch)
assert run._pid_is_studio_backend(1234, [111.5, 222.5]) is False
def test_recorded_records_keep_every_timestamp_for_a_pid(tmp_path):
(tmp_path / "studio-8888-1234.pid").write_text("1234\n111.5", encoding = "utf-8")
(tmp_path / "studio-9000-1234.pid").write_text("1234\n999.0", encoding = "utf-8")
def test_a_stale_record_on_another_port_does_not_hide_a_live_server(tmp_path, monkeypatch):
# 1234 was reused: the stale 8888 record must not stop us seeing 9000.
monkeypatch.setattr(run, "_pid_is_studio_backend", _REAL_IS_STUDIO_BACKEND)
monkeypatch.setattr(run, "_process_create_time", lambda pid: 999.0)
(tmp_path / "studio-8888-1234.pid").write_text("1234\n111.5\n", encoding = "utf-8")
(tmp_path / "studio-9000-1234.pid").write_text("1234\n999.0\n", encoding = "utf-8")
assert run._recorded_studio_records() == {1234: [111.5, 999.0]}
assert run._own_studio_on_port(8888, "127.0.0.1") is None
assert run._own_studio_on_port(9000, "127.0.0.1") == 1234
def test_legacy_records_match_an_in_process_studio(monkeypatch):
@ -224,46 +229,35 @@ def test_legacy_records_do_not_match_a_training_run(monkeypatch):
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")
# Our server holds ::1:8889; binding 127.0.0.1:8889 is not a conflict with us,
# so fall through to the next port instead of refusing.
(tmp_path / "studio-8889-8550.pid").write_text("8550\n\n::1", 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
assert run._own_studio_on_port(8889, "127.0.0.1") is None
assert run._own_studio_on_port(8889, "::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_address_matching(tmp_path):
assert run._addresses_collide("0.0.0.0", "127.0.0.1", 8889) is True
assert run._addresses_collide("127.0.0.1", "0.0.0.0", 8889) is True
assert run._addresses_collide("127.0.0.1", "127.0.0.1", 8889) is True
assert run._addresses_collide("::1", "127.0.0.1", 8889) is False
# An unrecorded address is unknown, so assume a conflict.
assert run._addresses_collide(None, "127.0.0.1", 8889) is True
def test_a_wildcard_bind_is_blocked_by_any_listener(tmp_path):
(tmp_path / "studio-8889-8550.pid").write_text("8550", encoding = "utf-8")
def test_a_hostname_resolves_the_same_way_the_bind_does(tmp_path):
# `localhost` and the address _is_port_free actually binds must agree, or a
# recorded server is missed and a duplicate starts.
recorded = run._bind_address("localhost", 8889)
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
assert run._addresses_collide(recorded, "localhost", 8889) is True
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")
(tmp_path / "studio-8889-8550.pid").write_text("8550\n\n127.0.0.1", encoding = "utf-8")
monkeypatch.setattr(run, "_is_port_free", lambda host, p: p >= 8890)
monkeypatch.setattr(
run,
"_get_pids_on_port",
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:
run._find_free_port("127.0.0.1", 8889, avoid_own_studio = True)
@ -271,8 +265,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(monkeypatch):
def test_fallback_still_skips_foreign_processes(tmp_path, monkeypatch):
# No record for 8889, so the blocker is not ours: keep falling back.
monkeypatch.setattr(run, "_is_port_free", lambda host, p: p >= 8890)
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