From 279afd1024390f915b034c85ce7133d47d8aaf52 Mon Sep 17 00:00:00 2001 From: Nilay Yadav Date: Wed, 29 Jul 2026 04:03:10 +0530 Subject: [PATCH] Detect our own server from our own records instead of a psutil listener scan --- studio/backend/run.py | 127 +++++++----------- studio/backend/tests/test_studio_pid_files.py | 124 ++++++++--------- 2 files changed, 106 insertions(+), 145 deletions(-) diff --git a/studio/backend/run.py b/studio/backend/run.py index c58e2381a1..b85b01a81b 100644 --- a/studio/backend/run.py +++ b/studio/backend/run.py @@ -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 = "" - 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) diff --git a/studio/backend/tests/test_studio_pid_files.py b/studio/backend/tests/test_studio_pid_files.py index 2df342f381..87e248c5b3 100644 --- a/studio/backend/tests/test_studio_pid_files.py +++ b/studio/backend/tests/test_studio_pid_files.py @@ -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