Keep every recorded start time per PID and accept in-process Studio servers
This commit is contained in:
Nilay Yadav
parent
3fbd60fa5f
commit
e897c025c9
4 changed files with 123 additions and 40 deletions
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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):
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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:
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue