From 4d472ed392ef3c9dd6508b06dd8fbdc46f0661c7 Mon Sep 17 00:00:00 2001 From: Nilay Yadav Date: Wed, 29 Jul 2026 05:37:01 +0530 Subject: [PATCH] Drop the command-line guess, fix Windows liveness, and free the PID record last --- studio/backend/run.py | 55 +++++++----- studio/backend/tests/test_studio_pid_files.py | 89 ++++++++++++++----- unsloth_cli/commands/studio.py | 27 +++--- unsloth_cli/tests/test_studio_stop.py | 43 ++++----- 4 files changed, 127 insertions(+), 87 deletions(-) diff --git a/studio/backend/run.py b/studio/backend/run.py index 8e28e0140a..e18c49e01a 100644 --- a/studio/backend/run.py +++ b/studio/backend/run.py @@ -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") diff --git a/studio/backend/tests/test_studio_pid_files.py b/studio/backend/tests/test_studio_pid_files.py index 6218cb4beb..4b35dae6d0 100644 --- a/studio/backend/tests/test_studio_pid_files.py +++ b/studio/backend/tests/test_studio_pid_files.py @@ -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): diff --git a/unsloth_cli/commands/studio.py b/unsloth_cli/commands/studio.py index e064ae8c89..7d0f8e4764 100644 --- a/unsloth_cli/commands/studio.py +++ b/unsloth_cli/commands/studio.py @@ -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": diff --git a/unsloth_cli/tests/test_studio_stop.py b/unsloth_cli/tests/test_studio_stop.py index 6a022be7e8..a50c0f5b5b 100644 --- a/unsloth_cli/tests/test_studio_stop.py +++ b/unsloth_cli/tests/test_studio_stop.py @@ -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()