diff --git a/studio/backend/run.py b/studio/backend/run.py index 5b1be1195b..4a8634da02 100644 --- a/studio/backend/run.py +++ b/studio/backend/run.py @@ -10,7 +10,7 @@ import os import sys import time from pathlib import Path -from typing import Optional, Tuple +from typing import Optional, Sequence, Tuple def _fix_torch_cuda_ld_path(): @@ -827,22 +827,25 @@ def _read_pid_record(path: Path) -> "tuple[int, float | None] | None": return int(lines[0].strip()), created -def _pid_is_studio_backend(pid: int, created: "float | None" = None) -> bool: +def _pid_is_studio_backend(pid: int, created_times: "Sequence[float | 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. + 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 ...`. """ - if created is not None: + known = [c for c in created_times if c is not None] + if known: actual = _process_create_time(pid) # Unknowable without psutil: trust the record rather than guess. - return actual is None or abs(actual - created) < 1.0 + return actual is None or any(abs(actual - c) < 1.0 for c in known) try: import psutil cmdline = " ".join(psutil.Process(pid).cmdline()).lower() except Exception: return True - return "run.py" in cmdline and "studio" in cmdline + 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": @@ -854,9 +857,12 @@ def _blocker_is_own_studio(blockers: "list[tuple[int, str]]") -> "int | None": return None -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]" = {} +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] except OSError: @@ -867,7 +873,7 @@ def _recorded_studio_records() -> "dict[int, float | None]": continue pid, created = record if _pid_alive(pid): - records.setdefault(pid, created) + records.setdefault(pid, []).append(created) elif path != _PID_FILE: path.unlink(missing_ok = True) return records diff --git a/studio/backend/tests/test_studio_pid_files.py b/studio/backend/tests/test_studio_pid_files.py index 2c6139b273..d25b1d5b9a 100644 --- a/studio/backend/tests/test_studio_pid_files.py +++ b/studio/backend/tests/test_studio_pid_files.py @@ -31,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, created = None: True) + monkeypatch.setattr(run, "_pid_is_studio_backend", lambda pid, created_times = (): True) yield @@ -110,7 +110,7 @@ def test_recorded_records_read_per_port_and_legacy_files(tmp_path): (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_records() == {8550: 111.5, 8600: None, 4242: None} + assert run._recorded_studio_records() == {8550: [111.5], 8600: [None], 4242: [None]} def test_recorded_records_ignore_corrupt_files(tmp_path): @@ -151,7 +151,7 @@ def test_our_studio_is_found_behind_a_foreign_listener(tmp_path): 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, created = None: False) + 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 @@ -161,8 +161,42 @@ 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 + assert run._pid_is_studio_backend(8550, [111.5]) is False + assert run._pid_is_studio_backend(8550, [999.0]) is True + + +def test_a_stale_record_does_not_veto_a_live_server_sharing_the_pid(monkeypatch): + # Crash leaves studio-8888-1234.pid, the OS reuses 1234 for a new server on + # another port. Keeping only the first timestamp would reject the live one. + 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(1234, [111.5, 999.0]) is True + 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") + + assert run._recorded_studio_records() == {1234: [111.5, 999.0]} + + +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. + monkeypatch.setattr(run, "_pid_is_studio_backend", _REAL_IS_STUDIO_BACKEND) + + class _FakeProcess: + def __init__(self, pid): + self.pid = pid + + def cmdline(self): + return ["/root/.unsloth/studio/unsloth_studio/bin/unsloth", "studio", "-p", "8901"] + + 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): diff --git a/unsloth_cli/commands/studio.py b/unsloth_cli/commands/studio.py index c55fde8ac9..64fc3da2bd 100644 --- a/unsloth_cli/commands/studio.py +++ b/unsloth_cli/commands/studio.py @@ -19,7 +19,7 @@ import urllib.error import urllib.request from datetime import datetime, timezone from pathlib import Path -from typing import List, Literal, Optional +from typing import List, Literal, Optional, Sequence import typer from unsloth_cli import _studio_deps @@ -2441,14 +2441,15 @@ def _read_pid_record(path: Path) -> "tuple[int, float | None] | 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. +def _pid_file_entries() -> "list[tuple[int, list[float | None], list[Path]]]": + """(pid, create_times, 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. + it mid-shutdown. Every recorded time is kept -- a stale file and a live server + can share a PID, and the stale one must not veto the live one. """ - by_pid: "dict[int, tuple[float | None, list[Path]]]" = {} + by_pid: "dict[int, tuple[list[float | None], list[Path]]]" = {} try: paths = sorted(STUDIO_HOME.glob(PID_FILE_GLOB)) + [_PID_FILE] except OSError: @@ -2464,28 +2465,30 @@ def _pid_file_entries() -> "list[tuple[int, float | None, list[Path]]]": path.unlink(missing_ok = True) continue pid, created = record - known_created, files = by_pid.setdefault(pid, (created, [])) + created_times, files = by_pid.setdefault(pid, ([], [])) + created_times.append(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()] + return [(pid, times, files) for pid, (times, files) in by_pid.items()] -def _pid_is_studio_server(pid: int, created: "float | None" = None) -> bool: +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. Start time pins it exactly; the cmdline check only covers legacy - records, and must not match `unsloth train` or a stray run.py.""" + 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.""" + known = [c for c in created_times if c is not None] try: import psutil proc = psutil.Process(pid) - if created is not None: - return abs(proc.create_time() - created) < 1.0 + if known: + actual = proc.create_time() + return any(abs(actual - c) < 1.0 for c in known) cmdline = " ".join(proc.cmdline()).lower() except Exception: # Unknowable without psutil: trust the record rather than never stopping. return True - return "run.py" in cmdline and "studio" in cmdline + return "studio" in cmdline and ("run.py" in cmdline or "unsloth" in cmdline) def _signal_stop(pid: int) -> "str | None": @@ -2517,8 +2520,8 @@ def stop(): raise typer.Exit(0) signalled, failed = [], [] - for pid, created, paths in entries: - if not _pid_alive(pid) or not _pid_is_studio_server(pid, created): + for pid, created_times, paths in entries: + if not _pid_alive(pid) or not _pid_is_studio_server(pid, created_times): for path in paths: path.unlink(missing_ok = True) continue diff --git a/unsloth_cli/tests/test_studio_stop.py b/unsloth_cli/tests/test_studio_stop.py index 9724f25504..e94f9b8b8e 100644 --- a/unsloth_cli/tests/test_studio_stop.py +++ b/unsloth_cli/tests/test_studio_stop.py @@ -48,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, created = None: True) + monkeypatch.setattr(studio_mod, "_pid_is_studio_server", lambda pid, created_times = (): True) def fake_kill(pid, _sig): killed.append(pid) @@ -104,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, created = None: True) + monkeypatch.setattr(studio_mod, "_pid_is_studio_server", lambda pid, created_times = (): True) killed = [] monkeypatch.setattr(studio_mod.os, "kill", lambda pid, _sig: killed.append(pid)) monkeypatch.setattr(sys, "platform", "linux") @@ -134,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, created = None: False) + monkeypatch.setattr(studio_mod, "_pid_is_studio_server", lambda pid, created_times = (): False) _write_pid(tmp_path, "studio-8901-8550.pid", 8550) result = _run_stop(studio_mod) @@ -144,6 +144,46 @@ def test_stop_does_not_signal_a_reused_pid(monkeypatch, tmp_path): assert not (tmp_path / "studio-8901-8550.pid").exists() +def test_stop_signals_a_live_server_whose_pid_has_a_stale_record(monkeypatch, tmp_path): + # Crash leaves studio-8888-8550.pid, the OS reuses 8550 for a new server on + # another port. The stale timestamp must not veto the live one. + 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-8888-8550.pid").write_text("8550\n111.5", encoding = "utf-8") + (tmp_path / "studio-9000-8550.pid").write_text("8550\n999.0", encoding = "utf-8") + + result = _run_stop(studio_mod) + + assert result.exit_code == 0, result.output + assert killed == [8550] + 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. + studio_mod = _studio() + + class _FakeProcess: + def __init__(self, pid): + self.pid = pid + + def cmdline(self): + return ["/root/.unsloth/studio/unsloth_studio/bin/unsloth", "studio", "-p", "8901"] + + monkeypatch.setitem(sys.modules, "psutil", SimpleNamespace(Process = _FakeProcess)) + + assert studio_mod._pid_is_studio_server(8550) is True + + 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() @@ -187,8 +227,8 @@ def test_pid_identity_check_uses_the_recorded_start_time(monkeypatch): 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 + assert studio_mod._pid_is_studio_server(8550, [111.5]) is True + assert studio_mod._pid_is_studio_server(8550, [999.0]) is False def test_stop_drops_a_record_whose_start_time_no_longer_matches(monkeypatch, tmp_path): @@ -254,7 +294,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, created = None: True) + monkeypatch.setattr(studio_mod, "_pid_is_studio_server", lambda pid, created_times = (): 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) @@ -274,7 +314,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, created = None: True) + monkeypatch.setattr(studio_mod, "_pid_is_studio_server", lambda pid, created_times = (): True) def fake_kill(pid, _sig): if pid == 8550: